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

Pathways 互動式工作負載是遠端 JAX 工作負載，在 VM 中執行，不屬於託管 Pathways 叢集的 GKE 叢集。與批次工作負載不同，互動式工作負載作業完成後，Pathways 叢集元件不會關閉，其他 JAX 用戶端仍可連線。本文以 Jupyter 筆記本為例，示範互動式工作負載。

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

事前準備

請確認您已備妥以下項目：

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

在互動模式中執行 Pathways

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} --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
   

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

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

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

使用通訊埠轉送功能連線

此時，您可以使用通訊埠轉送功能 (從任何可存取叢集控制層的主機)，存取 Proxy 伺服器：

請依據工作負載使用適當的指令：

Cluster Toolkit

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

kubectl

PROXY_POD=$(kubectl get pods | grep pathways-${USER}-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

在同一個主機上開啟新的終端機視窗。設定 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

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

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

使用 Cloud DNS 連線

在叢集中啟用 Cloud DNS，會將 Cloud DNS 供應商從 kube-dns 切換為 Cloud DNS。啟用後，系統會在虛擬私有雲中為 Cloud DNS 名稱建立私人 Cloud DNS 區域。詳情請參閱「將 Cloud DNS 用於 GKE」。

如果您啟用 Cloud DNS，並設定叢集、附加型虛擬私有雲或虛擬私有雲範圍，Kubernetes Cloud DNS 名稱就能從虛擬私有雲內部的非 GKE VM 解析。名稱格式為 <service_name>.<namespace>.svc.<custom_dns_domain>。Pathways head Pod 有一項名為 <jobset_name>-pathways-head-0-0.<jobset_name>.<namespace>.svc.<custom_dns_domain> 的服務。

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

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

   Cluster Toolkit

   host WORKLOAD-pathways-head-0-0.WORKLOAD.default.svc.USERNAME-test

   畫面會顯示類似以下的輸出：

   <WORKLOAD>-pathways-head-0-0.<WORKLOAD>.default.svc.<user>-test has address 10.0.2.75

   kubectl

   host pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test

   畫面會顯示類似以下的輸出：

   pathways-<user>-pathways-head-0-0.pathways-<user>.default.svc.<user>-test has address 10.0.2.75

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

   Cluster Toolkit

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

   kubectl

   JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://pathways-USERNAME-pathways-head-0-0.pathways-USERNAME.default.svc.USERNAME-test: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 部署作業，請建立由內部負載平衡器支援的服務。叢集不需啟用 Cloud DNS。

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

gcloud container clusters update ${CLUSTER} \
  --project=${PROJECT} \
  [--zone=${ZONE} | --region=${REGION}] \
  --enable-l4-ilb-subsetting

啟用 ILB 子集後，您可以使用下列 YAML 建立 LoadBalancer 類型的 Kubernetes 服務。這會導致 GKE 在叢集的虛擬私有雲中建立內部負載平衡器：

apiVersion: v1
kind: Service
metadata:
  name: pathways-USERNAME-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: pathways-USER
    jobset.sigs.k8s.io/replicatedjob-name: pathways-head
  ports:
  - name: tcp-port
    protocol: TCP
    port: 29000
    targetPort: 29000

將 USER 更新為您的 Google Cloud 使用者 ID，並將檔案儲存為 pathways-headless-ilb.yaml。

套用資訊清單：

kubectl apply -f pathways-headless-ilb.yaml

負載平衡器建立完成後 (約一分鐘後)，EXTERNAL-IP 欄就會有值：

kubectl get services
NAME                  TYPE           CLUSTER-IP      EXTERNAL-IP   PORT(S)        AGE
pathways-$USER       ClusterIP      None            <none>        <none>         30m
pathways-$USER-ilb   LoadBalancer   34.118.232.46   10.0.0.22     80:31246/TCP   2m41s

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

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

注意：請將上一個指令中的 JAX 後端 IP 位址 (10.0.0.22) 替換為 kubectl get services 顯示的值。

畫面會顯示類似以下的輸出：

[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 位址會與上述指令和輸出範例中顯示的不同。

Jupyter Notebook

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

建立 Vertex AI Workbench 執行個體

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

使用「Create new」(新建) 按鈕，建立新的 Workbench 執行個體 (從「Instances」(執行個體) 分頁)。確認網路與 GKE 叢集的網路相同。您可以使用指令列建立新的 Workbench 執行個體。

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

注意：--network 旗標為選用項目，如省略，系統會使用預設網路。

建立執行個體後，前往該執行個體並點選「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。詳情請參閱「使用 kubectl 在互動模式下執行 Pathways」。

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

前往筆記本執行個體

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

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 | grep jupyter-notebook-USERNAME | awk '{print $1}')
  kubectl port-forward pod/${MAIN_POD} 8888:8888

在本機瀏覽器中前往 http://localhost:8888?token=<var>your-token</var>。將 <your-token> 替換為 Jupyter 筆記本容器記錄中的權杖。

kubectl logs ${MAIN_POD}

輸出內容應如下所示：

...
Or copy and paste one of these URLs:
  http://jupyter-notebook-<user>-bbbdh:8888/lab?token=<token>
  http://127.0.0.1:8888/lab?token=<token>

Notebook 與 Pathways 叢集的連線

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

在筆記本中新增儲存格，用於安裝 pathwaysutils、將 JAX_PLATFORMS 設為 proxy，以及將 JAX_BACKEND_TARGET 設為 PROXY_ADDRESS。

注意：PROXY_ADDRESS 的值取決於您在「連線至 Pathways Proxy 伺服器」中選擇的連線選項。

!pip install pathwaysutils
%env JAX_PLATFORMS=proxy
  # Replace your proxy address below:
%env JAX_BACKEND_TARGET=PROXY_ADDRESS

新增第二個儲存格做為「hello world」類型檢查，並列印 Pathways 叢集中的裝置。

import pathwaysutils
import jax

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

如果一切運作正常，應該會看到訊息，指出系統偵測到 Pathways-on-Cloud 後端。

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

將程式碼新增至筆記本

新增自己的 JAX 程式碼，並在 Pathways 叢集的 TPU 上以互動方式執行。下列程式碼說明如何從單一筆記本對兩個切片執行運算。

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.
jax.jit(jnp.sin, device=jax.devices()[-1])(np.pi / 2.)

# 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 互動式工作負載

Cluster Toolkit

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

kubectl

kubectl delete -f pathways-headless-workload.yaml

後續步驟

• 建立採用 Pathways 的 GKE 叢集
• 使用 Pathways 執行批次工作負載
• 將 JAX 工作負載移植到 Pathways
• 使用 Pathways 進行彈性訓練
• 排解雲端 Pathways 問題