使用 Pathways 建立 GKE 叢集

本頁面說明平台管理員和機器學習 (ML) 工程師如何建立 Google Kubernetes Engine (GKE) 叢集，並設定為 Pathways 工作負載。您可以使用 Cluster Toolkit 自動佈建預先設定的叢集，或使用 Google Cloud CLI 手動建立叢集。

事前準備

繼續操作前，請先完成下列必要條件：

• 安裝 gcloud CLI。
• 安裝 kubectl 和 gke-gcloud-auth-plugin。
• 如果您使用 Cluster Toolkit，請設定 Cluster Toolkit 並更新至最新版本。
• 啟用 Kubernetes Engine API。
• 確認您有足夠的 TPU 配額或 TPU 預留項目，可建立所需的 TPU 配量。

注意：按照本頁步驟操作，會建立計費資源，包括 GKE 叢集、TPU v6e 切片節點集區和 CPU 節點集區。

設定本機環境

使用 Google Cloud 憑證登入 gcloud CLI：

gcloud auth login

如果您使用 Cluster Toolkit，也請產生 Terraform 使用的應用程式預設憑證 (ADC)：

gcloud auth application-default login

建立 GKE 叢集

在以下範例中，您將建立含有 TPU v6e 2x4 節點集區的叢集。您可以使用 Cluster Toolkit 或 gcloud CLI 建立叢集。

Cluster Toolkit

1. 按照「部署 GKE TPU v6e 叢集」中的操作說明，部署 GKE TPU 叢集。如要查看所有 TPU 部署選項的總覽，請參閱「Cloud TPU 部署總覽」。

   設定部署藍圖時，請按照下列步驟操作：

   • 將 enable_pathways_for_tpus 設為 true (Cluster Toolkit TPU 藍圖預設啟用)。這項設定會自動佈建 Pathways 控制器所需的 cpu-np CPU 節點集區 (n4-standard-64) 和 Kueue 設定。只有在未於 kueue_configuration_path 變數中指定自訂 Kueue 設定檔時，Cluster Toolkit 才會套用預設 Pathways Kueue 設定。
   • 將 num_slices 設為工作負載所需的 TPU 配量節點集區數量。
   • 將 tpu_topology 設為所選 TPU 切片拓撲 (例如 2x4 或 4x4)。
   • 選擇計費模式。根據預設，TPU v6e 部署作業檔案會使用特定的 TPU 預留項目，因此請在 reservation_affinity 變數中指定預留項目名稱。如要使用其他計費方案，請參閱「GKE 計費方案」。

2. 確認叢集節點的 STATUS 為 Ready：

   kubectl get nodes
   

   建立叢集後，您可以使用 gcluster job submit 指令搭配 --pathways 和 --pathways-gcs-location 旗標，提交 Pathways 工作負載。詳情請參閱「使用 Pathways 執行工作負載」。

gcloud

如要建立 GKE 叢集並設定為執行 Pathways 工作負載，請按照下列步驟操作：

1. 請設定下列環境變數：

   CLUSTER=GKE_CLUSTER_NAME
   PROJECT=PROJECT_ID
   ZONE=ZONE
   REGION=REGION
   CLUSTER_VERSION=GKE_CLUSTER_VERSION
   PW_CPU_MACHINE_TYPE="n2-standard-64"
   NETWORK=NETWORK
   SUBNETWORK=SUB_NETWORK
   CLUSTER_NODEPOOL_COUNT=3
   TPU_MACHINE_TYPE="ct6e-standard-4t"
   WORKERS_PER_SLICE=2
   TOPOLOGY="2x4"
   NUM_CPU_NODES=1

   更改下列內容：

   • GKE_CLUSTER_NAME：GKE 叢集名稱。
   • PROJECT_ID：您的 Google Cloud 專案 ID。
   • ZONE：建立叢集的可用區。請選擇提供 TPU v6e 的可用區。詳情請參閱「Cloud TPU 區域和可用區」。
   • REGION：選用：如果使用區域叢集，則為建立資源的區域。
   • GKE_CLUSTER_VERSION：選用：GKE 叢集版本。請使用 1.32.2-gke.1475000 以上版本，並選擇標準支援的版本。詳情請參閱 GKE 發布時間表。如未指定版本，請從下列 gcloud container clusters create 指令中移除 --cluster-version 旗標。
   • NETWORK：選用：虛擬私有雲 (VPC) 網路的名稱。您必須先建立網路，才能建立叢集。如未指定網路，請從下列 gcloud container clusters create 指令中移除 --network 旗標。
   • SUB_NETWORK：選用：子網路名稱。您必須先建立子網路，才能建立叢集。如未指定子網路，請從下列 gcloud container clusters create 指令中移除 --subnetwork 旗標。

   這個範例會為其餘變數設定預設值，您可以根據工作負載修改這些值：

   • PW_CPU_MACHINE_TYPE：叢集預設節點集區的機器類型，以及執行 Pathways head (Pathways 資源管理工具和 Proxy) 的 CPU 節點集區。
   • NUM_CPU_NODES：CPU 節點集區中的節點數量。
   • CLUSTER_NODEPOOL_COUNT：TPU 配量數量。每個切片都是獨立的節點集區。
   • TPU_MACHINE_TYPE 和 TOPOLOGY：TPU 機型和各個配量的 TPU 拓撲。
   • WORKERS_PER_SLICE：每個切片中的節點數量，也就是拓撲中的 TPU 晶片數量除以每個節點的晶片數量。舉例來說，2x4 拓撲有 8 個晶片，ct6e-standard-4t 節點有 4 個晶片，因此每個切片有 2 個節點。

2. 使用 gcloud container clusters create 指令建立 GKE 叢集：

   gcloud container clusters create ${CLUSTER} \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --cluster-version=${CLUSTER_VERSION} \
       --scopes=storage-full,gke-default,cloud-platform \
       --machine-type=${PW_CPU_MACHINE_TYPE} \
       --network=${NETWORK} \
       --subnetwork=${SUBNETWORK}
   

3. 使用 gcloud container node-pools create 指令建立 TPU 切片節點集區：

   for i in $(seq 1 ${CLUSTER_NODEPOOL_COUNT}); do
       gcloud container node-pools create "tpu-np-${i}" \
           --project=${PROJECT} \
           --zone=${ZONE} \
           --cluster=${CLUSTER} \
           --machine-type=${TPU_MACHINE_TYPE} \
           --num-nodes=${WORKERS_PER_SLICE} \
           --placement-type=COMPACT \
           --tpu-topology=${TOPOLOGY} \
           --scopes=storage-full,gke-default,cloud-platform \
           --workload-metadata=GCE_METADATA
   done
   

   這些指令不會指定特定預留項目。如要使用 TPU 預留項目，請新增 --reservation-affinity=specific 和 --reservation=RESERVATION_NAME 旗標，其中 RESERVATION_NAME 是 TPU 預留項目的名稱。如要使用 Spot VM，請新增 --spot 旗標。

4. 為 Pathways head 建立 CPU 節點集區，用於執行 Pathways 資源管理工具和 Proxy。將節點集區命名為 cpu-np，以便 gcluster job submit --pathways 指令偵測：

   gcloud container node-pools create cpu-np \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --cluster=${CLUSTER} \
       --machine-type=${PW_CPU_MACHINE_TYPE} \
       --num-nodes=${NUM_CPU_NODES} \
       --scopes=storage-full,gke-default,cloud-platform \
       --workload-metadata=GCE_METADATA
   

   注意：Pathways 工作負載會將暫存檔案和快取的編譯程式碼儲存在 Cloud Storage bucket。本範例中的節點集區使用 --workload-metadata=GCE_METADATA，因此 Pathways 容器會使用節點服務帳戶存取 Cloud Storage。根據預設，這個帳戶是 Compute Engine 預設服務帳戶。並授予這個服務帳戶 bucket 的 Storage 物件管理員角色 (roles/storage.objectAdmin)。

5. 取得叢集的憑證，並設定 kubectl 的預設命名空間：

   注意：如果您使用可用區叢集 (例如本範例中的叢集)，請指定 --zone。如果您使用的是區域叢集，請將 --zone=${ZONE} 替換為 --region=${REGION}。

   gcloud container clusters get-credentials ${CLUSTER} \
       --zone=${ZONE} \
       --project=${PROJECT} \
       && kubectl config set-context --current --namespace=default
   

6. 如要在 GKE 叢集上使用 Pathways 架構，請安裝 JobSet API：

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

7. 確認節點的 STATUS 為 Ready，且 jobset-controller-manager Pod 的 STATUS 為 Running：

   kubectl get nodes
   kubectl get pods --namespace=jobset-system
   

   輸出結果會與下列內容相似：

   NAME                                      STATUS   ROLES    AGE   VERSION
   gke-cluster-cpu-np-12345678-abcd          Ready    <none>   5m    v1.32.2-gke.1475000
   gke-cluster-tpu-np-1-12345678-efgh        Ready    <none>   5m    v1.32.2-gke.1475000
   ...
   NAME                                         READY   STATUS    RESTARTS   AGE
   jobset-controller-manager-6f7b8c9d4-xyz12    1/1     Running   0          2m
   

清除所用資源

為避免系統向您的 Google Cloud 帳戶收取本頁面所用資源的費用，請在完成後刪除叢集：

• Cluster Toolkit：執行下列指令，刪除部署作業：

  ./gcluster destroy DEPLOYMENT_FOLDER_NAME

  將 DEPLOYMENT_FOLDER_NAME 替換為部署資料夾的名稱。

• gcloud CLI：執行下列指令，刪除 GKE 叢集及其節點集區：

  gcloud container clusters delete ${CLUSTER} \
      --project=${PROJECT} \
      --zone=${ZONE}
  

後續步驟

• 使用 Pathways 執行批次工作負載
• Pathways 互動模式
• 使用 Pathways 進行彈性訓練
• 將 JAX 工作負載移植到 Pathways
• 排解雲端 Pathways 問題