在 AI for Science 平台落地异构调度：从「数据搬运工」到「算力调度器」的架构重构实录

上周有个需求砸下来：院里新建的 AI for Science 协作平台要上线对外开放，支撑分子动力学（MD）长周期模拟与大模型预训练混部。测试环境跑通的流程，一到生产环境就炸——不是 OOM 就是存储 IO 打满，研究员提个任务排队 6 小时起步，跑俩小时挂了还得从头来过。

项目背景简单交代：上海某研究院自建算力集群，约 200 张 H800 + 500 张 A800，存储端挂着 3PB CephFS（Ceph 18.2.2 Quincy）和 200TB Alluxio 缓存层。技术栈锁定在 Kubernetes 1.29.4、Slurm 23.11.8、KubeRay 1.1.0、Spring Boot 3.3.2（JDK 21.0.3）、Fluid 0.9.0（Alluxio Runtime 2.9.3）。核心痛点不在模型算法，而在异构负载下的数据亲和性调度与有状态长任务的生存保障。

背景：混部场景下的「三高」噩梦

先说业务画像。分子动力学典型作业：单任务跑 7×24 小时甚至数周，MPI 进程间通信极其敏感，中断一次等于白跑，Checkpoint 文件动辄 2TB+；大模型预训练则是 Gang Scheduling 启动，PyTorch DDP + NCCL，对带宽抖动零容忍，但支持 Tensor Parallel 切分、支持频繁torch.save。这俩负载混在同一套 K8s 集群里，还要套一层多租户配额（研究组级别），Slurm 管着裸金属 MPI 任务，K8s 管着容器化训练任务，中间隔着一层自研的统一提交网关。

测试阶段没暴露问题，因为测试数据集小、并发低、无抢占。真实上线第一周，三个指标扎眼：

1. 跨节点数据拉取占任务总耗时 68%（Alluxio 缓存命中率仅 12%），万兆网卡跑满却喂不饱 GPU。
2. Slurm 抢占 K8s 任务导致 PyTorchJob 重启率 43%，无法恢复，研究员直接在群里骂娘。
3. 配额失效：A 组提交 50 张卡的 RayJob，B 组 8 张卡的 MD 任务排队饿死，Kueue 的 ClusterQueue 配置在 Gang Scheduling 面前形同虚设。

过程：把「数据搬到算力」改成「算力贴着数据转」

1. 调度层拆解：别让 K8s 调度器干不擅长的事

最早想用 K8s 原生kube-scheduler+SchedulerPlugin解决数据亲和性，写了个PreFilter插件去查 FluidDatasetCR 的status.nodes字段，再给 Pod 打nodeAffinity。压测发现：调度延迟 P99 从 200ms 飙到 3.4s，因为每个 Pod 创建都要 List 所有 Dataset CR、计算拓扑距离、打分排序，调度器主线程被卡死。

换路子了。引入Scheduler Extender（HTTP 扩展器）做「粗筛」，只在Filter阶段调用，逻辑下沉到独立 Go 服务（scheduler-extender-v0.3.1），缓存 Dataset 节点分布在本地 Redis（Valkey 7.2.5），单次请求 15ms 以内。K8s 原生插件只保留Score阶段做微调。

```java
// Spring Boot 网关侧：任务提交时自动注入数据亲和性标签，不依赖用户手填 YAML
@RestController
@RequestMapping("/api/v1/jobs")
@RequiredArgsConstructor
public class JobSubmitController {

private final FluidDatasetClient fluidClient; // Feign 调用 Fluid Controller
private final SchedulerExtenderClient extenderClient;
private final KubernetesClient k8sClient;

@PostMapping
public 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 Operator；DDP/Ray -> Kueue + K8s JobSet
if (req.getWorkloadType() == WorkloadType.MPI_MD) {
return submitToSlurm(jobPod, req);
}
return submitToKueue(jobPod, req);
}

private Set computeOptimalNodes(Map> datasetNodeMap, int gpuCount) {
// 简单启发式：求交集，不足则按缓存完整度评分取 Top-N
return 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 接口，平台侧只负责「发信号、存制品、拉起新实例」。

```yaml

Kueue JobSet 模板：定义 PreStop Hook + 共享 PVC 挂载 Checkpoint 目录

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: ${{JOB_NAME}}
spec:
replicatedJobs:

• name: worker

replicas: ${{GPU_COUNT}}
template:
spec:
serviceAccountName: jobset-sa
restartPolicy: OnFailure
volumes:

• name: ckpt-volume

persistentVolumeClaim:
claimName: ${{PVC_NAME}}
containers:

• name: trainer

image: ${{IMAGE}}
command: ["bash", "-c", "torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py --ckpt-dir /ckpt"]
volumeMounts:

• mountPath: /ckpt

name: ckpt-volume
lifecycle:
preStop:
exec:
command: ["/bin/sh", "-c", "touch /ckpt/SIG_CHECKPOINT && sleep 120"] # 留 2 分钟刷盘
env:

• name: PLATFORM_CHECKPOINT_SIGNAL

value: "SIGUSR1" # 应用层监听信号触发 torch.save
```

平台侧控制器（job-lifecycle-controller v0.4.0）监听PodDeletionTimestamp，发现被抢占立即：

1. Patch PodterminationGracePeriodSeconds=180（防止默认 30s 截断）。
2. 向容器发SIGUSR1，业务代码捕获后执行torch.save({'model': ..., 'optimizer': ..., 'rng': ...}, f"/ckpt/ckpt_rank{dist.get_rank()}_step{step}.pt")。
3. 等待SIG_CHECKPOINT文件出现或超时，再允许 Pod 彻底终止。
4. 重建 JobSet 时，initContainer自动探测/ckpt最新 step，传入--resume-step启动。

效果：MD 任务虽不支持热迁移，但应用层每 30 分钟落盘一次（PLUMEDRESTART关键字），抢占恢复仅丢 ≤30 分钟算力，研究员能接受。

3. 多租户配额：Kueue + Slurm 统一视图的「伪全局调度」

Kueue 管不了 Slurm，Slurm 管不了 K8s。搞了个中间层quota-syncer（Go 协程每 10s 跑一次）：

• 从 Slurmsacctmgr show assoc format=account,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]float64

merged := 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.lua
local 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 then
return {0, used, limit} -- 拒绝
end
redis.call('INCRBY', used_key, request)
return {1, used + request, limit} -- 通过
```

方案对比：为何放弃「全 K8s」与「全 Slurm」

| 维度 | 纯 K8s (Kueue + JobSet + MPI Operator) | 纯 Slurm (Pyxis + Enroot) |混合模式（当前落地）|
| :--- | :--- | :--- | :--- |
|MPI 通信性能| 依赖 HostNetwork + 手动亲和性，NCCL 拓扑感知弱 | 原生srun --mpi=pmi2，拓扑感知最优 |Slurm 独占 MPI 池，零损耗|
|容器化生态兼容| 原生支持 Ray/PyTorchJob/KubeFlow | 需 Enroot 转镜像，调试链路长 |K8s 侧原生，Slurm 侧仅跑传统 MPI|
|抢占/回收粒度| Kueue Preemption 支持优先级，但 Gang Scheduling 易死锁 | Slurm 抢占成熟，但无感知容器状态 |分域抢占：Slurm 抢 Slurm，Kueue 抢 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_threads=8、bluestore_min_alloc_size=64K针对 2MB+ 大块顺序读，配合 Alluxioalluxio.user.file.readtype.default=CACHE_PROMOTE，单节点顺序读吞吐从 3.2 GB/s 稳在 9.8 GB/s（接近 100Gbps 网卡物理上限）。

总结：后端架构的「确定性」不在模型层，在调度层

这回重构的核心心得就三条：

1. 数据亲和性调度必须下沉到 Extender/Operator 层，别指望业务侧写nodeSelector，也别指望通用调度器插件能扛住高并发打分。
2. GPU 任务的 Checkpoint 是应用层契约，不是基础设施魔法。CRIU 在显存面前失效，强制约定SIGUSR1+torch.save/ PLUMED RESTART，平台只负责信号传递与制品版本管理。
3. 异构集群别强行合一，Slurm 管 MPI、K8s 管 Cloud-Native Training，配额层做最终一致性同步即可，别试图搞统一调度器，那是架构师的自嗨，不是工程师的交付物。

下一步打算把quota-syncer替换成基于 KueueAdmissionCheck的原生准入链路，把 Slurm 用量作为外部指标喂给 Prometheus Adapter，再砍掉一层 Redis 中间商。不过那是下个迭代的事了——当下集群跑稳了，研究员不骂娘了，这才是后端交付的硬指标。

#后端 #Java #SpringBoot #Kubernetes #Slurm #AI4S #异构调度

你在实际项目中有遇到类似问题吗？欢迎在评论区分享你的经验和解决方案。