从 Slack 抢 GPU 到自动排队:AI21 如何用 GKE 与 Kueue 缩短 AI 任务等待时间

2026-10-03 21 预计阅读时间: 1 分钟
来源: cloud.google.com AI 摘要 Original link

Disclaimer: This article is an AI-assisted summary. Read it together with the original source when precision matters. The summary may omit context, version differences, or edge cases and is not official documentation.

预计阅读时间:11 分钟

当 GPU 集群接近满载时,真正拖慢训练的往往不只是算力不足,而是调度失序:团队在聊天频道里争抢资源,空闲 GPU 散落在不同节点,大型训练任务即使看到“总量足够”,也迟迟无法启动。

AI21 在共享 GKE 集群中集中使用数千个 Google Cloud A3 与 A3 Ultra 实例,承载 Jamba 系列模型以及其他训练和智能体优化工作负载。引入 AI Hypercomputer、GKE 和 Kueue 后,其高优先级任务最长等待时间从 72 小时降到 12 小时,人工调度干预从每周约 20 次降到零,集群碎片率也从 15% 降到 8%。

接近 100% 利用率时,问题不只是“谁先用”

AI21 最初通过 Slack 的 #gpu-resources 频道协调容量。集群尚有余量时,这种方式足够灵活;但保留算力长期接近 100% 利用率后,每次申请都会变成跨团队谈判。

这里其实有两个不同问题。

1. 资源竞争:下一个轮到谁

训练团队、调试任务和推理服务都需要 GPU。如果没有统一队列、优先级和配额规则,团队负责人只能人工判断哪个项目更紧急。

这种方式难以审计,也容易出现两个极端:大型关键任务长期饥饿,或者高优先级任务频繁打断其他团队的工作。

2. 资源碎片:总量够,但形状不对

假设集群共有 8 张空闲 GPU,但分别以 1 + 1 + 4 + 2 的形式散落在四台机器上。对一个要求单节点 8 卡的任务来说,可用容量仍然是零。

多节点训练还会放大这个问题。如果任务的部分 Pod 先占住资源,其余 Pod 却无法调度,就可能形成“占着资源但无法工作”的僵局。人工重新分配优先级并不能解决这种装箱问题,调度系统需要理解任务的整体资源需求和集群拓扑。

为什么选择 Kueue,而不是替换 Kubernetes 调度器

AI21 评估了 Apache YuniKorn、Volcano 和 Kueue。最终选择 Kueue,核心原因不是它功能最多,而是它能够在保留标准 Kubernetes 调度链路的前提下增加批处理准入能力:

  • 不需要替换 Kubernetes 核心调度组件;
  • 不必重写现有 Job 规格;
  • 可以把一个训练任务作为整体排队和准入;
  • 能与 GKE 以及 Spot VM、Dynamic Workload Scheduler 等容量类型结合;
  • 可以通过队列、配额和优先级统一管理共享 GPU 池。

这里要区分“准入”和“放置”。Kueue 判断一个工作负载何时可以获得配额并进入集群,Kubernetes 调度器负责把具体 Pod 放到节点上。两者配合后,可以避免大型任务在资源尚未凑齐时过早占住部分 GPU。

AI21 还使用了两个重要能力:

  • Admission Fair Sharing(AFS):根据团队历史资源使用情况调整准入顺序,倾向于让使用较少的团队先获得机会,同时不抢占已经运行的任务。
  • Topology Aware Scheduling:让准入决策考虑节点和物理拓扑,避免批准一个实际上无法按所需形状落地的任务。

AFS 的价值在于把“公平”与“强制抢占”拆开。对昂贵且运行时间很长的训练任务来说,为了公平而中断已运行作业,可能浪费大量计算;调整尚未启动任务的顺序通常更稳妥。

可以这样实践:为共享 GPU 集群建立准入队列

下面是一个最小化示例,演示如何用 Kueue 管理 Kubernetes Job。它不是 AI21 生产配置的复刻,而是一套可以据此改造的起点。

运行前请确认:

  1. 集群已经安装 Kueue;
  2. GPU 节点已经安装 NVIDIA device plugin,能够报告 nvidia.com/gpu;
  3. 你有权限创建 ResourceFlavor、ClusterQueue 和 PriorityClass;
  4. 将示例中的 CPU、内存和 GPU 配额改成集群真实容量。
kubectl create namespace ai-training --dry-run=client -o yaml | kubectl apply -f -

kubectl apply -f - <<'EOF'
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: critical-training
value: 100000
preemptionPolicy: Never
globalDefault: false
description: High-priority model training without Kubernetes preemption
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: gpu-reserved
spec: {}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: shared-gpu
spec:
  namespaceSelector:
    matchLabels:
      kueue-access: shared-gpu
  resourceGroups:
    - coveredResources:
        - cpu
        - memory
        - nvidia.com/gpu
      flavors:
        - name: gpu-reserved
          resources:
            - name: cpu
              nominalQuota: 256
            - name: memory
              nominalQuota: 2Ti
            - name: nvidia.com/gpu
              nominalQuota: 32
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: training
  namespace: ai-training
spec:
  clusterQueue: shared-gpu
EOF

kubectl label namespace ai-training kueue-access=shared-gpu --overwrite

接着提交一个需要两个 Pod、每个 Pod 使用 8 张 GPU 的训练任务。suspend: true 很关键:任务先进入 Kueue 队列,只有获得整体配额后才会被解除挂起。

kubectl apply -f - <<'EOF'
apiVersion: batch/v1
kind: Job
metadata:
  name: distributed-training
  namespace: ai-training
  labels:
    kueue.x-k8s.io/queue-name: training
spec:
  suspend: true
  parallelism: 2
  completions: 2
  completionMode: Indexed
  template:
    spec:
      priorityClassName: critical-training
      restartPolicy: Never
      containers:
        - name: trainer
          image: nvidia/cuda:12.3.2-base-ubuntu22.04
          command:
            - bash
            - -lc
            - |
              echo "Starting worker ${JOB_COMPLETION_INDEX}"
              nvidia-smi
              sleep 300
          resources:
            requests:
              cpu: "8"
              memory: 32Gi
              nvidia.com/gpu: "8"
            limits:
              cpu: "8"
              memory: 32Gi
              nvidia.com/gpu: "8"
EOF

可以用以下命令观察队列和工作负载状态:

kubectl get clusterqueue shared-gpu
kubectl get localqueue -n ai-training
kubectl get workloads -n ai-training
kubectl describe workload -n ai-training
kubectl get jobs,pods -n ai-training -w

如果任务一直处于等待状态,优先检查三件事:配额是否足够、是否存在满足单 Pod GPU 请求的节点,以及节点标签、污点和亲和性规则是否把可用节点排除在外。

生产环境通常还应按团队创建不同的 LocalQueue,统一指向共享的 ClusterQueue,而不是给每个团队永久切割一块 GPU。这样既保留组织边界,也能让暂时空闲的容量服务其他团队。若要启用公平共享和拓扑感知,应根据正在使用的 Kueue 与 GKE 版本配置对应特性;这些 API 和开关可能随版本变化,不宜直接复制未经验证的旧配置。

成功指标不应只看 GPU 利用率

AI21 的总成本没有因为这次改造而下降。其保留算力原本就被有意维持在接近 100% 利用率,优化目标不是少买 GPU,而是减少等待和协调成本。

更有意义的指标包括:

  • 高优先级任务从提交到准入的 P50、P95 和最长等待时间;
  • 每周人工调整队列、删除 Pod或重新分配容量的次数;
  • 已分配但没有产生有效训练进度的 GPU 时间;
  • 因单节点或拓扑限制无法使用的碎片比例;
  • 各团队实际用量与其公平份额之间的偏差;
  • 任务因抢占、失败或部分调度而浪费的 GPU 小时数。

只观察平均利用率,可能掩盖严重问题:集群可以显示 98% 忙碌,但关键训练任务仍等待数天,甚至有一部分 GPU 被无法完成调度的 Pod 占住。

落地前的检查清单

将共享 GPU 集群从人工协调迁移到自动准入时,可以按以下顺序推进:

  1. 先统一入口:要求训练 Job 进入明确的 Kueue 队列,避免保留绕过队列的隐蔽通道。
  2. 区分优先级与公平性:业务关键程度决定优先级,团队历史用量决定公平顺序,两者不要混成一个字段。
  3. 默认采用整体准入:尤其是多 Pod、多节点训练,避免部分任务先占住 GPU。
  4. 把拓扑当成容量的一部分:32 张分散的 GPU 不一定等价于四台完整的 8-GPU 节点。
  5. 谨慎使用抢占:中断长时间训练的代价很高,可以先采用不抢占运行中任务的公平准入。
  6. 保留可观测性:记录排队原因、准入时间、配额消耗和失败事件,避免把 Slack 黑箱换成调度器黑箱。
  7. 再连接弹性容量:当保留容量用满后,再把相同准入策略扩展到 Spot VM 或动态申请的容量。

这次实践最值得借鉴的地方,并不是某个单独的调度功能,而是把 GPU 分配从“人与人协商”变成“策略驱动的整体准入”。当集群越接近满载,队列、公平性和拓扑意识就越不是附加功能,而是让昂贵算力真正转化为实验速度的基础设施。


相关推荐