Best Practices für die Priorisierung von KI-/ML-Arbeitslasten in GKE

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 Online-Inferenzdienst zusammen mit einer Reihe von unterbrechbaren Batch-Trainingsjobs ausführen.

Dieser Leitfaden enthält Empfehlungen für Plattformadministratoren und -operatoren sowie Daten- und KI-Spezialisten. Eine konsolidierte Übersicht über alle GKE-Best Practices finden Sie unter Best Practices für GKE.

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 HPC-Arbeitslasten (High Performance Computing, Hochleistungs-Computing) entwickelt wurde. Kueue kann erweitert werden, um andere Arbeitslasttypen zu verwalten, z. B. solche, die durch benutzerdefinierte Ressourcendefinitionen wie leaderworkerset definiert sind. Kueue verwaltet Kontingente und die Nutzung dieser Kontingente 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, die die durchschnittliche Wiederherstellungszeit (Mean Time To Recovery, MTTR) verkürzt. Hotswap ermöglicht die Preemption basierend auf der Priorität der Arbeitslast, wenn die Clusterressourcen vollständig ausgelastet sind und keine zusätzliche Kapazität verfügbar ist, weder von On-Demand-Instanzen noch von vorhandenen Reservierungen.

    • Wenn ein Knoten, auf dem eine Arbeitslast gehostet 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 vorzeitig beenden, 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 konfiguriert ist, eine laufende Arbeitslast mit niedriger Priorität entfernen, um deren Ressourcen zu erhalten. Dieser Vorgang 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:

  • Priorisieren Sie Jobs mit hoher Priorität, indem Sie ihnen eine höhere Kueue-WorkloadPriority zuweisen.
  • 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 bei einer Unterbrechung auf fehlerfreie Ressourcen umplanen 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.

Ein Beispiel für die Best Practice-Konfiguration finden Sie in diesem Dokument unter Hotswap-Beispiel.

Mehrere KI-Arbeitslasten konkurrieren um begrenzte Ressourcen Kombinieren Sie Kueue und Hotswap. 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 zu verkürzen, wenn eine Unterbrechung auftritt.

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

Beispiele für Best Practices bei der Implementierung

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, die es Arbeitslasten im Namespace default ermöglicht, die Cluster-Warteschlange tpu-training-cq zu verwenden.

Hot-Swap

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 ein Infrastrukturausfall die Arbeitslast mit hoher Priorität unterbricht, wird sie vom JobSet neu gestartet. Hotswap unterbricht die Arbeitslast mit niedriger Priorität, um die Arbeitslast mit hoher Priorität neu zu planen, bevor die Infrastruktur wiederhergestellt wird. Die Arbeitslast mit niedriger Priorität bleibt im Status „Fehlgeschlagen“. Dadurch wird die Leerlaufzeit der Arbeitslast 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 sowohl bei der ersten Planung als auch 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

Basierend auf 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 ihren 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 unterbrochen, 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