JobSet と Kueue を使用してマルチスライスワークロードをオーケストレートする

Autopilot Standard

このチュートリアルでは、リソース使用率を向上させるために、Google Kubernetes Engine（GKE）で複数のマルチスライスワークロードをオーケストレートする方法について説明します。例として Jax ワークロードをデプロイし、TPU マルチスライスで実行して、JobSet と Kueue を使用してジョブキューイングを実装します。Kueue は、使用可能なリソース、割り当て、チーム間での公平な共有のための階層に基づいて、Job を実行するタイミングを決定します。

このチュートリアルは、LLM のトレーニングに Kubernetes のコンテナオーケストレーション機能を使用することに関心のある ML エンジニア、プラットフォーム管理者、オペレーターを対象としています。 Google Cloud のコンテンツで使用されている一般的なロールとタスクの例の詳細については、一般的な GKE Enterprise ユーザーロールとタスクをご覧ください。

このページを読む前に、次のことをよく理解しておいてください。

Cloud TPU システムアーキテクチャでの現在の TPU バージョンの可用性
GKE での TPU マルチスライス

目標

3 つの v5e TPU スライスを含む GKE クラスタで環境を準備します。各 TPU スライスには、8 個のチップを含む 2x4 トポロジがあります。そのため、合計で 24 個の TPU v5e チップがあります。
Kueue リソースを作成して、ワークロード間で割り当てが公平に共有されるようにします。
マルチスライスワークロードを実行します。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

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 など）を防止できます。クラスタのロケーションが設定したデフォルトと異なる場合は、特定のコマンドでロケーションの指定が必要になることがあります。

v0.2.3 以降の JobSet をインストールします。
v0.4.1 以降の Kueue をインストールします。

環境を準備する

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 フラグは、クラスタのリリースチャンネルを指定します。この場合、Rapid チャンネルは GKE で使用可能な最新バージョンをサポートします。

Standard

gcloud container clusters create multislice-cluster \
    --location=LOCATION

LOCATION は、クラスタを作成するロケーションに置き換えます。ct5lp-hightpu-4t マシンタイプの容量があることを確認します。クラスタの作成には数分かかることもあります。

GKE Autopilot モードを使用している場合は、Kueue リソースを作成するセクションに進みます。バージョン 1.29.2-gke.1521000 以降を実行する Autopilot クラスタでは、デフォルトで TPU が有効になります。

3 つの Standard モード TPU スライスノードプールを作成する

このセクションでは、gcloud beta container node-pools create コマンドを使用して TPU ノードプールを作成します。

nodepool1 という名前で 1 つ目のノードプールを作成します。

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 は、ノードを作成するクラスタリージョンの 1 つ以上のゾーンに置き換えます。

nodepool2 という名前で 2 つ目のノードプールを作成します。

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 という名前の 3 つ目のノードプールを作成します。

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 は 3 つのノードプールを作成します。それぞれのノードプールは個別の TPU スライスです。

前の手順では、gcloud beta container node-pools create コマンドを使用してノードプールを作成しました。これらのコマンドでは、次のフラグを使用します。

--node-locations: GKE がノードプールを作成する 1 つ以上のゾーンのカンマ区切りリスト。
--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 トポロジを持つ 3 つの TPU スライスを作成します。各 TPU スライスには、8 個のチップ（合計 24 個の TPU チップ）を含む 2x4 トポロジがあります。
ClusterQueue: ワークロードとクラスタリソースを管理するグローバルキュー。
LocalQueue: 通常は単一のテナント（ユーザー）が実行する密接に関連したワークロードをグループ化します。各 LocalQueue は、ワークロードを実行するためにリソースが割り当てられる ClusterQueue を参照します。Kueue ワークロードはバッチワークロードを表す抽象化です。この場合の各ワークロードは JobSet です。

JobSet を使用してマルチスライスワークロードを定義する

このセクションでは、3 つの JobSet を作成します。Jobset は、Kubernetes Job のグループをユニットとして管理できるワークロード API です。JobSet の最も一般的なユースケースは分散トレーニングですが、バッチワークロードの実行にも使用できます。

次の JobSet は、スライス内の TPU チップの総数を出力する Jax ワークロードを実行し、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

Standard

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-1slice JobSet は、合計 1 つの TPU スライスを必要とする 1 つの Job を作成します。
multislice-2slice JobSet は、合計 2 つの TPU スライスを必要とする 2 つの Job を作成します。
multislice-3slice JobSet は、合計 3 つの TPU スライスを必要とする 3 つの Job を作成します。

クラスタには TPU スライスが 3 つしかないため、すべての JobSet を一度に実行できるわけではありません。Kueue が 3 つすべての multislice-3slice JobSet をキューに登録すると、その 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 リソースに応じて 1 つ以上のワークロードをキューに登録します。

ワークロードをモニタリングする

Google Cloud コンソールの JobSet とノードプールオブザーバビリティの指標とダッシュボードが一般提供されています。

ダッシュボード

GKE の TPU マルチホストノードプールのステータスを表示するには、Cloud Monitoring が提供する [GKE TPU Node Pool Status] ダッシュボードに移動します。

[GKE TPU Node Pool Status] に移動

詳細については、TPU ノードとノードプールの健全性指標をモニタリングするをご覧ください。

Google Cloud コンソールの [Kubernetes Engine AI/ML] ページの [AI デプロイ > ジョブ] タブには、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 の前回の中断の終了から現在の中断の開始までの間隔の分布。各サンプルは、前回の中断から現在の中断までの 1 つの期間を示します。

これらの指標は、GPU または TPU のレプリケートされたジョブが 1 つだけある 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

Standard

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 アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

プロジェクトを削除する

注意: プロジェクトを削除すると、次のような影響があります。

プロジェクト内のすべてのものが削除されます。このドキュメントのタスクで既存のプロジェクトを使用した場合、それを削除すると、そのプロジェクトで行った他の作業もすべて削除されます。
カスタムプロジェクト ID が失われます。このプロジェクトを作成したときに、将来使用するカスタムプロジェクト ID を作成した可能性があります。そのプロジェクト ID を使用した URL（たとえば、appspot.com）を保持するには、プロジェクト全体ではなくプロジェクト内の選択したリソースだけを削除します。

In the Google 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.

個々のリソースを削除する

Kueue リソースを削除します。

kubectl delete -f jobsets-multislice.yaml
kubectl delete -f kueue.yaml

クラスタを削除します。

gcloud container clusters delete multislice-cluster --region=LOCATION

次のステップ

Kueue の詳細を確認する。
GKE の Namespace 間で割り当てを共有するジョブキューイングシステムを実装する方法を学習する。

JobSet と Kueue を使用してマルチスライス ワークロードをオーケストレートする コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

目標

始める前に

環境を準備する

GKE クラスタを作成する

Autopilot

Standard

3 つの Standard モード TPU スライス ノードプールを作成する

Kueue リソースを作成する

JobSet を使用してマルチスライス ワークロードを定義する

Autopilot

Standard

Kueue がワークロードを承諾したことを確認する

ワークロードをモニタリングする

ダッシュボード

実行中の Pod をモニタリングする

指標を使用して JobSet の健全性をモニタリングする

JobSet の稼働時間、復旧時間（TTR）、中断間隔（TBI）をモニタリングする

Kueue ワークロードの優先度とプリエンプションを有効にする

Autopilot

Standard

クリーンアップ

プロジェクトを削除する

個々のリソースを削除する

次のステップ

JobSet と Kueue を使用してマルチスライスワークロードをオーケストレートする

3 つの Standard モード TPU スライスノードプールを作成する

JobSet を使用してマルチスライスワークロードを定義する