Pathways 互動式工作負載是遠端 JAX 工作負載,在 Google Kubernetes Engine (GKE) 叢集的 Pathways 控制器元件外部執行,例如在 VM 上。與批次工作負載不同,互動式工作負載完成後,不會關閉 Pathways 叢集元件,其他 JAX 用戶端仍可連線至這些元件。本文以 Jupyter 筆記本為例,示範互動式工作負載。
JAX 使用者可透過 IFRT 介面,將指令傳送至 Pathways 叢集。無論是從終端機、筆記本或任何與 Python 相容的環境執行,JAX 程式碼都能與 Pathways 資源互動。
事前準備
請完成下列前置作業:
- 建立 Pathways GKE 叢集。
- 安裝 Cluster Toolkit。
- 安裝 Kubernetes 工具。
- 安裝 Google Cloud CLI。
- 啟用 Kubernetes Engine API。
- 請確認您已為 Google Cloud 專案啟用 Pathways。
在互動模式中執行 Pathways
請設定下列環境變數:
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。
如果尚未建立 Pathways 構件的 Cloud Storage bucket,請建立一個:
gcloud storage buckets create gs://${BUCKET_NAME} \ --project=${PROJECT_ID} \ --location=${REGION}在叢集上建立 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-headlessgcluster 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 伺服器。
在互動模式下連線至 Pathways 叢集
連線至 Pathways 叢集時,可以選擇是否要轉送通訊埠。請使用下列其中一個部分,連線至 Pathways 叢集。
使用通訊埠轉送功能連線
從任何可存取叢集控制層的主機,將本機通訊埠 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
從虛擬私有雲中的主機連線,不使用通訊埠轉送
如果不想使用通訊埠轉送,可以透過 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 叢集:
確認可從非 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
使用 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
Jupyter Notebook
您可以使用 Vertex AI 建立 Jupyter 筆記本,也可以建立自行託管的 Jupyter 筆記本。
建立 Vertex AI Workbench 執行個體
設定並驗證 Pathways 叢集後,即可透過 Vertex AI Workbench Jupyter 筆記本存取 GKE TPU VM。下列設定操作說明假設 GKE Pathways 叢集位於同一個虛擬私有雲網路 (除非您另行設定,否則為預設網路)。
- 前往 Google Cloud 控制台的「Vertex AI Workbench」頁面。
在「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網路。
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 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)。
Notebook 與 Pathways 叢集的連線
- 在 JupyterLab 中建立新的 Python 3 筆記本。
- 連線至 Pathways Proxy 伺服器。
在筆記本中新增儲存格,安裝 pathwaysutils,將 JAX_PLATFORMS 設為 proxy,並將 JAX_BACKEND_TARGET 設為 grpc://PROXY_ADDRESS:29000:
!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 位址)。如果您使用 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 工作負載。如果只有一個切片,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 互動式工作負載
Cluster Toolkit
gcluster job cancel WORKLOAD \ --cluster=CLUSTER \ --project=PROJECT_ID \ --location=ZONE
kubectl
kubectl delete job WORKLOAD