使用 Pathways 运行交互式工作负载

Pathways 交互式工作负载是在 Google Kubernetes Engine (GKE) 集群的 Pathways 控制器组件之外(例如在虚拟机上)运行的远程 JAX 工作负载。与批处理工作负载不同,交互式工作负载完成后,不会关闭 Pathways 集群组件,这些组件仍可供其他 JAX 客户端连接。本文档以 Jupyter 笔记本为例来演示交互式工作负载。

JAX 用户使用 IFRT 接口向 Pathways 集群发送命令。JAX 代码(无论是从终端、笔记本还是任何与 Python 兼容的环境执行)都可以与 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:存储桶和集群所在的区域。
    • ZONE:您计划运行工作负载的可用区。
    • CLUSTER:GKE 集群的名称
    • BUCKET_NAME:用于 Pathways 制品的 Cloud Storage 存储桶。
  2. 如果您尚未为 Pathways 工件创建 Cloud Storage 存储桶,请创建一个:

    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 头 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 代理服务器。

以交互模式连接到 Pathways 集群

无论是否使用端口转发,您都可以连接到 Pathways 集群。使用以下部分之一连接到 Pathways 集群。

使用端口转发进行连接

从有权访问集群控制平面的任何主机,将本地端口 29000 转发到 Pathways 头 Pod 以访问代理服务器:

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

从 VPC 中的主机连接,而不使用端口转发

如果您不想使用端口转发,可以使用 Cloud DNS 或内部负载均衡器连接到 Pathways 集群。

使用 Cloud DNS 进行连接

在集群中启用 Cloud DNS 会将 DNS 提供商从 kube-dns 切换到 Cloud DNS。启用后,GKE 会在您的虚拟私有云 (VPC) 网络中为 DNS 名称创建专用 Cloud DNS 区域。如需了解详情,请参阅使用 Cloud DNS for GKE。

如果您启用具有 VPC 范围或额外 VPC 范围的 Cloud DNS,则可以从 VPC 网络内的非 GKE 虚拟机解析 Kubernetes DNS 名称。在集群范围内,只有集群中的节点才能解析这些名称。名称的格式为 SERVICE_NAME.NAMESPACE.svc.CUSTOM_DNS_DOMAIN。 Pathways 头 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
    

使用内部负载均衡器进行连接

对于 VPC 网络中指向 Pathways 部署的专用 IP 地址,请创建一个由内部负载均衡器 (ILB) 支持的 Service。这不需要您的集群启用 Cloud DNS。

对于具有许多虚拟机的集群,如果您创建内部负载均衡器,请启用 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 在集群的 VPC 网络内创建内部负载均衡器:

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

您无需进行端口转发,即可从与集群位于同一 VPC 网络的宿主机访问 Pathways Deployment:

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 笔记本

您可以使用 Vertex AI 创建 Jupyter 笔记本,也可以创建自托管 Jupyter 笔记本。

创建 Vertex AI Workbench 实例

设置并验证 Pathways 集群后,您可以从 Vertex AI Workbench Jupyter 笔记本访问 GKE TPU 虚拟机。以下设置说明假定您的 GKE Pathways 集群位于同一 VPC 网络中(除非您另有配置,否则该网络为默认网络)。

  1. 前往 Google Cloud 控制台中的 Vertex AI Workbench 页面。
  2. 在实例标签页上,点击新建,并确保网络与 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 集群的 VPC 网络的名称。--network 标志是可选的;如果省略该标志,该命令将使用 default 网络。
  3. Vertex AI 创建实例后,前往该实例并点击打开 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 代理服务器完成应用层传输安全 (ALTS) 握手。如果没有此字段,则在使用 Workload Identity Federation for GKE 的节点池上,连接会失败。dnsPolicy: ClusterFirstWithHostNet 字段可让笔记本解析集群内 DNS 名称,例如 Pathways 头 Pod 的名称。如需了解详情,请参阅排查 Pathways on Cloud 问题。

连接到笔记本实例

使用端口转发从本地机器连接到笔记本:

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)。

笔记本与 Pathways 集群的连接

  1. 在 JupyterLab 中,创建一个新的 Python 3 笔记本。
  2. 连接到 Pathways 代理服务器。

在笔记本中,添加一个单元以安装 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 代理服务器的主机或 IP 地址(例如 127.0.0.1、Cloud DNS 主机名或内部负载均衡器 IP 地址)。如果您使用 kubectl 部署了自托管笔记本,请将 PROXY_ADDRESS 替换为 Pathways 头 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

后续步骤