本教學課程示範如何在 Google Kubernetes Engine (GKE) 上協調多個多切片工作負載，以提升資源使用率。您會部署 Jax 工作負載做為範例，在 TPU Multislice 上執行，並使用 JobSet 和 Kueue 實作工作佇列。Kueue 會根據可用資源、配額和團隊間的公平分享階層，決定何時應執行 Job。
本教學課程的適用對象為機器學習 (ML) 工程師、平台管理員和操作員，他們對 Kubernetes 的容器自動化調度管理功能有興趣，並想瞭解如何使用這項功能訓練 LLM。如要進一步瞭解我們在 Google Cloud 內容中提及的常見角色和範例工作，請參閱「常見的 GKE 使用者角色和工作」。
閱讀本頁面之前，請先熟悉下列概念：
- 目前可用的 TPU 版本 (Cloud TPU 系統架構)
- GKE 中的 TPU Multislice
目標
- 準備環境，其中包含三個 v5e TPU 節點的 GKE 叢集。每個 TPU 節點都有
2x4拓撲，內含 8 個晶片。因此總共有 24 個 TPU v5e TPU 晶片。 - 建立 Kueue 資源，確保工作負載之間公平分配配額。
- 執行 Multislice 工作負載。
事前準備
開始之前，請確認您已完成下列工作：
- 啟用 Google Kubernetes Engine API。 啟用 Google Kubernetes Engine API
- 如要使用 Google Cloud CLI 執行這項工作，請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI，請執行
gcloud components update，取得最新版本。
準備環境
在 Google Cloud 控制台中啟動 Cloud Shell 執行個體：
開啟 Cloud Shell使用
gcloud config set指令設定預設環境變數：gcloud config set project PROJECT_ID將 PROJECT_ID 替換為 Google Cloud 專案 ID。
執行 1.29.2-gke.1521000 以上版本的 Autopilot 叢集預設會啟用 TPU。Autopilot 叢集上的 TPU 會在工作負載規格中設定。詳情請參閱「使用 JobSet 定義多切片工作負載」一節。
建立 GKE 叢集
在 Cloud Shell 中建立 GKE 叢集：
Autopilot
gcloud container clusters create-auto multislice-cluster \
--location=LOCATION \
--cluster-version 1.29.2-gke.1521000 \
--release-channel rapid
在這個指令中：
--location標記用於指定叢集的 Compute Engine 位置。--cluster-version標記會指定叢集的 Kubernetes 版本。--release-channel標記會指定叢集的發布版本。在這種情況下，快速通道支援 GKE 的最新版本。
標準
gcloud container clusters create multislice-cluster \
--location=LOCATION
將 LOCATION 替換為要建立叢集的位置。確認該區域有 ct5lp-hightpu-4t 機型容量。建立叢集可能需要幾分鐘的時間。
如果您使用 GKE Autopilot 模式，請跳至「建立 Kueue 資源」一節。執行 1.29.2-gke.1521000 以上版本的 Autopilot 叢集預設會啟用 TPU。
建立三個標準模式 TPU 節點集區
在本節中，您將使用 gcloud beta container node-pools create 指令建立 TPU 節點集區。
建立名為
nodepool1的第一個節點集區：gcloud beta container node-pools create nodepool1 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID將 NODE_LOCATION 替換為叢集區域中的一或多個區域，您要在這些區域中建立節點。
建立名為
nodepool2的第二個節點集區：gcloud beta container node-pools create nodepool2 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID建立名為
nodepool3的第三個節點集區：gcloud beta container node-pools create nodepool3 \ --location=LOCATION \ --cluster=multislice-cluster \ --node-locations=NODE_LOCATION \ --machine-type=ct5lp-hightpu-4t \ --tpu-topology=2x4 \ --project=PROJECT_ID
GKE 會建立三個節點集區。每個節點集區都是獨立的 TPU 配量。
在先前的步驟中，您使用 gcloud beta container node-pools create 指令建立節點集區。這些指令會使用下列標記：
--node-locations：以半形逗號分隔的清單，當中列出 GKE 建立節點集區的一或多個區域。--machine-type：節點使用的機器類型。在本例中，您使用了ct5lp-hightpu-4t。如要進一步瞭解與 TPU 相容的機器類型，請參閱「選擇 TPU 版本」一文中的表格。--tpu-topology：節點集區要使用的 TPU 拓撲。在本例中，您使用了2x4。如要進一步瞭解 TPU 拓撲，請參閱「選擇 TPU 拓撲」。
建立 Kueue 資源
建立下列
kueue.yaml資訊清單：apiVersion: kueue.x-k8s.io/v1beta1 kind: ResourceFlavor metadata: name: "vlp-24" spec: nodeLabels: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "cluster-queue" spec: namespaceSelector: {} queueingStrategy: BestEffortFIFO resourceGroups: - coveredResources: ["google.com/tpu"] flavors: - name: "vlp-24" resources: - name: "google.com/tpu" nominalQuota: 24 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: default name: multislice-queue spec: clusterQueue: cluster-queue套用
kueue.yaml資訊清單：kubectl apply -f kueue.yaml
GKE 會建立下列 Kueue 資源：
- ResourceFlavor：
叢集中資源的抽象概念。在本範例中，GKE 會建立三個具有
2x4拓撲的 TPU 節點。每個 TPU 配量都有2x4拓撲，內含 8 個晶片 (總共 24 個 TPU 晶片)。 - ClusterQueue： 管理工作負載和叢集資源的全域佇列。
- LocalQueue： 將密切相關的工作負載分組，這些工作負載通常由單一租戶 (使用者) 執行。每個 LocalQueue 都會指向 ClusterQueue，資源會從該 ClusterQueue 分配，以執行工作負載。Kueue 工作負載是代表批次工作負載的抽象概念，在本例中，每個工作負載都是 JobSet。
使用 JobSet 定義 Multislice 工作負載
在本節中，您將建立三個 JobSet。Jobset 是工作負載 API，可讓您以單元形式管理一組 Kubernetes Job。JobSet 最常見的用途是分散式訓練，但您也可以用來執行批次工作負載。
下列 JobSet 會執行 Jax 工作負載，輸出切片中的 TPU 晶片總數，然後休眠 60 秒來模擬模型訓練時間，最後結束。
在叢集中安裝 JobSet API：
VERSION=v0.8.1 kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml建立下列
jobsets-multislice.yaml資訊清單：Autopilot
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-1slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-2slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 2 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-3slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 3 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4標準
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-1slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 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("Global device count:", jax.device_count())' resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-2slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 2 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 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("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: 4 --- apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-3slice labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 3 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4套用
jobsets-multislice.yaml資訊清單：kubectl apply -f jobsets-multislice.yaml
GKE 會使用下列資源要求建立 Job：
multislice-1sliceJobSet 會建立一個 Job，總共需要一個 TPU 配量。multislice-2sliceJobSet 會建立兩項工作，總共需要兩個 TPU 配量。multislice-3sliceJobSet 會建立三個工作，總共需要三個 TPU 配量。
由於叢集只有三個 TPU 切片，因此並非所有 JobSet 都能同時執行。
Kueue 將所有三個 multislice-3sliceJobSet 加入佇列後，這些 Job 會單獨執行，直到完成為止。multislice-1slice 和 multislice-2slice 會等待，然後一起執行。
確認 Kueue 是否已接受工作負載
在 Kueue 中檢查已加入佇列的工作負載：
kubectl get workloads輸出結果會與下列內容相似：
NAME QUEUE ADMITTED BY AGE jobset-multislice-1slice-2530a multislice-queue 3s jobset-multislice-2slice-ffb02 multislice-queue 4s jobset-multislice-3slice-8c695 multislice-queue cluster-queue 10s
Kueue 會根據工作負載所需的 TPU 資源，將一或多個工作負載加入佇列。
監控工作負載
Google Cloud 控制台中的 JobSet 和節點集區可觀測性指標與資訊主頁，現已全面開放使用。
資訊主頁
如要查看 GKE 上 TPU 多主機節點集區的狀態，請前往 Cloud Monitoring 提供的 GKE TPU 節點集區狀態資訊主頁：
詳情請參閱「監控 TPU 節點和節點集區的健康狀態指標」。
在Google Cloud 控制台的「Kubernetes Engine AI/ML」頁面中，「AI deployment > Jobs」分頁會顯示 JobSet 監控資訊主頁，其中包含 JobSet 和基礎架構的健康狀態與效能等全面資訊，例如 JobSet 狀態、副本就緒狀態和副本狀態。資訊主頁也會顯示基礎架構指標，包括 CPU、GPU、TPU、記憶體和儲存空間指標。 詳情請參閱「使用指標監控 JobSet 狀態」。
監控哪些 Pod 正在執行
kubectl get pods
輸出結果會與下列內容相似：
NAME READY STATUS RESTARTS AGE
multislice-1slice-slice-0-0-pf2ll 1/1 Running 0 1s
multislice-1slice-slice-0-1-55g62 1/1 Running 0 1s
multislice-2slice-slice-0-0-f4hf7 1/1 Running 0 3s
multislice-2slice-slice-0-1-c8kv7 1/1 Running 0 3s
multislice-2slice-slice-1-0-7h46t 1/1 Running 0 3s
multislice-2slice-slice-1-1-lj9hb 1/1 Running 0 3s
multislice-3slice-slice-0-0-wzq9t 0/1 Completed 0 2m31s
multislice-3slice-slice-0-1-zf4dp 0/1 Completed 0 2m30s
multislice-3slice-slice-1-0-hbfn5 0/1 Completed 0 2m31s
multislice-3slice-slice-1-1-45fgl 0/1 Completed 0 2m30s
multislice-3slice-slice-2-0-wjbp4 0/1 Completed 0 2m30s
multislice-3slice-slice-2-1-lwnvs 0/1 Completed 0 2m30s
您會看到 GKE 已排定、建立及執行 multislice-3slice 的 Pod。接著，GKE 會從 multislice-1slice 和 multislice-2slice JobSet 執行 Pod。
使用指標監控 JobSet 健康狀態
如要瞭解 JobSet 是否正常執行，或推斷是否中斷，可以使用 JobSet 指標套件中的 Prometheus 指標，例如 kube_jobset_succeeded_replicas。
請注意，Jobset 健康狀態指標僅適用於 GKE 1.32.1-gke.135700 以上版本。使用支援版本建立的新叢集，預設會啟用 JobSet 健康狀態指標。如果現有叢集升級至支援的版本，客戶必須手動啟用 JobSet 指標套件。詳情請參閱說明文件。
在本教學課程中，請使用下列 PromQL 查詢檢查 JobSet 是否完成：
kube_jobset_succeeded_replicas{
cluster="multislice-cluster",
jobset_name=~"mulitslice-.*"}
監控 JobSet 運作時間、復原時間 (TTR) 和中斷間隔時間 (TBI)
下列指標有助於監控 JobSet 的可用性：
kubernetes.io/jobset/uptime：JobSet 的總可用時間。kubernetes.io/jobset/times_to_recover：JobSet 的復原期分配。每個樣本都代表 JobSet 停機期間的單一復原事件。kubernetes.io/jobset/times_between_interruptions：JobSet 中，前一次中斷結束到目前中斷開始之間的間隔時間分布。每個樣本都代表前一次和目前中斷之間的單一時間長度。
這些指標適用於只有一個 GPU 或 TPU 複製工作的 JobSet。指標的計算依據僅為該單一複製工作的可用性。 所有 GKE 版本都支援這些指標。
如要查看本教學課程中使用的 JobSet 正常運作時間，請執行下列 PromQL 查詢：
avg_over_time(
kubernetes_io:jobset_uptime{
monitored_resource="k8s_entity", entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}])
如要查看本教學課程中 JobSet 的 TBI 分布情形，請執行下列 PromQL 查詢：
histogram_quantile(0.50,
sum_over_time(
kubernetes_io:jobset_times_between_interruptions_bucket{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
您可以將查詢間隔延長至較長的時間範圍 (例如 7 天)，並計算這段期間的中斷情形平均間隔時間 (MTBI)：
sum(sum_over_time(
kubernetes_io:jobset_times_between_interruptions_sum{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
kubernetes_io:jobset_times_between_interruptions_count{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
如要查看 TTR 分布情形，可以執行下列 PromQL 查詢：
histogram_quantile(0.50,
sum_over_time(
kubernetes_io:jobset_times_to_recover_bucket{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
將查詢間隔延長至較長的時間範圍 (例如 7 天) 後，您就可以計算這段期間的平均修復時間 (MTTR)：
sum(sum_over_time(
kubernetes_io:jobset_times_to_recover_sum{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
kubernetes_io:jobset_times_to_recover_count{
monitored_resource="k8s_entity",entity_type="jobset",
entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
啟用 Kueue 工作負載優先順序和搶占功能
您可以視需要指派 Kueue 工作負載優先順序，決定 Kueue 允許佇列工作負載的順序。
更新
ClusterQueue，加入搶占政策：apiVersion: kueue.x-k8s.io/v1beta1 kind: ResourceFlavor metadata: name: "vlp-24" spec: nodeLabels: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "cluster-queue" spec: namespaceSelector: {} resourceGroups: - coveredResources: ["google.com/tpu"] flavors: - name: "vlp-24" resources: - name: "google.com/tpu" nominalQuota: 24 preemption: reclaimWithinCohort: Any withinClusterQueue: LowerPriority --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: default name: multislice-queue spec: clusterQueue: cluster-queue為要指派給工作負載的每個不同優先順序等級建立
PriorityClass：apiVersion: scheduling.k8s.io/v1 kind: PriorityClass metadata: name: low-priority value: 100 globalDefault: false description: "This low priority class should be used for some Pods only."將
priorityClassName指派給 JobSet：Autopilot
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: low-priority labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 priorityClassName: low-priority containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4 # Number of TPU chips per worker標準
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: low-priority labels: kueue.x-k8s.io/queue-name: multislice-queue annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: 1 template: spec: parallelism: 2 completions: 2 backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 priorityClassName: low-priority containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 securityContext: privileged: true command: - bash - -c - | sleep 60 resources: limits: google.com/tpu: 4 # Number of TPU chips per worker
GKE 包含搶占政策，可定義 Kueue 如何指派可用資源。這項政策規定，如果優先順序較高的工作負載需要資源，系統可以搶占工作負載的資源。優先順序值較低的工作負載，較有可能遭到優先順序較高的工作負載搶占。
清除所用資源
如要避免系統向您的 Google Cloud 帳戶收取本教學課程中所用資源的相關費用，請刪除含有該項資源的專案，或者保留專案但刪除個別資源。
刪除專案
- In the Google Cloud console, go to the Manage resources page.
- 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.
刪除個別資源
刪除 Kueue 資源：
kubectl delete -f jobsets-multislice.yaml kubectl delete -f kueue.yaml刪除叢集：
gcloud container clusters delete multislice-cluster --region=LOCATION
後續步驟
- 進一步瞭解 Kueue。
- 瞭解如何在 GKE 上實作工作佇列系統，並在命名空間之間共用配額。