使用 DWS 和 Kueue 最佳化 TPU 上的 AI 訓練

本文說明如何設定 GKE 叢集，盡可能提高 Tensor Processing Unit (TPU) 的 AI 訓練可用性。您可以使用名為 Kueue 的開放原始碼工作佇列工具，以及 Google Cloud的 Dynamic Workload Scheduler (DWS)，設定自動備援系統。

您在本文件中定義的設定會建立主要節點集區和備份節點集區：

  • 隨選 (方案 A)：這個節點集區是節點的首選。系統會先嘗試在這些隨選機器上排定工作。「隨選」是指這些機器開始運作後，您就能穩定不間斷地存取。
  • DWS 彈性啟動 (方案 B)：這是備份節點集區。如果 Plan A 的機器無法使用，Kueue 排程程式會自動將工作指派給這個 Plan 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 角色，只要您已獲授角色，即可選取任何專案。
    • 建立專案：如要建立專案，您需要具備專案建立者角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

    前往專案選取器

  2. 確認專案已啟用計費功能 Google Cloud 。

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

    啟用 API 時所需的角色

    如要啟用 API，您需要服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

    啟用 API

  4. 在 Google Cloud 控制台中啟用 Cloud Shell。

    啟用 Cloud Shell

  5. 請確認您有足夠的先占配額，可以使用 TPU 彈性啟動 VM。如果預設配額不符合需求，請申請提高配額。詳情請參閱「Cloud TPU 配額」和「設定 Cloud TPU 環境」。

定義環境變數

如要簡化本文中執行的指令，可以在 Cloud Shell 中設定環境變數。這些變數會儲存專案 ID、節點集區名稱和 GKE 叢集位置等值。 Google Cloud

定義這些變數後，您就可以重複使用這些變數，方法是參照變數名稱 (例如 $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：專案的專屬 ID 編號 (例如 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 彈性啟動節點集區的名稱。(這是您的 Plan 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 運作資源。您需要這個運算子，才能管理稍後在本文件中提交的 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 彈性啟動機器 (您的 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，將 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 區段指定指令碼所需的硬體 (在本例中為 tpu-v6e-slice，拓撲為 4x4)。
  • 選取佇列：kueue.x-k8s.io/queue-name 標籤會將工作直接傳送至 Kueue。這樣一來，系統就會啟用自動備援邏輯。
  • 容許 DWS 彈性啟動節點：tolerations 區段可讓工作在 B 方案節點集區上執行。由於 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 指派至該處。
    • 由於工作 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 控制台的「Manage resources」(管理資源) 頁面。

    前往「Manage resources」(管理資源)

  • 在專案清單中選取要刪除的專案，然後點選「Delete」(刪除)。
  • 在對話方塊中輸入專案 ID，然後按一下 [Shut down] (關閉) 以刪除專案。

    • 刪除個別資源

    如要保留您在本文件中使用的 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 資訊主頁，並確認兩項工作都已順利執行。

    後續步驟