Mengorkestrasi workload Multislice menggunakan JobSet dan Kueue

Autopilot Standard

Tutorial ini menunjukkan cara menjalankan beban kerja Jax menggunakan TPU Multislice di Google Kubernetes Engine (GKE) dan Kueue. Kueue menerapkan antrean tugas, memutuskan kapan dan kapan Tugas harus dimulai, berdasarkan kuota dan hierarki agar dapat berbagi resource secara adil di antara tim.

Tutorial ini menunjukkan cara mengorkestrasi beberapa workload Multislice yang memerlukan resource TPU agar berjalan serentak.

Sebelum menggunakan TPU di GKE, sebaiknya selesaikan jalur pembelajaran berikut:

Pelajari ketersediaan versi TPU saat ini dengan arsitektur sistem Cloud TPU.
Pelajari TPU Multislice di GKE.

Tujuan

Tutorial ini ditujukan bagi administrator GKE yang sudah memiliki cluster GKE dan ingin menjalankan workload Multislice untuk pertama kalinya.

Tutorial ini membahas langkah-langkah berikut:

Siapkan lingkungan Anda menggunakan cluster GKE dengan tiga irisan TPU v5e. Setiap slice TPU memiliki topologi 2x4 dan empat chip per host. Oleh karena itu, total 24 chip TPU v5e.
Membuat resource Kueue untuk memastikan bahwa kuota dibagi secara adil di antara beban kerja.
Jalankan workload Multislice Anda.

Sebelum memulai

Sebelum memulai, pastikan Anda telah menjalankan tugas berikut:

Aktifkan Google Kubernetes Engine API.

Aktifkan Google Kubernetes Engine API

Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu initialize gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.
Catatan: Untuk penginstalan gcloud CLI yang ada, pastikan untuk menyetel properti compute/region dan compute/zone. Dengan menyetel lokasi default, Anda dapat menghindari error di gcloud CLI yang seperti ini: One of [--zone, --region] must be supplied: Please specify location.

Menginstal JobSet v0.2.3 atau yang lebih baru.
Install Kueue v0.4.1 atau yang lebih baru.

Menyiapkan lingkungan

Di konsol Google Cloud, mulai instance Cloud Shell:
Buka Cloud Shell
Tetapkan variabel lingkungan default:
```
gcloud config set project PROJECT_ID
gcloud config set compute/region COMPUTE_REGION
```
Ganti nilai berikut:
- PROJECT_ID: project ID Google Cloud Anda.
- COMPUTE_REGION: region Compute Engine.

Cluster Autopilot yang menjalankan versi 1.29.2-gke.1521000 atau yang lebih baru mengaktifkan TPU secara default. TPU pada cluster Autopilot dikonfigurasi dalam spesifikasi beban kerja. Untuk mengetahui informasi selengkapnya, lihat bagian Menentukan workload Multislice dengan JobSets.

Membuat cluster GKE

Di Cloud Shell, buat cluster GKE:

Autopilot

gcloud container clusters create-auto multislice-cluster \
    --location=LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

Standar

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

Ganti LOCATION dengan lokasi tempat Anda ingin membuat cluster. Pastikan kapasitas untuk jenis mesin ct5lp-hightpu-4t. Pembuatan cluster mungkin memerlukan waktu beberapa menit.

Jika Anda menggunakan mode GKE Autopilot, lanjutkan ke bagian Membuat resource Kueue. Cluster autopilot yang menjalankan versi 1.29.2-gke.1521000 atau yang lebih baru mengaktifkan TPU secara default.

Membuat tiga kumpulan node TPU mode Standar

Buat kumpulan node pertama bernama 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 \
    --num-nodes=2 \
    --project=PROJECT_ID

Ganti NODE_LOCATION dengan satu atau beberapa zona di region cluster tempat Anda ingin membuat node.

Buat kumpulan node kedua bernama 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 \
    --num-nodes=2 \
    --project=PROJECT_ID

Buat kumpulan node ketiga bernama 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 \
    --num-nodes=2 \
    --project=PROJECT_ID

GKE membuat tiga node pool. Setiap kumpulan node adalah irisan TPU terpisah.

Membuat sumber daya Kueue

Buat manifes kueue.yaml berikut:

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

Terapkan manifes kueue.yaml:
```
kubectl apply -f kueue.yaml
```
GKE membuat resource Kueue berikut:

ResourceFlavor: Abstraksi resource dalam cluster. Dalam contoh ini, tiga irisan TPU dengan topologi 2x4 dan empat chip per host, sehingga menghasilkan 24 chip TPU.
ClusterQueue: Antrean global yang mengelola beban kerja dan resource cluster.
LocalQueue: Mengelompokkan beban kerja yang terkait erat yang biasanya dijalankan oleh satu tenant (pengguna). Setiap LocalQueue mengarah ke ClusterQueue tempat resource dialokasikan untuk menjalankan beban kerjanya. Beban Kerja Kueue adalah abstraksi yang mewakili beban kerja batch, dalam hal ini, setiap beban kerja adalah JobSet.

Menentukan workload Multislice dengan JobSets

Di bagian ini, Anda akan membuat tiga JobSet. JobSet ini menjalankan workload Jax yang menghasilkan jumlah global chip TPU dalam slice, lalu tidur selama 60 detik untuk menyimulasikan waktu pelatihan model, lalu keluar.

Buat manifes jobsets-multislice.yaml berikut:

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

Standar

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

Terapkan manifes jobsets-multislice.yaml:

kubectl apply -f jobsets-multislice.yaml

GKE membuat Tugas dengan permintaan resource berikut:

multislice-1slice JobSet membuat satu Tugas yang memerlukan total satu irisan TPU.
multislice-2slice JobSet membuat dua Tugas yang memerlukan total dua irisan TPU.
multislice-3slice JobSet membuat tiga Tugas yang memerlukan total tiga slice TPU.

Karena cluster hanya memiliki tiga irisan TPU, tidak semua JobSets dapat dijalankan sekaligus. Saat Kueue mengantrekan ketiga multislice-3slice JobSet, Tugasnya berjalan sendiri hingga selesai. multislice-1slice dan multislice-2slice akan menunggu dan berjalan bersama setelahnya.

Memverifikasi bahwa Kueue menerima workload

Memeriksa beban kerja dalam antrean di Kueue:

kubectl get workloads

Outputnya mirip dengan hal berikut ini:

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 mengantrekan satu atau beberapa workload, bergantung pada resource TPU yang dibutuhkan.

Memantau beban kerja

Pantau pod mana yang sedang berjalan:

kubectl get pods

Outputnya mirip dengan hal berikut ini:

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

Lihat bahwa GKE menjadwalkan, membuat, dan menjalankan Pod untuk multislice-3slice terlebih dahulu. Kemudian, GKE menjalankan Pod dari multislice-1slice dan multislice-2slice JobSets.

Memungkinkan prioritas dan preemption workload Kueue

Jika ingin, Anda dapat menetapkan prioritas beban kerja Kueue yang menentukan urutan beban kerja dalam antrean diterima oleh Kueue.

Update ClusterQueue Anda agar memiliki kebijakan preemption:

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

Buat PriorityClass untuk setiap tingkat prioritas berbeda yang ingin Anda tetapkan ke workload:

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."

Tetapkan priorityClassName ke JobSet Anda:

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

Standar

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
  ```

Pembersihan

Agar tidak perlu membayar biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

Menghapus project

Di konsol Google Cloud, buka halaman Manage resource.
Buka Manage resource
Pada daftar project, pilih project yang ingin Anda hapus, lalu klik Delete.
Pada dialog, ketik project ID, lalu klik Shut down untuk menghapus project.

Menghapus resource satu per satu

Hapus sistem kuota Kueue:

kubectl delete -n team-a localqueue
kubectl delete -n team-b localqueue
kubectl delete clusterqueue
kubectl delete clusterqueue
kubectl delete clusterqueue
kubectl delete resourceflavor
kubectl delete resourceflavor
kubectl delete resourceflavor

Hapus manifes Kueue:

VERSION=kueue.x-k8s.io/v1beta1
kubectl delete -f \
    https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml

Hapus kluster:

gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION

Langkah selanjutnya

Pelajari Kueue lebih lanjut
Pelajari cara Mengimplementasikan sistem antrean Tugas dengan berbagi kuota antar-namespace di GKE.