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

预览版

此产品或功能 受服务专用条款的约束。 非正式版产品和功能“按原样”提供，可能只能获得有限的支持。 如需了解详情，请参阅发布阶段说明。

Pathway 交互式工作负载是在虚拟机（不属于托管 Pathway 集群的 GKE 集群）内运行的远程 JAX 工作负载。与批处理工作负载不同，交互式工作负载操作完成后，系统不会关闭 Pathways 集群组件，这些组件仍可供其他 JAX 客户端连接。本文档使用 Jupyter 笔记本作为示例来演示交互式工作负载。

通过 IFRT 接口，JAX 用户可以向 Pathways 集群发送命令。JAX 代码（无论是在终端、笔记本还是任何与 Python 兼容的环境中执行）都可以与 Pathways 资源无缝交互。

准备工作

请确保您已备妥：

• 使用 XPK 创建了 GKE 集群
• 已安装的 XPK
• 已安装的 Kubernetes 工具
• 安装了 gcloud CLI
• 已启用 TPU API
• 已启用 Google Kubernetes Engine API
• 确保已为您的 Google Cloud 项目启用 Pathways

以互动模式运行 Pathway

您可以使用 xpk 或 kubectl 在互动模式下运行 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：Pathways 工作负载使用的节点池数量
   • 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 代理服务器。

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_RESTARTS：PathwaysJob 可重启的最大次数
   • TPU_MACHINE_TYPE：您要使用的 TPU 机器类型，支持的值示例：“ct6e-standard-8t”“ct5p-hightpu-4t”
   • TOPOLOGY：TPU 拓扑
   • WORKLOAD_NODEPOOL_COUNT：Pathways 工作负载使用的节点池数量
   • BUCKET_NAME：用于存储临时文件的 Cloud Storage 存储桶
   如需更改上一个 YAML 中由 WORKLOAD_NODEPOOL_COUNT 指定的节点池数量（pathways-worker 副本），您需要删除此 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 集群。

使用端口转发进行连接

此时，您可以使用端口转发（从任何有权访问集群控制平面的主机）来访问代理服务器：

使用适合您工作负载的命令：

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

从 VPC 中的主机连接，无需使用端口转发

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

使用 Cloud DNS 进行连接

在集群中启用 Cloud DNS 会将 Cloud DNS 提供商从 kube-dns 切换到 Cloud DNS。启用后，系统会在您的 Virtual Private Cloud 中为 Cloud DNS 名称创建一个专用 Cloud DNS 区域。如需了解详情，请参阅使用 Cloud DNS for GKE。

如果您启用 Cloud DNS 并指定集群范围、附加 VPC 范围或 VPC 范围，则 Kubernetes Cloud DNS 名称可从 Virtual Private Cloud 内的非 GKE 虚拟机解析。名称的格式为 <service_name>.<namespace>.svc.<custom_dns_domain>。Pathways 头 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

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

对于 VPC 中指向您的 Pathways 部署的专用 IP 地址，请创建一个由内部负载均衡器提供支持的服务。这不需要您的集群启用 Cloud DNS。

对于具有大量虚拟机的集群，如果您要创建内部负载平衡器，建议您启用 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

负载均衡器创建后（大约 1 分钟后），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

您无需在与集群位于同一 VPC 中的主机上进行端口转发，即可访问 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 地址将与上述命令和输出示例中显示的 IP 地址不同。

Jupyter 笔记本

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

创建 Vertex AI Workbench 实例

设置并验证 Pathways 集群后，您可以通过 Vertex AI Jupyter 笔记本访问 GKE TPU 虚拟机。以下设置说明假设您的 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]

注意： --network 标志是可选的，如果您省略该标志，系统会使用默认网络。

创建实例后，前往该实例，然后点击打开 Jupyterlab。

创建自托管的 Jupyter 笔记本实例

以下命令展示了如何使用 XPK 创建自托管 Jupyter 笔记本实例：

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 笔记本实例。在创建无头 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

前往笔记本实例

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

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

在笔记本中，添加一个用于安装 pathwaysutils 的单元，将 JAX_PLATFORMS 设置为 proxy，并将 JAX_BACKEND_TARGET 设置为 PROXY_ADDRESS。

注意： PROXY_ADDRESS 的值取决于您在连接到 Pathways 代理服务器中选择的连接选项。

!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 交互式集群

XPK

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

后续步骤

• 通过学习路线创建 GKE 集群
• 使用 Pathways 进行多主机推理
• 使用 Pathways 的批量工作负载
• Pathway 互动模式
• 将 JAX 工作负载移植到 Pathways
• 通过 Pathways 进行弹性训练
• 问题排查途径