Mengoptimalkan prioritas workload AI/ML GKE

Dokumen ini menjelaskan alat dan praktik terbaik untuk memaksimalkan pemanfaatan resource dan meminimalkan waktu henti workload AI/ML heterogen di Google Kubernetes Engine (GKE), terutama saat tidak ada kapasitas dalam reservasi atau melalui resource on-demand. Workload heterogen mengacu pada berbagai jenis workload AI/ML yang berjalan secara bersamaan dalam cluster GKE yang sama. Misalnya, Anda dapat menjalankan layanan inferensi online yang sensitif terhadap latensi bersama serangkaian tugas pelatihan batch yang dapat diinterupsi.

Panduan ini memberikan rekomendasi untuk admin dan operator Platform serta spesialis Data dan AI.

Manfaat penentuan prioritas workload AI/ML

Workload heterogen memiliki prioritas yang berbeda dan berbagi kapasitas serta resource yang terbatas. Praktik terbaik di halaman ini menjelaskan cara mengonfigurasi GKE dan alat open source untuk membantu Anda mendapatkan manfaat berikut:

  • Meminimalkan waktu non-operasional untuk workload prioritas tinggi.
  • Menjalankan beban kerja prioritas tinggi dengan cepat.
  • Mengoptimalkan konsumsi resource.

Latar belakang

GKE mendukung alat open source berikut untuk mengoptimalkan pemanfaatan resource.

  • Kueue: sistem antrean workload berbasis Kubernetes yang dirancang untuk workload batch, AI, dan komputasi berperforma tinggi. Kueue dapat diperluas untuk mengelola jenis workload lain, seperti yang ditentukan oleh Definisi Resource Kustom seperti leaderworkerset. Kueue mengelola kuota dan cara workload menggunakannya di cluster Kubernetes. Kueue membuat keputusan tentang kapan workload menunggu, kapan workload dimulai (misalnya, dengan membuat Pod), dan kapan Pod milik workload di-preempt.

    Untuk mengetahui informasi selengkapnya tentang Kueue, lihat dokumentasi Konsep Kueue.

  • Hot swap: teknik yang mengurangi rata-rata waktu pemulihan (MTTR). Hotswap memungkinkan penghentian sementara berdasarkan prioritas workload saat resource cluster digunakan sepenuhnya dan tidak ada kapasitas tambahan yang tersedia, baik dari instance sesuai permintaan maupun reservasi yang ada.

    • Jika node yang menghosting workload menjadi tidak sehat, workload akan dijadwalkan ulang pada node cadangan yang memenuhi syarat. Jika tidak ada node pengganti yang tersedia, Hotswap dapat menghentikan sementara workload dengan prioritas lebih rendah untuk memberi ruang bagi workload yang sedang dipulihkan.
    • Jika Anda mengonfigurasi Pod dengan PriorityClass, workload yang dikonfigurasi dengan prioritas lebih tinggi akan mengeluarkan workload prioritas rendah yang sedang berjalan untuk mendapatkan sumber dayanya. Proses pengusiran ini dikenal sebagai pengambilalihan.

Kasus penggunaan

Gunakan tabel berikut untuk memahami praktik terbaik untuk setiap kasus penggunaan:

Kasus penggunaan Praktik terbaik Deskripsi
Beberapa workload dengan prioritas yang berbeda Gunakan Kueue untuk menentukan antrean dan menetapkan prioritas ke workload berdasarkan kepentingannya. Kueue dapat mengelola kuota sehingga tim atau project tertentu memiliki akses ke sejumlah resource.

Kueue memungkinkan Anda menerapkan konfigurasi berikut:

  • Prioritaskan Pekerjaan dengan prioritas tinggi dengan menetapkan WorkloadPriority Kueue yang lebih tinggi untuk Pekerjaan tersebut.
  • Aktifkan antrean berbagi yang adil Kueue sehingga semua beban kerja akhirnya menerima resource, bahkan yang berprioritas rendah.

Untuk menguji konfigurasi praktik terbaik, lihat contoh Kueue dalam dokumen ini.

Anda harus mengurangi MTTR saat ini. Gunakan Hotswap untuk menjadwalkan ulang workload di resource yang responsif saat terjadi gangguan, dan menggantikan workload berprioritas rendah dengan workload berprioritas tinggi.

Hotswap memungkinkan Anda menerapkan konfigurasi berikut:

  • Konfigurasi PriorityClasses untuk menentukan tingkat prioritas bagi workload Anda.
  • Tetapkan PriorityClasses yang lebih tinggi untuk workload penting.
  • Secara otomatis menjadwalkan ulang workload di node yang berfungsi dengan baik saat terjadi gangguan.

Untuk menguji konfigurasi praktik terbaik, lihat Contoh penggantian cepat dalam dokumen ini.

Beberapa workload AI bersaing untuk mendapatkan sumber daya terbatas Gabungkan Kueue dan Hotswap. Kombinasi ini menyediakan sistem yang andal yang memprioritaskan workload penting baik selama penjadwalan awal maupun selama runtime.

Kueue dan Hotswap memungkinkan Anda menerapkan konfigurasi berikut:

  • Gunakan Kueue untuk mengelola penjadwalan dan penerimaan awal workload berdasarkan prioritas.
  • Gunakan Hotswap untuk menangani gangguan beban kerja dan memungkinkan pemulihan cepat. Hot swap membantu mengurangi waktu pemulihan workload prioritas tinggi saat terjadi gangguan.

Untuk menguji konfigurasi praktik terbaik, lihat Contoh Kueue dan Hotswap dalam dokumen ini.

Contoh penerapan praktik terbaik

Contoh berikut menunjukkan cara menerapkan Kueue dan Hotswap, serta cara menggabungkannya untuk praktik terbaik yang dijelaskan di bagian sebelumnya.

Kueue

Contoh manifes berikut menunjukkan konfigurasi Kueue:

  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq

Manifes ini melakukan hal berikut:

  • Menentukan ResourceFlavor bernama tpu-v6e-slice yang menentukan label node untuk slice TPU v6e.
  • Menentukan ClusterQueue bernama tpu-training-cq yang mengelola kuota untuk resource TPU.
  • Menentukan LocalQueue bernama default-queue yang memungkinkan workload di namespace default menggunakan antrean cluster tpu-training-cq.

Hotswap

Contoh berikut menunjukkan konfigurasi Hotswap yang menentukan dua Class Prioritas, low-priority-job dan high-priority-job. Konfigurasi Hotswap ini membuat workload JobSet prioritas tinggi dan menggunakan MaxText.

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                -   python3
                -   MaxText/train.py
                -   MaxText/configs/base.yml
                -   model_name=llama2-7b
                -   run_name=<UNIQUE RUN NAME>
                -   steps=300
                -   base_output_directory=gs://<OUTPUT BUCKET>
                -   dataset_path=gs://max-datasets-rogue
                -   max_target_length=4096
                -   dataset_type=synthetic
                -   enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

Berdasarkan konfigurasi ini, Hotswap akan melakukan tindakan berikut:

  • Jika kegagalan infrastruktur mengganggu beban kerja berprioritas tinggi, JobSet akan memulainya kembali. Hot-swap mendahului workload berprioritas rendah untuk menjadwalkan ulang workload berprioritas tinggi sebelum infrastruktur dipulihkan. Beban kerja prioritas rendah tetap dalam status gagal. Proses ini mengurangi waktu tunggu workload secara signifikan.
  • Saat infrastruktur pulih, Hotswap akan menjadwalkan ulang beban kerja berprioritas rendah di node pool yang pulih.

Kueue dan Hotswap

Gabungkan Kueue dan Hotswap saat Anda beroperasi di lingkungan yang kompleks dengan sumber daya terbatas. Kombinasi ini menyediakan sistem yang andal yang memprioritaskan beban kerja penting selama penjadwalan awal dan selama runtime.

Contoh berikut menunjukkan konfigurasi gabungan Kueue dan Hotswap. Contoh ini menggunakan MaxText:

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: low-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: low-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: low-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=low-priority-run
                - steps=30000
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: high-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=high-priority-run
                - steps=300
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

Berdasarkan konfigurasi ini, Kueue digabungkan dengan Hotswap, dan melakukan tindakan berikut:

  • Kueue mengelola penerimaan JobSet low-jax-trillium dan high-jax-trillium ke dalam antrean cluster berdasarkan prioritas yang ditentukan dan resource yang tersedia.
  • Jika JobSet high-jax-trillium terganggu oleh kegagalan infrastruktur, Hotswap akan mendahului JobSet low-jax-trillium untuk menjadwalkan ulang JobSet berprioritas tinggi.
  • Hotswap memastikan JobSet prioritas tinggi dimulai ulang dengan cepat, sehingga meminimalkan waktu tidak aktifnya.
  • Saat infrastruktur pulih, Hotswap akan menjadwalkan ulang JobSet prioritas rendah di node pool yang dipulihkan.

Langkah berikutnya