使用 DWS 和 Kueue 在 TPU 上优化 AI 训练

本文档将介绍如何配置 GKE 集群，以最大限度地提高张量处理单元 (TPU) 上 AI 训练的可用性。您可以使用名为 Kueue 的开源作业排队工具和 Google Cloud的动态工作负载调度器 (DWS) 来配置自动回退系统。

您在此文档中定义的配置会设置一个主节点池和一个备份节点池：

  • 按需 (方案 A)：此节点池是您的首选节点。系统会先尝试在这些按需机器上调度作业。“按需”是指这些机器一旦开始运行，您就可以可靠地、不间断地访问它们。
  • DWS 灵活启动（方案 B）：这是您的备用节点池。当 A 方案机器不可用时，Kueue 调度程序会自动将作业分配给此 B 方案池。然后，DWS 会搜索备选硬件，但由于该硬件也可能无法使用，因此无法保证立即访问。不过，DWS 不会放弃：它会将您的请求排队最多 7 天，并在机器可用时立即自动提供。

这种方法可以最大限度地减少作业在队列中等待的时间。这意味着，您不必手动检查可用资源，也不必针对不同的机器重写脚本。

配置步骤概览

如需配置自动回退系统，您必须完成多个配置步骤。这有助于将此配置分为两类：

  • 集群管理员任务：一次性基础架构配置，例如创建 GKE 集群、预配节点池和安装 Kueue 调度控制器。
  • AI 开发者任务：重复性的日常工作流程，例如定义训练作业要求和提交工作负载。

即使您自己执行所有这些步骤，也要牢记这一区别，这有助于明确整个流程。

在配置系统之前，请查看您将执行的配置步骤。

主题 任务
配置基础架构（集群管理员） 1. 创建 GKE 集群
2. 创建节点池
3. 验证节点池中灵活启动的状态
安装和配置 Kueue（集群管理员） 1. 安装 Kueue
2. 定义配置规则
运行训练作业 (AI 开发者) 1. 创建 ConfigMap
2. 定义 RayJob 清单
3. 提交工作负载
4. 连接到 RayJob
5. 查看日志

主要概念

  • 按需付费（方案 A）节点池：主要的高优先级节点池。作业始终会尝试先使用此池。
  • DWS 灵活启动（方案 B）节点池：备用节点池。如果主池中的机器不可用，系统会自动使用此池来搜索可用的硬件。
  • Kueue：用于管理作业队列的调度程序。它会拦截您的作业请求，并决定使用哪个节点池（方案 A 或方案 B）。
  • 作业：您要运行的 AI 训练工作负载。在本文档中，您将使用 RayJob 清单来定义它。

准备工作

  1. 在 Google Cloud 控制台的项目选择器页面上，选择或创建 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator)，该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。

    转到“项目选择器”

  2. 验证是否已为您的 Google Cloud 项目启用结算功能。

  3. 启用 Google Kubernetes Engine API 和 Cloud TPU API。

    启用 API 所需的角色

    如需启用 API，您需要拥有 Service Usage Admin IAM 角色 (roles/serviceusage.serviceUsageAdmin)，该角色包含 serviceusage.services.enable 权限。了解如何授予角色。

    启用 API

  4. 在 Google Cloud 控制台中，激活 Cloud Shell。

    激活 Cloud Shell

  5. 确保您有足够的抢占式配额来使用 TPU 灵活启动虚拟机。如果默认配额不足以满足您的需求，请申请更高的分配额。 如需了解详情，请参阅 Cloud TPU 配额。

定义环境变量

为了简化您在本文档中运行的命令，您可以在 Cloud Shell 中设置环境变量。这些变量用于存储 Google Cloud 项目 ID、节点池名称和 GKE 集群位置等值。

定义这些变量后，您可以通过引用变量名称（例如 $CLUSTER_NAME）在多个命令中重复使用它们，而无需每次都重新输入或替换值。这种方法可让本教程更易于理解，并降低出错风险。

如需在 Cloud Shell 中定义以下有用的环境变量，请运行以下命令：

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

以下是对这些环境变量的说明：

  • PROJECT_ID：您的 Google Cloud 项目的 ID。
  • PROJECT_NUMBER：项目的唯一标识符编号（例如 123456789012）。
  • ZONE：集群的计算可用区（例如 us-east5-b）。请选择提供所选加速器类型的可用区。如需了解可用性信息，请参阅 Cloud TPU 配额或 GPU 配额。
  • REGION：您创建集群资源的区域（例如 us-east5）。
  • CLUSTER_NAME：您为 GKE 集群选择的名称。
  • GKE_VERSION：集群的 GKE 版本。使用 1.34 或更高版本。
  • ONDEMAND_NODEPOOL：标准按需付费节点池的名称。（这是您的方案 A 节点池）。
  • DWS_NODEPOOL：DWS 灵活启动节点池的名称。（这是您的备选方案节点池）。

配置基础架构（集群管理员）

作为集群管理员，您需要配置 GKE 集群和节点池以支持回退机制。

创建 GKE 集群

首先，创建 GKE 集群。此集群是您在其中安装 Kueue 控制器、配置节点池并运行 AI 训练作业的环境。如需创建并连接到集群，请按以下步骤操作：

  1. 创建集群：

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    此命令使用以下关键标志：

    • --addons=RayOperator：在集群上安装 Ray Operator。您需要此运算符来管理您在本文档后面部分提交的 RayJob 工作负载。
    • --enable-image-streaming：使集群能够更快地拉取容器映像。此功能可显著缩短大型 AI 容器映像开始运行所需的时间。

  2. 检索集群的凭证，以便 kubectl CLI 可以连接到该集群。此命令会更新您的 Kubernetes 配置文件，该文件默认存储在 ~/.kube/config 目录中：

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

创建节点池

为您的环境创建主节点池和备份节点池：按需节点池（方案 A）和 DWS 灵活启动节点池（方案 B）：

  1. 创建按需节点池：此池将作为训练作业的主要资源：

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    在此示例中，您的首选机器是 TPU v6e 加速器。您可以使用 --machine-type=ct6e-standard-4t 标志指定此硬件。您可以更改此机器类型，以匹配您希望为 AI 模型使用的硬件（例如 GPU 或不同的 TPU）。

  2. 创建 DWS 灵活启动节点池：在此示例中，您选择与为主要按需节点池选择的机器类型 (--machine-type=ct6e-standard-4t) 相同的机器类型。备选方案 B 节点池不必使用其他机器类型。如果您非常想要这款特定硬件，那么将其作为备选方案意味着，如果这款硬件无法立即获得，您会改用其他方法来获取。此替代方法使用 DWS，该服务会持续搜索可用硬件，最长可达 7 天：

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    这些命令使用以下关键标志：

    • --num-nodes=0、--min-nodes=0、--max-nodes=4 和 --enable-autoscaling：此组合可让节点池在作业需要时从零个节点扩容，并在空闲时缩容，从而有助于节省费用。
    • --tpu-topology：定义 TPU 芯片的物理排列方式。您之所以要指定此布局，是因为芯片的物理排列会影响分布式训练作业的运行速度。
    • --reservation-affinity=none：有助于确保节点池不会占用预先预留的硬件。 Google Cloud 可让您预留特定机器，以帮助确保可用性。将此标志设置为 none 会告知系统绕过这些预留，而是动态请求未预留的机器。
    • --enable-queued-provisioning 和 --flex-start：（仅限方案 B 池）这些标志可让 DWS 在灵活容量可用时，从该容量中为方案 B 池预配节点。

验证节点池中灵活启动的状态

检查 DWS 灵活启动节点池，并验证是否已启用灵活启动：

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

如果启用了灵活启动，则输出为 True。

安装和配置 Kueue（集群管理员）

在本部分中，您将在集群上安装 Kueue 控制器。请注意，Kueue 是一种用于管理作业队列的调度程序。它会拦截您的作业请求，决定使用哪个节点池（按需或 DWS 灵活启动），然后分配作业。

安装 Kueue

运行以下命令以安装 Kueue。此命令会从官方代码库下载安装清单，并将其应用于您的集群：

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

定义配置规则

创建用于定义优先级规则的 YAML 清单。这些规则会告知 Kueue 先使用按需付费池，再使用 DWS 灵活启动池：

  1. 创建一个名为 dws-tpu-queue.yaml 的文件，其中包含以下内容。此文件定义了两种资源变种（按需和 DWS 灵活启动）以及一个用于确定这些资源变种优先级的集群队列。此配置文件定义了 Kueue 用于处理作业的逻辑：

    • ResourceFlavor：在本文档开头，您创建了两个节点池，并使用环境变量 ${ONDEMAND_NODEPOOL} 和 ${DWS_NODEPOOL} 为它们分配了名称。在创建这些节点池时，GKE 会自动为这些节点池中的每个节点添加您为这些环境变量选择的名称作为标签。ResourceFlavor 部分会告知 Kueue 查找具有这些标签的节点。
    • ClusterQueue：清单的这一部分定义了优先级规则。它会先列出按需实例，因此 Kueue 会尝试先预配按需机器。如果 Kueue 无法获取这些机器，则会尝试改而预配 DWS 灵活启动机器。
    • Quotas：该文件设置了配额，即作业在按需节点池中任意给定时间可使用的总资源（例如 CPU、内存和 TPU 芯片）的限制。当作业达到此限额时，Kueue 会自动尝试预配您在 dws-tpu-queue.yaml 中配置的 DWS 灵活启动机器（您的备选机器），这些机器的配额限额要高得多。 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. 将配置应用到集群。以下命令使用名为 envsubst 的命令行工具来替换文件 dws-tpu-queue.yaml 中显示的占位变量。envsubst 将占位符替换为您之前定义的环境变量的值：

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

运行训练作业（AI 开发者）

作为 AI 开发者，您可以通过创建 RayJob 清单来定义和提交训练工作负载。您可以在此清单中指定资源需求，而集群管理员之前使用 Kueue 和 DWS 配置的自动化回退系统会为您处理底层节点池。

在本部分中，您将执行以下步骤：

  • 创建 Python 训练脚本。
  • 将该脚本存储在 Kubernetes ConfigMap 中。
  • 部署将 ConfigMap 作为卷装载的 RayJob，以便在节点上执行训练脚本。

执行这些步骤后，Ray Train 会自动将 JAX 工作负载分配到各个节点，而 Kueue 会负责获取所需的机器。

训练脚本

将以下 Python 脚本复制并粘贴到名为 train.py 的文件中：

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

训练脚本使用 JAX（一种用于高性能数值计算的 Python 库）来训练线性回归模型。此脚本是一个简化的示例，旨在演示如何使用 DWS 和 Kueue 实现自动回退，不执行数据并行或模型并行。

请注意，训练脚本的 ScalingConfig 部分定义了训练作业的硬件要求。该部分请求 4x4 TPU 拓扑，这与您之前配置的节点池的物理布局相匹配。

创建 ConfigMap

将 train.py 脚本的内容上传到 Kubernetes ConfigMap 对象中。这样，集群就可以存储脚本并使其可供 RayJob 使用：

kubectl create configmap jax-train-script --from-file=train.py

您在下一部分中定义的 RayJob 会将此 ConfigMap 装载为卷。这样一来，脚本文件就会显示在 Ray 容器内，以便 Ray 软件能够找到并执行该文件。

应用 RayJob 清单

创建一个名为 rayjob-tpu-v6e-dws.yaml 的文件，其中包含以下内容。此清单定义了训练作业，并告知系统如何路由该作业：

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

此清单包含三个可让回退系统正常运行的配置：

  • 请求特定硬件：nodeSelector 部分指定了脚本所需的硬件（在此示例中，为具有 4x4 拓扑的 tpu-v6e-slice）。
  • 选择队列：kueue.x-k8s.io/queue-name 标签会将作业直接路由到 Kueue。这会启用自动回退逻辑。
  • 容忍 DWS 灵活启动节点：tolerations 部分允许作业在备选方案节点池运行。由于 GKE 会专门标记（添加污点）DWS 灵活启动节点，以防止正常工作负载意外在这些节点上运行，因此您的作业必须明确容忍 cloud.google.com/gke-queued 污点。

提交工作负载

为了证明回退系统正常运行，您需要提交两个作业。第一个作业会消耗方案 A 的按需容量，这会迫使第二个作业回退到方案 B 的 DWS 灵活启动容量。

运行以下命令以提交这两个作业。该命令使用 for 循环和 envsubst 将唯一的作业 ID 注入到每次运行的清单中：

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

提交作业后，系统会按如下方式处理工作负载：

  1. 拦截：Kueue 使用队列标签检测作业，并暂时暂停这些作业。
  2. 决策：Kueue 会根据管理员的规则评估资源可用性。它会先检查方案 A 池。
  3. 作业：
    • 由于方案 A 的资源可供第一个作业使用，因此 Kueue 会将作业 1 分配给方案 A。
    • 由于作业 1 消耗了方案 A 的资源，因此 Kueue 会自动将作业 2 分配给方案 B（DWS 灵活启动）池。
  4. 启动：Kueue 会取消暂停作业。此操作会触发 GKE 集群自动扩缩器来预配节点并启动训练脚本。

连接到 RayJob

作为最后的验证步骤，您可以使用 kubectl port-forward 命令连接到 Ray 信息中心，并观看作业的运行情况。

如需检查第一个作业的状态，请运行以下命令：

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

运行此命令后，打开网络浏览器并前往 http://localhost:8265。在 Ray 信息中心内，您可以查看作业状态和报告的指标，以验证这两个作业是否已在各自的节点池中成功完成。

您还可以通过运行以下命令来查看第一个作业的日志：

kubectl logs job/rayjob-tpu-v6e-dws-1

训练脚本的截断输出应类似于以下内容。您应该会在输出内容的末尾附近看到消息 Training Complete! 和 Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded：

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

清理

为避免因本文档中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留该项目但删除各个资源。

删除项目

  • 在 Google Cloud 控制台中，前往管理资源页面。

    转到“管理资源”

  • 在项目列表中，选择要删除的项目，然后点击删除。
  • 在对话框中输入项目 ID，然后点击关闭以删除项目。

    • 逐个删除资源

    如果您希望保留本文档中使用的 GGoogle Cloud 项目，请运行以下命令来删除集群：

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    摘要

    在本文档中，您配置并测试了 Ray 训练环境。此环境使用主节点池和备份 DWS 池来最大限度提高硬件可用性。通过在主机器不可用时自动回退到 DWS，您可以最大限度地缩短训练作业排队等待的时间。

    为了让此功能正常运行，您执行了以下步骤：

    1. 创建了 GKE 集群：建立了用于托管节点池和调度工具的环境。
    2. 配置了节点池：创建了一个按需节点池（方案 A）和一个 DWS 节点池（方案 B）。
    3. 安装并配置了 Kueue：部署了 Kueue 控制器，并应用了优先级规则，指示系统先尝试方案 A，然后再回退到方案 B。
    4. 创建了 ConfigMap：将简化的 JAX 训练脚本部署到集群，以充当测试工作负载。
    5. 定义了 RayJob 清单：配置了作业以请求特定硬件、路由到 Kueue 控制器并容忍 DWS 节点。
    6. 提交了工作负载：提交了两个作业，以强制 Kueue 在方案 A 的资源被消耗时自动将第二个作业路由到方案 B。
    7. 验证结果：使用端口转发连接到 Ray 控制面板，并确认这两个作业都已成功运行。

    后续步骤