Planifier des charges de travail GKE avec la planification sensible à la topologie

Les charges de travail d'IA et de ML nécessitent une communication importante entre les pods. En raison de cette exigence, la bande passante réseau entre les pods a un impact direct sur le temps d'exécution et le coût de la charge de travail. Cette bande passante dépend de l'emplacement des instances de machine virtuelle (VM) dans le cluster.

Ce document explique comment optimiser la planification de vos charges de travail d'IA ou de ML à grande échelle sur un cluster Google Kubernetes Engine (GKE) pour améliorer les performances et la fiabilité. Plus précisément, vous configurez votre cluster pour qu'il utilise la planification basée sur la topologie (TAS, Topology Aware Scheduling) pour une communication à faible latence. Cette approche minimise la surcharge de communication et permet d'optimiser les performances de vos charges de travail.

Qu'est-ce que la planification tenant compte de la topologie (TAS) ?

TAS peut améliorer considérablement l'efficacité de l'entraînement des grands modèles de langage (LLM). TAS place stratégiquement les nœuds de calcul sur la topologie du réseau pour minimiser la surcharge de communication lors de l'agrégation des gradients, qui nécessite que les nœuds de calcul communiquent dans un ordre de classement spécifique. En minimisant les sauts de réseau entre les nœuds de calcul qui communiquent séquentiellement, TAS réduit les conflits réseau et optimise l'utilisation de la bande passante, ce qui permet une convergence plus rapide et des temps d'entraînement plus courts. Avec des modèles LLM de plus en plus volumineux, TAS est essentiel pour maximiser les performances et l'évolutivité de l'entraînement distribué.

TAS fonctionne mieux avec une capacité dense, qui peut être obtenue grâce aux réservations. Avec les VM à démarrage flexible ou les VM Spot, il est moins probable que votre capacité soit allouée à proximité. Il est donc possible que TAS ne fonctionne pas bien dans ce scénario.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
  • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser la Google Cloud CLI pour cette tâche, installez et initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.
  • Pour vous connecter à votre cluster, exécutez la commande suivante :

    gcloud container clusters get-credentials CLUSTER_NAME
    

    Remplacez CLUSTER_NAME par le nom de votre cluster.

Préparer votre cluster GKE

Pour préparer votre cluster GKE à exécuter des charges de travail avec TAS, procédez comme suit :

  1. Installer Kueue avec TAS activé

  2. Afficher la topologie de votre cluster GKE

  3. Configurer Kueue

Installer Kueue avec TAS activé

Nous vous recommandons d'utiliser la planification basée sur la topologie avec Kueue, un système Kubernetes natif qui gère les quotas et la manière dont les jobs doivent les consommer. TAS nécessite Kueue version 0.10.0 ou ultérieure, et vous devez l'activer explicitement.

Pour installer Kueue et activer TAS, sélectionnez l'une des options suivantes :

Fichier manifeste Kueue

  1. Installez Kueue :

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml
    
  2. Activez TAS dans Kueue :

    kubectl -n kueue-system patch deployment kueue-controller-manager \
        --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'
    

Chart Helm

Installez Kueue avec TAS activé à l'aide d'un chart Helm :

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

Après avoir installé Kueke, vous devez le configurer pour qu'il comprenne l'infrastructure qu'il gère, comme expliqué dans la section suivante.

Afficher la topologie de votre cluster GKE

Avant d'afficher la topologie des nœuds A4X, A4, A3 Ultra, A3 Mega et A3 High (8 GPU) provisionnés en tant que VM Spot, vous devez définir un emplacement compact sur les nœuds GKE pour exposer leur topologie physique pour TAS. Sinon, vous rencontrerez des erreurs.

Pour afficher la topologie des nœuds de votre cluster GKE dans un pool de nœuds spécifique, exécutez la commande suivante :

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

Remplacez NODE_POOL_NAME par le nom du pool de nœuds.

Pour comprendre la topologie physique des nœuds GKE sur vos VM dans la sortie, reportez-vous aux libellés de nœuds suivants :

  • cloud.google.com/gce-topology-block : ID spécifique à l'organisation du bloc réservé dans lequel se trouve la VM.

  • cloud.google.com/gce-topology-subblock : ID spécifique à l'organisation du sous-bloc dans lequel se trouve la VM.

  • cloud.google.com/gce-topology-host : ID de l'hôte sur lequel se trouve la VM.

  • kubernetes.io/hostname : nom d'hôte du nœud Kubernetes. Ce nom d'hôte est généralement également le nom du nœud GKE.

Plus deux VM partagent de valeurs de libellé, plus elles sont physiquement proches l'une de l'autre. Pour en savoir plus sur ces termes, consultez Terminologie.

Configurer Kueue

Après avoir installé Kueue, vous devez le configurer pour spécifier l'infrastructure qu'il gère. En règle générale, Kueue nécessite une définition de quota de ressource ClusterQueue pour une infrastructure statique ou une infrastructure dynamique avec l'autoscaling de cluster activé. La ressource ClusterQueue n'accepte une charge de travail que si les ressources qu'elle demande sont inférieures ou égales au pool de ressources défini dans ClusterQueue. Une fois que vous avez configuré Kueue comme décrit dans cette section, Kueue accepte les charges de travail à l'aide du TAS comme suit :

  • Charges de travail TAS : Kueue vérifie à la fois la topologie de l'infrastructure physique et son utilisation actuelle.

  • Charges de travail non TAS : Kueue ne vérifie pas la topologie de l'infrastructure physique. Kueue gère l'intégralité du quota défini dans la configuration et laisse l'attribution des nœuds à kube-scheduler.

Pour comprendre comment fournir une définition de quota de ressources ClusterQueue à Kueue, consultez les exemples suivants :

  • Quota très élevé : Kueue n'arrête pratiquement jamais l'admission d'une charge de travail en fonction des ressources demandées. En fonction des définitions TAS, Kueue peut accepter ou non les charges de travail en fonction de la topologie de l'infrastructure. Pour en savoir plus, consultez Quota de ressources très élevé.

  • Quota réaliste : Kueue n'accepte la charge de travail que si les ressources qu'elle demande respectent les limites de quota de ressources. En fonction des définitions TAS, Kueue vérifie ensuite la topologie de l'infrastructure avant d'admettre la charge de travail. Pour en savoir plus, consultez Quotas de ressources réalistes.

Toutes les références aux quotas de ressources dans les sections suivantes concernent les quotas de ressources ClusterQueue.

Quota de ressources très élevé

L'exemple suivant utilise un quota de ressources très élevé, de sorte que Kueue n'arrête jamais une charge de travail en fonction du quota de ressources disponible. Kueue utilise plutôt les informations de topologie des nœuds disponibles pour essayer de faire correspondre la topologie aux exigences de la charge de travail.

Pour utiliser la définition de quota de ressources suivante, procédez comme suit :

  1. Ouvrez l'éditeur de fichiers de votre choix. Ensuite, incluez la définition de quota suivante dans un fichier YAML nommé kueue-tas-config-very-high-quota.yaml :

      apiVersion: kueue.x-k8s.io/v1alpha1
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
    ---
      kind: ResourceFlavor
      apiVersion: kueue.x-k8s.io/v1beta1
      metadata:
        name: "tas-flavor"
      spec:
        nodeLabels:
          cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
        topologyName: "gke-default"
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: NoSchedule
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: ClusterQueue
      metadata:
        name: "tas-cluster-queue"
      spec:
        namespaceSelector: {}
        resourceGroups:
        - coveredResources: ["nvidia.com/gpu"]
          flavors:
          - name: "tas-flavor"
            resources:
            - name: "nvidia.com/gpu"
              nominalQuota: 10000000
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: LocalQueue
      metadata:
        namespace: "default"
        name: "tas-user-queue"
      spec:
        clusterQueue: "tas-cluster-queue"
    

    Remplacez NODE_POOL_NAME par le nom du pool de nœuds.

  2. Créez et appliquez la configuration du quota de ressources pour le système de mise en file d'attente des tâches Kueue :

    kubectl create -f kueue-tas-config-very-high-quota.yaml
    

Quota de ressources réaliste

L'exemple précédent ne configurait que les ressources GPU. Toutefois, Kueue peut gérer toutes les ressources compatibles avec Kubernetes.

L'exemple suivant définit un quota de ressources plus réaliste, incluant le processeur, la mémoire et le GPU. Ceci concerne 100 machines a3-ultragpu-8g. Une seule machine dispose de 224 processeurs virtuels, de 2 944 Go de mémoire et de 8 GPU.

Pour utiliser la définition de quota de ressources suivante, procédez comme suit :

  1. Ouvrez l'éditeur de fichiers de votre choix. Ensuite, incluez la définition de quota suivante dans un fichier YAML nommé kueue-tas-config-real-quota.yaml :

      apiVersion: kueue.x-k8s.io/v1alpha1
      kind: Topology
      metadata:
        name: "gke-default"
      spec:
        levels:
        - nodeLabel: "cloud.google.com/gce-topology-block"
        - nodeLabel: "cloud.google.com/gce-topology-subblock"
        - nodeLabel: "cloud.google.com/gce-topology-host"
        - nodeLabel: "kubernetes.io/hostname"
    ---
      kind: ResourceFlavor
      apiVersion: kueue.x-k8s.io/v1beta1
      metadata:
        name: "tas-flavor"
      spec:
        nodeLabels:
          cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
        topologyName: "gke-default"
        tolerations:
        - key: "nvidia.com/gpu"
          operator: "Exists"
          effect: NoSchedule
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: ClusterQueue
      metadata:
        name: "tas-cluster-queue"
      spec:
        namespaceSelector: {} # match all
        resourceGroups:
        - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
          flavors:
          - name: "tas-flavor"
            resources:
            # numbers below represent quota of 100 a3-ultragpu-8g machines
            - name: "cpu"
              nominalQuota: 22400
            - name: "memory"
              nominalQuota: 294400Gi
            - name: "nvidia.com/gpu"
              nominalQuota: 800
    ---
      apiVersion: kueue.x-k8s.io/v1beta1
      kind: LocalQueue
      metadata:
        namespace: "default"
        name: "tas-user-queue"
      spec:
        clusterQueue: "tas-cluster-queue"
    

    Remplacez NODE_POOL_NAME par le nom du pool de nœuds.

  2. Créez et appliquez une configuration de quota de ressources pour le système de mise en file d'attente de tâches Kueue :

    kubectl create -f kueue-tas-config-real-quota.yaml
    

    Le résultat ressemble à ce qui suit :

    topology.kueue.x-k8s.io/gke-default created
    resourceflavor.kueue.x-k8s.io/tas-flavor created
    clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
    localqueue.kueue.x-k8s.io/tas-user-queue created
    

Planifier des charges de travail avec TAS à l'aide de Kueue

Les scénarios suivants montrent comment demander à Kueue et à TAS de gérer les combinaisons courantes de charges de travail et d'infrastructures à l'aide des types et des niveaux de requêtes de topologie :

  • Voici les types de demandes de topologie disponibles (préférées ou requises) :

    • kueue.x-k8s.io/podset-preferred-topology : Kueue donne la priorité à la planification de l'ensemble de la charge de travail dans un niveau de topologie donné, mais admet tout de même une charge de travail qui ne correspond pas à ce niveau de topologie. Pour une charge de travail qui aurait pu tenir dans un seul niveau de topologie, Kueue peut planifier cette charge de travail sur plusieurs instances de ce niveau de topologie.

    • kueue.x-k8s.io/podset-required-topology : Kueue continue d'essayer d'admettre cette charge de travail jusqu'à ce que la charge de travail entière puisse s'adapter au niveau de topologie choisi.

  • Voici les niveaux de requête de topologie disponibles, qui vous permettent d'être plus ou moins précis sur l'infrastructure physique sur laquelle vous préférez ou devez exécuter votre job :

    • cloud.google.com/gce-topology-block

    • cloud.google.com/gce-topology-subblock

    • cloud.google.com/gce-topology-host

    • kubernetes.io/hostname

Pour planifier des charges de travail à l'aide de ces valeurs, utilisez le fichier YAML de job suivant :

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

Remplacez les variables suivantes :

  • JOB_NAME : nom du job.

  • NUMBER_OF_REPLICAS : nombre de pods en cours d'exécution en parallèle.

  • ANNOTATIONS_STRING : consultez le tableau suivant :

    Type et niveau de topologie demandés Description ANNOTATIONS_STRING
    Préférentiellement exécuté dans un nom d'hôte (recommandé) Cette configuration acceptera votre charge de travail tant qu'il y aura suffisamment de ressources disponibles pour répondre à ses besoins, même si la capacité est fragmentée. Kueue planifie vos pods de la manière la plus compacte possible. kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    Obligatoire pour s'exécuter dans un hôte

    Cette configuration n'acceptera votre charge de travail que s'il existe un hôte disposant de suffisamment de ressources pour répondre aux exigences de votre charge de travail.

    Cela est utile lorsqu'il existe plusieurs VM par hôte (par exemple, des types de machines plus petits) ou que plusieurs pods peuvent s'exécuter sur un même nœud. Dans ce cas, si la charge de travail est acceptée, elle s'exécutera sur un seul hôte.

    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    Préféré pour s'exécuter dans un hôte Cette configuration acceptera votre charge de travail tant qu'il y aura suffisamment de ressources disponibles pour répondre à ses besoins, même si la capacité est fragmentée. Kueue tentera de planifier vos pods dans un hôte et utilisera des hôtes supplémentaires si nécessaire. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    Obligatoire pour l'exécution dans un sous-bloc Cette configuration n'acceptera votre charge de travail que s'il existe un sous-bloc avec suffisamment de ressources pour répondre aux exigences de ressources de votre charge de travail. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    Préféré pour l'exécution dans un sous-bloc Cette configuration acceptera votre charge de travail tant qu'il y aura suffisamment de ressources disponibles pour répondre à ses besoins, même si la capacité est fragmentée. Kueue tentera de planifier vos pods dans un sous-bloc et utilisera des sous-blocs supplémentaires si nécessaire. Dans ce cas, Kueue classera plus haut un sous-bloc avec plus de capacité disponible, même s'il est fragmenté, par rapport à un sous-bloc avec juste assez de capacité pour répondre aux exigences. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    Obligatoire pour l'exécution dans un bloc Cette configuration n'autorise votre charge de travail que si les ressources disponibles dans un bloc répondent à ses exigences en termes de ressources. Si elle est acceptée, Kueue réduit au minimum le nombre de sous-blocs et d'hôtes pour planifier la charge de travail. Cela peut entraîner une fragmentation de votre capacité disponible. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    Préféré pour s'exécuter dans un bloc Cette configuration acceptera votre charge de travail tant qu'il y aura suffisamment de ressources disponibles pour répondre à ses besoins, même si la capacité est fragmentée. Kueue tentera de planifier vos pods dans un bloc et utilisera des blocs supplémentaires si nécessaire. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Planifier des charges de travail à l'aide de PodGroup avec TAS et Kueue

Lorsque vous utilisez des PodGroups, vous devez spécifier trois champs supplémentaires pour chaque pod d'un PodGroup :

Selon le framework de ML que vous utilisez, un leader de PodGroup peut nécessiter ou non un GPU. En raison d'une limitation de Kueue, ces cas doivent être traités différemment. Les exemples suivants montrent comment créer un PodGroup de trois pods avec un leader et deux workers.

Cas 1 : Le leader est également un nœud de calcul et nécessite un GPU

Si le leader est l'un des nœuds de calcul et qu'il nécessite également un GPU, il peut avoir n'importe quel nombre dans le PodGroup. Par souci de simplicité, dans l'exemple suivant, l'index du leader est 0 :

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Cas 2 : Le leader n'est pas un nœud de calcul et ne nécessite pas de GPU

Si le leader n'est pas l'un des workers en raison de la limitation de Kueue, il doit avoir le dernier index du PodGroup, en raison de la façon dont Kueue crée les PodSets. Si le leader ne possède pas le dernier index et que le premier nœud de calcul n'utilise pas le premier index, Kueue n'appliquera pas d'attribution de rang.

Consultez l'exemple ci-dessous :

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Étapes suivantes