使用 Pathways 執行互動式工作負載

Pathways 互動式工作負載是遠端 JAX 工作負載，在叢集的 Pathways 控制器元件外執行，例如在 VM 上。與批次工作負載不同，互動式工作負載完成後，不會關閉 Pathways 叢集元件，這些元件仍可供其他 JAX 用戶端連線。本文以 Jupyter 筆記本為例，示範互動式工作負載。您可以在 Google Kubernetes Engine (GKE) 叢集或採用 Slurm 的 Cluster Director 叢集上，透過 Pathways 執行互動式工作負載。

本頁面提供這兩個平台的操作說明。在每個部分中，選取「GKE」或「Cluster Director」分頁標籤。

JAX 使用者可透過 IFRT 介面，將指令傳送至 Pathways 叢集。無論是從終端機、筆記本或任何與 Python 相容的環境執行，JAX 程式碼都能與 Pathways 資源互動。

事前準備

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

  gcloud

  安裝 Google Cloud CLI，然後 使用聯合身分登入 gcloud CLI。 登入後，執行下列指令來初始化 Google Cloud CLI：

  gcloud init

  REST

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

  安裝 Google Cloud CLI，然後 使用聯合身分登入 gcloud CLI。

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

事前準備

GKE

請完成下列前置作業：

• 建立 Pathways GKE 叢集。
• 安裝 Cluster Toolkit。
• 安裝 Kubernetes 工具。
• 安裝 Google Cloud CLI。
• 啟用 Kubernetes Engine API。
• 請確認您已為 Google Cloud 專案啟用 Pathways。

Cluster Director

請完成下列前置作業：

• 建立 Cluster Director 叢集並設定 Pathways 環境。
• 安裝 gcloud CLI。
• 確認您具備叢集登入節點的 SSH 存取權。
• 在執行 JAX 用戶端的主機上，安裝 Python 3 和 venv 模組。在 Debian 和 Ubuntu 上，安裝 python3-venv 套件。

在互動模式中執行 Pathways

GKE

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

   export WORKLOAD=WORKLOAD
   export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT
   export COMPUTE_TYPE="COMPUTE_TYPE"
   export TOPOLOGY="TOPOLOGY"
   export PROJECT_ID=PROJECT_ID
   export REGION=REGION
   export ZONE=ZONE
   export CLUSTER=CLUSTER
   export BUCKET_NAME=BUCKET_NAME

   更改下列內容：

   • WORKLOAD：用於識別工作負載的專屬名稱 (由於 DNS 標籤限制，長度不得超過 22 個字元)。
   • WORKLOAD_NODEPOOL_COUNT：Pathways 工作負載使用的節點集區數量。
   • COMPUTE_TYPE：TPU 機型，例如 ct6e-standard-4t。如要進一步瞭解各 TPU 版本支援的 TPU 類型，請參閱「TPU 版本」。
   • TOPOLOGY：TPU 放置拓撲，例如 2x4。
   • PROJECT_ID：Google Cloud 專案 ID。
   • REGION：bucket 和叢集所在的區域。
   • ZONE：您打算執行工作負載的可用區。
   • CLUSTER：GKE 叢集的名稱。
   • BUCKET_NAME：Pathways 構件的 Cloud Storage bucket。

2. 如果尚未建立 Pathways 構件的 Cloud Storage bucket，請建立一個：

   gcloud storage buckets create gs://${BUCKET_NAME} \
       --project=${PROJECT_ID} \
       --location=${REGION}
   

3. 在叢集上建立 Pathways 容器。如要執行無頭工作負載，請使用 gcluster job submit 指令 (詳情請參閱「Cluster Toolkit 工作提交指南」)：

   gcluster job submit \
       --pathways \
       --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
       --name=${WORKLOAD} \
       --cluster=${CLUSTER} \
       --project=${PROJECT_ID} \
       --location=${ZONE} \
       --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
       --compute-type=${COMPUTE_TYPE} \
       --topology=${TOPOLOGY} \
       --pathways-headless
   

   gcluster job submit 指令會尋找名為 cpu-np 或 pathways-np 的 CPU 節點集區，以便執行 Pathways head Pod。如果 CPU 節點集區有其他名稱 (例如叢集建立指南中的 cpu-pathways-np)，請將 --pathways-head-np=NODE_POOL_NAME 旗標新增至本頁的 gcluster job submit 指令。否則指令會失敗，並顯示 failed to auto-detect a suitable CPU node pool 錯誤。

此時，JAX 工作負載可以連線至 IFRT Proxy 伺服器。

Cluster Director

1. 連線至叢集的登入節點：

   gcloud compute ssh CLUSTER_ID-login-001 \
       --project=PROJECT_ID \
       --zone=ZONE \
       --tunnel-through-iap

   更改下列內容：

   • CLUSTER_ID：Cluster Director 叢集的名稱。
   • PROJECT_ID：Google Cloud 專案 ID。
   • ZONE：叢集所在的可用區。

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

   export SLICE_SHAPE=TOPOLOGY
   export MAX_SLICES=NUM_SLICES
   export MIN_SLICES=NUM_SLICES
   export GCS_BUCKET=gs://BUCKET_NAME

   更改下列內容：

   • TOPOLOGY：每個切片的 TPU 拓撲，例如 2x4。必須與 TPU 節點集的拓撲相符。
   • NUM_SLICES：TPU 配量數量。將 MIN_SLICES 和 MAX_SLICES 設為相同的值，這樣 Pathways 叢集就會在所有配量都執行時才準備就緒。
   • BUCKET_NAME：現有的 Pathways 構件 Cloud Storage bucket。

   如果 TPU 分區不是叢集的預設分區，請將 PARTITION 設為該分區的名稱。

3. 以無頭模式啟動 Pathways 叢集：

   cd ~/pathways
   sbatch start_pathways.sh
   

4. 等待服務記錄 pathways-service-JOB_ID.out (其中 JOB_ID 是 Slurm 工作 ID) 顯示類似下方的輸出內容：

   Pathways Runtime Service is READY!
    Discovery file    : /home/USER/pathways/cluster.env
    IFRT Proxy Target : grpc://RM_IP:29000
   

   RM_IP 是 Proxy 伺服器的內部 IP 位址。檔案 ~/pathways/cluster.env 包含相同的位址。

此時，JAX 工作負載可以連線至 IFRT Proxy 伺服器。

在互動模式下連線至 Pathways 叢集

連線至 Pathways 叢集時，可以選擇是否要轉送通訊埠。請使用下列其中一個部分，連線至 Pathways 叢集。

在 Cluster Director 上，JAX 用戶端必須使用與 Pathways 伺服器相同的 JAX 版本 (0.11.1)。

使用通訊埠轉送功能連線

GKE

從任何可存取叢集控制層的主機，將本機通訊埠 29000 轉送至 Pathways head Pod，以便存取 Proxy 伺服器：

PROXY_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
PROXY_PORT=29000
kubectl port-forward ${PROXY_POD} ${PROXY_PORT}:${PROXY_PORT}

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

Forwarding from 127.0.0.1:29000 -> 29000
Forwarding from [::1]:29000 -> 29000

在同一個主機上開啟新的終端機視窗。使用 Python 3.11 以上版本建立虛擬環境、安裝 pathwaysutils 和 jax 套件、設定 JAX_PLATFORMS 和 JAX_BACKEND_TARGET 環境變數，然後執行匯入 pathwaysutils 和 jax 的 Python 指令碼：

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(144,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(145,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(146,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(147,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
device(148,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(149,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(150,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(151,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
device(162,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(163,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(164,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(165,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
device(166,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(167,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(168,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
device(169,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-13 21:38:51.267523: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

Cluster Director

從任何可透過 SSH 連線至叢集登入節點的主機 (例如工作站或叢集子網路外的 VM)，透過登入節點轉送 Proxy 通訊埠：

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

注意：如果您在 Compute Engine VM 上使用預設服務帳戶執行 gcloud CLI，VM 必須具備 cloud-platform 存取權範圍，否則 gcloud compute ssh 會失敗並顯示 insufficient authentication scopes。或者，您也可以在 VM 上執行 gcloud auth login，使用自己的帳戶。

在同一個主機上開啟新的終端機視窗。設定 JAX_PLATFORMS 和 JAX_BACKEND_TARGET 環境變數，並執行匯入 pathwaysutils 和 jax 的 Python 指令碼。主機沒有通往 RM_IP 的路徑，因此用戶端會連線至通道的本機端 127.0.0.1:29000：

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

當指令碼結束並關閉連線時，系統會顯示 IFRT proxy server disconnected 警告。這是正常情況。

注意：每個 JAX 呼叫都會透過 SSH 通道進行。使用通訊埠轉送功能驗證連線。如果是筆記本，請在叢集的子網路中執行 JAX 用戶端。

從虛擬私有雲中的主機連線，不使用通訊埠轉送

GKE

如果不想使用通訊埠轉送，可以透過 Cloud DNS 或內部負載平衡器連線至 Pathways 叢集。

使用 Cloud DNS 連線

在叢集中啟用 Cloud DNS，會將 DNS 供應商從 kube-dns 變更為 Cloud DNS。啟用後，GKE 會在虛擬私有雲 (VPC) 網路中，為 DNS 名稱建立私人 Cloud DNS 區域。詳情請參閱「將 Cloud DNS 用於 GKE」一文。

如果啟用 Cloud DNS，並設定虛擬私有雲範圍或附加虛擬私有雲範圍，則可從虛擬私有雲網路中的非 GKE VM 解析 Kubernetes DNS 名稱。如果範圍是叢集，則只有叢集中的節點可以解析這些名稱。名稱格式為 SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN。Pathways head Pod 有一項名為 WORKLOAD-pathways-head-0-0.WORKLOAD.NAMESPACE.svc.CUSTOM_DNS_DOMAIN 的服務。

下列指令說明如何使用 Cloud DNS 連線至 Pathways 叢集：

1. 確認可從非 GKE 主機解析 leader DNS 項目：

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN

   更改下列內容：

   • WORKLOAD：Pathways 工作負載的名稱。
   • CUSTOM_DNS_DOMAIN：為叢集設定的自訂 DNS 網域。

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

   WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN has address 10.0.2.75

2. 使用 DNS 名稱連線至 Pathways 叢集：

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.CUSTOM_DNS_DOMAIN:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

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

   [device(216,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(217,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(218,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(219,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
   device(220,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(221,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(222,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(223,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
   device(234,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(235,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(236,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(237,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
   device(238,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(239,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(240,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
   device(241,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
   Waiting up to 5 seconds.
   Sent all pending logs.
   2024-11-14 00:02:49.882044: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled
   

使用內部負載平衡器連線

如要將虛擬私有雲網路中的私人 IP 位址指向 Pathways 部署作業，請建立由內部負載平衡器 (ILB) 支援的 Service。叢集不必啟用 Cloud DNS。

如果叢集有多個 VM，請在建立內部負載平衡器時啟用 ILB 子設定。詳情請參閱「在現有叢集中啟用 GKE 子集」。如果未啟用 ILB 子設定，叢集中的所有節點都會是所有內部負載平衡器的後端執行個體群組的一部分，且無法擴充至超過 250 個節點。啟用 ILB 子集設定後，GKE 會建立網路端點群組，而非執行個體群組，且只會納入執行 Service 服務 Pod 的節點。啟用 ILB 子集功能時，需要約 15 分鐘的一次性設定延遲時間。下列指令會啟用 ILB 子集：

gcloud container clusters update ${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=CLUSTER_LOCATION \
    --enable-l4-ilb-subsetting

將 CLUSTER_LOCATION 替換為叢集所在的區域或可用區。

啟用 ILB 子集後，請使用下列 YAML 資訊清單，建立 LoadBalancer 類型的 Kubernetes Service。這會導致 GKE 在叢集的虛擬私有雲網路中建立內部負載平衡器：

apiVersion: v1
kind: Service
metadata:
  name: WORKLOAD-ilb
  annotations:
    networking.gke.io/load-balancer-type: "Internal"
    networking.gke.io/internal-load-balancer-allow-global-access: "true"
spec:
  type: LoadBalancer
  externalTrafficPolicy: Local
  selector:
    jobset.sigs.k8s.io/jobset-name: WORKLOAD
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

將 WORKLOAD 替換為 Pathways 工作負載的名稱，並將檔案儲存為 pathways-headless-ilb.yaml。

套用資訊清單：

kubectl apply -f pathways-headless-ilb.yaml

GKE 建立負載平衡器後 (約一分鐘後)，請確認 EXTERNAL-IP 欄顯示 IP 位址：

kubectl get services

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

NAME             TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)           AGE
WORKLOAD       ClusterIP      None            <none>        <none>            30m
WORKLOAD-ilb   LoadBalancer   34.118.232.46   10.0.0.22     29000:31246/TCP   2m41s

您可以從與叢集位於相同虛擬私有雲網路的主機，存取 Pathways 部署作業，而不必轉送通訊埠：

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://EXTERNAL_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

將 EXTERNAL_IP 替換為 kubectl get services 顯示的 EXTERNAL-IP 值，例如 10.0.0.22。

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

[device(288,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(290,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(291,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3),
 device(292,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(293,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(294,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(295,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=0,default_mem=device,mem_spaces=3),
 device(306,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(307,TPU_DEVICE,coords=[1,0,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(308,TPU_DEVICE,coords=[0,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(309,TPU_DEVICE,coords=[1,1,0,0],vtask=0,slice=1,default_mem=device,mem_spaces=3),
 device(310,TPU_DEVICE,coords=[0,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(311,TPU_DEVICE,coords=[1,0,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(312,TPU_DEVICE,coords=[0,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3),
 device(313,TPU_DEVICE,coords=[1,1,1,0],vtask=1,slice=1,default_mem=device,mem_spaces=3)]
Waiting up to 5 seconds.
Sent all pending logs.
2024-11-14 00:30:07.296917: W external/xla/xla/python/ifrt_proxy/client/grpc_client.cc:63] IFRT proxy server disconnected: CANCELLED: Cancelled

注意：您的 IP 位址與上述指令和輸出範例中顯示的位址不同。

Cluster Director

叢集子網路中的主機 (例如 Compute Engine VM) 可以直接連線至 Proxy 伺服器。叢集導向器的防火牆允許來自叢集子網路的流量，但不允許來自相同虛擬私有雲網路其他子網路的流量：

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://RM_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

將 RM_IP 替換為服務記錄中的 IFRT Proxy 位址。

注意：IFRT Proxy 伺服器不會驗證用戶端。任何可連上 RM_IP:29000 的主機，都能使用 Pathways 叢集。

Jupyter Notebook

GKE

您可以使用 Vertex AI 建立 Jupyter 筆記本，也可以建立自行託管的 Jupyter 筆記本。

建立 Vertex AI Workbench 執行個體

設定並驗證 Pathways 叢集後，即可透過 Vertex AI Workbench Jupyter 筆記本存取 GKE TPU VM。下列設定操作說明假設 GKE Pathways 叢集位於同一個虛擬私有雲網路 (除非您另行設定，否則為預設網路)。

1. 前往 Google Cloud 控制台的「Vertex AI Workbench」頁面。

2. 在「Instances」(執行個體) 分頁中，按一下「Create new」(新建)，並確認網路與 GKE 叢集的網路相符。

   或者，您也可以使用指令列建立新的 Vertex AI Workbench 執行個體：

   gcloud workbench instances create INSTANCE_NAME \
       --project=PROJECT_ID \
       --machine-type=e2-standard-4 \
       --data-disk-size=100 \
       --location=ZONE \
       --network=NETWORK

   更改下列內容：

   • INSTANCE_NAME：Vertex AI Workbench 執行個體的名稱。
   • PROJECT_ID：Google Cloud 專案 ID。
   • ZONE：要建立執行個體的可用區。
   • NETWORK：GKE 叢集的虛擬私有雲網路名稱。--network 旗標為選用項目。如果省略，指令會使用 default 網路。

3. Vertex AI 建立執行個體後，請前往該執行個體，然後按一下「Open JupyterLab」。

建立自行託管的 Jupyter 筆記本執行個體

下列指令說明如何使用 gcluster job submit 指令，建立自行託管的 Jupyter 筆記本執行個體：

gcluster job submit \
    --pathways \
    --name=${WORKLOAD} \
    --cluster=${CLUSTER} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --num-slices=${WORKLOAD_NODEPOOL_COUNT} \
    --compute-type=${COMPUTE_TYPE} \
    --topology=${TOPOLOGY} \
    --pathways-gcs-location="gs://${BUCKET_NAME}/pathways-artifacts" \
    --image=jupyter/base-notebook \
    --command="start-notebook.sh"

下列 YAML 資訊清單說明如何使用 kubectl 建立自行託管的 Jupyter 筆記本執行個體。建立無頭 Pathways 叢集後，請套用下列 YAML 資訊清單。詳情請參閱「在互動模式下執行 Pathways」：

apiVersion: batch/v1
kind: Job
metadata:
  name: WORKLOAD
spec:
  template:
    spec:
      hostNetwork: true
      dnsPolicy: ClusterFirstWithHostNet
      restartPolicy: OnFailure
      containers:
      - name: jupyter-notebook
        image: jupyter/base-notebook  # Use the appropriate Jupyter image
        ports:
        - containerPort: 8888

hostNetwork: true 欄位可讓筆記本中的 JAX 用戶端，與 Pathways Proxy 伺服器完成應用程式層傳輸安全性 (ALTS) 握手。如果沒有這個欄位，在採用 Workload Identity Federation for GKE 的節點集區上，連線會失敗。dnsPolicy: ClusterFirstWithHostNet 欄位可讓筆記本解析叢集內 DNS 名稱，例如 Pathways head Pod 的名稱。詳情請參閱「排解 Cloud 上的 Pathways 問題」。

Cluster Director

您可以在叢集子網路中的現有 Compute Engine VM 上執行 Jupyter 筆記本，也可以在叢集的登入節點上建立自行託管的 Jupyter 筆記本。在這兩種情況下，筆記型電腦都會連線至 Proxy 伺服器，不需通訊埠轉送。

使用叢集子網路中的現有 VM

在 VM 的虛擬環境中安裝 JupyterLab 並啟動：

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

建立自行託管的 Jupyter 筆記本執行個體

在叢集的登入節點上，安裝並啟動 JupyterLab。使用叢集的 Python 在主目錄中建立虛擬環境。登入節點的系統 Python 版本過舊，不適用於目前的 JAX 版本，因此 pip 在該處找不到 jax==0.11.1：

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

連線至筆記本執行個體

GKE

使用通訊埠轉送功能，從本機電腦連線至筆記本：

Cluster Toolkit

MAIN_POD=$(kubectl get pods | grep ${WORKLOAD}-pathways-head | awk '{print $1}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

kubectl

MAIN_POD=$(kubectl get pods \
    --selector=batch.kubernetes.io/job-name=WORKLOAD \
    --output=jsonpath='{.items[0].metadata.name}')
kubectl port-forward pod/${MAIN_POD} 8888:8888

在本機瀏覽器中，前往 http://localhost:8888?token=TOKEN。將 TOKEN 替換為 Jupyter 筆記本容器記錄中的權杖：

kubectl logs ${MAIN_POD}

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

...
Or copy and paste one of these URLs:
  http://NODE_NAME:8888/lab?token=TOKEN
  http://127.0.0.1:8888/lab?token=TOKEN

NODE_NAME 是執行筆記本 Pod 的 GKE 節點主機名稱 (例如 gke-my-cluster-cpu-np-1a2b3c4d-ef56)。

Cluster Director

如果筆記本位於 VM 或登入節點，請使用通訊埠轉送功能，從本機電腦連線至該主機。例如，登入節點：

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

在本機瀏覽器中，前往 http://localhost:8888?token=TOKEN。將 TOKEN 替換為 JupyterLab 啟動時輸出的權杖。

Notebook 與 Pathways 叢集的連線

1. 在 JupyterLab 中建立新的 Python 3 筆記本。
2. 連線至 Pathways Proxy 伺服器。

在筆記本中新增儲存格，安裝 pathwaysutils，將 JAX_PLATFORMS 設為 proxy，並將 JAX_BACKEND_TARGET 設為 grpc://PROXY_ADDRESS:29000：

注意：PROXY_ADDRESS 的值取決於您在「以互動模式連線至 Pathways 叢集」中選擇的連線選項。在 Cluster Director 上，將 jax==0.11.1 安裝在同一個儲存格中。在 source ~/pathways/cluster.env 後啟動的自行託管筆記本中，變數已設定完成。

!pip install pathwaysutils
# Set JAX_PLATFORMS and JAX_BACKEND_TARGET if you're not using `gcluster job submit` to deploy the notebook
%env JAX_PLATFORMS=proxy
  # Replace your proxy address in the following line:
%env JAX_BACKEND_TARGET=grpc://PROXY_ADDRESS:29000

將 PROXY_ADDRESS 替換為「以互動模式連線至 Pathways 叢集」中的 Pathways Proxy 伺服器主機或 IP 位址 (例如 127.0.0.1、Cloud DNS 主機名稱或內部負載平衡器 IP 位址)。在 Cluster Director 上，如果您使用通訊埠轉送，請使用 127.0.0.1，或使用服務記錄中的 IFRT Proxy 位址。如果您使用 kubectl 部署自行託管的筆記本，請將 PROXY_ADDRESS 替換為 Pathways head Pod 的叢集內 DNS 名稱，讓 JAX_BACKEND_TARGET 值為 grpc://WORKLOAD-pathways-head-0-0.WORKLOAD:29000。

新增第二個儲存格，驗證連線能力並輸出 Pathways 叢集中的裝置：

import pathwaysutils
import jax

pathwaysutils.initialize()
print(jax.devices())

如果連線成功，輸出內容會包含一則訊息，指出 JAX 偵測到 Pathways-on-Cloud 後端。

列出的 JAX 裝置數量，與您建立 Pathways 叢集時指定的 TPU 晶片數量和配量數量相符。

將程式碼新增至筆記本

新增自己的 JAX 程式碼，並在 Pathways 叢集的 TPU 上以互動方式執行。下列程式碼示範如何從單一筆記本對兩個切片執行運算。如要執行已完成樣本，請將 WORKLOAD_NODEPOOL_COUNT 設為 2 以上，建立至少有兩個配量的 Pathways 工作負載。在 Cluster Director 上，將 MIN_SLICES 和 MAX_SLICES 設為 2 以上。如果只有一個切片，slice1_devices 會是空白，範例的最後一部分會失敗，並顯示 ValueError：

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np

# You can use JAX APIs as usual across any of the devices.
x = jax.device_put(np.pi / 2., jax.devices()[-1])
jax.jit(jnp.sin)(x)

# pmap can run across all devices on all slices
num_tpus = jax.device_count()
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i')
x = jnp.arange(num_tpus)
y = f(x)
print(y)

# You can also target devices from a specific slice
slice0_devices = [d for d in jax.devices() if d.slice_index == 0]
f = jax.pmap(lambda x: lax.psum(1, 'i'), 'i', devices=slice0_devices)
x = jnp.arange(len(slice0_devices))
y = f(x)
print(y)
print(y.global_shards)

# You can send data produced on one slice to another slice
slice1_devices = [d for d in jax.devices() if d.slice_index == 1]
g = jax.pmap(lambda x: x + lax.axis_index('i'), 'i', devices=slice1_devices)
z = g(y)
print(z)
print(z.global_shards)

刪除 Pathways 互動式工作負載

GKE

Cluster Toolkit

gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE

kubectl

kubectl delete job WORKLOAD

Cluster Director

在叢集的登入節點上，取消 Pathways 服務工作：

scancel -u ${USER} -n pathways-service

這個指令也會取消 Pathways worker 工作。TPU VM 是靜態節點，因此會持續運作，並可供其他 Slurm 工作使用。

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

後續步驟

• 建立採用 Pathways 的 GKE 叢集
• 建立採用 Pathways 的 Cluster Director 叢集
• 使用 Pathways 執行批次工作負載
• 將 JAX 工作負載移植到 Pathways
• 使用 Pathways 進行彈性訓練
• Troubleshoot Pathways on Cloud