AI 和 ML 工作负载需要大量的 Pod 到 Pod 通信。由于此要求，Pod 之间的网络带宽会直接影响工作负载执行时间和费用。此带宽取决于集群中虚拟机 (VM) 实例的放置位置。
本文档介绍如何在 Google Kubernetes Engine (GKE) 集群上优化大规模 AI 或机器学习工作负载的调度，以提高性能和可靠性。具体来说，您将配置集群以使用拓扑感知调度 (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 的工作负载，请完成以下步骤：
安装启用了 TAS 的 Kueue
我们建议您搭配使用 TAS 与 Kueue。Kueue 是一个 Kubernetes 原生系统，可用于管理配额以及作业使用配额的方式。TAS 需要使用 Kueue 版本 0.10.0 或更高版本，并且您必须明确启用它。
如需安装 Kueue 并启用 TAS，请选择以下选项之一：
Kueue 清单
安装 Kueue：
kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml在 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 图表安装启用了 TAS 的 Kueue：
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"
安装 Kueke 后，您必须对其进行配置，使其了解其管理的相应基础架构，如下一部分中所述。
查看 GKE 集群的拓扑
在查看以抢占式虚拟机形式预配的 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 仅在工作负载请求的资源量不超过这些资源配额限制时才允许该工作负载运行。然后，Kueue 会根据 TAS 定义检查基础设施拓扑，然后再允许工作负载。如需了解详情，请参阅实际资源配额。
以下各部分中对资源配额的所有引用均指 ClusterQueue 资源配额。
资源配额非常高
以下示例使用了非常高的资源配额，这样一来，Kueue 就永远不会根据可用资源配额停止工作负载。相反，Kueue 会使用可用节点的拓扑信息，尝试将拓扑与工作负载的要求相匹配。
如需使用以下资源配额定义，请完成以下步骤：
打开您选择的文件编辑器。然后，在名为
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替换为节点池的名称。为 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。
如需使用以下资源配额定义，请完成以下步骤：
打开您选择的文件编辑器。然后，在名为
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替换为节点池的名称。为 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-blockcloud.google.com/gce-topology-subblockcloud.google.com/gce-topology-hostkubernetes.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 指定三个额外的字段：
标签：
kueue.x-k8s.io/pod-group-name：用于聚合的 PodGroup 的名称。
kueue.x-k8s.io/pod-group-pod-index：PodGroup 中每个 Pod 的索引。
注释：
- kueue.x-k8s.io/pod-group-total-count：PodGroup 中的 Pod 总数。
根据您使用的机器学习框架，PodGroup 的领导者可能需要 GPU，也可能不需要。由于 Kueue 的限制，这些情况需要以不同的方式处理。以下示例演示了如何创建包含 3 个 Pod 的 PodGroup，其中 1 个是 leader，2 个是 worker。
情形 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 的方式。如果 leader 没有最后一个索引，而第一个 worker 没有使用第一个索引，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"
后续步骤
如需详细了解如何在 GKE 集群中启用节点健康状况预测，请参阅启用节点健康状况预测。
如需了解如何管理与 GKE 集群和 AI 工作负载相关的常见事件，请参阅管理 AI 优化型 GKE 集群。
如需详细了解如何使用 Kueue 在 GKE 上调度作业，请参阅使用 Kueue 部署批处理系统。