Mengoptimalkan pelatihan AI di TPU dengan DWS dan Kueue

Dokumen ini mengajarkan cara mengonfigurasi cluster GKE untuk memaksimalkan ketersediaan pelatihan AI di Tensor Processing Unit (TPU). Anda mengonfigurasi sistem penggantian otomatis menggunakan alat antrean pekerjaan open source yang disebut Kueue, dan Google Cloud's Dynamic Workload Scheduler (DWS).

Konfigurasi yang Anda tentukan dalam dokumen ini menyiapkan kumpulan node utama dan kumpulan node cadangan:

  • Sesuai permintaan (Paket A): Node pool ini adalah pilihan pertama Anda untuk node. Sistem akan mencoba menjadwalkan tugas di mesin on-demand ini terlebih dahulu. Istilah on-demand berarti bahwa setelah mesin ini mulai berjalan, Anda memiliki akses yang andal dan tanpa gangguan ke mesin tersebut.
  • DWS flex-start (Plan B): Ini adalah kumpulan node cadangan Anda. Jika mesin Plan A tidak tersedia, program penjadwalan Kueue akan otomatis menetapkan tugas Anda ke pool Plan B ini. Kemudian, DWS akan mencari hardware Plan B, tetapi tidak menjamin akses langsung karena hardware tersebut mungkin juga tidak tersedia. Namun, DWS tidak menyerah: DWS akan menahan permintaan Anda dalam antrean hingga 7 hari dan otomatis menyediakan mesin segera setelah tersedia.

Pendekatan ini meminimalkan waktu tunggu pekerjaan dalam antrean. Artinya, Anda tidak perlu memeriksa ketersediaan resource secara manual atau menulis ulang skrip untuk mesin yang berbeda.

Ringkasan langkah-langkah konfigurasi

Untuk mengonfigurasi sistem penggantian otomatis, Anda harus menyelesaikan beberapa langkah konfigurasi. Konfigurasi ini dapat dibagi menjadi dua kategori:

  • Tugas Administrator Cluster: Konfigurasi infrastruktur satu kali, seperti membuat cluster GKE, menyediakan node pool, dan menginstal pengontrol penjadwalan Kueue.
  • Tugas Developer AI: Alur kerja harian yang berulang, seperti menentukan persyaratan tugas pelatihan dan mengirimkan workload.

Meskipun Anda melakukan semua langkah ini sendiri, perlu diingat bahwa perbedaan ini akan membantu memperjelas keseluruhan proses.

Sebelum mengonfigurasi sistem, tinjau langkah-langkah konfigurasi yang akan Anda lakukan.

Topik Tugas
Mengonfigurasi infrastruktur (Administrator Cluster) 1. Buat cluster GKE
2. Buat node pool
3. Memverifikasi status mulai fleksibel di node pool
Menginstal dan mengonfigurasi Kueue (Administrator Cluster) 1. Instal Kueue
2. Tentukan aturan konfigurasi
Menjalankan tugas pelatihan (Developer AI) 1. Buat ConfigMap
2. Tentukan manifes RayJob
3. Kirimkan workload
4. Hubungkan ke RayJob
5. Periksa log

Konsep utama

  • Node pool on-demand (Paket A): Node pool utama dengan prioritas tinggi. Pekerjaan Anda selalu mencoba menggunakan kumpulan ini terlebih dahulu.
  • Kumpulan node mulai fleksibel DWS (Plan B): Kumpulan node cadangan. Jika mesin di kumpulan primer tidak tersedia, sistem akan otomatis menggunakan kumpulan ini untuk menelusuri hardware yang tersedia.
  • Kueue: Program penjadwalan yang mengelola antrean tugas. Fitur ini mencegat permintaan tugas Anda dan memutuskan kumpulan node mana yang akan digunakan (Paket A atau Paket B).
  • Tugas: Workload pelatihan AI yang ingin Anda jalankan. Dalam dokumen ini, Anda menentukannya menggunakan manifes RayJob.

Sebelum memulai

  1. Di konsol Google Cloud , pada halaman pemilih project, pilih atau buat project Google Cloud .

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.

    Buka pemilih project

  2. Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  3. Aktifkan Google Kubernetes Engine, Cloud TPU API.

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran.

    Aktifkan API

  4. Di konsol Google Cloud , aktifkan Cloud Shell.

    Aktifkan Cloud Shell

  5. Pastikan Anda memiliki kuota preemptible yang cukup untuk menggunakan VM flex-start TPU. Jika kuota default tidak cukup untuk kebutuhan Anda, minta alokasi yang lebih tinggi. Untuk mengetahui detailnya, lihat kuota Cloud TPU dan Menyiapkan lingkungan Cloud TPU.

Menentukan variabel lingkungan

Untuk menyederhanakan perintah yang Anda jalankan dalam dokumen ini, Anda dapat menetapkan variabel lingkungan di Cloud Shell. Variabel ini menyimpan nilai seperti ID project Google Cloud , nama kumpulan node, dan lokasi cluster GKE Anda.

Setelah menentukan variabel ini, Anda dapat menggunakannya kembali di beberapa perintah dengan mereferensikan nama variabel (misalnya, $CLUSTER_NAME) daripada mengetik ulang atau mengganti nilai setiap kali. Pendekatan ini mempermudah proses untuk diikuti dan mengurangi risiko error.

Untuk menentukan variabel lingkungan berguna berikut di Cloud Shell, jalankan perintah berikut:

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

Berikut penjelasan variabel lingkungan ini:

  • PROJECT_ID: ID Google Cloud project Anda.
  • PROJECT_NUMBER: Nomor ID unik untuk project Anda (misalnya, 123456789012).
  • ZONE: Zona komputasi untuk cluster Anda (misalnya, us-east5-b). Pilih zona dengan ketersediaan untuk jenis akselerator yang Anda pilih. Lihat kuota Cloud TPU atau kuota GPU untuk mengetahui info ketersediaan.
  • REGION: Region tempat Anda membuat resource cluster (misalnya, us-east5).
  • CLUSTER_NAME: Nama yang Anda pilih untuk cluster GKE Anda.
  • GKE_VERSION: Versi GKE cluster Anda. Gunakan versi 1.34 atau yang lebih baru.
  • ONDEMAND_NODEPOOL: Nama node pool on-demand standar Anda. (Ini adalah node pool Plan A Anda).
  • DWS_NODEPOOL: Nama node pool flex-start DWS Anda. (Ini adalah node pool Plan B Anda).

Mengonfigurasi infrastruktur (Administrator Cluster)

Sebagai Administrator Cluster, Anda mengonfigurasi cluster GKE dan kumpulan node untuk mendukung mekanisme penggantian.

Membuat cluster GKE

Pertama, buat cluster GKE. Cluster ini adalah lingkungan tempat Anda menginstal pengontrol Kueue, mengonfigurasi node pool, dan menjalankan tugas pelatihan AI. Untuk membuat dan terhubung ke cluster, lakukan langkah-langkah berikut:

  1. Membuat cluster:

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    Perintah ini menggunakan flag utama berikut:

    • --addons=RayOperator: Menginstal Ray Operator di cluster Anda. Anda memerlukan operator ini untuk mengelola workload RayJob yang Anda kirimkan nanti dalam dokumen ini.
    • --enable-image-streaming: Memungkinkan cluster Anda menarik image container lebih cepat. Fitur ini secara signifikan mengurangi waktu yang diperlukan agar image container AI berukuran besar mulai berjalan.

  2. Ambil kredensial cluster agar kubectl CLI dapat terhubung ke cluster tersebut. Perintah ini memperbarui file konfigurasi Kubernetes Anda, yang disimpan secara default di direktori ~/.kube/config:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

Buat node pool

Buat node pool utama dan cadangan untuk lingkungan Anda: node pool on-demand (Plan A) dan node pool flex-start DWS (Plan B):

  1. Buat node pool sesuai permintaan: Kumpulan ini berfungsi sebagai resource utama untuk tugas pelatihan:

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Dalam contoh ini, mesin pilihan pertama Anda adalah akselerator TPU v6e. Anda menentukan hardware ini menggunakan flag --machine-type=ct6e-standard-4t. Anda dapat mengubah jenis mesin ini agar sesuai dengan hardware, seperti GPU atau TPU yang berbeda, yang Anda inginkan untuk model AI Anda.

  2. Buat node pool flex-start DWS: Dalam contoh ini, Anda memilih jenis mesin yang sama (--machine-type=ct6e-standard-4t) dengan yang Anda pilih untuk pool on-demand utama. Node pool Plan B tidak harus menggunakan jenis mesin yang berbeda. Karena Anda benar-benar menginginkan hardware tertentu ini, menjadikannya pilihan Rencana B hanya berarti Anda beralih ke metode lain untuk mendapatkannya jika tidak tersedia secara langsung. Metode alternatif ini menggunakan DWS, yang terus-menerus mencari hardware yang tersedia hingga 7 hari:

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Perintah ini menggunakan flag kunci berikut:

    • --num-nodes=0, --min-nodes=0, --max-nodes=4, dan --enable-autoscaling: Kombinasi ini memungkinkan node pool di-scaling naik dari nol node saat tugas membutuhkannya dan di-scaling kembali ke bawah saat tidak ada aktivitas, yang membantu menghemat biaya.
    • --tpu-topology: Menentukan pengaturan fisik chip TPU. Anda menentukan tata letak ini karena pengaturan fisik chip memengaruhi seberapa cepat tugas pelatihan terdistribusi Anda berjalan.
    • --reservation-affinity=none: Membantu memastikan bahwa node pool tidak menggunakan hardware yang telah Anda pesan sebelumnya. Google Cloud Memungkinkan Anda memesan mesin tertentu untuk membantu memastikan ketersediaan. Menyetel tanda ini ke none akan memberi tahu sistem untuk melewati reservasi tersebut dan meminta mesin yang tidak direservasi secara dinamis.
    • --enable-queued-provisioning dan --flex-start: (Khusus pool Plan B) Flag ini memungkinkan DWS menyediakan node untuk pool Plan B Anda dari kapasitas fleksibel saat tersedia.

Memverifikasi status mulai fleksibel di node pool

Periksa node pool flex-start DWS dan pastikan flex-start diaktifkan:

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

Jika flex-start diaktifkan, outputnya adalah True.

Menginstal dan mengonfigurasi Kueue (Administrator Cluster)

Di bagian ini, Anda akan menginstal pengontrol Kueue di cluster Anda. Ingat bahwa Kueue adalah program penjadwalan yang mengelola antrean tugas. Layanan ini mencegat permintaan tugas Anda, memutuskan node pool mana yang akan digunakan (mulai fleksibel DWS atau sesuai permintaan), lalu menetapkan tugas.

Menginstal Kueue

Jalankan perintah berikut untuk menginstal Kueue. Perintah ini akan mendownload manifes penginstalan dari repositori resmi dan menerapkannya ke cluster Anda:

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

Menentukan aturan konfigurasi

Buat manifes YAML yang menentukan aturan prioritas. Aturan ini memberi tahu Kueue untuk menggunakan kumpulan on-demand terlebih dahulu dan kumpulan flex-start DWS kedua:

  1. Buat file bernama dws-tpu-queue.yaml dengan konten berikut. File ini menentukan dua ragam resource (mulai fleksibel DWS dan sesuai permintaan) serta antrean cluster yang memprioritaskannya. File konfigurasi ini menentukan logika yang digunakan Kueue untuk menangani tugas Anda:

    • ResourceFlavor: Di awal dokumen ini, Anda membuat dua node pool dan menetapkan nama menggunakan variabel lingkungan ${ONDEMAND_NODEPOOL} dan ${DWS_NODEPOOL}. Saat membuat node pool tersebut, GKE secara otomatis memberi label pada setiap node di node pool tersebut dengan nama yang Anda pilih untuk variabel lingkungan tersebut. Bagian ResourceFlavor memberi tahu Kueue untuk mencari node dengan label tersebut.
    • ClusterQueue: Bagian manifes ini menentukan aturan prioritas. Bagian ini mencantumkan jenis instance sesuai permintaan terlebih dahulu, sehingga Kueue mencoba menyediakan mesin sesuai permintaan terlebih dahulu. Jika Kueue tidak dapat memperoleh mesin tersebut, Kueue akan mencoba menyediakan mesin flex-start DWS.
    • Quotas: File menetapkan kuota, yang merupakan batas total resource (seperti CPU, memori, dan chip TPU) yang dapat digunakan tugas Anda kapan saja di node pool sesuai permintaan. Saat tugas Anda mencapai batas ini, Kueue akan otomatis mencoba menyediakan mesin flex-start DWS (mesin Rencana B Anda), yang Anda konfigurasi di dws-tpu-queue.yaml dengan batas kuota yang jauh lebih tinggi. 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. Terapkan konfigurasi ke cluster Anda. Perintah berikut menggunakan alat command line yang disebut envsubst untuk mengganti variabel placeholder yang muncul dalam file dws-tpu-queue.yaml. envsubst menggantikan placeholder dengan nilai variabel lingkungan yang Anda tentukan sebelumnya:

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

Menjalankan tugas pelatihan (Developer AI)

Sebagai Developer AI, Anda menentukan dan mengirimkan workload pelatihan dengan membuat manifes RayJob. Anda menentukan persyaratan resource dalam manifes ini, dan sistem penggantian otomatis, yang dikonfigurasi oleh administrator Cluster sebelumnya dengan Kueue dan DWS, akan menangani node pool yang mendasarinya untuk Anda.

Di bagian ini, Anda akan melakukan langkah-langkah berikut:

  • Buat skrip pelatihan Python.
  • Simpan skrip tersebut di ConfigMap Kubernetes.
  • Deploy RayJob yang memasang ConfigMap sebagai volume sehingga skrip pelatihan dapat dieksekusi di node.

Setelah melakukan langkah-langkah ini, Ray Train akan otomatis mendistribusikan beban kerja JAX di seluruh node, dan Kueue akan menangani perolehan mesin yang Anda butuhkan.

Skrip pelatihan

Salin dan tempel skrip Python berikut ke dalam file bernama train.py:

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

Skrip pelatihan menggunakan JAX, library Python untuk komputasi numerik berperforma tinggi, guna melatih model regresi linier. Skrip ini adalah contoh sederhana yang dirancang untuk menunjukkan cara menggunakan DWS dan Kueue untuk penggantian otomatis, dan tidak melakukan paralelisme data atau paralelisme model.

Perhatikan bahwa bagian ScalingConfig dari skrip pelatihan menentukan persyaratan hardware untuk tugas pelatihan. Bagian tersebut meminta topologi TPU 4x4, yang cocok dengan tata letak fisik node pool yang Anda konfigurasi sebelumnya.

Buat ConfigMap

Upload konten skrip train.py Anda ke dalam objek ConfigMap Kubernetes. Hal ini memungkinkan cluster menyimpan skrip dan membuatnya tersedia untuk RayJob Anda:

kubectl create configmap jax-train-script --from-file=train.py

RayJob yang Anda tentukan di bagian berikutnya akan memasang ConfigMap ini sebagai volume. Hal ini membuat file skrip muncul di dalam container Ray sehingga software Ray dapat menemukan dan mengeksekusinya.

Terapkan manifes RayJob

Buat file bernama rayjob-tpu-v6e-dws.yaml dengan konten berikut. Manifes ini menentukan tugas pelatihan Anda dan memberi tahu sistem cara merutekannya:

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

Manifes ini mencakup tiga konfigurasi yang membuat sistem penggantian berfungsi:

  • Meminta hardware tertentu: Bagian nodeSelector menentukan hardware yang diperlukan skrip Anda (dalam contoh ini, tpu-v6e-slice dengan topologi 4x4).
  • Memilih antrean: Label kueue.x-k8s.io/queue-name mengarahkan tugas Anda langsung ke Kueue. Tindakan ini mengaktifkan logika penggantian otomatis.
  • Mentoleransi node flex-start DWS: Bagian tolerations memungkinkan tugas berjalan di node pool Plan B Anda. Karena node flex-start DWS ditandai (diberi taint) secara khusus oleh GKE agar workload normal tidak berjalan secara tidak sengaja di node tersebut, tugas Anda harus secara eksplisit mentoleransi taint cloud.google.com/gke-queued.

Kirimkan workload

Untuk membuktikan bahwa sistem penggantian berfungsi, Anda harus mengirimkan dua tugas. Job pertama menggunakan kapasitas on-demand Paket A, yang memaksa job kedua melakukan failover ke kapasitas fleksibel mulai DWS Paket B.

Jalankan perintah berikut untuk mengirimkan dua tugas. Perintah ini menggunakan loop for dan envsubst untuk menyuntikkan ID tugas unik ke dalam manifes untuk setiap proses:

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

Setelah Anda mengirimkan tugas, sistem akan menangani workload sebagai berikut:

  1. Penyadapan: Kueue mendeteksi tugas menggunakan label antrean dan menangguhkannya untuk sementara.
  2. Keputusan: Kueue mengevaluasi ketersediaan resource berdasarkan aturan administrator. Sistem akan memeriksa kumpulan Plan A terlebih dahulu.
  3. Tugas:
    • Karena resource Plan A tersedia untuk tugas pertama, Kueue menetapkan Tugas 1 di sana.
    • Karena Tugas 1 menggunakan resource Paket A, Kueue secara otomatis menetapkan Tugas 2 ke kumpulan Paket B (DWS flex-start).
  4. Peluncuran: Kueue melanjutkan tugas. Tindakan ini memicu autoscaler cluster GKE untuk menyediakan node dan memulai skrip pelatihan.

Menghubungkan ke RayJob

Sebagai langkah verifikasi terakhir, Anda dapat menggunakan perintah kubectl port-forward untuk terhubung ke Dasbor Ray dan melihat tugas Anda berjalan.

Untuk memeriksa status tugas pertama, jalankan perintah berikut:

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

Setelah menjalankan perintah ini, buka browser web dan buka http://localhost:8265. Di Dasbor Ray, Anda dapat melihat status tugas dan metrik yang dilaporkan untuk memverifikasi bahwa kedua tugas berhasil diselesaikan di kumpulan node masing-masing.

Anda juga dapat melihat log untuk tugas pertama dengan menjalankan perintah berikut:

kubectl logs job/rayjob-tpu-v6e-dws-1

Output skrip pelatihan yang dipangkas akan terlihat seperti berikut. Anda akan melihat pesan Training Complete! dan Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded di dekat akhir output:

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

Pembersihan

Agar akun Google Cloud Anda tidak dikenai biaya untuk resource yang digunakan dalam dokumen ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

Menghapus project

  • Di Konsol Google Cloud , buka halaman Manage resources.

    Buka Kelola 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

    Jika Anda ingin menyimpan project GGoogle Cloud yang Anda gunakan dalam dokumen ini, jalankan perintah berikut untuk menghapus cluster:

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    Ringkasan

    Dalam dokumen ini, Anda telah mengonfigurasi dan menguji lingkungan pelatihan Ray. Lingkungan ini menggunakan kumpulan node utama dan kumpulan DWS cadangan untuk memaksimalkan ketersediaan hardware. Dengan otomatis melakukan failover ke DWS saat mesin utama tidak tersedia, Anda meminimalkan waktu yang dihabiskan tugas pelatihan Anda untuk menunggu dalam antrean.

    Agar ini berfungsi, Anda telah melakukan langkah-langkah berikut:

    1. Membuat cluster GKE: Menyiapkan lingkungan untuk menghosting node pool dan alat penjadwalan.
    2. Mengonfigurasi node pool: Membuat node pool sesuai permintaan (Paket A) dan node pool DWS (Paket B).
    3. Menginstal dan mengonfigurasi Kueue: Men-deploy pengontrol Kueue dan menerapkan aturan prioritas yang menginstruksikan sistem untuk mencoba Rencana A terlebih dahulu dan melakukan failover ke Rencana B.
    4. Membuat ConfigMap: Men-deploy skrip pelatihan JAX yang disederhanakan ke cluster untuk berfungsi sebagai beban kerja pengujian.
    5. Menentukan manifes RayJob: Mengonfigurasi tugas untuk meminta hardware tertentu, merutekan ke pengontrol Kueue, dan mentoleransi node DWS.
    6. Mengirimkan workload: Mengirimkan dua tugas untuk memaksa Kueue merutekan tugas kedua secara otomatis ke Rencana B saat resource Rencana A digunakan.
    7. Memverifikasi hasil: Menggunakan penerusan port untuk terhubung ke Dasbor Ray dan mengonfirmasi bahwa kedua tugas berhasil dijalankan.

    Langkah berikutnya