创建具有 Pathways 的 Cluster Director 集群

本页面向平台管理员和机器学习 (ML) 工程师介绍了如何创建 Slurm 编排的 Cluster Director 集群，并将其配置为运行 Pathways 工作负载。完成这些步骤后，您将预配一个包含 TPU v6e 切片、CPU 控制器节点、共享 Google Cloud Managed Lustre 文件系统和 /home/pathways 运行时环境的集群。Pathways 未预安装在 Cluster Director 集群上：创建集群后，您需要将 Pathways 脚本下载到登录节点，并运行一次设置脚本。

准备工作

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

    gcloud

    在 Google Cloud 控制台中，激活 Cloud Shell。

    激活 Cloud Shell

    Cloud Shell 会话随即会在 Google Cloud 控制台的底部启动，并显示命令行提示符。Cloud Shell 是一个已安装 Google Cloud CLI 且已为当前项目设置值的 Shell 环境。该会话可能需要几秒钟时间来完成初始化。

    REST

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

      安装 Google Cloud CLI。

      如果您使用的是外部身份提供方 (IdP)，则必须先使用联合身份登录 gcloud CLI。

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

准备工作

在开始之前，请完成以下前提条件：

  • 安装 Google Cloud CLI。
  • 启用 Hypercompute Cluster API。
  • 确保您的用户账号和 Compute Engine 默认服务账号具有创建和连接到 Cluster Director 集群所需的 Identity and Access Management (IAM) 角色。如需查看角色列表，请参阅 Cluster Director 快速入门中的准备工作。
  • 在您计划创建集群的可用区中创建 TPU v6e 预留（例如，为两个 2x4 切片预留至少四个 ct6e-standard-4t 虚拟机）。
  • 在集群节点上配置无密码 sudo，以便自动恢复切片。请让您的管理员为您授予项目的 Compute OS Admin Login (roles/compute.osAdminLogin) 角色。
  • 为 Pathways 暂存数据和工作负载输出创建 Cloud Storage 存储桶，并向您的用户账号和 Compute Engine 默认服务账号授予项目的 Storage Object User (roles/storage.objectUser) 角色。

设置本地环境

使用您的 Google Cloud 凭证登录：

gcloud auth login

使用适合您的工作负载的值定义以下环境变量：

export PROJECT=PROJECT_ID
export REGION=REGION
export ZONE=ZONE
export CLUSTER_ID=CLUSTER_ID
export RESERVATION_NAME=RESERVATION_NAME

替换以下内容：

  • PROJECT_ID：您的Google Cloud 项目 ID。
  • REGION：您在其中创建资源的区域。
  • ZONE：您在其中创建资源的可用区。
  • CLUSTER_ID：集群名称。 名称最多可包含 10 个字符，并且只能使用数字和小写字母。
  • RESERVATION_NAME：ZONE 中 TPU v6e 预留的名称。

创建 Cluster Director 集群

在以下示例中，您将使用预留创建包含两个 TPU v6e 2x4 切片（四个 ct6e-standard-4t 虚拟机）的集群，以及一个用于 Pathways 控制器的 CPU 节点。所有节点都会在 /home 处装载共享的 Managed Lustre 文件系统。如需详细了解集群定义中的字段，请参阅创建自定义集群。

  1. 使用 REST API 创建集群：

    创建请求正文：

    cat <<EOF > "${CLUSTER_ID}_create.json"
{
  "networkResources": {
    "net": {
      "config": {
        "newNetwork": {
          "network": "projects/${PROJECT}/global/networks/${CLUSTER_ID}-net"
        }
      }
    }
  },
  "storageResources": {
    "home": {
      "config": {
        "newLustre": {
          "lustre": "projects/${PROJECT}/locations/${ZONE}/instances/${CLUSTER_ID}-lustre",
          "filesystem": "lustre",
          "capacityGb": "18000",
          "perUnitStorageThroughput": "500"
        }
      }
    }
  },
  "computeResources": {
    "crtpu": {
      "config": {
        "newReservedInstances": {
          "reservation": "projects/${PROJECT}/zones/${ZONE}/reservations/${RESERVATION_NAME}"
        }
      }
    },
    "crcpu": {
      "config": {
        "newOnDemandInstances": {
          "zone": "${ZONE}",
          "machineType": "n2-highmem-32"
        }
      }
    }
  },
  "orchestrator": {
    "slurm": {
      "loginNodes": {
        "count": "1",
        "zone": "${ZONE}",
        "machineType": "n2-standard-16",
        "enableOsLogin": true,
        "enablePublicIps": true,
        "bootDisk": {
          "type": "pd-balanced",
          "sizeGb": "100"
        },
        "storageConfigs": [
          {
            "id": "home",
            "localMount": "/home"
          }
        ]
      },
      "nodeSets": [
        {
          "id": "nstpu",
          "computeId": "crtpu",
          "staticNodeCount": "4",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "resourcePolicyConfig": {
            "acceleratorTopology": "2x4"
          },
          "computeInstance": {
            "bootDisk": {
              "type": "hyperdisk-balanced",
              "sizeGb": "100"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        },
        {
          "id": "nscpu",
          "computeId": "crcpu",
          "staticNodeCount": "1",
          "maxDynamicNodeCount": "0",
          "enablePublicIps": true,
          "computeInstance": {
            "bootDisk": {
              "type": "pd-balanced",
              "sizeGb": "100",
              "image": "projects/clusterdirector-public-images/global/images/family/tpu-v6e-v5p-ubuntu2204-slurm2605-v20260902"
            }
          },
          "storageConfigs": [
            {
              "id": "home",
              "localMount": "/home"
            }
          ]
        }
      ],
      "partitions": [
        {
          "id": "tpu",
          "nodeSetIds": ["nstpu"]
        },
        {
          "id": "cpu",
          "nodeSetIds": ["nscpu"]
        }
      ],
      "defaultPartition": "tpu",
      "config": {
        "requeueExitCodes": [1, 2, 134, 137, 139, 143]
      }
    }
  }
}
EOF

    然后发送请求：

    curl -s -X POST \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -d @"${CLUSTER_ID}_create.json" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters?clusterId=${CLUSTER_ID}"

    该集群包含以下内容：

    • tpu 分区：预留 TPU v6e 虚拟机的静态节点集。staticNodeCount 是虚拟机数量，acceleratorTopology 是每个切片的拓扑。每个 ct6e-standard-4t 虚拟机都有四个芯片，因此 2x4 切片有两个虚拟机。如需更多切片，请乘以：两个 2x4 切片需要 4 的 staticNodeCount。
    • cpu 分区：一个运行 Pathways 资源管理器、IFRT 代理服务器和 MaxText 客户端的按需虚拟机。Pathways 脚本会将作业提交到名为 cpu 的分区。虚拟机使用 n2-highmem-32 机器类型，因为 Pathways 服务和 MaxText 在保存检查点时总共需要多达 240 GB 的内存。它还使用具有 Slurm 26.05 的映像，该映像支持 TPU 切片，因为 Pathways 服务会从此虚拟机提交 TPU 工作进程作业。
    • requeueExitCodes：Slurm 会将退出时返回其中一个 代码的作业重新排入队列。此设置可让 Slurm 在切片中的任何虚拟机（而不仅仅是第一个虚拟机）被抢占时重启 TPU 工作进程作业。

  2. 等待集群准备就绪。您可以在 Google Cloud 控制台的 Cluster Director 页面上跟踪进度，也可以在以下命令的输出中检查 reconciling 是否为 false：

    curl -s \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

设置 Pathways 环境

Pathways 环境位于共享文件系统上的 /home/pathways 中，因此您只需为每个集群设置一次。它包含 Pathways 容器映像、具有 JAX 和 pathwaysutils 的 Python 环境以及 Pathways 脚本。

  1. 连接到登录节点。如需了解其他连接方式，请参阅连接到集群。

    gcloud compute ssh ${CLUSTER_ID}-login-001 \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --tunnel-through-iap

  2. 验证 cpu 和 tpu 分区是否存在且处于 idle 状态：

    sinfo -o "%P %a %D %T %N %f"

    输出类似于以下内容：

    PARTITION AVAIL NODES STATE NODELIST FEATURES
tpu* up 4 idle ${CLUSTER_ID}-nstpu-[0-3]
cpu up 1 idle ${CLUSTER_ID}-nscpu-0

  3. 将 Pathways 脚本下载到登录节点：

    mkdir -p ~/pathways-scripts && cd ~/pathways-scripts
for f in setup_env.sh start_pathways.sh run_worker_slice.sh run_maxtext.sh; do
  curl -fsSLO "https://raw.githubusercontent.com/GoogleCloudPlatform/slurm-gcp/tpu-2605/pathways/scripts/${f}"
done

    此命令会下载以下脚本：

    • setup_env.sh：对 Pathways 环境执行一次性设置。
    • start_pathways.sh：启动 Pathways 服务。
    • run_worker_slice.sh：为一个 TPU 切片启动 Pathways 工作器（由 start_pathways.sh 提交）。
    • run_maxtext.sh：运行 MaxText 训练工作负载示例。

  4. 运行设置脚本。大约需要 5 分钟：

    chmod +x *.sh
./setup_env.sh

    该脚本会拉取 Pathways 服务器和代理映像，使用 JAX、pathwaysutils 和 MaxText 构建 Python 环境，将脚本复制到 /home/pathways，并将 ~/pathways 关联到 /home/pathways。最后会输出如下所示的内容：

    Environment ready: JAX 0.11.1

  5. 可选：如需向其他集群用户授予对 Pathways 环境的访问权限，请让他们将 /home/pathways 关联到其主目录：

    ln -sfn /home/pathways ~/pathways

此时，集群已准备好使用 Pathways。您在运行工作负载时启动 Pathways 服务，因为批处理工作负载和交互式工作负载会使用不同的设置启动该服务。如需了解详情，请参阅使用 Pathways 运行批量工作负载和使用 Pathways 运行交互式工作负载。

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

清理

为避免因本页中使用的资源导致您的 Google Cloud 账号产生费用，请在完成后删除 Cluster Director 集群：

curl -s -X DELETE \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://hypercomputecluster.googleapis.com/v1alpha/projects/${PROJECT}/locations/${REGION}/clusters/${CLUSTER_ID}"

后续步骤