Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

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

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

通过 IFRT 接口，JAX 用户可以向 Pathways 集群发送命令。JAX 代码（无论是在终端、笔记本还是任何兼容 Python 的环境中执行）都可以与 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 placement 拓扑（例如 2x4）
- PROJECT_ID：您的 Google Cloud 项目 ID
- REGION：您的存储桶和集群所在的区域
- ZONE：您计划运行工作负载的可用区
- CLUSTER：GKE 集群的名称
- BUCKET_NAME：用于存放 Pathways 制品的 Cloud Storage 存储桶
如果您尚未为 Pathways 制品创建 Cloud Storage 存储桶，请创建一个：
```
gcloud storage buckets create gs://${BUCKET_NAME} --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-headless

此时，您的 JAX 工作负载可以连接到 IFRT 代理服务器。

以互动模式连接到 Pathways 集群

您可以通过端口转发或不通过端口转发来连接到 Pathways 集群。请根据以下部分之一连接到 Pathways 集群。

使用端口转发进行连接

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

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

Cluster Toolkit

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。

如果您启用 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 集群：

确认可从非 GKE 主机解析主 Cloud DNS 条目：

Cluster Toolkit

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

使用 Cloud DNS 名称连接到 Pathways 集群：

Cluster Toolkit

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())'

您看到的输出结果应该类似于以下内容：

[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 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]

创建实例后，前往该实例并点击打开 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。如需了解详情，请参阅使用 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

前往笔记本实例

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

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 | 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 集群的连接

在 Jupyterlab 中，创建一个新的 Python 3 笔记本
连接到 Pathways 代理服务器

在笔记本中，添加一个用于安装 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”类型的检查，并打印 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 互动式工作负载

Cluster Toolkit

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

kubectl

kubectl delete -f pathways-headless-workload.yaml

使用 Pathways 运行交互式工作负载 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

准备工作

以互动模式运行 Pathways

以互动模式连接到 Pathways 集群

使用端口转发进行连接

Cluster Toolkit

kubectl

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

使用 Cloud DNS 进行连接

Cluster Toolkit

kubectl

Cluster Toolkit

kubectl

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

Jupyter 笔记本

创建 Vertex AI Workbench 实例

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

前往笔记本实例

Cluster Toolkit

kubectl

笔记本电脑与 Pathways 集群的连接

将代码添加到笔记本

删除 Pathways 互动式工作负载

Cluster Toolkit

kubectl

后续步骤

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