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

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

您在本文档中定义的配置将设置主节点池和备份节点池：

  • 按需 (方案 A)： 此节点池是您的首选节点。系统会先尝试在这些按需机器上调度作业。“按需”一词是指，一旦这些机器开始运行，您就可以可靠地不间断地访问它们。
  • DWS 灵活启动（方案 B）： 这是您的备份节点池。当方案 A 机器不可用时，Kueue 调度程序会自动将您的作业分配给此方案 B 池。然后，DWS 会搜索方案 B 硬件，但它不保证立即访问，因为该硬件也可能不可用。不过，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、Cloud TPU API API。

    启用 API 所需的角色

    如需启用 API，您需要拥有 serviceusage.services.enable 权限。如果您已创建项目，则可能已通过 Owner 角色 (roles/owner) 拥有此权限。否则，您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。了解如何授予角色。

    启用 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 灵活启动节点池的名称。（这是您的方案 B 节点池）。

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

作为集群管理员，您可以配置 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 节点池不必使用不同的机器类型。由于您非常需要此特定硬件，因此将其作为方案 B 选择意味着，如果该硬件无法立即使用，您将切换到其他获取方法。此替代方法使用 DWS，该方法会持续搜索可用硬件，最多 7 天：

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --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：（仅限方案 A 池）此标志可防止按需节点池使用预留容量。DWS 灵活启动节点池不会使用预留。
    • --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 灵活启动机器（方案 B 机器），您在 dws-tpu-queue.yaml 中配置了这些机器，其配额限制要高得多。 
      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 中。
  • 部署 RayJob，该 RayJob 会将 ConfigMap 作为卷装载，以便可以在节点上执行训练脚本。

执行这些步骤后，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 部分允许作业在方案 B 节点池运行。由于 DWS 灵活启动节点由 GKE 特殊标记（污点），因此普通工作负载不会意外地在这些节点上运行，因此您的作业必须明确容忍 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 分配给该资源。
    • 由于作业 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 信息中心，并确认这两个作业都成功运行。

    后续步骤