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 in Reservierungen oder über On-Demand-Ressourcen keine Kapazität verfügbar ist. Heterogene Arbeitslasten sind verschiedene Arten von KI/ML-Arbeitslasten, die gleichzeitig im selben GKE-Cluster ausgeführt werden. Beispielsweise können Sie einen latenzempfindlichen Online-Inferenzdienst zusammen mit einer Reihe von unterbrechbaren Batch-Trainingsjobs ausführen.

Dieser Leitfaden enthält Empfehlungen für Plattformadministratoren und -operatoren sowie für Daten- und KI-Experten. Eine zusammenfassende Ü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. In den Best Practices auf dieser Seite wird beschrieben, wie Sie GKE und Open-Source-Tools konfigurieren, um folgende Vorteile zu erzielen:

  • Ausfallzeiten für Arbeitslasten mit hoher Priorität minimieren
  • Arbeitslasten mit hoher Priorität schnell ausführen
  • 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 die Nutzung von Kontingenten durch Arbeitslasten in einem Kubernetes-Cluster. Kueue entscheidet, wann eine Arbeitslast warten muss, 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 Kueue Konzepten Dokumentation.

  • Hotswap:Eine Technik, mit der die durchschnittliche Wiederherstellungszeit (Mean Time To Recovery, MTTR) verkürzt wird. Hotswap ermöglicht das vorzeitige Beenden basierend auf der Arbeitslastpriorität, 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, beendet die mit höherer Priorität konfigurierte Arbeitslast eine laufende Arbeitslast mit niedriger Priorität, um ihre Ressourcen zu nutzen. Dieser Vorgang wird als vorzeitiges Beenden 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 Verwenden Sie Kueue, um Warteschlangen zu definieren und Arbeitslasten basierend auf ihrer Wichtigkeit Prioritäten zuzuweisen. 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-Warteschlangenverwaltung von Kueue, damit alle Arbeitslasten Ressourcen erhalten, auch solche mit niedriger Priorität.

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

Sie müssen die aktuelle MTTR reduzieren. Verwenden Sie Hotswap, um Arbeitslasten bei einer Unterbrechung auf fehlerfreien Ressourcen neu zu planen und Arbeitslasten mit niedriger Priorität zugunsten von Arbeitslasten mit hoher Priorität vorzeitig zu beenden.

Mit Hotswap können Sie die folgenden Konfigurationen anwenden:

  • Konfigurieren Sie PriorityClasses zur Definition von Prioritätsstufen für Ihre Arbeitslasten.
  • Weisen Sie kritischen Arbeitslasten höhere PriorityClasses zu.
  • Planen Sie Arbeitslasten bei Unterbrechungen automatisch auf fehlerfreien Knoten neu.

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

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:

  • Verwenden Sie Kueue, um die erste Planung und Zulassung von Arbeitslasten basierend auf der Priorität zu verwalten.
  • Verwenden Sie Hotswap, um Arbeitslast Unterbrechungen zu verarbeiten und eine schnelle Wiederherstellung zu ermöglichen. Hotswap trägt dazu bei, die Wiederherstellungszeit einer Arbeitslast mit hoher Priorität bei einer Unterbrechung zu verkürzen.

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

Beispiele für Best-Practice-Implementierungen

In den folgenden Beispielen wird gezeigt, 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 ein ResourceFlavor mit dem Namen tpu-v6e-slice, das 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 in dem default-Namespace die Clusterwarteschlange tpu-training-cq verwenden können.

Hotswap

Das folgende Beispiel zeigt eine Hotswap-Konfiguration, die zwei PriorityClasses definiert: low-priority-job und high-priority-job. Diese Hotswap-Konfiguration erstellt eine JobSet-Arbeitslast mit hoher Priorität und verwendet 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

Basierend auf 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 beendet die Arbeitslast mit niedriger Priorität vorzeitig, um die Arbeitslast mit hoher Priorität neu zu planen, bevor die Infrastruktur wiederhergestellt ist. Die Arbeitslast mit niedriger Priorität bleibt im Status „Fehlgeschlagen“. Dieser Vorgang verkürzt die Leerlaufzeit der Arbeitslast erheblich.
  • 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 Zulassung der JobSets low-jax-trillium und high-jax-trillium in die Clusterwarteschlange basierend auf ihren definierten Prioritäten und den verfügbaren Ressourcen.
  • Wenn das JobSet high-jax-trillium durch einen Infrastrukturausfall unterbrochen wird, beendet Hotswap das JobSet low-jax-trillium vorzeitig, um das JobSet mit hoher Priorität neu zu planen.
  • Hotswap sorgt dafür, 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