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

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

本页面提供了适用于这两个平台的说明。在每个部分中，选择 GKE 或 Cluster Director 标签页。

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

准备工作

• 选择标签页以了解您打算如何使用本页面上的示例：

  gcloud

  安装 Google Cloud CLI，然后 使用联合身份登录 gcloud CLI。 登录后，运行以下命令来初始化 Google Cloud CLI：

  gcloud init

  REST

  如需在本地开发环境中使用本页面上的 REST API 示例，请使用您提供给 gcloud CLI 的凭证。

  安装 Google Cloud CLI，然后 使用联合身份登录 gcloud CLI。

  如需了解详情，请参阅 Google Cloud 身份验证文档中的使用 REST 时进行身份验证。

准备工作

GKE

满足以下前提条件：

• 创建具有 Pathways 的 GKE 集群。
• 安装 Cluster Toolkit。
• 安装 Kubernetes 工具。
• 安装 Google Cloud CLI。
• 启用 Kubernetes Engine API。
• 确保您的 Google Cloud 项目已启用 Pathways。

Cluster Director

满足以下前提条件：

• 创建 Cluster Director 集群并设置 Pathways 环境。
• 安装 gcloud CLI。
• 验证您是否拥有对集群登录节点的 SSH 访问权限。
• 在运行 JAX 客户端的主机上安装 Python 3 和 venv 模块。在 Debian 和 Ubuntu 上，安装 python3-venv 软件包。

在交互模式下运行 Pathways

GKE

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

Cluster Director

1. 连接到集群的登录节点：

   gcloud compute ssh CLUSTER_ID-login-001 \
       --project=PROJECT_ID \
       --zone=ZONE \
       --tunnel-through-iap

   替换以下内容：

   • CLUSTER_ID：Cluster Director 集群的名称。
   • PROJECT_ID：您的Google Cloud 项目 ID。
   • ZONE：集群所在的可用区。

2. 设置以下环境变量：

   export SLICE_SHAPE=TOPOLOGY
   export MAX_SLICES=NUM_SLICES
   export MIN_SLICES=NUM_SLICES
   export GCS_BUCKET=gs://BUCKET_NAME

   替换以下内容：

   • TOPOLOGY：每个切片的 TPU 拓扑，例如 2x4。它必须与 TPU 节点集的拓扑匹配。
   • NUM_SLICES：TPU 切片的数量。将 MIN_SLICES 和 MAX_SLICES 设置为相同的值，以便仅在所有切片都运行时 Pathways 集群才准备就绪。
   • BUCKET_NAME：用于存储 Pathways 制品的现有 Cloud Storage 存储桶。

   如果您的 TPU 分区不是集群的默认分区，请将 PARTITION 设置为其名称。

3. 在无头模式下启动 Pathways 集群：

   cd ~/pathways
   sbatch start_pathways.sh
   

4. 等待服务日志 pathways-service-JOB_ID.out（其中 JOB_ID 是 Slurm 任务 ID）显示类似于以下内容的输出：

   Pathways Runtime Service is READY!
    Discovery file    : /home/USER/pathways/cluster.env
    IFRT Proxy Target : grpc://RM_IP:29000
   

   RM_IP 是代理服务器的内部 IP 地址。文件 ~/pathways/cluster.env 包含相同的地址。

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

以交互模式连接到 Pathways 集群

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

在 Cluster Director 上，JAX 客户端必须使用与 Pathways 服务器相同的 JAX 版本，即 0.11.1。

使用端口转发进行连接

GKE

从有权访问集群控制平面的任何主机，将本地端口 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

Cluster Director

从可通过 SSH 连接到集群登录节点的任何主机（例如您的工作站或集群子网外部的虚拟机），通过登录节点转发代理端口：

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 29000:RM_IP:29000

注意：如果您在具有默认服务账号的 Compute Engine 虚拟机上运行 gcloud CLI，则该虚拟机需要 cloud-platform 访问权限范围，否则 gcloud compute ssh 会失败并显示 insufficient authentication scopes。 或者，在虚拟机上运行 gcloud auth login 以使用您自己的账号。

在同一主机上，打开一个新的终端窗口。设置 JAX_PLATFORMS 和 JAX_BACKEND_TARGET 环境变量，然后运行导入 pathwaysutils 和 jax 的 Python 脚本。主机没有通往 RM_IP 的路由，因此客户端会连接到隧道的本地端 127.0.0.1:29000：

python3 -m venv .venv
source .venv/bin/activate
pip install pathwaysutils jax==0.11.1

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(288,TPU_DEVICE,coords=[0,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 device(289,TPU_DEVICE,coords=[1,0,0,0],logical_task=0,slice=0,default_mem=device,mem_spaces=3),
 ...
 device(313,TPU_DEVICE,coords=[1,3,0,0],logical_task=1,slice=1,default_mem=device,mem_spaces=3)]
W1029 22:52:22.897483    2361 grpc_client.cc:91] IFRT proxy server disconnected: CANCELLED: Stream removed (CANCELLED)

当脚本退出并关闭其连接时，系统会显示 IFRT proxy server disconnected 警告。这是正常现象。

注意：每个 JAX 调用都会通过 SSH 隧道。使用端口转发来验证连接。对于笔记本，请在集群的子网中运行 JAX 客户端。

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

GKE

如果您不想使用端口转发，可以使用 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

注意：您的 IP 地址与上述命令和输出示例中显示的地址不同。

Cluster Director

集群子网中的主机（例如 Compute Engine 虚拟机）可以直接连接到代理服务器。Cluster Director 的防火墙允许来自集群子网的流量，但不允许来自同一 VPC 网络的其他子网的流量：

JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://RM_IP:29000 python -c 'import pathwaysutils; import jax; import pprint; pathwaysutils.initialize(); pprint.pprint(jax.devices())'

将 RM_IP 替换为服务日志中的 IFRT 代理地址。

注意：IFRT 代理服务器不会对客户端进行身份验证。任何可以访问 RM_IP:29000 的主机都可以使用 Pathways 集群。

Jupyter 笔记本

GKE

您可以使用 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 Director

您可以在集群子网中的现有 Compute Engine 虚拟机上运行 Jupyter 笔记本，也可以在集群的登录节点上创建自托管式 Jupyter 笔记本。在这两种情况下，笔记本都会连接到代理服务器，而无需端口转发。

使用集群子网中的现有虚拟机

在虚拟机上，在虚拟环境中安装 JupyterLab 并启动它：

python3 -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
jupyter lab --no-browser --port=8888

创建自托管 Jupyter 笔记本实例

在集群的登录节点上，安装 JupyterLab 并启动它。使用集群的 Python 在主目录中创建虚拟环境。登录节点的系统 Python 版本太旧，无法支持当前的 JAX 版本，因此 pip 无法在其中找到 jax==0.11.1：

cd ~
/home/pathways/venv/bin/python -m venv .venv
source .venv/bin/activate
pip install jupyterlab pathwaysutils jax==0.11.1
source ~/pathways/cluster.env
jupyter lab --no-browser --port=8888

连接到笔记本实例

GKE

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

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

Cluster Director

对于虚拟机或登录节点上的笔记本，请使用端口转发从本地机器连接到该主机。例如，对于登录节点：

gcloud compute ssh CLUSTER_ID-login-001 \
    --project=PROJECT_ID \
    --zone=ZONE \
    --tunnel-through-iap \
    -- -N -L 8888:localhost:8888

在本地浏览器中，前往 http://localhost:8888?token=TOKEN。将 TOKEN 替换为 JupyterLab 启动时输出的令牌。

笔记本与 Pathways 集群的连接

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

在笔记本中，添加一个单元以安装 pathwaysutils，将 JAX_PLATFORMS 设置为 proxy，并将 JAX_BACKEND_TARGET 设置为 grpc://PROXY_ADDRESS:29000：

注意：PROXY_ADDRESS 的值取决于您在以交互模式连接到 Pathways 集群中选择的连接选项。在 Cluster Director 上，在同一单元中安装 jax==0.11.1。在您于 source ~/pathways/cluster.env 之后启动的自托管笔记本上，这些变量已设置。

!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 地址）。在 Cluster Director 上，如果您使用端口转发，请使用 127.0.0.1，或者使用服务日志中的 IFRT 代理地址。如果您使用 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 工作负载。在 Cluster Director 上，将 MIN_SLICES 和 MAX_SLICES 设置为 2 或更大值。如果只有一个切片，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 交互式工作负载

GKE

Cluster Toolkit

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

kubectl

kubectl delete job WORKLOAD

Cluster Director

在集群的登录节点上，取消 Pathways 服务作业：

scancel -u ${USER} -n pathways-service

此命令还会取消 Pathways 工作器作业。TPU 虚拟机是静态节点，因此它们会继续运行并可供其他 Slurm 作业使用。

如果您遇到问题，请参阅排查 Pathways on Cloud 问题。

后续步骤

• 创建具有 Pathways 的 GKE 集群
• 创建具有 Pathways 的 Cluster Director 集群
• 使用 Pathways 运行批量工作负载
• 将 JAX 工作负载移植到 Pathways
• 使用 Pathways 进行弹性训练
• 排查 Pathways on Cloud 问题