Bonnes pratiques pour la priorisation des charges de travail d'IA/de ML GKE

Autopilot Standard

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/de 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 des 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 tâches d'entraînement par lot pouvant être interrompues.

Ce guide fournit des recommandations aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA. Pour obtenir une vue d'ensemble de toutes les bonnes pratiques GKE, consultez Bonnes pratiques pour GKE.

Avantages de la hiérarchisation 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 des outils Open Source pour vous aider à bénéficier des avantages suivants :

Minimiser les temps d'arrêt pour les charges de travail à priorité élevée.
Exécuter rapidement les charges de travail à priorité élevée.
Optimiser 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 Kubernetes natif 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, tels que ceux définis par des définitions de ressources personnalisées comme 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 attend, quand elle démarre (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 Kueue.
Hotswap : technique qui réduit le temps moyen de récupération (MTTR). Hotswap 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 défectueux, la charge de travail est replanifié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 libérer de l'espace pour 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 en cours d'exécution de faible priorité 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.

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 : Hiérarchisez les tâches à priorité élevée en leur attribuant une valeur `WorkloadPriority` Kueue plus élevée. Activez la mise en file d'attente é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éempter les charges de travail de faible priorité au profit des charges de travail à priorité élevée.	Hotswap vous permet d'appliquer les configurations suivantes : Configure PriorityClasses to define priority levels for your workloads. Attribuez des `PriorityClasses` plus élevées aux charges de travail critiques. Replanifiez automatiquement les charges de travail sur des nœuds sains en cas d'interruption. Pour tester la configuration des bonnes pratiques, consultez l'exemple Hotswap dans ce document.
Plusieurs charges de travail d'IA en concurrence pour des ressources limitées	Associez Kueue et Hotswap. Cette combinaison fournit un système robuste qui hiérarchise les 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 initiale et l'admission 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. Hotswap permet de réduire le temps de récupération d'une charge de travail à priorité élevée en cas d'interruption. Pour tester la configuration des bonnes pratiques, consultez l'exemple Kueue et Hotswap dans ce document.

Kueue vous permet d'appliquer les configurations suivantes :

Hiérarchisez les tâches à priorité élevée en leur attribuant une valeur WorkloadPriority Kueue plus élevée.
Activez la mise en file d'attente é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éempter les charges de travail de faible priorité au profit des charges de travail à priorité élevée.

Hotswap vous permet d'appliquer les configurations suivantes :

Configure PriorityClasses to define priority levels for your workloads.
Attribuez des PriorityClasses plus élevées aux charges de travail critiques.
Replanifiez automatiquement les charges de travail sur des nœuds sains en cas d'interruption.

Pour tester la configuration des bonnes pratiques, consultez l'exemple Hotswap dans ce document.

Plusieurs charges de travail d'IA en concurrence pour des ressources limitées

Associez Kueue et Hotswap. Cette combinaison fournit un système robuste qui hiérarchise les 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 initiale et l'admission 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. Hotswap permet de réduire le temps de récupération d'une charge de travail à priorité élevée en cas d'interruption.

Pour tester la configuration des bonnes pratiques, 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 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 une ResourceFlavor nommée tpu-v6e-slice qui spécifie les libellés de nœud pour les tranches TPU v6e.
Définit une ClusterQueue nommée tpu-training-cq qui gère le quota pour les ressources TPU.
Définit un LocalQueue nommé default-queue qui permet aux charges de travail dans l'espace de noms default d'utiliser la file d'attente de cluster tpu-training-cq.

Hotswap

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 à priorité élevée 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

En fonction de cette configuration, Hotswap effectue les actions suivantes :

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

Kueue et Hotswap

Associez Kueue et Hotswap lorsque vous travaillez dans un environnement complexe avec des ressources limitées. Cette combinaison fournit un système robuste qui hiérarchise les charges de travail critiques lors de la planification initiale et de 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

En fonction de cette configuration, Kueue est associé à 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 à priorité élevée.
Hotswap s'assure que le JobSet à priorité élevée redémarre rapidement, ce qui minimise son temps d'inactivité.
Lorsque l'infrastructure est récupérée, Hotswap replanifie le JobSet de faible priorité dans le pool de nœuds récupéré.

Étape suivante

Découvrez comment déployer des charges de travail GPU dans GKE.
Découvrez comment déployer des charges de travail TPU dans GKE.

Bonnes pratiques pour la priorisation des charges de travail d'IA/de ML GKE Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.