Optimiser la priorisation des charges de travail d'IA/de ML GKE

Ce document décrit les outils et les bonnes pratiques permettant de maximiser l'utilisation des ressources et de minimiser les temps d'arrêt des charges de travail d'IA/ML hétérogènes dans Google Kubernetes Engine (GKE), en particulier lorsqu'il n'y a pas de capacité dans les réservations ou via les ressources à la demande. Les charges de travail hétérogènes font référence à différents types de charges de travail d'IA/de ML qui s'exécutent simultanément dans le même cluster GKE. Par exemple, vous pouvez exécuter un service d'inférence en ligne sensible à la latence en parallèle d'une série de jobs d'entraînement par lot interruptibles.

Ce guide fournit des recommandations pour les administrateurs et opérateurs de plate-forme, ainsi que pour les spécialistes des données et de l'IA.

Avantages de la priorisation des charges de travail d'IA/de ML

Les charges de travail hétérogènes ont des priorités différentes et partagent une capacité et des ressources limitées. Les bonnes pratiques de cette page décrivent comment configurer GKE et les outils Open Source pour vous aider à bénéficier des avantages suivants :

  • Minimisez les temps d'arrêt pour les charges de travail hautement prioritaires.
  • Exécutez rapidement les charges de travail à priorité élevée.
  • Optimisez la consommation des ressources.

Arrière-plan

GKE est compatible avec les outils Open Source suivants pour optimiser l'utilisation des ressources.

  • Kueue : système de mise en file d'attente des charges de travail natif à Kubernetes, conçu pour les charges de travail par lot, d'IA et de calcul hautes performances. Kueue peut être étendu pour gérer d'autres types de charges de travail, comme celles définies par des définitions de ressources personnalisées telles que leaderworkerset. Kueue gère les quotas et la manière dont les charges de travail les consomment dans un cluster Kubernetes. Kueue décide quand une charge de travail doit attendre, quand elle doit démarrer (par exemple, en créant le pod) et quand un pod appartenant à une charge de travail est préempté.

    Pour en savoir plus sur Kueue, consultez la documentation sur les concepts de Kueue.

  • Échange à chaud : technique qui réduit le temps moyen de récupération (MTTR). L'échange à chaud permet la préemption en fonction de la priorité de la charge de travail lorsque les ressources du cluster sont entièrement utilisées et qu'aucune capacité supplémentaire n'est disponible, que ce soit à partir d'instances à la demande ou de réservations existantes.

    • Lorsqu'un nœud hébergeant une charge de travail devient non opérationnel, la charge de travail est reprogrammée sur des nœuds de secours éligibles. Si aucun nœud de secours n'est disponible, Hotswap peut préempter une charge de travail de priorité inférieure pour faire de la place à la charge de travail en cours de récupération.
    • Si vous configurez vos pods avec PriorityClass, la charge de travail configurée avec une priorité plus élevée évince une charge de travail de faible priorité en cours d'exécution pour acquérir ses ressources. Ce processus d'éviction est appelé "préemption".

Cas d'utilisation

Utilisez le tableau suivant pour comprendre les bonnes pratiques pour chaque cas d'utilisation :

Cas d'utilisation Bonne pratique Description
Plusieurs charges de travail avec des priorités différentes Utilisez Kueue pour définir des files d'attente et attribuer des priorités aux charges de travail en fonction de leur importance. Kueue peut gérer les quotas afin que certaines équipes ou certains projets aient accès à une quantité définie de ressources.

Kueue vous permet d'appliquer les configurations suivantes :

  • Priorisez les jobs à priorité élevée en leur attribuant un WorkloadPriority Kueue plus élevé.
  • Activez la mise en file d'attente avec partage équitable de Kueue afin que toutes les charges de travail reçoivent des ressources, même celles de faible priorité.

Pour tester la configuration des bonnes pratiques, consultez l'exemple Kueue dans ce document.
Vous devez réduire le MTTR actuel. Utilisez Hotswap pour replanifier les charges de travail dans des ressources saines en cas d'interruption, et préemptez les charges de travail de faible priorité au profit de celles de priorité élevée.

L'échange à chaud vous permet d'appliquer les configurations suivantes :

  • Configurez des PriorityClasses pour définir des niveaux de priorité pour vos charges de travail.
  • Attribuez un PriorityClasses plus élevé aux charges de travail critiques.
  • Replanifiez automatiquement les charges de travail sur des nœuds opérationnels en cas d'interruption.

Pour tester la configuration recommandée, consultez l'exemple de remplacement à chaud dans ce document.
Plusieurs charges de travail d'IA se disputent des ressources limitées. Combiner Kueue et Hotswap Cette combinaison fournit un système robuste qui donne la priorité aux charges de travail critiques lors de la planification initiale et de l'exécution.

Kueue et Hotswap vous permettent d'appliquer les configurations suivantes :

  • Utilisez Kueue pour gérer la planification et l'admission initiales des charges de travail en fonction de leur priorité.
  • Utilisez Hotswap pour gérer les interruptions de charge de travail et permettre une récupération rapide. L'échange à chaud permet de réduire le temps de récupération d'une charge de travail à haute priorité en cas d'interruption.

Pour tester la configuration recommandée, consultez l'exemple Kueue et Hotswap dans ce document.

Exemples d'implémentations de bonnes pratiques

Les exemples suivants montrent comment implémenter Kueue et Hotswap, et comment les combiner pour appliquer les bonnes pratiques décrites dans la section précédente.

Kueue

L'exemple de fichier manifeste suivant montre une configuration 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

Ce fichier manifeste effectue les opérations suivantes :

  • Définit un ResourceFlavor nommé tpu-v6e-slice qui spécifie les libellés de nœud pour les tranches TPU v6e.
  • Définit un ClusterQueue nommé tpu-training-cq qui gère le quota des ressources TPU.
  • Définit un LocalQueue nommé default-queue qui permet aux charges de travail de l'espace de noms default d'utiliser la file d'attente de cluster tpu-training-cq.

Échange à chaud

L'exemple suivant montre une configuration Hotswap qui définit deux classes de priorité, low-priority-job et high-priority-job. Cette configuration Hotswap crée une charge de travail JobSet à haute priorité et utilise 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

Sur la base de cette configuration, Hotswap effectue les actions suivantes :

  • Si une défaillance de l'infrastructure interrompt la charge de travail haute priorité, JobSet la redémarre. L'échange à chaud préempte la charge de travail à faible priorité pour replanifier la charge de travail à priorité élevée avant la récupération de l'infrastructure. La charge de travail à faible priorité reste à l'état "Échec". Ce processus réduit considérablement le temps d'inactivité de la charge de travail.
  • Lorsque l'infrastructure est rétablie, Hotswap replanifie la charge de travail de faible priorité dans le pool de nœuds rétabli.

Kueue et Hotswap

Combinez Kueue et Hotswap lorsque vous travaillez dans un environnement complexe avec des ressources limitées. Cette combinaison fournit un système robuste qui donne la priorité aux charges de travail critiques lors de la planification initiale et pendant l'exécution.

L'exemple suivant montre une configuration combinée de Kueue et Hotswap. Cet exemple utilise 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

Avec cette configuration, Kueue est combiné à Hotswap et effectue les actions suivantes :

  • Kueue gère l'admission des JobSets low-jax-trillium et high-jax-trillium dans la file d'attente du cluster en fonction de leurs priorités définies et des ressources disponibles.
  • Si le JobSet high-jax-trillium est interrompu par une défaillance de l'infrastructure, Hotswap préempte le JobSet low-jax-trillium pour replanifier le JobSet de haute priorité.
  • L'échange à chaud garantit que le JobSet à priorité élevée redémarre rapidement, ce qui minimise son temps d'inactivité.
  • Lorsque l'infrastructure est rétablie, Hotswap reprogramme le JobSet de faible priorité dans le pool de nœuds rétabli.

Étapes suivantes