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

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

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

事前準備

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

在互動模式中執行 Pathways

您可以使用 xpkkubectl,以互動模式執行 Pathways。

XPK

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

    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 專案 ID
    • ZONE:您打算執行工作負載的可用區
    • CLUSTER:GKE 叢集名稱
  2. 在叢集上建立 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 容器。

  1. 替換預留位置、複製下列 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_RESTARTSPathwaysJob 最多可重新啟動的次數
    • TPU_MACHINE_TYPE:要使用的 TPU 機型,例如支援的值:「ct6e-standard-8t」、「ct5p-hightpu-4t」
    • TOPOLOGY:TPU 拓撲
    • WORKLOAD_NODEPOOL_COUNT:路徑工作負載使用的節點集區數量
    • BUCKET_NAME:用於儲存暫存檔案的 Cloud Storage bucket
    如要變更先前 YAML 檔案中 WORKLOAD_NODEPOOL_COUNT 指定的節點集區數量 (路徑工作站副本),請刪除 PathwaysJob 並建立新的 PathwaysJob,其中包含更新後的節點集區數量。此外,您也需要重新啟動所有已連線的筆記本,才能與新的 Pathways 叢集建立連線。
  2. 套用 pathways-headless-workload.yaml 檔案:
      kubectl apply -f pathways-headless-workload.yaml
      
  3. 執行 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_PLATFORMSJAX_BACKEND_TARGET 環境變數,然後執行匯入 pathwaysutilsjax 的 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 叢集:

  1. 確認可從非 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
  2. 使用 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 叢集的連線

  1. 在 Jupyterlab 中,建立新的 Python 3 筆記本
  2. 連線至 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

後續步驟