通过 Pathways 进行弹性训练

本页面介绍了如何为 Google Kubernetes Engine (GKE) 上的 Pathways 工作负载配置弹性功能，以便大规模训练作业可以从计划内中断和计划外硬件故障中恢复。

Pathways 通过以下方式提供弹性优势：

  • 暂停-恢复：让您的工作负载能够容忍计划内中断（例如抢占通知），而无需自定义抢占处理代码。
  • 弹性训练：通过在训练代码中使用恢复逻辑，让您的工作负载能够容忍意外硬件故障，而不会导致客户端崩溃。
  • 多层级检查点设置：让您的工作负载能够使用共置的 Python 在 TPU 工作进程之间保存高频本地和对等复制检查点，并通过异步 Cloud Storage 备份来缩短恢复时间。
  • 热插拔：让优先级较高的 GKE JobSet 工作负载在发生硬件故障时快速接管备用或优先级较低的 TPU 切片。

准备工作

在开始之前，请完成以下任务：

如果您使用弹性训练，请在工作负载容器映像中安装以下软件包：

  • pathwaysutils
  • Orbax (orbax-checkpoint)：>= 0.12.6（使用 ocp.v1.training.pathways.Snapshotter 的内存中快照需要）
  • Grain (grain)：如果您使用 grain.experimental.ElasticIterator，则此字段为必填字段

如果您使用多层级检查点，还需满足以下前提条件：

  • 按照在 GKE 上使用多层级检查点机制训练大规模机器学习模型中的说明，为多层级检查点机制配置 GKE 集群：
    • 在集群上启用多层检查点和 Cloud Storage FUSE (GcsFuseCsiDriver) 插件，以及 Workload Identity Federation for GKE。
    • 创建启用了分层命名空间的 Cloud Storage 存储桶。
    • 向多层级检查点 Kubernetes 服务账号授予 roles/storage.objectUser 和 roles/storage.insightsCollectorService Identity and Access Management (IAM) 角色。
    • 部署 CheckpointConfiguration 自定义资源，指定 Cloud Storage 存储桶、TPU 节点选择器和内存中卷大小 (inMemoryVolumeSize)。
  • 构建并推送共置的 Python 边车容器映像，例如使用边车 Dockerfile 示例。在工作负载和边车映像中，使用与 Pathways 版本一致且满足以下最低版本要求的匹配 jax 和 jaxlib 版本：
    • Orbax 版本：>= 0.12.4
    • JAX 版本（包括共置的 Python）：>= 0.10.1

暂停-恢复

GKE 通常会在抢占加速器 Pod 之前向其发送抢占通知。Pathways 默认在所有 Cloud Deploy 部署中启用抢占容忍度，并且 Pathways 加速器作业会监听这些通知。

收到抢占通知后，Pathways 会首先确定是否可以透明地保存和恢复工作负载。如果可以，Pathways 会尝试在 GKE 逐出加速器作业之前，通过将机器学习工作负载的当前状态写入永久性存储（例如 Cloud Storage）来透明地暂停该工作负载。当 GKE 稍后重新安排作业时，Pathways 会通过读回其持久化状态来恢复机器学习工作负载。

如果工作负载无法恢复，Pathways 会关闭加速器作业，并在配置了弹性训练的情况下将故障转发到工作负载。如果未配置弹性训练，GKE 会根据 JobSet 重启政策重启整个工作负载。

使用 JAX 定义的典型机器学习工作负载依赖于无状态 Pathways XLA 组件，这些组件可以使用高带宽内存 (HBM) 快照进行恢复。 某些机器学习工作负载（例如使用 JAX 共置 Python API 定义的工作负载）依赖于有状态 Pathways 组件，无法恢复。

弹性训练

弹性训练可让您的训练作业在发生硬件故障或切片抢占时继续运行，而无需重启整个工作负载。在多切片训练作业中，当单个切片不可用时，Pathways 会使单控制器客户端进程和健康的 TPU 切片保持活跃状态。

弹性训练将 Pathways 系统功能与训练代码中的恢复逻辑相结合：

  • 故障检测：当 TPU 切片发生硬件故障或抢占时，Pathways 会取消分配受影响的切片，并在您的工作负载下次访问该切片上的数据时引发 jax.errors.JaxRuntimeError。由于客户端进程保持活跃状态，因此您的代码可以捕获此错误、重新配置资源并继续训练。
  • 恢复模式：pathwaysutils.elastic.manager.Manager 类及其 elastic_retry 修饰器支持两种恢复策略：
    • 暂停和恢复：如果您将 minimum_slice_count 设置为 None 或切片总数，则当切片发生故障时，训练会暂停，并在所有切片恢复正常后继续。由于切片数量在多次尝试中不会发生变化，因此您的工作负载可以重复使用其现有设备网格和已编译的可执行文件。
    • 副本大小调整：如果您将 minimum_slice_count 设置为小于总切片数，则当切片失败时，系统会通过重新构建设备网格并针对较小的拓扑重新编译，在剩余的活跃切片上继续训练。后台监控器线程会定期检查非活跃切片，并在 elastic_manager.available_inactive_slices 中记录新恢复健康的切片。然后，训练循环可以使用替换切片重建设备网格，针对扩展的拓扑重新初始化，并将当前状态放置到新网格上。
  • 状态恢复：如需在弹性事件后恢复训练，您的工作负载会使用内存中快照恢复模型参数、优化器状态和数据迭代器状态。orbax.checkpoint 中的 Snapshotter 类会将复制的状态从 TPU HBM 异步复制到每个切片上的 pinned_host CPU 内存。在恢复期间，Snapshotter 会根据幸存的主机副本重构状态，并将其直接重新分片到活跃的 TPU 网格上，而无需从远程存储空间读取。

在 Pathways 代理服务器上启用弹性

如需在 GKE 中的 Pathways 上启用弹性设备集，请将 --num_elastic_instances 标志添加到 JobSet 清单中 pathways-proxy 容器的 args 字段，并指定工作负载中的弹性 TPU 切片数量：

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

替换以下内容：

  • BUCKET_NAME：用于 Pathways 临时文件的 Cloud Storage 存储桶。
  • NUM_SLICES：工作负载中的 TPU 切片数量。

过滤活跃设备并构建设备网格

在 Pathways 工作负载中，jax.devices() 会在客户端进程的生命周期内返回初始分配中的所有设备。使用弹性训练时，请使用 elastic_manager.active_devices，以便设备网格和批次大小仅使用正常运行的活跃切片：

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

在副本大小调整模式下，当活跃切片数量发生变化时，请使用弹性数据加载器，例如 grain.experimental.ElasticIterator。标准每主机数据迭代器会保存每主机的状态，并且无法在不同数量的主机上恢复。相比之下，ElasticIterator 会跟踪单个全局数据集位置，您可以在任意数量的活跃切片中恢复该位置。

弹性训练循环

弹性训练循环可协调切片监控、故障恢复和内存中快照，如以下代码示例所示：

  • train 函数会在重试封装容器之外创建一个 manager.Manager() 实例和一个 Orbax Snapshotter，以便在多次尝试之间保留切片跟踪和主机内存快照。
  • 将 @elastic_manager.elastic_retry 应用于 run_training_loop 会将函数配置为等待健康切片、将 jax.default_device 设置为活跃切片、使后台监控器线程保持活跃状态，并在 jax.errors.JaxRuntimeError 上重试。
  • 在尝试恢复之初，您的代码必须从 elastic_manager.active_devices 重新构建设备网格，并调用 snapshotter.load 将最新快照恢复到更新后的网格分片上。让失败尝试中的设备缓冲区超出范围，以便 JAX 释放其内存。
  • 在循环的每个步骤中，您的代码会定期调用 snapshotter.save 来备份训练状态，并检查 elastic_manager.available_inactive_slices 以检测运行状况良好的替换切片。当替换切片可用时，您的代码会将其添加到 elastic_manager.active_slice_indices，重新构建网格，重新初始化训练循环，并将状态放置在扩展网格上。

以下示例展示了如何通过结合使用 manager.Manager、elastic_retry 和 Orbax 的内存中 Snapshotter 来实现弹性训练循环：

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

配置弹性管理器和重试行为

您需要为每个进程实例化一次 manager.Manager(devices=None)。如果您省略 devices，管理器会使用 jax.devices() 并按设备的 slice_index 属性对设备进行分组。

Manager 实例公开了以下属性，以便在训练期间检查集群拓扑：

  • slice_to_devices：从每个 slice_index 到属于该切片的 jax.Device 对象序列的映射。
  • all_slice_indices：分配中所有切片索引的 frozenset。
  • active_slice_indices：参与活跃计算网格的切片索引的 frozenset。
  • active_devices：属于活跃切片的 jax.Device 对象列表，按 slice_index 排序。
  • inactive_slice_indices：分配中未参与计算网格的切片索引的 frozenset。
  • available_inactive_slices：后台监控线程检测到健康且准备好重新加入的非活跃切片索引的 frozenset。
  • total_slice_count 和 active_slice_count：分配的切片总数和活跃切片数。
  • default_device：来自活跃切片的设备 0，elastic_retry 会自动将其设置为每次尝试的默认 JAX 设备。

您可以通过将以下参数传递给 @elastic_manager.elastic_retry 来自定义重试和伸缩行为：

  • minimum_slice_count：在开始初始执行或重试尝试之前，所需的健康切片数量下限。如果设置为 None（默认值），elastic_retry 会等待分配中的所有切片 (total_slice_count)，并在暂停和恢复模式下运行。如果设置为较小的整数（例如 1），elastic_retry 会在副本大小调整模式下运行。
  • poll_interval：弹性训练管理器 (Manager) 在训练期间执行切片健康检查的时间间隔（以秒为单位）。默认值为 10。
  • timeout：在引发 TimeoutError 之前，等待 minimum_slice_count 切片变为活跃状态的最长秒数。默认值为 None（无超时）。
  • pre_callback：一个可选回调，在管理器更新 active_slice_indices 之后、封装的训练循环函数运行之前，在每次尝试开始时运行。
  • on_elastic_event_callback：一个可选回调，在 elastic_retry 捕获到切片 jax.errors.JaxRuntimeError 后、下一次重试尝试开始前运行。您可以使用此回调来清理不完整的检查点或重置数据迭代器。
  • retry_policy：一个可调用的 (attempt: int, error: Exception) -> bool，用于确定是否在弹性事件后继续重试。使用 manager.ElasticRetryLimit(max_attempts) 限制总尝试次数。如果省略，elastic_retry 将无限期重试。

在弹性事件中，一旦堆栈帧展开且没有变量引用这些数组，JAX 就会释放失败尝试中的设备缓冲区。避免在全局变量或长期回调中保留对失败尝试中的设备数组的引用，以便在下次尝试开始之前释放内存。

将内存中快照与 Orbax 结合使用

ocp.v1.training.pathways.Snapshotter 类（在 orbax-checkpoint 0.12.6 版及更高版本中提供）在切片之间的 pinned_host 内存中维护训练状态的内存中备份：

  • Snapshotter(replica_axis_index=0)：初始化快照工具，并指定哪个网格轴表示数据并行切片副本。保存状态中的所有 JAX 数组都必须沿此网格轴进行复制，以便在切片发生故障时，幸存的切片仍保留模型和优化器状态的完整副本。如果您的状态包含单设备数组，请在调用 save 之前在网格中复制这些数组。
  • save(step, state)：将 state 异步传输到 pinned_host 内存，并沿 replica_axis_index 将其拆分为每个副本树。如果之前的快照传输仍在进行中，save 会跳过新请求，以避免训练循环停滞。
  • load(abstract_state, reset_snapshot_state=True)：探测已保存的主机副本以识别幸存切片，串联正常副本，并将状态重新分片到 abstract_state 指定的目标分片。
  • join()：在任何已加入队列的快照操作完成之前，阻止其他操作。
  • latest：返回最近完成的快照的元数据（包括 step），如果没有任何快照完成，则返回 None。

减少切片失败时的分片

当 TPU 切片在执行期间失败或被抢占时，Pathways 会引发 jax.errors.JaxRuntimeError（例如 DATA_LOSS、DEADLINE_EXCEEDED、NOT_FOUND 或 INTERNAL）。在用 @elastic_manager.elastic_retry 修饰的函数中，elastic_retry 会捕获此异常并检查 elastic.is_error_due_to_slice_down(error)：

  • 如果错误不是由切片事件引起的，elastic_retry 会立即重新引发异常。
  • 如果错误是由切片缩减事件引起的，elastic_retry 会停止所有活跃的 JAX 性能分析器跟踪记录，运行 on_elastic_event_callback，并评估 retry_policy(attempt, error)。
  • 在开始下一次尝试之前，elastic_retry 会调用 elastic.wait_for_slices 一次，以轮询切片运行状况，直到至少有 minimum_slice_count 个切片处于活跃状态。然后，它会更新 elastic_manager.active_slice_indices 并在幸存的切片上运行训练循环函数。

在切片重新加入时增加分片

当工作负载在副本大小调整模式下以较少的切片数量运行时，后台监控线程会定期对 elastic_manager.inactive_slice_indices 运行健康检查。当修复或重新安排的切片通过健康检查时，监控线程会将其切片索引添加到 elastic_manager.available_inactive_slices。

如果 elastic_manager.available_inactive_slices 在训练步骤结束时非空，则工作负载会就地扩容，而不会丢失训练进度：

  • 将 elastic_manager.available_inactive_slices 中的切片索引添加到 elastic_manager.active_slice_indices 会将重新联接的切片纳入 elastic_manager.active_devices，并将其从 available_inactive_slices 中清除。
  • 通过 elastic_manager.active_devices 重新构建设备网格，并重新初始化训练步骤和数据迭代器，使计算适应扩展后的拓扑。
  • 调用 jax.device_put 会将当前训练状态放置到新的网格分片上，而无需从快照恢复。

多层级检查点

大规模训练作业经常会中断，直接向 Cloud Storage 写入检查点或从中恢复检查点可能会成为训练进度的瓶颈。多层级检查点机制通过在 GKE 上协调三个存储层级，减少了检查点开销和恢复时间：

  • 本地 RAM 磁盘：TPU 工作进程会将高频检查点写入每个工作器节点上装载的内存中 CSI 卷 (multitier-checkpoint.csi.storage.gke.io)。此卷在同一节点上的 Pod 重启期间会保留。
  • 集群内对等复制：GKE 复制器服务会在集群中的对等 TPU 工作进程之间复制本地检查点分片。
  • Cloud Storage 备份：复制器服务会将本地检查点异步备份到启用了分层命名空间的 Cloud Storage 存储桶，而不会阻塞 TPU 训练步骤。

在多控制器 JAX 中，每个 TPU 主机都会运行自己的训练脚本，而 Pathways 则不同，它会在 pathways-head Pod 上运行单个控制器，并将工作分派给 pathways-worker Pod。为了直接在每个 TPU 主机的本地 RAM 磁盘上读取和写入检查点分片，而无需通过单个控制器路由检查点数据，Pathways 使用共置 Python。每个 worker Pod 都会运行一个 colocated-python-sidecar 容器和一个 pathways-worker 容器，并且这两个容器都会装载多层级检查点 CSI 卷和共享内存目录 (/tmp/sidecar)。在初始化和训练期间，Orbax 会将检查点设置、保存和恢复从 Pathways 控制器分派到每个 TPU 主机上的工作器边车。在缩减或扩容活跃 TPU 切片的数量时，Pathways 上的多层级检查点还支持弹性训练。

Pathways 支持通过 MaxText 使用多层级检查点。如需了解设置和配置说明，请参阅 MaxText 多层级检查点指南。

配置 Pathways 工作负载

您可以使用 Cluster Toolkit (gcluster job submit) 或配置自定义 Pathways JobSet 清单，为 Pathways 工作负载配置多层级检查点。

使用 Cluster Toolkit

使用 gcluster job submit --pathways 提交作业时，请传递 --gke-mtc-enabled、--gke-mtc-ramdisk-dir 和 --pathways-colocated-python-sidecar-image 标志：

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

将 SIDECAR_IMAGE 替换为共置的 Python 边车容器映像的 URI。

设置这些标志后，gcluster 会自动配置以下内容：

  • 将 --sidecar_name=external 传递给 pathways-proxy 容器。
  • 在每个 worker Pod 中运行 colocated-python-sidecar 容器，并在 pathways-worker 和 colocated-python-sidecar 上装载共享 /tmp/sidecar 目录 (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar)。
  • 在 worker Pod 上的 pathways-worker 和 colocated-python-sidecar 中，将 multitier-checkpoint.csi.storage.gke.io CSI 卷装载到 --gke-mtc-ramdisk-dir 指定的路径，但不在 pathways-head Pod 上装载。

使用自定义 Pathways JobSet 清单

如果您自行管理 Pathways JobSet 清单，而不是使用 gcluster，请将 --sidecar_name=external 传递给 pathways-proxy。在 worker Pod 规范中，将共享边车内存卷和 multitier-checkpoint.csi.storage.gke.io CSI 卷装载到 pathways-worker 和 colocated-python-sidecar 容器上：

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

在自定义 JAX 工作负载中使用 Orbax 多层检查点

在 Pathways 上运行的自定义 JAX 训练脚本中，使用 orbax.checkpoint.experimental.emergency.multi_tier_checkpointing 软件包，并在 initialize_multi_tier_checkpointing 和 ReplicatorCheckpointManagerOptions 中设置 use_colocated_python=True：

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

热插拔

热插拔是 GKE JobSet API 的一项功能，可让高优先级作业快速接管低优先级作业的资源，从而最大限度地减少停机时间并加快恢复速度。

创建 JobSet 时，GKE 会按照 JobSet 配置中的指定，在多个切片上调度工作负载。如果一个或多个切片发生硬件故障，GKE 会将受影响的 Pod 标记为失败。如果您的 GKE 集群有一个运行较低优先级作业的备用切片，JobSet 控制器会将故障切片的工作负载从较高优先级作业重新映射到该备用切片。此重新映射通常用时不超过 1 分钟。

JobSet 重启时，在以下情况下可能会发生热插拔：

  • 默认模式：如果同一集群中有空闲的备用 TPU 切片，Kubernetes 调度器会优先将重启的作业调度到这些切片上，而不是等待故障切片修复。这种调度行为可加快恢复速度。
  • 异构工作负载：在运行多个工作负载且配置了 Kubernetes PriorityClass 资源的集群中，重启的 JobSet 可能会触发热插拔。如果重启的作业的亲和性与优先级较低的作业的资源匹配，则 Kubernetes 会抢占优先级较低的作业，以便优先级较高的作业可以启动，而无需等待切片修复。例如，您可以使用 PriorityClass 资源为 Pathways 工作器 Pod 配置不同的优先级。

如需在集群中使用优先级，请定义 PriorityClass 清单并将其保存为 high-prior-job.yaml：

apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

将此 YAML 应用到您的 GKE 集群：

kubectl apply -f high-prior-job.yaml

接下来，将新的优先级类附加到 Pathways 工作负载。如果您使用 Cluster Toolkit 提交工作负载，请将 --priority=high-prior-job 传递给 gcluster job submit 命令，该命令会在 pathways-head 和 worker Pod 上设置 priorityClassName。如果您编写自己的 JobSet 清单，请在运行 pathways-worker 容器的 worker 复制作业的 Pod 模板中设置 priorityClassName：

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

如果您使用启用了 Pathways 的 Cluster Toolkit TPU 蓝图创建了集群，则该集群已包含 very-low、low、medium、high 和 very-high 优先级类别，因此您可以将其中一个传递给 --priority，而无需自行创建。

后续步骤