通过拓扑感知调度 (TAS) 安排 GKE 工作负载

AI 和 ML 工作负载需要大量的 Pod 间通信。由于此要求，Pod 之间的网络带宽直接影响工作负载执行时间和费用。此带宽取决于集群中虚拟机 (VM) 实例的放置位置。

本文档介绍了如何在 Google Kubernetes Engine (GKE) 集群上优化大规模 AI 或 ML 工作负载的调度，以提高性能和可靠性。具体来说，您可以将集群配置为使用拓扑感知调度 (TAS) 进行低延迟通信。这种方法可以最大限度地减少通信开销，并有助于最大限度地提高工作负载的性能。

什么是拓扑感知调度 (TAS)？

TAS 可以 显著提高大语言模型 (LLM) 训练的效率。TAS 会在网络拓扑上策略性地放置工作器，以最大限度地减少梯度聚合期间的通信开销，这需要工作器按特定排名顺序进行通信。通过最大限度地减少按顺序通信的工作器之间的网络跃点，TAS 可以减少网络争用并优化带宽利用率，从而加快收敛速度并缩短训练时间。随着 LLM 模型越来越大，TAS 对于最大限度地提高分布式训练的性能和可伸缩性至关重要。

TAS 最适合密集放置的容量，这可以通过预留获得。使用灵活启动虚拟机或 Spot 虚拟机时，您的容量不太可能紧密分配在一起，因此 TAS 在这种情况下可能无法正常工作。

准备工作

在开始之前，请确保您已执行以下任务：

  • 启用 Google Kubernetes Engine API。
    • 启用 Google Kubernetes Engine API
  • 如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。 如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。

  • 如需连接到集群，请运行以下命令：

    gcloud container clusters get-credentials CLUSTER_NAME

    将 CLUSTER_NAME 替换为您的集群名称。

准备 GKE 集群

如需准备 GKE 集群以运行启用 TAS 的工作负载，请完成以下步骤：

  1. 安装 Kueue 并启用 TAS

  2. 查看 GKE 集群的拓扑

  3. 配置 Kueue

安装 Kueue 并启用 TAS

我们建议您将 TAS 与 Kueue搭配使用。Kueue 是一个 Kubernetes 原生系统 ，可用于管理配额以及作业使用配额的方式。TAS 需要 Kueue 版本 0.10.0 或更高版本，并且您必须明确启用它。

如需安装 Kueue 并启用 TAS，请选择以下选项之一：

Kueue 清单

  1. 安装 Kueue：

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml

  2. 在 Kueue 中启用 TAS：

    kubectl -n kueue-system patch deployment kueue-controller-manager \
    --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'

Helm 图表

使用 Helm 图表安装 Kueue 并启用 TAS：

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

安装 Kueue 后，您必须对其进行配置，使其了解它所管理的基础架构，如下一部分中所述。

查看 GKE 集群的拓扑

在查看以 Spot 虚拟机形式预配的 A4X、A4、A3 Ultra、A3 Mega 和 A3 High（8 个 GPU）节点 的拓扑之前，您必须在 GKE 节点上定义 紧凑放置，以向 TAS 公开其物理拓扑。否则，您会遇到错误。

如需查看特定节点池中 GKE 集群节点的拓扑，请运行以下命令：

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

将 NODE_POOL_NAME 替换为节点池的名称。

如需了解输出中虚拟机上 GKE 节点的物理拓扑，请参阅以下节点标签：

  • cloud.google.com/gce-topology-block：虚拟机所在的预留块的组织专用 ID。

  • cloud.google.com/gce-topology-subblock：虚拟机所在的子块的组织专用 ID。

  • cloud.google.com/gce-topology-host：虚拟机所在的主机的 ID。

  • kubernetes.io/hostname：Kubernetes 节点的主机名。此主机名通常也是 GKE 节点名称。

两个虚拟机共享的标签值越多，这两个虚拟机在物理上就越靠近。如需详细了解这些术语，请参阅 术语。

配置 Kueue

安装 Kueue 后，您必须配置 Kueue 以指定它所管理的基础架构。通常，Kueue 需要 ClusterQueue资源 配额定义，可以是静态基础架构，也可以是启用了 集群自动扩缩的动态基础架构。仅当 工作负载请求的 资源小于或等于 ClusterQueue 中定义的 资源池时，ClusterQueue 才会允许工作负载。按照本部分中的说明配置 Kueue 后，Kueue 会使用 TAS 允许工作负载，如下所示：

  • TAS 工作负载：Kueue 会检查物理 基础架构的拓扑及其当前使用情况。

  • 非 TAS 工作负载：Kueue 不会检查物理 基础架构的拓扑。Kueue 会管理配置中定义的整个配额，并将节点分配留给 kube-scheduler。

如需了解如何向 Kueue 提供 ClusterQueue 资源配额定义，请查看以下示例：

  • 非常高的配额：Kueue 实际上永远不会根据请求的资源停止允许工作负载 。根据 TAS 定义，Kueue 可能会或可能不会根据基础架构拓扑允许工作负载。如需了解更多 信息，请参阅非常高的资源配额。

  • 实际配额：仅当 工作负载请求的资源在这些资源配额限制范围内时，Kueue 才会允许工作负载。根据 TAS 定义，Kueue 随后会检查基础架构拓扑，然后再允许工作负载。如需了解详情，请参阅 实际资源配额。

以下部分中对资源配额的所有引用均指 ClusterQueue 资源配额。

非常高的资源配额

以下示例使用非常高的资源配额，因此 Kueue 永远不会根据可用资源配额停止工作负载。相反，Kueue 会使用可用节点的拓扑信息来尝试将拓扑与工作负载的要求相匹配。

如需使用以下资源配额定义，请完成以下步骤：

  1. 打开您选择的文件编辑器。然后，在名为 kueue-tas-config-very-high-quota.yaml 的 YAML 文件中添加以下配额定义：

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {}
    resourceGroups:
    - coveredResources: ["nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        - name: "nvidia.com/gpu"
          nominalQuota: 10000000
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    将 NODE_POOL_NAME 替换为节点池的名称。

  2. 为 Kueue 作业排队系统创建并应用资源配额配置：

    kubectl create -f kueue-tas-config-very-high-quota.yaml

实际资源配额

上一个示例仅配置了 GPU 资源。不过，Kueue 可以管理所有与 Kubernetes 兼容的资源。

以下示例定义了更实际的资源配额，包括 CPU、内存和 GPU。这是针对 100 台 a3-ultragpu-8g 机器的。一台机器具有 224 个 vCPU、2944 GB 内存和 8 个 GPU。

如需使用以下资源配额定义，请完成以下步骤：

  1. 打开您选择的文件编辑器。然后，在名为 kueue-tas-config-real-quota.yaml 的 YAML 文件中添加以下配额定义：

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {} # match all
    resourceGroups:
    - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        # numbers below represent quota of 100 a3-ultragpu-8g machines
        - name: "cpu"
          nominalQuota: 22400
        - name: "memory"
          nominalQuota: 294400Gi
        - name: "nvidia.com/gpu"
          nominalQuota: 800
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    将 NODE_POOL_NAME 替换为节点池的名称。

  2. 为 Kueue 作业排队系统创建并应用资源配额配置：

    kubectl create -f kueue-tas-config-real-quota.yaml

    输出类似于以下内容：

    topology.kueue.x-k8s.io/gke-default created
resourceflavor.kueue.x-k8s.io/tas-flavor created
clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
localqueue.kueue.x-k8s.io/tas-user-queue created

使用 Kueue 通过 TAS 调度工作负载

以下场景展示了如何使用拓扑请求类型和拓扑请求级别指示 Kueue 和 TAS 管理常见的工作负载和基础架构组合：

  • 以下是可用的拓扑请求类型 （首选或 必需）：

    • kueue.x-k8s.io/podset-preferred-topology：Kueue 会优先在给定的拓扑级别内调度整个工作负载，但仍会允许不适合此拓扑级别的工作负载。对于可能适合单个拓扑级别的工作负载，Kueue 可能会跨该拓扑级别的多个实例调度该工作负载。

    • kueue.x-k8s.io/podset-required-topology：Kueue 会继续尝试允许此工作负载，直到整个工作负载可以适合所选的拓扑级别。

  • 以下是可用的拓扑请求级别，可让您 更具体地说明您希望或 要求作业运行的物理基础架构：

    • cloud.google.com/gce-topology-block

    • cloud.google.com/gce-topology-subblock

    • cloud.google.com/gce-topology-host

    • kubernetes.io/hostname

如需使用这些值调度工作负载，请使用以下作业 YAML 文件：

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

执行以下变量替换操作：

  • JOB_NAME：作业的名称。

  • NUMBER_OF_REPLICAS：并行运行的 Pod 数量。

  • ANNOTATIONS_STRING：请参阅下表：

    请求的拓扑类型和级别 说明 ANNOTATIONS_STRING
    首选 在 主机名 内运行（推荐） 只要有足够的资源来满足工作负载的资源要求，即使容量分散，此配置也会允许您的工作负载。Kueue 会尽可能紧凑地调度您的 Pod。 kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    必需 在 主机 内运行

    当且仅当有主机具有足够的资源来满足工作负载的资源要求时，此配置才会允许您的工作负载。

    当每个主机有多个虚拟机（例如，较小的机器类型）或单个节点上可以运行多个 Pod 时，此配置非常有用。在这种情况下，如果允许工作负载，它将在单个主机上运行。

    		 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    首选 在 主机 内运行 只要有足够的资源来满足工作负载的资源要求，即使容量分散，此配置也会允许您的工作负载。Kueue 会尝试在主机内调度您的 Pod，并根据需要使用其他主机。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    必需 在 子块 内运行 当且仅当有子块具有足够的资源来满足工作负载的资源要求时，此配置才会允许您的工作负载。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    首选 在 子块 内运行 只要有足够的资源来满足工作负载的资源要求，即使容量分散，此配置也会允许您的工作负载。Kueue 会尝试在子块内调度您的 Pod，并根据需要使用其他子块。在这种情况下，与仅具有满足要求的容量的子块相比，Kueue 会将具有更多可用容量的子块（即使容量分散）的排名更高。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    必需 在 块 内运行 当且仅当块中可用的资源满足工作负载的资源要求时，此配置才会允许您的工作负载。 如果允许，Kueue 会最大限度地减少用于调度工作负载的子块和主机数量。这可能会导致可用容量分散。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    首选 在 块 内运行 只要有足够的资源来满足工作负载的资源要求，即使容量分散，此配置也会允许您的工作负载。Kueue 会尝试在块内调度您的 Pod，并根据需要使用其他块。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

使用 Kueue 通过 TAS 使用 PodGroup 调度工作负载

使用 PodGroup 时，您必须为 PodGroup 中的每个 Pod 指定三个额外的字段：

根据您使用的 ML 框架，PodGroup 的领导者可能需要 GPU，也可能不需要。由于 Kueue 的限制，这些情况需要以不同的方式处理。以下示例演示了如何创建包含三个 Pod（一个领导者和两个工作器）的 PodGroup。

情况 1：领导者也是工作器，并且需要 GPU

如果领导者是工作器之一，并且还需要 GPU，则领导者可以在 PodGroup 中拥有任意数量。为简单起见，在以下示例中，领导者的索引为 0：

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

情况 2：领导者不是工作器，并且不需要 GPU

如果领导者不是工作器之一（由于 Kueue 限制），则领导者必须在 PodGroup 中具有最后一个索引，这是因为 Kueue 创建 PodSet 的方式。如果领导者没有最后一个索引，并且第一个工作器没有使用第一个索引，则 Kueue 不会应用排名分配。

请参阅以下示例：

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

后续步骤