在 AI for Science 平台落地异构调度从「数据搬运工」到「算力调度器」的架构重构实录上周有个需求砸下来院里新建的 AI for Science 协作平台要上线对外开放支撑分子动力学MD长周期模拟与大模型预训练混部。测试环境跑通的流程一到生产环境就炸——不是 OOM 就是存储 IO 打满研究员提个任务排队 6 小时起步跑俩小时挂了还得从头来过。项目背景简单交代上海某研究院自建算力集群约 200 张 H800 500 张 A800存储端挂着 3PB CephFSCeph 18.2.2 Quincy和 200TB Alluxio 缓存层。技术栈锁定在 Kubernetes 1.29.4、Slurm 23.11.8、KubeRay 1.1.0、Spring Boot 3.3.2JDK 21.0.3、Fluid 0.9.0Alluxio Runtime 2.9.3。核心痛点不在模型算法而在异构负载下的数据亲和性调度与有状态长任务的生存保障。背景混部场景下的「三高」噩梦先说业务画像。分子动力学典型作业单任务跑 7×24 小时甚至数周MPI 进程间通信极其敏感中断一次等于白跑Checkpoint 文件动辄 2TB大模型预训练则是 Gang Scheduling 启动PyTorch DDP NCCL对带宽抖动零容忍但支持 Tensor Parallel 切分、支持频繁torch.save。这俩负载混在同一套 K8s 集群里还要套一层多租户配额研究组级别Slurm 管着裸金属 MPI 任务K8s 管着容器化训练任务中间隔着一层自研的统一提交网关。测试阶段没暴露问题因为测试数据集小、并发低、无抢占。真实上线第一周三个指标扎眼跨节点数据拉取占任务总耗时 68%Alluxio 缓存命中率仅 12%万兆网卡跑满却喂不饱 GPU。Slurm 抢占 K8s 任务导致 PyTorchJob 重启率 43%无法恢复研究员直接在群里骂娘。配额失效A 组提交 50 张卡的 RayJobB 组 8 张卡的 MD 任务排队饿死Kueue 的 ClusterQueue 配置在 Gang Scheduling 面前形同虚设。过程把「数据搬到算力」改成「算力贴着数据转」1. 调度层拆解别让 K8s 调度器干不擅长的事最早想用 K8s 原生kube-schedulerSchedulerPlugin解决数据亲和性写了个PreFilter插件去查 FluidDatasetCR 的status.nodes字段再给 Pod 打nodeAffinity。压测发现调度延迟 P99 从 200ms 飙到 3.4s因为每个 Pod 创建都要 List 所有 Dataset CR、计算拓扑距离、打分排序调度器主线程被卡死。换路子了。引入Scheduler ExtenderHTTP 扩展器做「粗筛」只在Filter阶段调用逻辑下沉到独立 Go 服务scheduler-extender-v0.3.1缓存 Dataset 节点分布在本地 RedisValkey 7.2.5单次请求 15ms 以内。K8s 原生插件只保留Score阶段做微调。java// Spring Boot 网关侧任务提交时自动注入数据亲和性标签不依赖用户手填 YAMLRestControllerRequestMapping(/api/v1/jobs)RequiredArgsConstructorpublic class JobSubmitController {private final FluidDatasetClient fluidClient; // Feign 调用 Fluid Controllerprivate final SchedulerExtenderClient extenderClient;private final KubernetesClient k8sClient;PostMappingpublic ResponseEntity submit(Valid RequestBody JobSubmitRequest req) {// 1. 解析数据集引用查询 Fluid 缓存分布List datasetNames req.getDatasetRefs();Map datasetNodeMap fluidClient.getCachedNodes(datasetNames);if (datasetNodeMap.isEmpty()) {throw new BusinessException(DATASET_NOT_CACHED, 目标数据集未预热至任何节点请先提交预热任务);}// 2. 计算最优节点集合交集优先并集兜底Set candidateNodes computeOptimalNodes(datasetNodeMap, req.getGpuCount());// 3. 构建 PodSpec注入 nodeAffinity tolerations专用节点池污点Pod jobPod buildPodSpec(req, candidateNodes);// 4. 根据负载画像分发MPI - Slurm OperatorDDP/Ray - Kueue K8s JobSetif (req.getWorkloadType() WorkloadType.MPI_MD) {return submitToSlurm(jobPod, req);}return submitToKueue(jobPod, req);}private Set computeOptimalNodes(Map datasetNodeMap, int gpuCount) {// 简单启发式求交集不足则按缓存完整度评分取 Top-Nreturn datasetNodeMap.values().stream().reduce((a, b) - a.stream().filter(b::contains).collect(Collectors.toList())).orElseGet(() - datasetNodeMap.values().stream().flatMap(List::stream).collect(Collectors.groupingBy(Function.identity(), Collectors.counting())).entrySet().stream().sorted(Map.Entry.comparingByValue().reversed()).limit(gpuCount / 8 2) // 每节点 8 卡冗余 2 个节点.map(Map.Entry::getKey).collect(Collectors.toSet()));}} 这个方案虽然官方文档不推荐 Extender维护成本高但在我们「数据集固定、节点拓扑变化慢」的场景下比插件模式稳得多。插件模式适合通用调度器开发我们是业务侧深度定制别被「最佳实践」绑架。2. 抢占与生存CRIU 在 GPU 场景是个坑应用层 Checkpoint 才是正解Slurm 抢占 K8s 任务最初想用 CRIU (3.17.1) 做容器级热迁移/检查点。实测发现CRIU 根本 Dump 不下来显存状态cuda-checkpoint只能存 host 侧内存恢复时 NCCL Communicator 直接炸报ncclInternalError。折腾两周放弃。转而强制要求业务侧实现应用层 Checkpoint 接口平台侧只负责「发信号、存制品、拉起新实例」。yamlKueue JobSet 模板定义 PreStop Hook 共享 PVC 挂载 Checkpoint 目录apiVersion: jobset.x-k8s.io/v1alpha2kind: JobSetmetadata:name: ${{JOB_NAME}}spec:replicatedJobs:name: workerreplicas: ${{GPU_COUNT}}template:spec:serviceAccountName: jobset-sarestartPolicy: OnFailurevolumes:name: ckpt-volumepersistentVolumeClaim:claimName: ${{PVC_NAME}}containers:name: trainerimage: ${{IMAGE}}command: [bash, -c, torchrun --standalone --nnodes1 --nproc_per_node8 train.py --ckpt-dir /ckpt]volumeMounts:mountPath: /ckptname: ckpt-volumelifecycle:preStop:exec:command: [/bin/sh, -c, touch /ckpt/SIG_CHECKPOINT sleep 120] # 留 2 分钟刷盘env:name: PLATFORM_CHECKPOINT_SIGNALvalue: SIGUSR1 # 应用层监听信号触发 torch.save平台侧控制器job-lifecycle-controller v0.4.0监听PodDeletionTimestamp发现被抢占立即Patch PodterminationGracePeriodSeconds180防止默认 30s 截断。向容器发SIGUSR1业务代码捕获后执行torch.save({model: ..., optimizer: ..., rng: ...}, f/ckpt/ckpt_rank{dist.get_rank()}_step{step}.pt)。等待SIG_CHECKPOINT文件出现或超时再允许 Pod 彻底终止。重建 JobSet 时initContainer自动探测/ckpt最新 step传入--resume-step启动。效果MD 任务虽不支持热迁移但应用层每 30 分钟落盘一次PLUMEDRESTART关键字抢占恢复仅丢 ≤30 分钟算力研究员能接受。3. 多租户配额Kueue Slurm 统一视图的「伪全局调度」Kueue 管不了 SlurmSlurm 管不了 K8s。搞了个中间层quota-syncerGo 协程每 10s 跑一次从 Slurmsacctmgr show assoc formataccount,GrpTRESMins拿各账户已用 GPU 分钟数。从 KueueClusterQueue.status.flavors[].used拿 K8s 侧占用。汇总写入 Redisquota:{org_id}:used_gpu_minutes提交网关准入时只看这键。go// quota-syncer 核心同步逻辑片段func (s *Syncer) syncLoop() {ticker : time.NewTicker(10 * time.Second)for range ticker.C {slurmUsage, _ : s.slurmClient.GetAccountUsage() // map[org]float64 (gpu-min)kueueUsage, _ : s.kueueClient.GetClusterQueueUsage() // map[org]float64merged : make(map[string]float64)for k, v : range slurmUsage { merged[k] v }for k, v : range kueueUsage { merged[k] v }pipe : s.redis.Pipeline()for org, used : range merged {limit : s.getQuotaLimit(org) // 从 CM 读配额上限pipe.Set(ctx, fmt.Sprintf(quota:%s:used, org), used, 0)pipe.Set(ctx, fmt.Sprintf(quota:%s:limit, org), limit, 0)// 提交网关 Lua 脚本原子检查if used request limit then reject}_, _ pipe.Exec(ctx)}}提交网关侧 Lua 脚本保证原子性lua-- check_and_consume_quota.lualocal used_key KEYS[1]local limit_key KEYS[2]local request tonumber(ARGV[1])local used tonumber(redis.call(GET, used_key) or 0)local limit tonumber(redis.call(GET, limit_key) or 0)if used request limit thenreturn {0, used, limit} -- 拒绝endredis.call(INCRBY, used_key, request)return {1, used request, limit} -- 通过方案对比为何放弃「全 K8s」与「全 Slurm」| 维度 | 纯 K8s (Kueue JobSet MPI Operator) | 纯 Slurm (Pyxis Enroot) |混合模式当前落地|| :--- | :--- | :--- | :--- ||MPI 通信性能| 依赖 HostNetwork 手动亲和性NCCL 拓扑感知弱 | 原生srun --mpipmi2拓扑感知最优 |Slurm 独占 MPI 池零损耗||容器化生态兼容| 原生支持 Ray/PyTorchJob/KubeFlow | 需 Enroot 转镜像调试链路长 |K8s 侧原生Slurm 侧仅跑传统 MPI||抢占/回收粒度| Kueue Preemption 支持优先级但 Gang Scheduling 易死锁 | Slurm 抢占成熟但无感知容器状态 |分域抢占Slurm 抢 SlurmKueue 抢 K8s跨域仅配额联动||数据亲和性调度| Scheduler Plugin/Extender 灵活 |sbatch --constraint表达力弱 |Extender Fluid 缓存感知跨域数据预热走统一 DataLoad Job||运维复杂度| 单一栈但 MPI Operator 仍不稳定 (v0.7.0) | 成熟稳定但容器化体验差 |双栈运维成本 30%换取业务零改造上线| 别信「一套调度器走天下」。MPI 死磕 K8s 是给运维找罪受PyTorch 死磕 Slurm 是给研发找罪受。分域治理、配额联动、数据解耦才是混部的务实解。效果数字说话上线 3 周观测窗口2026-07-01 至 2026-07-21对比重构前基线| 指标 | 重构前 | 重构后 | 变化幅度 || :--- | :--- | :--- | :--- ||平均数据加载耗时首轮 Epoch/Step| 2h 14min |18min|↓ 86.5%||Alluxio 缓存命中率| 12% |91%|↑ 79pp||任务排队 P99 时长| 4h 12min |19min|↓ 92.4%||抢占导致的有效算力损失| 100%全量重跑 |≤ 30min/次|↓ 99%||集群 GPU 平均利用率| 45% |78%|↑ 33pp||跨租户配额争抢投诉| 日均 6 单 |0 单|清零|存储侧顺带优化了 Cephosd_op_threads8、bluestore_min_alloc_size64K针对 2MB 大块顺序读配合 Alluxioalluxio.user.file.readtype.defaultCACHE_PROMOTE单节点顺序读吞吐从 3.2 GB/s 稳在 9.8 GB/s接近 100Gbps 网卡物理上限。总结后端架构的「确定性」不在模型层在调度层这回重构的核心心得就三条数据亲和性调度必须下沉到 Extender/Operator 层别指望业务侧写nodeSelector也别指望通用调度器插件能扛住高并发打分。GPU 任务的 Checkpoint 是应用层契约不是基础设施魔法。CRIU 在显存面前失效强制约定SIGUSR1torch.save/ PLUMED RESTART平台只负责信号传递与制品版本管理。异构集群别强行合一Slurm 管 MPI、K8s 管 Cloud-Native Training配额层做最终一致性同步即可别试图搞统一调度器那是架构师的自嗨不是工程师的交付物。下一步打算把quota-syncer替换成基于 KueueAdmissionCheck的原生准入链路把 Slurm 用量作为外部指标喂给 Prometheus Adapter再砍掉一层 Redis 中间商。不过那是下个迭代的事了——当下集群跑稳了研究员不骂娘了这才是后端交付的硬指标。#后端 #Java #SpringBoot #Kubernetes #Slurm #AI4S #异构调度你在实际项目中有遇到类似问题吗欢迎在评论区分享你的经验和解决方案。
