概述

xpk (加速处理套件，发音为 x-p-k) 是一款软件工具，帮助 Cloud 开发者在 GKE 上的 TPUs 和 GPUs 等加速器上编排训练作业。xpk 将 TPUs、GPU (HGX H100) 和 CPU (n2-standard-32) 的“多主机 Pod”视为一等公民。

xpk 将资源分配与运行作业解耦。有两种结构：集群（已分配的虚拟机）和工作负载（训练作业）。集群表示您可用的物理资源。工作负载表示训练作业——在任何时候，其中一些将完成，一些正在运行，一些将排队，等待集群资源可用。

理想的工作流程首先为已预留的所有 ML 硬件分配集群。然后，无需重新分配，按需提交作业。通过消除作业之间的重新分配需求，使用预安装依赖项的 Docker 容器和跨编译，这些排队作业可以以最短启动时间运行。此外，因为工作负载在完成时将硬件返回到共享池，因此开发人员可以实现更高效的有限硬件资源利用。并且可以在资源利用率较低的夜间运行自动化测试。

xpk 支持以下 TPU 类型

• v4
• v5e
• v5p

以及以下 GPU 类型

• a100
• h100

以及以下 CPU 类型

• n2-standard-32

安装

要安装 xpk，请运行以下命令

pip install xpk

如果您通过克隆 GitHub 存储库运行 XPK，请首先运行以下命令以开始使用 XPK 命令

git clone https://github.com/google/xpk.git
cd xpk
# Install dependencies such as cloud-accelerator-diagnostics
pip install .

如果您看到错误消息：“此环境由外部管理”，请使用虚拟环境。

示例

  ## One time step of creating the venv
  VENV_DIR=~/venvp3
  python3 -m venv $VENV_DIR
  ## Enter your venv.
  source $VENV_DIR/bin/activate
  ## Clone the repository and installing dependencies.
  git clone https://github.com/google/xpk.git
  cd xpk
  # Install dependencies such as cloud-accelerator-diagnostics
  pip install .

XPK 大规模（>1k VM）

按照xpk-large-scale-guide.sh中的用户指南操作，使用xpk管理超过1000个虚拟机的GKE集群。执行以下步骤以使用XPK设置支持大规模训练和高吞吐量的GKE集群，并运行作业。我们建议您手动复制每一步的命令并验证每一步的输出。

示例用法

要开始，请确保通过gcloud config set设置您的GCP项目和区域，就像平常一样。

以下是参考命令。典型的流程是从Cluster Create开始，然后进行多个Workload Create。为了了解系统状态，您可能需要使用Cluster List或Workload List命令。最后，您可以使用Cluster Delete进行清理。

如果您的工作负载运行失败，请使用xpk inspector进行进一步调查。

集群创建

首先通过gcloud config或xpk参数设置项目和区域。

PROJECT_ID=my-project-id
ZONE=us-east5-b
# gcloud config:
gcloud config set project $PROJECT_ID
gcloud config set compute/zone $ZONE
# xpk arguments
xpk .. --zone $ZONE --project $PROJECT_ID

创建的集群是区域集群，以支持跨所有区域的GKE控制平面。

• 集群创建（预配预留容量）

  # Find your reservations
  gcloud compute reservations list --project=$PROJECT_ID
  # Run cluster create with reservation.
  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-256 \
  --num-slices=2 \
  --reservation=$RESERVATION_ID
  

• 集群创建（预配按需容量）

  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=4 --on-demand
  

• 集群创建（预配spot / 可抢占容量）

  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=4 --spot
  

• 为Pathways创建集群：可以使用--enable-pathways创建与Pathways兼容的集群。

  python3 xpk.py cluster create \
  --cluster xpk-pw-test \
  --num-slices=4 --on-demand \
  --tpu-type=v5litepod-16 \
  --enable-pathways
  

• 可以使用相同的--cluster name再次调用集群创建来修改切片数量或重试失败的步骤。

  例如，如果用户创建了一个包含4个切片的集群

  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=4  --reservation=$RESERVATION_ID
  

  然后重新创建一个包含8个切片的集群。命令将重新运行以创建4个新切片

  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=8  --reservation=$RESERVATION_ID
  

  然后重新创建一个包含6个切片的集群。命令将重新运行以删除2个切片。在删除切片时，命令将警告用户。使用--force跳过提示。

  python3 xpk.py cluster create \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=6  --reservation=$RESERVATION_ID
  
  # Skip delete prompts using --force.
  
  python3 xpk.py cluster create --force \
  --cluster xpk-test --tpu-type=v5litepod-16 \
  --num-slices=6  --reservation=$RESERVATION_ID
  

创建Vertex AI Tensorboard

注意：此功能在XPK >= 0.4.0中可用。在您的Google Cloud控制台中启用Vertex AI API以使用此功能。确保您的用户账户已分配Vertex AI管理员角色。

Vertex AI Tensorboard是开源Tensorboard的完全托管版本。有关Vertex AI Tensorboard的更多信息，请访问此处。请注意，Vertex AI Tensorboard仅在这些区域中可用。

您可以使用Cluster Create命令为您集群创建Vertex AI Tensorboard。XPK将为每个集群创建单个Vertex AI Tensorboard实例。

• 在默认区域使用默认Tensorboard名称创建Vertex AI Tensorboard

python3 xpk.py cluster create \
--cluster xpk-test --num-slices=1 --tpu-type=v4-8 \
--create-vertex-tensorboard

将在us-central1（默认区域）中创建名为xpk-test-tb-instance（<args.cluster>-tb-instance）的Vertex AI Tensorboard。

• 在用户指定的区域使用默认Tensorboard名称创建Vertex AI Tensorboard

python3 xpk.py cluster create \
--cluster xpk-test --num-slices=1 --tpu-type=v4-8 \
--create-vertex-tensorboard --tensorboard-region=us-west1

将在us-west1中创建名为xpk-test-tb-instance（<args.cluster>-tb-instance）的Vertex AI Tensorboard。

• 在默认区域使用用户指定的Tensorboard名称创建Vertex AI Tensorboard

python3 xpk.py cluster create \
--cluster xpk-test --num-slices=1 --tpu-type=v4-8 \
--create-vertex-tensorboard --tensorboard-name=tb-testing

将在us-central1中创建名为tb-testing的Vertex AI Tensorboard。

• 在用户指定的区域使用用户指定的Tensorboard名称创建Vertex AI Tensorboard

python3 xpk.py cluster create \
--cluster xpk-test --num-slices=1 --tpu-type=v4-8 \
--create-vertex-tensorboard --tensorboard-region=us-west1 --tensorboard-name=tb-testing

将在us-west1中创建名为tb-testing的Vertex AI Tensorboard实例。

• 在不受支持的区域创建Vertex AI Tensorboard

python3 xpk.py cluster create \
--cluster xpk-test --num-slices=1 --tpu-type=v4-8 \
--create-vertex-tensorboard --tensorboard-region=us-central2

将失败集群创建过程，因为Vertex AI Tensorboard在us-central2中不受支持。

集群删除

• 集群删除（取消预配容量）

  python3 xpk.py cluster delete \
  --cluster xpk-test
  

集群列表

• 集群列表（查看预配容量）

  python3 xpk.py cluster list
  

集群描述

• 集群描述（查看容量）

  python3 xpk.py cluster describe \
  --cluster xpk-test
  

集群缓存镜像

• 集群缓存镜像（启用更快的启动时间）

  python3 xpk.py cluster cacheimage \
  --cluster xpk-test --docker-image gcr.io/your_docker_image \
  --tpu-type=v5litepod-16
  

工作负载创建

• 工作负载创建（提交训练作业）

  python3 xpk.py workload create \
  --workload xpk-test-workload --command "echo goodbye" \
  --cluster xpk-test \
  --tpu-type=v5litepod-16
  

• 针对路径创建工作负载：在启用路径的集群（使用--enable-pathways创建）上，可以使用--use-pathways提交路径工作负载

  路径工作负载示例

  python3 xpk.py workload create \
  --workload xpk-pw-test \
  --num-slices=1 \
  --tpu-type=v5litepod-16 \
  --use-pathways \
  --cluster xpk-pw-test \
  --docker-name='user-workload' \
  --docker-image=<maxtext docker image> \
  --command='python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=<output directory> dataset_path=<dataset path> per_device_batch_size=1 enable_checkpointing=false enable_profiler=false remat_policy=full global_parameter_scale=4 steps=300 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True run_name=$(USER)-pw-xpk-test-1'
  

  在启用路径的集群（使用--enable-pathways创建）上，也可以提交常规工作负载

  路径工作负载示例

  python3 xpk.py workload create \
  --workload xpk-regular-test \
  --num-slices=1 \
  --tpu-type=v5litepod-16 \
  --cluster xpk-pw-test \
  --docker-name='user-workload' \
  --docker-image=<maxtext docker image> \
  --command='python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=<output directory> dataset_path=<dataset path> per_device_batch_size=1 enable_checkpointing=false enable_profiler=false remat_policy=full global_parameter_scale=4 steps=300 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True run_name=$(USER)-pw-xpk-test-1'
  

为生产作业设置max-restarts

• --max-restarts <值>：默认情况下，此值为0。当作业终止时，此作业将重启0次。对于生产作业，建议将其增加到较大数值，例如50。真实作业可能会因硬件故障和软件更新而中断。我们假设您的作业已实现检查点，以便作业可以在中断附近重新启动。

工作负载优先级和抢占

• 使用--priority=优先级设置您的工作负载优先级

  我们定义了五个优先级：非常低、低、中、高、非常高。默认优先级是中。

  优先级决定

  1. 排队作业的顺序。

     排队作业按以下顺序排列：非常低 < 低 < 中 < 高 < 非常高

  2. 低优先级工作负载的抢占。

     高优先级作业将驱逐低优先级作业。被驱逐的作业将返回队列，并相应地重新激活。

  一般示例

  python3 xpk.py workload create \
  --workload xpk-test-medium-workload --command "echo goodbye" --cluster \
  xpk-test --tpu-type=v5litepod-16 --priority=medium
  

创建Vertex AI实验以将数据上传到Vertex AI Tensorboard

注意：此功能在XPK >= 0.4.0中可用。在您的Google Cloud控制台中启用Vertex AI API以使用此功能。请确保您的用户账户和附加到集群节点池的Compute Engine服务帐户已分配Vertex AI管理员角色。

Vertex AI实验是一种帮助跟踪和分析在Vertex AI Tensorboard上运行的实验的工具。要了解更多关于Vertex AI实验的信息，请访问此处。

XPK将在workload create命令中创建一个Vertex AI实验，并将其附加到在cluster create期间为集群创建的Vertex AI Tensorboard。如果在此功能发布之前已创建集群，则不会为该集群创建Vertex AI Tensorboard，并且workload create将失败。重新运行cluster create以创建Vertex AI Tensorboard，然后再次运行workload create以安排您的工作负载。

• 使用默认实验名称创建Vertex AI实验

python3 xpk.py workload create \
--cluster xpk-test --workload xpk-workload \
--use-vertex-tensorboard

将创建一个名为xpk-test-xpk-workload（<args.cluster>-<args.workload>）的Vertex AI实验。

• 使用用户指定的实验名称创建Vertex AI实验

python3 xpk.py workload create \
--cluster xpk-test --workload xpk-workload \
--use-vertex-tensorboard --experiment-name=test-experiment

将创建一个名为test-experiment的Vertex AI实验。

查看MaxText示例，了解如何更新您的工作负载，以便自动将收集在Tensorboard目录中的日志上传到由workload create创建的Vertex AI实验。

工作负载删除

• 工作负载删除（删除训练作业）

  python3 xpk.py workload delete \
  --workload xpk-test-workload --cluster xpk-test
  

  这将仅删除位于xpk-test集群中的xpk-test-workload工作负载。

• 工作负载删除（删除集群中的所有训练作业）

  python3 xpk.py workload delete \
  --cluster xpk-test
  

  这将删除xpk-test集群中的所有工作负载。如果在此提示下键入y或yes，则开始删除。多个工作负载删除将分批处理以优化处理。

• 工作负载删除支持过滤。删除符合用户标准的一部分作业。多个工作负载删除将分批处理以优化处理。

  • 按作业过滤：filter-by-job

  python3 xpk.py workload delete \
  --cluster xpk-test --filter-by-job=$USER
  

  这将删除xpk-test集群中所有名称以$USER开头的作业。如果在此提示下键入y或yes，则开始删除。

  • 按状态过滤：filter-by-status

  python3 xpk.py workload delete \
  --cluster xpk-test --filter-by-status=QUEUED
  

  这将删除所有在xpk-test集群中的工作负载，这些工作负载的状态为“已接受”或“已驱逐”，并且正在运行的虚拟机数量为0。只有在您在提示符处输入y或yes时，删除才会开始。状态可以是：EVERYTHING、FINISHED、RUNNING、QUEUED、FAILED、SUCCESSFUL。

工作负载列表

• 工作负载列表（参见训练作业）

  python3 xpk.py workload list \
  --cluster xpk-test
  

• 示例工作负载列表输出

  以下示例显示了四个不同状态的工作

  • user-first-job-failed：filter-status为FINISHED和FAILED。
  • user-second-job-success：filter-status为FINISHED和SUCCESSFUL。
  • user-third-job-running：filter-status为RUNNING。
  • user-forth-job-in-queue：filter-status为QUEUED。
  • user-fifth-job-in-queue-preempted：filter-status为QUEUED。

  Jobset Name                     Created Time           Priority   TPU VMs Needed   TPU VMs Running/Ran   TPU VMs Done      Status     Status Message                                                  Status Time
  user-first-job-failed           2023-1-1T1:00:00Z      medium     4                4                     <none>            Finished   JobSet failed                                                   2023-1-1T1:05:00Z
  user-second-job-success         2023-1-1T1:10:00Z      medium     4                4                     4                 Finished   JobSet finished successfully                                    2023-1-1T1:14:00Z
  user-third-job-running          2023-1-1T1:15:00Z      medium     4                4                     <none>            Admitted   Admitted by ClusterQueue cluster-queue                          2023-1-1T1:16:00Z
  user-forth-job-in-queue         2023-1-1T1:16:05Z      medium     4                <none>                <none>            Admitted   couldn't assign flavors to pod set slice-job: insufficient unused quota for google.com/tpu in flavor 2xv4-8, 4 more need   2023-1-1T1:16:10Z
  user-fifth-job-preempted        2023-1-1T1:10:05Z      low        4                <none>                <none>            Evicted    Preempted to accommodate a higher priority Workload             2023-1-1T1:10:00Z
  

• 工作负载列表支持筛选。观察匹配用户标准的作业的一部分。

  • 按状态过滤：filter-by-status

  通过相应作业的状态筛选工作负载列表。状态可以是：EVERYTHING、FINISHED、RUNNING、QUEUED、FAILED、SUCCESSFUL

  • 按作业过滤：filter-by-job

  通过作业的名称筛选工作负载列表。

  python3 xpk.py workload list \
  --cluster xpk-test --filter-by-job=$USER
  

• 工作负载列表支持等待特定作业的完成。XPK将跟踪现有作业，直到其完成或达到timeout（如果提供），然后列出作业。如果没有指定timeout，则默认值设置为最大值，即1周。您还可以将timeout=0设置为轮询作业一次。
  （注意：在创建工作负载时必须设置restart-on-user-code-failure，否则工作负载将始终以Completed状态完成。）

  等待作业完成。

  python3 xpk.py workload list \
  --cluster xpk-test --wait-for-job-completion=xpk-test-workload
  

  等待作业完成，超时时间为300秒。

  python3 xpk.py workload list \
  --cluster xpk-test --wait-for-job-completion=xpk-test-workload \
  --timeout=300
  

  返回代码
  0：工作负载完成并成功完成。
  124：在工作负载完成之前达到了超时。
  125：工作负载已完成，但未成功完成。
  1：其他失败。

检查器

• 检查器提供调试信息，以了解集群健康状态以及为什么工作负载没有运行。检查器输出将保存到文件中。

  python3 xpk.py inspector \
    --cluster $CLUSTER_NAME \
    --project $PROJECT_ID \
    --zone $ZONE
  

• 可选参数

  • --print-to-terminal：将命令输出打印到终端和文件。
  • --workload $WORKLOAD_NAME 检查器将写入与工作负载：$WORKLOAD_NAME相关的调试信息

• 示例输出

  在本例中，xpk检查器的输出保存在/tmp/tmp0pd6_k1o中。

  [XPK] Starting xpk
  [XPK] Task: `Set Cluster` succeeded.
  [XPK] Task: `Local Setup: gcloud version` is implemented by `gcloud version`, hiding output unless there is an error.
  [XPK] Task: `Local Setup: Project / Zone / Region` is implemented by `gcloud config get project; gcloud config get compute/zone; gcloud config get compute/region`, hiding output unless there is an error.
  [XPK] Task: `GKE: Cluster Details` is implemented by `gcloud beta container clusters list --project $PROJECT --region $REGION | grep -e NAME -e $CLUSTER_NAME`, hiding output unless there is an error.
  [XPK] Task: `GKE: Node pool Details` is implemented by `gcloud beta container node-pools list --cluster $CLUSTER_NAME  --project=$PROJECT --region=$REGION`, hiding output unless there is an error.
  [XPK] Task: `Kubectl: All Nodes` is implemented by `kubectl get node -o custom-columns='NODE_NAME:metadata.name, READY_STATUS:.status.conditions[?(@.type=="Ready")].status, NODEPOOL:metadata.labels.cloud\.google\.com/gke-nodepool'`, hiding output unless there is an error.
  [XPK] Task: `Kubectl: Number of Nodes per Node Pool` is implemented by `kubectl get node -o custom-columns=':metadata.labels.cloud\.google\.com/gke-nodepool' | sort | uniq -c`, hiding output unless there is an error.
  [XPK] Task: `Kubectl: Healthy Node Count Per Node Pool` is implemented by `kubectl get node -o custom-columns='NODE_NAME:metadata.name, READY_STATUS:.status.conditions[?(@.type=="Ready")].status, NODEPOOL:metadata.labels.cloud\.google\.com/gke-nodepool' | grep -w True | awk {'print $3'} | sort | uniq -c`, hiding output unless there is an error.
  [XPK] Task: `Kueue: ClusterQueue Details` is implemented by `kubectl describe ClusterQueue cluster-queue`, hiding output unless there is an error.
  [XPK] Task: `Kueue: LocalQueue Details` is implemented by `kubectl describe LocalQueue multislice-queue`, hiding output unless there is an error.
  [XPK] Task: `Kueue: Kueue Deployment Details` is implemented by `kubectl describe Deployment kueue-controller-manager -n kueue-system`, hiding output unless there is an error.
  [XPK] Task: `Jobset: Deployment Details` is implemented by `kubectl describe Deployment jobset-controller-manager -n jobset-system`, hiding output unless there is an error.
  [XPK] Task: `Kueue Manager Logs` is implemented by `kubectl logs deployment/kueue-controller-manager -n kueue-system --tail=100 --prefix=True`, hiding output unless there is an error.
  [XPK] Task: `Jobset Manager Logs` is implemented by `kubectl logs deployment/jobset-controller-manager -n jobset-system --tail=100 --prefix=True`, hiding output unless there is an error.
  [XPK] Task: `List Jobs with filter-by-status=EVERYTHING with filter-by-jobs=None` is implemented by `kubectl get workloads -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime"  `, hiding output unless there is an error.
  [XPK] Task: `List Jobs with filter-by-status=QUEUED with filter-by-jobs=None` is implemented by `kubectl get workloads -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime"  | awk -e 'NR == 1 || ($7 ~ "Admitted|Evicted|QuotaReserved" && ($5 ~ "<none>" || $5 == 0)) {print $0}' `, hiding output unless there is an error.
  [XPK] Task: `List Jobs with filter-by-status=RUNNING with filter-by-jobs=None` is implemented by `kubectl get workloads -o=custom-columns="Jobset Name:.metadata.ownerReferences[0].name,Created Time:.metadata.creationTimestamp,Priority:.spec.priorityClassName,TPU VMs Needed:.spec.podSets[0].count,TPU VMs Running/Ran:.status.admission.podSetAssignments[-1].count,TPU VMs Done:.status.reclaimablePods[0].count,Status:.status.conditions[-1].type,Status Message:.status.conditions[-1].message,Status Time:.status.conditions[-1].lastTransitionTime"  | awk -e 'NR == 1 || ($7 ~ "Admitted|Evicted" && $5 ~ /^[0-9]+$/ && $5 > 0) {print $0}' `, hiding output unless there is an error.
  [XPK] Find xpk inspector output file: /tmp/tmp0pd6_k1o
  [XPK] Exiting XPK cleanly
  

GPU使用情况

为了使用XPK进行GPU，您可以通过使用device-type标志来实现。

• 集群创建（预配预留容量）

  # Find your reservations
  gcloud compute reservations list --project=$PROJECT_ID
  
  # Run cluster create with reservation.
  python3 xpk.py cluster create \
  --cluster xpk-test --device-type=h100-80gb-8 \
  --num-nodes=2 \
  --reservation=$RESERVATION_ID
  

• 集群删除（取消预配容量）

  python3 xpk.py cluster delete \
  --cluster xpk-test
  

• 集群列表（查看预配容量）

  python3 xpk.py cluster list
  

• 集群描述（查看容量）

  python3 xpk.py cluster describe \
  --cluster xpk-test
  

• 集群缓存镜像（启用更快的启动时间）

  python3 xpk.py cluster cacheimage \
  --cluster xpk-test --docker-image gcr.io/your_docker_image \
  --device-type=h100-80gb-8
  

• 安装NVIDIA GPU设备驱动程序

  # List available driver versions
  gcloud compute ssh $NODE_NAME --command "sudo cos-extensions list"
  
  # Install the default driver
  gcloud compute ssh $NODE_NAME --command "sudo cos-extensions install gpu"
  # OR install a specific version of the driver
  gcloud compute ssh $NODE_NAME --command "sudo cos-extensions install gpu -- -version=DRIVER_VERSION"
  

• 运行工作负载

  # Submit a workload
  python3 xpk.py workload create \
  --cluster xpk-test --device-type h100-80gb-8 \
  --workload xpk-test-workload \
  --command="echo hello world"
  

• 工作负载删除（删除训练作业）

  python3 xpk.py workload delete \
  --workload xpk-test-workload --cluster xpk-test
  

  这将仅删除位于xpk-test集群中的xpk-test-workload工作负载。

• 工作负载删除（删除集群中的所有训练作业）

  python3 xpk.py workload delete \
  --cluster xpk-test
  

  这将删除xpk-test集群中的所有工作负载。只有在您在提示符处输入y或yes时，删除才会开始。

• 工作负载删除支持筛选。删除匹配用户标准的作业的一部分。

  • 按作业过滤：filter-by-job

  python3 xpk.py workload delete \
  --cluster xpk-test --filter-by-job=$USER
  

  这将删除xpk-test集群中所有名称以$USER开头的作业。如果在此提示下键入y或yes，则开始删除。

  • 按状态过滤：filter-by-status

  python3 xpk.py workload delete \
  --cluster xpk-test --filter-by-status=QUEUED
  

  这将删除所有在xpk-test集群中的工作负载，这些工作负载的状态为“已接受”或“已驱逐”，并且正在运行的虚拟机数量为0。只有在您在提示符处输入y或yes时，删除才会开始。状态可以是：EVERYTHING、FINISHED、RUNNING、QUEUED、FAILED、SUCCESSFUL。

CPU使用情况

为了使用XPK进行CPU，您可以通过使用device-type标志来实现。

• 集群创建（预配按需容量）

  # Run cluster create with on demand capacity.
  python3 xpk.py cluster create \
  --cluster xpk-test \
  --device-type=n2-standard-32-256 \
  --num-slices=1 \
  --default-pool-cpu-machine-type=n2-standard-32 \
  --on-demand
  

  请注意，CPUs的device-type格式为“-”，因此在上述示例中，用户请求256个n2-standard-32类型的虚拟机。目前支持使用小于1000个虚拟机的工作负载。

• 运行工作负载

  # Submit a workload
  python3 xpk.py workload create \
  --cluster xpk-test \
  --num-slices=1 \
  --device-type=n2-standard-32-256 \
  --workload xpk-test-workload \
  --command="echo hello world"
  

使用XPK进行自动配置

XPK可以使用节点自动配置（NAP）支持动态分配集群容量。

支持从单个XPK集群支持多个拓扑大小，并根据传入的工作负载请求动态配置。XPK用户不需要手动重新配置集群。

启用自动配置将使集群最初大约需要30分钟进行升级。

创建具有自动配置的集群

以下集群将启用自动配置，该集群具有[0, 8]个v4 TPU芯片（最高1xv4-16）以进行扩展。

XPK目前不支持同一集群中不同代加速器（如v4和v5p TPU）。

CLUSTER_NAME=my_cluster
NUM_SLICES=2
DEVICE_TYPE=v4-8
RESERVATION=reservation_id
PROJECT=my_project
ZONE=us-east5-b

python3 xpk.py cluster create \
  --cluster $CLUSTER_NAME \
  --num-slices=$NUM_SLICES \
    --device-type=$DEVICE_TYPE \
  --zone=$ZONE \
  --project=$PROJECT \
  --reservation=$RESERVATION \
  --enable-autoprovisioning

1. 定义起始加速器配置和容量类型。

   --device-type=$DEVICE_TYPE \
   --num-slice=$NUM_SLICES
   

2. 可选设置自定义的 最小 / 最大 芯片数。NAP 将根据 最大 - 最小 芯片数调整集群大小。默认情况下，最大 设置为当前集群配置的大小，最小 设置为 0。这允许 NAP 使用所有资源进行缩放。

   --autoprovisioning-min-chips=$MIN_CHIPS \
   --autoprovisioning-max-chips=$MAX_CHIPS
   

3. 即将推出功能： 设置节点池在无工作负载运行时自动删除的超时时间。目前为 10 分钟。

4. 即将推出功能： 设置超时时间为无限。这将使空闲节点池配置持续运行，直到由新工作负载更新。

使用自动配置更新集群

CLUSTER_NAME=my_cluster
NUM_SLICES=2
DEVICE_TYPE=v4-8
RESERVATION=reservation_id
PROJECT=my_project
ZONE=us-east5-b

python3 xpk.py cluster create \
  --cluster $CLUSTER_NAME \
  --num-slices=$NUM_SLICES \
    --device-type=$DEVICE_TYPE \
  --zone=$ZONE \
  --project=$PROJECT \
  --reservation=$RESERVATION \
  --enable-autoprovisioning

使用不同数量的芯片更新之前已自动配置的集群

• 选项 1：通过创建新的集群节点池配置。

CLUSTER_NAME=my_cluster
NUM_SLICES=2
DEVICE_TYPE=v4-16
RESERVATION=reservation_id
PROJECT=my_project
ZONE=us-east5-b

# This will create 2x v4-16 node pools and set the max autoprovisioned chips to 16.
python3 xpk.py cluster create \
  --cluster $CLUSTER_NAME \
  --num-slices=$NUM_SLICES \
    --device-type=$DEVICE_TYPE \
  --zone=$ZONE \
  --project=$PROJECT \
  --reservation=$RESERVATION \
  --enable-autoprovisioning

• 选项 2：通过增加 --autoprovisioning-max-chips。

CLUSTER_NAME=my_cluster
NUM_SLICES=0
DEVICE_TYPE=v4-16
RESERVATION=reservation_id
PROJECT=my_project
ZONE=us-east5-b

# This will clear the node pools if they exist in the cluster and set the max autoprovisioned chips to 16
python3 xpk.py cluster create \
  --cluster $CLUSTER_NAME \
  --num-slices=$NUM_SLICES \
    --device-type=$DEVICE_TYPE \
  --zone=$ZONE \
  --project=$PROJECT \
  --reservation=$RESERVATION \
  --enable-autoprovisioning \
  --autoprovisioning-max-chips 16

在具有自动配置的集群上运行工作负载

重新配置 --device-type 和 --num-slices

CLUSTER_NAME=my_cluster
NUM_SLICES=2
DEVICE_TYPE=v4-8
NEW_RESERVATION=new_reservation_id
PROJECT=my_project
ZONE=us-east5-b
# Create a 2x v4-8 TPU workload.
python3 xpk.py workload create \
  --cluster $CLUSTER \
  --workload ${USER}-nap-${NUM_SLICES}x${DEVICE_TYPE}_$(date +%H-%M-%S) \
  --command "echo hello world from $NUM_SLICES $DEVICE_TYPE" \
  --device-type=$DEVICE_TYPE \
  --num-slices=$NUM_SLICES \
  --zone=$ZONE \
  --project=$PROJECT

NUM_SLICES=1
DEVICE_TYPE=v4-16

# Create a 1x v4-16 TPU workload.
python3 xpk.py workload create \
  --cluster $CLUSTER \
  --workload ${USER}-nap-${NUM_SLICES}x${DEVICE_TYPE}_$(date +%H-%M-%S) \
  --command "echo hello world from $NUM_SLICES $DEVICE_TYPE" \
  --device-type=$DEVICE_TYPE \
  --num-slices=$NUM_SLICES \
  --zone=$ZONE \
  --project=$PROJECT

# Use a different reservation from what the cluster was created with.
python3 xpk.py workload create \
  --cluster $CLUSTER \
  --workload ${USER}-nap-${NUM_SLICES}x${DEVICE_TYPE}_$(date +%H-%M-%S) \
  --command "echo hello world from $NUM_SLICES $DEVICE_TYPE" \
  --device-type=$DEVICE_TYPE \
  --num-slices=$NUM_SLICES \
  --zone=$ZONE \
  --project=$PROJECT \
  --reservation=$NEW_RESERVATION

1. (可选) 定义容量类型。默认情况下，容量类型将与集群创建时匹配。

   --reservation=my-reservation-id | --on-demand | --spot
   

2. 使用 --device-type 设置工作负载的拓扑。

   NUM_SLICES=1
   DEVICE_TYPE=v4-8
   --device-type=$DEVICE_TYPE \
   --num-slices=$NUM_SLICES \
   

如何将 Docker 镜像添加到 xpk 工作负载

默认行为是 xpk workload create 将本地目录 (--script-dir) 作为层添加到基本 Docker 镜像 (--base-docker-image) 并运行工作负载命令。如果您不希望这种行为，可以直接使用 --docker-image。不要在同一个命令中混合两个流程的参数。

推荐/默认 Docker 流程：--base-docker-image 和 --script-dir

此流程将 --script-dir 拉取到 --base-docker-image 并运行新的 Docker 镜像。

• 以下参数默认为可选。xpk 将使用通用基本 Docker 镜像拉取本地目录。

  • --base-docker-image 设置 xpk 将开始的基础镜像。

  • --script-dir 设置要拉取到镜像中的目录。默认情况下，为当前工作目录。

  有关更多信息，请参阅 python3 xpk.py workload create --help。

• 默认情况下，将本地目录拉取到基础镜像的示例

  echo -e '#!/bin/bash \n echo "Hello world from a test script!"' > test.sh
  python3 xpk.py workload create --cluster xpk-test \
  --workload xpk-test-workload-base-image --command "bash test.sh" \
  --tpu-type=v5litepod-16 --num-slices=1
  

• 适用于正常大小作业（少于 10k 加速器）的推荐流程

  python3 xpk.py workload create --cluster xpk-test \
  --workload xpk-test-workload-base-image --command "bash custom_script.sh" \
  --base-docker-image=gcr.io/your_dependencies_docker_image \
  --tpu-type=v5litepod-16 --num-slices=1
  

可选直接 Docker 镜像配置：--docker-image

如果用户想直接设置用于工作负载的 Docker 镜像而不是将当前工作目录添加为层，则将 --docker-image 设置为用于工作负载的镜像。

• 使用 --docker-image 运行

  python3 xpk.py workload create --cluster xpk-test \
  --workload xpk-test-workload-base-image --command "bash test.sh" \
  --tpu-type=v5litepod-16 --num-slices=1 --docker-image=gcr.io/your_docker_image
  

• 适用于大型作业（多于 10k 加速器）的推荐流程

  python3 xpk.py cluster cacheimage \
  --cluster xpk-test --docker-image gcr.io/your_docker_image
  # Run workload create with the same image.
  python3 xpk.py workload create --cluster xpk-test \
  --workload xpk-test-workload-base-image --command "bash test.sh" \
  --tpu-type=v5litepod-16 --num-slices=1 --docker-image=gcr.io/your_docker_image
  

更多高级事实

• 工作负载创建有两种相互排斥的方式来覆盖工作负载的环境

  • a --env 标志，用于单独指定每个环境变量。格式为

    --env VARIABLE1=value --env VARIABLE2=value

  • a --env-file 标志，允许从文件中指定容器的环境。用法与 Docker 的 --env-file 标志 相同

  示例环境文件

  LIBTPU_INIT_ARGS=--my-flag=true --performance=high
  MY_ENV_VAR=hello
  

• 工作负载创建接受一个 --debug-dump-gcs 标志，该标志是 GCS 存储桶的路径。传递此标志将设置 XLA_FLAGS='--xla_dump_to=/tmp/xla_dump/' 并将每个工作者的 hlo 溢出上传到指定的 GCS 存储桶。

集成测试流程

通过 Github Workflows 和 Actions 对存储库代码进行测试。目前执行了三种类型的测试

• 每 24 小时运行一次的夜间构建
• 在将代码推送到 main 分支时运行的构建
• 在 PR 获得批准时运行的构建

更多信息请参阅 此处

故障排除

CPUs 的 无效机器类型。

XPK 将创建一个区域 GKE 集群。如果您看到以下问题

Invalid machine type e2-standard-32 in zone $ZONE_NAME

请选择一个在该区域所有区域中存在的 CPU 类型。

# Find CPU Types supported in zones.
gcloud compute machine-types list --zones=$ZONE_LIST
# Adjust default cpu machine type.
python3 xpk.py cluster create --default-pool-cpu-machine-type=CPU_TYPE ...

权限问题：需要以下其中一个 ["permission_name"] 权限

1. 根据权限错误确定所需的角色

   # For example: `requires one of ["container.*"] permission(s)`
   # Add [Kubernetes Engine Admin](https://cloud.google.com/iam/docs/understanding-roles#kubernetes-engine-roles) to your user.
   

2. 将角色添加到您项目中的用户。

   访问 iam-admin 或使用 gcloud 命令行工具

   PROJECT_ID=my-project-id
   CURRENT_GKE_USER=$(gcloud config get account)
   ROLE=roles/container.admin  # container.admin is the role needed for Kubernetes Engine Admin
   gcloud projects add-iam-policy-binding $PROJECT_ID --member user:$CURRENT_GKE_USER --role=$ROLE
   

3. 检查用户权限是否正确。

   访问 iam-admin 或使用 gcloud 命令行工具

   PROJECT_ID=my-project-id
   CURRENT_GKE_USER=$(gcloud config get account)
   gcloud projects get-iam-policy $PROJECT_ID --filter="bindings.members:$CURRENT_GKE_USER" --flatten="bindings[].members"
   

4. 确认您已使用正确的用户本地登录。

   gcloud auth login
   

基于权限错误所需的角色

• 需要 ["container.*"] 中的一个权限

  将 Kubernetes Engine Admin 角色添加到您的用户。

• ERROR: (gcloud.monitoring.dashboards.list) 用户没有权限访问项目实例（或它可能不存在）

  将 Monitoring Viewer 角色添加到您的用户。

预订故障排除

如何确定您的预订及其大小/利用率

PROJECT_ID=my-project
ZONE=us-east5-b
RESERVATION=my-reservation-name
# Find the reservations in your project
gcloud beta compute reservations list --project=$PROJECT_ID
# Find the tpu machine type and current utilization of a reservation.
gcloud beta compute reservations describe $RESERVATION --project=$PROJECT_ID --zone=$ZONE

TPU 工作负载调试

详细日志

如果您的工作负载存在问题，请在安排时尝试设置 --enable-debug-logs。这将为您提供更多详细的日志来帮助定位问题。例如

python3 xpk.py workload create \
--cluster --workload xpk-test-workload \
--command="echo hello world" --enable-debug-logs

请查看 libtpu 日志 和 Tensorflow 日志 以获取有关启用日志记录标志的更多信息。

收集堆栈跟踪

cloud-tpu-diagnostics PyPI 软件包可用于生成在 GKE 中运行的工作负载的堆栈跟踪。此软件包在程序发生段错误、浮点异常或非法操作异常等故障时将转储 Python 跟踪。此外，它还将定期收集堆栈跟踪，以帮助您调试程序无响应的情况。您必须对在 Kubernetes 主容器中运行的 Docker 镜像进行以下更改，以启用定期堆栈跟踪收集。

# main.py

from cloud_tpu_diagnostics import diagnostic
from cloud_tpu_diagnostics.configuration import debug_configuration
from cloud_tpu_diagnostics.configuration import diagnostic_configuration
from cloud_tpu_diagnostics.configuration import stack_trace_configuration

stack_trace_config = stack_trace_configuration.StackTraceConfig(
                      collect_stack_trace = True,
                      stack_trace_to_cloud = True)
debug_config = debug_configuration.DebugConfig(
                stack_trace_config = stack_trace_config)
diagnostic_config = diagnostic_configuration.DiagnosticConfig(
                      debug_config = debug_config)

with diagnostic.diagnose(diagnostic_config):
	main_method()  # this is the main method to run

此配置将在每个 Kubernetes Pod 的 /tmp/debugging 目录内开始收集堆栈跟踪。

探索堆栈跟踪

要探索在 Kubernetes Pod 的临时目录中收集的堆栈跟踪，您可以使用以下命令配置一个侧边容器，该容器将读取 /tmp/debugging 目录中的跟踪。

python3 xpk.py workload create \
 --workload xpk-test-workload --command "python3 main.py" --cluster \
 xpk-test --tpu-type=v5litepod-16 --deploy-stacktrace-sidecar