本页面上的部分或全部信息可能不适用于 Trusted Cloud by S3NS。

此页面由 Cloud Translation API 翻译。

使用 TPU 和灵活启动预配模式运行小批量工作负载

Autopilot Standard

本指南介绍如何使用灵活启动预配模式优化中小规模训练工作负载的 TPU 预配。在本指南中，您将使用灵活启动来部署包含 TPU 切片节点池的工作负载。

本指南适用于机器学习 (ML) 工程师、平台管理员和运维人员，以及对使用 Kubernetes 容器编排功能运行批量工作负载感兴趣的数据和 AI 专家。如需详细了解我们在 Trusted Cloud by S3NS 内容中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。

灵活启动价格

如果您的工作负载需要根据需要动态配置资源（最多 7 天，短期预留）、无需复杂的配额管理且可经济高效地访问，建议使用灵活启动。灵活启动由动态工作负载调度器提供支持，并按照动态工作负载调度器价格计费：

vCPU、GPU 和 TPU 可享受折扣（最高 53%）。
随用随付。

准备工作

在开始之前，请确保您已执行以下任务：

启用 Google Kubernetes Engine API。

启用 Google Kubernetes Engine API

如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请运行 gcloud components update 以获取最新版本。
注意：对于现有 gcloud CLI 安装，请务必设置 compute/region 属性。如果您主要使用可用区级集群，请改为设置 compute/zone。通过设置默认位置，您可以避免 gcloud CLI 中出现如下错误：One of [--zone, --region] must be supplied: Please specify location。如果集群的位置与您设置的默认位置不同，您可能需要在某些命令中指定位置。

验证您是否拥有运行 1.33.0-gke.1712000 版或更高版本的 Autopilot 集群或标准集群。
确认您已熟悉灵活启动的限制。
使用 Standard 集群时，验证您是否至少维护了一个未启用灵活启动的节点池，以便集群能够正常运行。
验证您在节点位置是否拥有抢占式 TPU 的配额。

创建具有灵活启动的节点池

如需在现有 Standard 集群上创建启用了灵活启动的节点池，请使用 gcloud CLI。

如果您在 Autopilot 模式下使用集群，请跳过本部分，然后前往运行批量工作负载部分。

您可以创建启用了灵活启动的单主机或多主机 TPU 切片节点池：

创建具有灵活启动的节点池：
单个主机
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --reservation-affinity=none \
    --enable-autoscaling \
    --flex-start \
    --num-nodes 0 \
    --min-nodes=0 \
    --max-nodes=1
```
替换以下内容：
- NODE_POOL_NAME：您为节点池选择的名称。
- CLUSTER_NAME：集群的名称。
- LOCATION_NAME：集群控制平面的计算区域。
- NODE_ZONES：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
- MACHINE_TYPE：用于节点的机器类型。如需详细了解与 TPU 兼容的机器类型，请使用选择 TPU 版本中的表。
多主机
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --flex-start \
    --enable-autoscaling \
    --num-nodes=0 \
    --max-nodes=2 \
    --reservation-affinity=none \
    --no-enable-autorepair
```
替换以下内容：
- NODE_POOL_NAME：您为节点池选择的名称。
- CLUSTER_NAME：集群的名称。
- LOCATION_NAME：集群控制平面的计算区域。
- NODE_ZONES：GKE 在其中创建节点池的一个或多个可用区的英文逗号分隔列表。
- MACHINE_TYPE：用于节点的机器类型。例如，您可以为 TPU Trillium 使用 ct6e-standard-4t。如需详细了解可用的机器类型，请参阅选择 TPU 版本。
- TPU_TOPOLOGY：TPU 切片的物理拓扑。拓扑格式取决于 TPU 版本。如需详细了解 TPU 拓扑，请使用选择拓扑中的表。
上述命令在创建启用了灵活启动的节点池时会使用以下必要标志：
- --enable-autoscaling：灵活启动在工作负载运行时仅预配必要的计算资源。您必须设置以下参数：
  - --num-nodes=0
  - --min-nodes=0
  - 将 --max-nodes 设置为 TPU 切片所需的虚拟机数量。
    
    例如，您的节点池创建命令可以包含以下参数：
```
...
--machine-type=ct6e-standard-4t \
--tpu-topology=4x4 \
--enable-autoscaling \
--num-nodes=0 \
--max-nodes=4 \
```
    此命令将 --max-nodes 字段设置为 4，因为 4x4 拓扑包含 16 个芯片，而每个 ct6e-standard-4t 虚拟机有 4 个芯片。
  集群自动扩缩器会扩容到工作负载所需的节点数量。工作负载完成后，集群自动扩缩器会将集群缩减至零个节点。
- --reservation-affinity=none：灵活启动不使用或不需要预留。

验证节点池中灵活启动的状态：

gcloud container node-pools describe NODE_POOL_NAME \
    --cluster CLUSTER_NAME \
    --location LOCATION_NAME \
    --format="get(config.flexStart)"

如果节点池中启用了灵活启动，则 flexStart 字段会设置为 True。

运行批量工作负载

在本部分中，您将创建一个作业来调度启用了灵活启动的 TPU 节点。Kubernetes 中的 Job 控制器会创建一个或多个 Pod，并确保它们成功执行特定任务。

在 Trusted Cloud 控制台中，点击 激活 Cloud Shell 以启动 Cloud Shell 会话。系统会在 Trusted Cloud 控制台的底部窗格中打开会话。

创建一个名为 dws-flex-start.yaml 的文件：

单个主机

针对 dws-flex-start.yaml 文件使用以下清单：

apiVersion: batch/v1
kind: Job
metadata:
  name: job-1
spec:
  template:
    spec:
      nodeSelector:
        cloud.google.com/gke-flex-start: "true"
        cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
        cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
      containers:
      - name: container-1
        image: gcr.io/k8s-staging-perf-tests/sleep:latest
        args: ["3600s"] # Sleep for 1 hour
        resources:
          requests:
              google.com/tpu: NUM_CHIPS
          limits:
              google.com/tpu: NUM_CHIPS
      restartPolicy: OnFailure

多主机

针对 dws-flex-start.yaml 文件使用以下清单：

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: job-1
---
apiVersion: batch/v1
kind: Job
metadata:
  name: job-1
spec:
  backoffLimit: 0
  completions: 2
  parallelism: 2
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
          cloud.google.com/gke-flex-start: "true"
          cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
          cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: NUM_CHIPS
          limits:
            google.com/tpu: NUM_CHIPS

替换以下内容：

ACCELERATOR_TYPE：创建节点池时使用的 TPU 加速器的类型。例如 tpu-v4-podslice 或 tpu-v5-lite-podslice。
TPU_TOPOLOGY：TPU 切片的物理拓扑。例如，该值可以是 4x4x4 或 2x2，具体取决于 TPU 版本。
NUM_CHIPS：每个虚拟机中的 TPU 芯片数量为 1、4 或 8。如需了解详情，请参阅 TPU 版本。

应用 dws-flex-start.yaml 清单：
```
kubectl apply -f dws-flex-start.yaml
```

验证作业是否在同一节点上运行：

kubectl get pods

输出类似于以下内容：

NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>

清理

为避免因本页面中使用的资源导致您的 Trusted Cloud by S3NS 账号产生费用，请删除包含这些资源的项目，或者保留该项目但删除各个资源。

删除项目

In the Trusted Cloud console, go to the Manage resources page.
Go to Manage resources
In the project list, select the project that you want to delete, and then click Delete.
In the dialog, type the project ID, and then click Shut down to delete the project.

逐个删除资源

删除作业：

kubectl delete job -l "job-name in (job-1,job-2)"

删除节点池：

gcloud container node-pools delete NODE_POOL_NAME \
      --location LOCATION_NAME

删除集群：

gcloud container clusters delete CLUSTER_NAME

使用 TPU 和灵活启动预配模式运行小批量工作负载

灵活启动价格

准备工作

创建具有灵活启动的节点池

单个主机

多主机

运行批量工作负载

单个主机

多主机

清理

删除项目

逐个删除资源

后续步骤