使用 Pathways 建立 Cluster Director 叢集

本頁面說明平台管理員和機器學習 (ML) 工程師如何建立 Slurm 協調的 Cluster Director 叢集，並設定 Pathways 工作負載。完成這些步驟後，您會佈建一個叢集，其中包含 TPU v6e 切片、CPU 控制器節點、共用 Google Cloud Managed Lustre 檔案系統，以及 /home/pathways 執行階段環境。Pathways 不會預先安裝在 Cluster Director 叢集上：建立叢集後，您需要將 Pathways 指令碼下載到登入節點，並執行設定指令碼一次。

事前準備

  • 請根據您打算如何使用本頁範例，選取對應的分頁標籤：

    gcloud

    在 Google Cloud 控制台啟用 Cloud Shell。

    啟用 Cloud Shell

    Google Cloud 控制台底部會開啟 Cloud Shell 工作階段，並顯示指令列提示。Cloud Shell 是已安裝 Google Cloud CLI 的殼層環境，並已設定適用於您目前專案的值。工作階段可能要幾秒鐘的時間才能初始化。

    REST

    如要在本機開發環境中使用本頁的 REST API 範例，請使用您提供給 gcloud CLI 的憑證。

      安裝 Google Cloud CLI。

      若您採用的是外部識別資訊提供者 (IdP)，請先使用聯合身分登入 gcloud CLI。

    詳情請參閱 Google Cloud 驗證說明文件中的「使用 REST 進行驗證」一節。

事前準備

開始之前，請先完成下列必要條件：

  • 安裝 Google Cloud CLI。
  • 啟用 Hypercompute Cluster API。
  • 確認您的使用者帳戶和 Compute Engine 預設服務帳戶具備必要的身分與存取權管理 (IAM) 角色，可建立及連線至 Cluster Director 叢集。如需角色清單，請參閱 Cluster Director 快速入門導覽課程中的「事前準備」一節。
  • 在您打算建立叢集的可用區中，建立 TPU v6e 預留項目 (例如，兩個 2x4 切片至少需要四個 ct6e-standard-4t VM)。
  • 在叢集節點上設定無密碼 sudo，以便自動復原切片。請要求管理員授予您專案的 Compute OS 管理員登入 (roles/compute.osAdminLogin) 角色。
  • 建立 Cloud Storage bucket，用於 Pathways 暫存資料和工作負載輸出內容，並將專案的 Storage 物件使用者 (roles/storage.objectUser) 角色授予您的使用者帳戶和 Compute Engine 預設服務帳戶。

設定本機環境

使用 Google Cloud 憑證登入：

gcloud auth login

定義下列環境變數，並根據工作負載設定適當的值：

export PROJECT=PROJECT_ID
export REGION=REGION
export ZONE=ZONE
export CLUSTER_ID=CLUSTER_ID
export RESERVATION_NAME=RESERVATION_NAME

更改下列內容：

  • PROJECT_ID：Google Cloud 專案 ID。
  • REGION：建立資源的區域。
  • ZONE：建立資源的可用區。
  • CLUSTER_ID：叢集名稱。名稱最多可包含 10 個字元，且只能使用數字和小寫英文字母。
  • RESERVATION_NAME：ZONE 中 TPU v6e 預留項目的名稱。

建立 Cluster Director 叢集

在以下範例中，您將使用預留項目建立一個叢集，其中包含兩個 TPU v6e 2x4 切片 (四個 ct6e-standard-4t VM)，以及一個用於 Pathways 控制器的 CPU 節點。所有節點都會在 /home 掛接共用的 Managed Lustre 檔案系統。如要進一步瞭解叢集定義中的欄位，請參閱「建立自訂叢集」。

  1. 使用 REST API 建立叢集：

    建立要求主體：

    cat <<EOF > "${CLUSTER_ID}_create.json"
{
  "networkResources": {
    "net": {
      "config": {
        "newNetwork": {
          "network": "projects/${PROJECT}/global/networks/${CLUSTER_ID}-net"
        }
      }
    }
  },
  "storageResources": {
    "home": {
      "config": {
        "newLustre": {
          "lustre": "projects/${PROJECT}/locations/${ZONE}/instances/${CLUSTER_ID}-lustre",
          "filesystem": "lustre",
          "capacityGb": "18000",
          "perUnitStorageThroughput": "500"
        }
      }
    }
  },
  "computeResources": {
    "crtpu": {
      "config": {
        "newReservedInstances": {
          "reservation": "projects/${PROJECT}/zones/${ZONE}/reservations/${RESERVATION_NAME}"
        }
      }
    },
    "crcpu": {
      "config": {
        "newOnDemandInstances": {
          "zone": "${ZONE}",
          "machineType": "n2-highmem-32"
        }
      }
    }
  },
  "orchestrator": {
    "slurm": {
      "loginNodes": {
        "count": "1",
        "zone": "${ZONE}",
        "machineType": "n2-standard-16",
        "enableOsLogin": true,
        "enablePublicIps": true,
        "bootDisk": {
          "type": "pd-balanced",
          "sizeGb": "100"
        },
        "storageConfigs": [
          {
            "id": "home",
            "localMount": "/home"
          }
        ]
      },
      "nodeSets": [
        {
          "id": "nstpu",
          "computeId": "crtpu",
          "staticNodeCount": "4",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "resourcePolicyConfig": {
            "acceleratorTopology": "2x4"
          },
          "computeInstance": {
            "bootDisk": {
              "type": "hyperdisk-balanced",
              "sizeGb": "100"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        },
        {
          "id": "nscpu",
          "computeId": "crcpu",
          "staticNodeCount": "1",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "computeInstance": {
            "bootDisk": {
              "type": "pd-balanced",
              "sizeGb": "100",
              "image": "projects/clusterdirector-public-images/global/images/family/tpu-v6e-v5p-ubuntu2204-slurm2605-v20260902"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        }
      ],
      "partitions": [
        {
          "id": "tpu",
          "nodeSetIds": ["nstpu"]
        },
        {
          "id": "cpu",
          "nodeSetIds": ["nscpu"]
        }
      ],
      "defaultPartition": "tpu",
      "config": {
        "requeueExitCodes": [1, 2, 134, 137, 139, 143]
      }
    }
  }
}
EOF

    然後傳送要求：

    curl -s -X POST \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -d @"${CLUSTER_ID}_create.json" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters?clusterId=${CLUSTER_ID}"

    叢集包含下列項目：

    • tpu 分區：預留 TPU v6e VM 的靜態節點集。staticNodeCount 是 VM 數量，acceleratorTopology 則是每個切片的拓撲。每個 ct6e-standard-4t VM 有四個晶片，因此 2x4 切片有兩個 VM。如需更多切片，請乘以：兩個 2x4 切片需要 4 的 staticNodeCount。
    • cpu 分區：一個隨選 VM，執行 Pathways 資源管理工具、IFRT Proxy 伺服器和 MaxText 用戶端。Pathways 指令碼會將工作提交至名為 cpu 的分區。VM 會使用 n2-highmem-32 機型，因為 Pathways 服務和 MaxText 在儲存檢查點時，最多需要 240 GB 的記憶體。此外，由於 Pathways 服務會從這個 VM 提交 TPU 工作站工作，因此也會使用 Slurm 26.05 的映像檔，支援 TPU 配量。
    • requeueExitCodes：Slurm 會重新排入以其中一個代碼退出的工作。這項設定可讓 Slurm 在 slice 中的任何 VM (不只是第一個) 遭到先占時，重新啟動 TPU 工作站工作。

  2. 等待叢集準備就緒。您可以在 Google Cloud 控制台的「Cluster Director」頁面追蹤進度，或在下列指令的輸出內容中，確認 reconciling 為 false：

    curl -s \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

設定 Pathways 環境

Pathways 環境位於共用檔案系統的 /home/pathways 中，因此每個叢集只需設定一次。其中包含 Pathways 容器映像檔、具備 JAX 和 pathwaysutils 的 Python 環境，以及 Pathways 指令碼。

  1. 連線至登入節點。如要瞭解其他連線方式，請參閱「連線至叢集」。

    gcloud compute ssh ${CLUSTER_ID}-login-001 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --tunnel-through-iap

  2. 確認 cpu 和 tpu 分區存在，且處於 idle 狀態：

    sinfo -o "%P %a %D %T %N %f"

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

    PARTITION AVAIL NODES STATE NODELIST FEATURES
tpu* up 4 idle ${CLUSTER_ID}-nstpu-[0-3]
cpu up 1 idle ${CLUSTER_ID}-nscpu-0

  3. 將 Pathways 指令碼下載至登入節點：

    mkdir -p ~/pathways-scripts && cd ~/pathways-scripts
for f in setup_env.sh start_pathways.sh run_worker_slice.sh run_maxtext.sh; do
  curl -fsSLO "https://raw.githubusercontent.com/GoogleCloudPlatform/slurm-gcp/tpu-2605/pathways/scripts/${f}"
done

    這個指令會下載下列指令碼：

    • setup_env.sh：執行 Pathways 環境的一次性設定。
    • start_pathways.sh：啟動 Pathways 服務。
    • run_worker_slice.sh：啟動一個 TPU 切片的 Pathways worker (由 start_pathways.sh 提交)。
    • run_maxtext.sh：執行 MaxText 訓練工作負載範例。

  4. 執行設定指令碼。大約需要 5 分鐘：

    chmod +x *.sh
./setup_env.sh

    指令碼會提取 Pathways 伺服器和 Proxy 映像檔，使用 JAX、pathwaysutils 和 MaxText 建構 Python 環境，將指令碼複製到 /home/pathways，並將 ~/pathways 連結至 /home/pathways。結尾會顯示類似下列內容的輸出：

    Environment ready: JAX 0.11.1

  5. 選用：如要將 Pathways 環境的存取權授予其他叢集使用者，請讓他們將 /home/pathways 連結至主目錄：

    ln -sfn /home/pathways ~/pathways

此時，叢集已可使用 Pathways。執行工作負載時，您會啟動 Pathways 服務，因為批次和互動式工作負載會以不同的設定啟動該服務。詳情請參閱「使用 Pathways 執行批次工作負載」和「使用 Pathways 執行互動式工作負載」。

如果遇到問題，請參閱「排解 Cloud 上的 Pathways 問題」。

清除所用資源

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

curl -s -X DELETE \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

後續步驟