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

预览版

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

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

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

准备工作

请确保您已备妥：

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

以互动模式运行 Pathways

您可以使用 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 工作器副本），您需要删除 此 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。启用后，系统会在虚拟私有云中为 Cloud DNS 名称创建专用 Cloud DNS 区域。如需了解详情，请参阅 使用 Cloud DNS for GKE。

如果您使用集群、附加 VPC 或 VPC 范围启用 Cloud DNS，则 Kubernetes Cloud DNS 名称可从 Virtual Private Cloud 内的非 GKE 虚拟机解析。这些名称的 格式为 <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. 确认领导者 Cloud DNS 条目可从非 GKE 主机解析：

   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

使用您的 Google Cloud 用户 ID 更新 USER，并将该文件另存为 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 开发者在线课程集群位于同一虚拟私有云网络中（除非您另行配置，否则这是默认网络）。前往 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. 连接到开发者在线课程代理服务器

在笔记本中，添加一个单元以安装 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

后续步骤

• 使用 Pathways 创建 GKE 集群
• 使用 Pathways 进行多主机推理
• 使用 Pathways 处理批量工作负载
• Pathways 互动模式
• 将 JAX 工作负载移植到 Pathways
• 使用 Pathways 进行弹性训练
• 排查 Pathways 问题