推理工作负载的 TPU 收集调度

Trillium (v6e) 包含一项名为“集合调度”的功能，可让您对一组 TPU 切片进行分组，这些切片可以是单主机或多主机，用于提供同一模型的副本。此功能适用于 Cloud TPU 和 GKE 配置。

本文档介绍了如何将收集作业调度与 Cloud TPU API 搭配使用。如需详细了解如何将收集作业调度功能与 GKE 搭配使用，请参阅 GKE 文档。

通过为推理工作负载创建集合，Google Cloud 可以限制和简化对推理工作负载操作的干扰。这对于需要高可用性的推理工作负载非常有用。 Google Cloud 可确保集合能够高可用地管理传入流量。集合中的部分 slice 始终可用于处理传入流量。

集合中的每个 TPU 切片都将采用相同的加速器类型和拓扑。

收集时间安排仅适用于 v6e。

通过 Cloud TPU API 创建集合

使用 Cloud TPU API 请求已排队的资源时，您可以使用 --workload-type = AVAILABILITY-OPTIMIZED 标志创建集合。此标志会向 Cloud TPU 基础架构指明，它旨在用于注重可用性的工作负载。

以下命令使用 Cloud TPU API 预配集合：

gcloud alpha compute tpus queued-resources create serving-QR \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --node-count ${NODE_COUNT} \
   --node-prefix "servingTPU" \
   --workload-type = AVAILABILITY-OPTIMIZED

--node-count 标志用于指定队列中资源的 slice 数量。这会创建一组 TPU 切片。

可选：--node-prefix 标志用于为 slice 名称指定前缀。

注意 ：v6e 支持的配置中介绍了支持的加速器类型。