Priorisierung von KI-/ML-Arbeitslasten in GKE optimieren

In diesem Dokument werden Tools und Best Practices beschrieben, mit denen Sie die Ressourcennutzung maximieren und die Ausfallzeiten heterogener KI-/ML-Arbeitslasten in Google Kubernetes Engine (GKE) minimieren können, insbesondere wenn keine Kapazität in Reservierungen oder über On-Demand-Ressourcen verfügbar ist. Heterogene Arbeitslasten beziehen sich auf verschiedene Arten von KI-/ML-Arbeitslasten, die gleichzeitig im selben GKE-Cluster ausgeführt werden. Sie können beispielsweise einen latenzsensitiven Onlinedienst für die Inferenz zusammen mit einer Reihe von unterbrechbaren Batchtrainingsjobs ausführen.

Dieser Leitfaden enthält Empfehlungen für Plattformadministratoren und ‑operatoren sowie für Daten- und KI-Spezialisten.

Vorteile der Priorisierung von KI-/ML-Arbeitslasten

Heterogene Arbeitslasten haben unterschiedliche Prioritäten und teilen sich begrenzte Kapazitäten und Ressourcen. Die Best Practices auf dieser Seite beschreiben, wie Sie GKE und Open-Source-Tools konfigurieren, um die folgenden Vorteile zu erzielen:

  • Ausfallzeiten für Arbeitslasten mit hoher Priorität minimieren.
  • Schnelle Ausführung von Arbeitslasten mit hoher Priorität
  • Ressourcenverbrauch optimieren

Hintergrund

GKE unterstützt die folgenden Open-Source-Tools zur Optimierung der Ressourcennutzung.

  • Kueue:Ein Kubernetes-natives System zur Warteschlangenverwaltung von Arbeitslasten, das für Batch-, KI- und Hochleistungs-Computing-Arbeitslasten entwickelt wurde. Kueue kann erweitert werden, um andere Arbeitslasttypen zu verwalten, z. B. solche, die durch benutzerdefinierte Ressourcendefinitionen wie leaderworkerset definiert werden. Kueue verwaltet Kontingente und deren Nutzung durch Arbeitslasten in einem Kubernetes-Cluster. Kueue entscheidet, wann eine Arbeitslast wartet, wann eine Arbeitslast gestartet wird (z. B. durch Erstellen des Pods) und wann ein Pod, der zu einer Arbeitslast gehört, vorzeitig beendet wird.

    Weitere Informationen zu Kueue finden Sie in der Dokumentation zu Kueue-Konzepten.

  • Hotswap:Eine Technik, mit der die durchschnittliche Wiederherstellungszeit (Mean Time To Recovery, MTTR) verkürzt wird. Hotswap ermöglicht die Präemption basierend auf der Arbeitslastpriorität, wenn die Clusterressourcen vollständig ausgelastet sind und keine zusätzliche Kapazität verfügbar ist, weder durch On-Demand-Instanzen noch durch vorhandene Reservierungen.

    • Wenn ein Knoten, auf dem eine Arbeitslast ausgeführt wird, fehlerhaft wird, wird die Arbeitslast auf geeigneten Ersatzknoten neu geplant. Wenn keine Ersatzknoten verfügbar sind, kann Hotswap eine Arbeitslast mit niedrigerer Priorität unterbrechen, um Platz für die wiederherzustellende Arbeitslast zu schaffen.
    • Wenn Sie Ihre Pods mit PriorityClass konfigurieren, wird eine Arbeitslast mit höherer Priorität, die mit PriorityClass konfiguriert ist, eine laufende Arbeitslast mit niedriger Priorität entfernen, um deren Ressourcen zu erhalten. Dieser Prozess wird als Preemption bezeichnet.

Anwendungsfälle

In der folgenden Tabelle finden Sie die Best Practices für die einzelnen Anwendungsfälle:

Anwendungsfall Best Practice Beschreibung
Mehrere Arbeitslasten mit unterschiedlichen Prioritäten Mit Kueue können Sie Warteschlangen definieren und Arbeitslasten basierend auf ihrer Wichtigkeit Prioritäten zuweisen. Kueue kann Kontingente verwalten, sodass bestimmte Teams oder Projekte Zugriff auf eine bestimmte Menge an Ressourcen haben.

Mit Kueue können Sie die folgenden Konfigurationen anwenden:

  • Weisen Sie Jobs mit hoher Priorität eine höhere Kueue-WorkloadPriority zu, um ihnen Vorrang einzuräumen.
  • Aktivieren Sie die Fair-Share-Warteschlange von Kueue, damit alle Arbeitslasten, auch solche mit niedriger Priorität, Ressourcen erhalten.

Informationen zum Testen der Best Practice-Konfiguration finden Sie im Kueue-Beispiel in diesem Dokument.

Sie müssen die aktuelle MTTR reduzieren. Mit Hotswap können Sie Arbeitslasten auf fehlerfreie Ressourcen umplanen, wenn eine Unterbrechung auftritt, und Arbeitslasten mit niedriger Priorität zugunsten von Arbeitslasten mit hoher Priorität unterbrechen.

Mit Hotswap können Sie die folgenden Konfigurationen anwenden:

  • Konfigurieren Sie PriorityClasses, um Prioritätsstufen für Ihre Arbeitslasten zu definieren.
  • Weisen Sie kritischen Arbeitslasten einen höheren PriorityClasses zu.
  • Arbeitslasten bei Unterbrechungen automatisch auf fehlerfreien Knoten neu planen.

Informationen zum Testen der Best Practice-Konfiguration finden Sie im Hotswap-Beispiel in diesem Dokument.

Mehrere KI-Arbeitslasten konkurrieren um begrenzte Ressourcen. Kueue und Hotswap kombinieren Diese Kombination bietet ein robustes System, das kritische Arbeitslasten sowohl bei der ersten Planung als auch während der Laufzeit priorisiert.

Mit Kueue und Hotswap können Sie die folgenden Konfigurationen anwenden:

  • Mit Kueue können Sie die anfängliche Planung und Zulassung von Arbeitslasten basierend auf der Priorität verwalten.
  • Mit Hotswap können Sie Unterbrechungen von Arbeitslasten bewältigen und eine schnelle Wiederherstellung ermöglichen. Hotswap trägt dazu bei, die Wiederherstellungszeit einer Arbeitslast mit hoher Priorität nach einer Unterbrechung zu verkürzen.

Ein Beispiel zum Testen der Best Practice-Konfiguration finden Sie in diesem Dokument unter Beispiel für Kueue und Hotswap.

Beispiele für Best-Practice-Implementierungen

Die folgenden Beispiele zeigen, wie Sie Kueue und Hotswap implementieren und wie Sie sie für die im vorherigen Abschnitt beschriebenen Best Practices kombinieren.

Kueue

Das folgende Beispielmanifest zeigt eine Kueue-Konfiguration:

  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

Das Manifest tut Folgendes:

  • Definiert eine ResourceFlavor mit dem Namen tpu-v6e-slice, die die Knotenlabels für TPU v6e-Slices angibt.
  • Definiert eine ClusterQueue mit dem Namen tpu-training-cq, die das Kontingent für TPU-Ressourcen verwaltet.
  • Definiert eine LocalQueue mit dem Namen default-queue, mit der Arbeitslasten im Namespace default die Clusterwarteschlange tpu-training-cq verwenden können.

Hotswap

Das folgende Beispiel zeigt eine Hotswap-Konfiguration, in der zwei Prioritätsklassen definiert sind: low-priority-job und high-priority-job. Bei dieser Hotswap-Konfiguration wird eine JobSet-Arbeitslast mit hoher Priorität erstellt und MaxText verwendet.

  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

Auf Grundlage dieser Konfiguration führt Hotswap die folgenden Aktionen aus:

  • Wenn die Arbeitslast mit hoher Priorität durch einen Infrastrukturfehler unterbrochen wird, wird sie vom JobSet neu gestartet. Hotswap beendet die Arbeitslast mit niedriger Priorität vorzeitig, um die Arbeitslast mit hoher Priorität neu zu planen, bevor die Infrastruktur wiederhergestellt wird. Die Arbeitslast mit niedriger Priorität bleibt im Status „Fehler“. Dadurch wird die Leerlaufzeit von Arbeitslasten erheblich reduziert.
  • Wenn die Infrastruktur wiederhergestellt ist, plant Hotswap die Arbeitslast mit niedriger Priorität im wiederhergestellten Knotenpool neu.

Kueue und Hotswap

Kombinieren Sie Kueue und Hotswap, wenn Sie in einer komplexen Umgebung mit begrenzten Ressourcen arbeiten. Diese Kombination bietet ein robustes System, das kritische Arbeitslasten bei der ersten Planung und während der Laufzeit priorisiert.

Das folgende Beispiel zeigt eine kombinierte Kueue- und Hotswap-Konfiguration. In diesem Beispiel wird MaxText verwendet:

  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

Auf Grundlage dieser Konfiguration wird Kueue mit Hotswap kombiniert und führt die folgenden Aktionen aus:

  • Kueue verwaltet die Aufnahme von low-jax-trillium- und high-jax-trillium-JobSets in die Clusterwarteschlange basierend auf den definierten Prioritäten und verfügbaren Ressourcen.
  • Wenn der high-jax-trillium-JobSet durch einen Infrastrukturfehler unterbrochen wird, wird der low-jax-trillium-JobSet durch Hotswap vorzeitig beendet, um den JobSet mit hoher Priorität neu zu planen.
  • Durch Hotswap wird sichergestellt, dass das JobSet mit hoher Priorität schnell neu gestartet wird, wodurch die Leerlaufzeit minimiert wird.
  • Wenn die Infrastruktur wiederhergestellt ist, plant Hotswap das JobSet mit niedriger Priorität im wiederhergestellten Knotenpool neu.

Nächste Schritte