Pathways 互動式工作負載是遠端 JAX 工作負載,在不屬於 Pathways 叢集主機的 GKE 叢集的 VM 中執行。與批次工作負載不同,互動式工作負載作業完成後,不會關閉 Pathways 叢集元件,其他 JAX 用戶端仍可連線。本文以 Jupyter 筆記本為例,示範如何使用互動式工作負載。
JAX 使用者可透過 IFRT 介面將指令傳送至 Pathways 叢集。無論是從終端機、筆記本或任何與 Python 相容的環境執行,JAX 程式碼都能與 Pathways 資源順暢互動。
事前準備
請確認您已備妥以下項目:
- 使用 XPK 建立 GKE 叢集
- 已安裝的 XPK
- 已安裝的 Kubernetes 工具
- 已安裝 gcloud CLI
- 已啟用 TPU API
- 已啟用 Google Kubernetes Engine API
- 確認專案已啟用 Pathways Google Cloud
在互動模式中執行 Pathways
您可以使用 xpk 或 kubectl,以互動模式執行 Pathways。
XPK
請設定下列環境變數:
export WORKLOAD=WORKLOAD export WORKLOAD_NODEPOOL_COUNT=WORKLOAD_NODEPOOL_COUNT export TPU_TYPE=TPU_TYPE export PROJECT_ID=PROJECT export ZONE=ZONE \ export CLUSTER=CLUSTER
更改下列內容:
WORKLOAD:將此項設為專屬名稱,以識別工作負載WORKLOAD_NODEPOOL_COUNT:路徑工作負載使用的節點集區數量TPU_TYPE:TPU 類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的 TPU 類型,請參閱「TPU 版本」一節。PROJECT:您的 Google Cloud 專案 IDZONE:您打算執行工作負載的可用區CLUSTER:GKE 叢集名稱
在叢集上建立 Pathways 容器。如要執行無頭工作負載,請執行下列指令:
xpk workload create-pathways \ --headless \ --workload=${WORKLOAD} \ --num-slices=${WORKLOAD_NODEPOOL_COUNT} \ --tpu-type=${TPU_TYPE} \ --project=${PROJECT} \ --zone=${ZONE} \ --cluster=${CLUSTER}
此時,JAX 工作負載即可連線至 IFRT Proxy 伺服器。
kubectl
下列 YAML 與批次工作負載 YAML 相同,但未指定 main 容器。
- 替換預留位置、複製下列 YAML,然後貼到名為
pathways-headless-workload.yaml的檔案中。 取代下列項目:apiVersion: pathways-job.pathways.domain/v1 kind: PathwaysJob metadata: name: pathways-USERNAME spec: maxRestarts: MAX_RESTARTS workers: - type: TPU_MACHINE_TYPE topology: TOPOLOGY numSlices: WORKLOAD_NODEPOOL_COUNT pathwaysDir: gs://BUCKET_NAME controller: deploymentMode: default
USERNAME:您的使用者名稱MAX_RESTARTS:PathwaysJob最多可重新啟動的次數TPU_MACHINE_TYPE:要使用的 TPU 機型,例如支援的值:「ct6e-standard-8t」、「ct5p-hightpu-4t」TOPOLOGY:TPU 拓撲WORKLOAD_NODEPOOL_COUNT:路徑工作負載使用的節點集區數量BUCKET_NAME:用於儲存暫存檔案的 Cloud Storage bucket
WORKLOAD_NODEPOOL_COUNT指定的節點集區數量 (路徑工作站副本),請刪除PathwaysJob並建立新的PathwaysJob,其中包含更新後的節點集區數量。此外,您也需要重新啟動所有已連線的筆記本,才能與新的 Pathways 叢集建立連線。 - 套用
pathways-headless-workload.yaml檔案:kubectl apply -f pathways-headless-workload.yaml
- 執行
kubectl get pods,確認 Pod 中的所有容器都在執行中。以下輸出內容適用於 2 個配量的 v5p 2x2x2,其中USER是執行指令的使用者 ID:NAME READY STATUS RESTARTS AGE pathways-USER-pathways-head-0-0-n848j 2/2 Running 0 49s pathways-USER-pathways-workers-0-0-jxt2z 1/1 Running 0 71s pathways-USER-pathways-workers-0-1-cxmhc 1/1 Running 0 70s pathways-USER-pathways-workers-1-0-5kmz9 1/1 Running 0 71s pathways-USER-pathways-workers-1-1-vg5n4 1/1 Running 0 71s
在互動模式中連線至 Pathways 叢集
您可以連線至 Pathways 叢集,無論是否要進行連接埠轉送都可以。請使用下列其中一個章節,連線至 Pathways 叢集。
使用通訊埠轉送連線
此時,您可以使用通訊埠轉送 (從任何可存取叢集控制層的主機),存取 Proxy 伺服器:
請使用適合工作負載的指令:
XPK
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 區域。詳情請參閱「使用 GKE 適用的 Cloud DNS」。
如果您啟用叢集、附加型 VPC 或 VPC 範圍的 Cloud DNS,虛擬私有雲內的非 GKE VM 就能解析 Kubernetes Cloud DNS 名稱。名稱格式為 <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 叢集:
確認可從非 GKE 主機解析領導者 Cloud DNS 項目:
XPK
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
使用 Cloud DNS 名稱連線至 Pathways 叢集:
XPK
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 位址,請建立由內部負載平衡器支援的服務。這項操作不需要叢集啟用 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 在叢集的 VPC 內建立內部負載平衡器:
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
您可以在與叢集相同的虛擬私有雲中,存取路徑部署作業,不必在主機上轉送連接埠:
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())'
畫面會顯示類似以下的輸出:
[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
Jupyter Notebook
您可以使用 Vertex AI 建立 Jupyter 筆記本,也可以建立自行代管的 Jupyter 筆記本。
建立 Vertex AI Workbench 執行個體
設定並驗證 Pathways 叢集後,即可透過 Vertex AI Jupyter 筆記本存取 GKE TPU VM。下列設定操作說明假設您的 GKE Pathways 叢集位於相同的虛擬私有雲網路 (除非您已設定其他網路,否則這是預設網路)。前往 Vertex AI Workbench 控制台。
使用「建立新執行個體」按鈕,在「執行個體」分頁中建立新的 Workbench 執行個體。確認網路與 GKE 叢集的網路相同。您可以使用指令列建立新的 Workbench 執行個體。
gcloud workbench instances create INSTANCE_NAME \ --machine-type=e2-standard-4 \ --data-disk-size=100 \ --location=ZONE \ [--network=NETWORK]
執行個體建立完成後,請前往該執行個體並點選「Open Jupyterlab」。
建立自行代管的 Jupyter 筆記本執行個體
以下指令說明如何使用 XPK 建立自行代管的 Jupyter Notebook 執行個體:
xpk workload create-pathways \
--workload=${WORKLOAD} \
--num-slices=${WORKLOAD_NODEPOOL_COUNT} \
--tpu-type=${TPU_TYPE} \
--project=${PROJECT} \
--zone=${ZONE} \
--cluster=${CLUSTER} \
--docker-image=jupyter/base-notebook \
--command "start-notebook.sh"
下列 YAML 顯示如何使用 kubectl 建立自行代管的 Jupyter Notebook 執行個體。建立無頭 Pathways 叢集後,請套用下列 YAML。詳情請參閱「使用 kubectl 以互動模式執行路徑」。
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
前往筆記本執行個體
使用通訊埠轉送功能,從本機電腦連線至筆記本:
XPK
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>
筆記本與 Pathways 叢集的連線
- 在 Jupyterlab 中,建立新的 Python 3 筆記本
- 連線至 Pathways Proxy 伺服器
在筆記本中新增儲存格,安裝 pathwaysutils、將 JAX_PLATFORMS 設為 proxy,並將 JAX_BACKEND_TARGET 設為 PROXY_ADDRESS。
!pip install pathwaysutils %env JAX_PLATFORMS=proxy # Replace your proxy address below: %env JAX_BACKEND_TARGET=PROXY_ADDRESS
新增第二個儲存格做為「Hello World」型別檢查,並在路徑叢集中列印裝置。
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 互動式叢集
XPK
xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE
kubectl
kubectl delete -f pathways-headless-workload.yaml
後續步驟
- 使用 Pathways 建立 GKE 叢集
- 使用 Pathways 進行多主機推論
- 使用 Pathways 批次處理工作負載
- 路徑互動模式
- 將 JAX 工作負載移植到 Pathways
- 透過 Pathways 進行彈性訓練
- 疑難排解路徑