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. Pour cette raison, la bande passante réseau entre les pods a un impact direct sur le temps d'exécution et le coût des charges 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 réduit la surcharge de communication et permet d'optimiser les performances de vos charges de travail.

Qu'est-ce que la planification basée sur la topologie (TAS) ?

TAS peut améliorer considérablement l'efficacité de l'entraînement des grands modèles de langage (LLM). Elle place stratégiquement les nœuds de calcul sur la topologie du réseau afin de réduire 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 réduisant le nombre de sauts réseau entre les nœuds de calcul qui communiquent de manière séquentielle, la planification basée sur la topologie réduit la contention 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 de LLM de plus en plus volumineux, la planification basée sur la topologie est essentielle pour optimiser les performances et l'évolutivité de l'entraînement distribué.

La planification basée sur la topologie fonctionne mieux avec une capacité dense, qui peut être obtenue grâce à des réservations. Avec les VM à démarrage flexible ou les VM Spot, il est moins probable que votre capacité soit allouée à proximité. La planification basée sur la topologie peut donc ne pas fonctionner correctement 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 la planification basée sur la topologie, procédez comme suit :

  1. Installer Kueue avec la planification basée sur la topologie activée

  2. Afficher la topologie de votre cluster GKE

  3. Configurer Kueue

Installer Kueue avec la planification basée sur la topologie activée

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. La planification basée sur la topologie nécessite Kueue version 0.10.0 ou ultérieure, et vous devez l'activer explicitement.

Pour installer Kueue et activer la planification basée sur la topologie, 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 la planification basée sur la topologie 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 la planification basée sur la topologie activée à 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"

Une fois Kueue installé, 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 à la planification basée sur la topologie. 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 les valeurs de libellé partagées par deux VM sont nombreuses, plus les VM sont physiquement proches les unes des autres. Pour en savoir plus sur ces termes, consultez la section Terminologie.

Configurer Kueue

Une fois Kueue installé, vous devez le configurer pour spécifier l'infrastructure qu'il gère. En règle générale, Kueue nécessite une ClusterQueue définition de quota de ressources 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 demandées par la charge de travail sont inférieures ou égales au pool de ressources défini dans la ressource ClusterQueue. Une fois que vous avez configuré Kueue comme décrit dans cette section, Kueue accepte les charges de travail à l'aide de la planification basée sur la topologie comme suit :

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

  • Charges de travail sans planification basée sur la topologie : Kueue ne vérifie pas la topologie de l'infrastructure physique. Kueue gère l'ensemble du quota défini dans la configuration et laisse l'attribution des nœuds à kube-scheduler.

Pour savoir 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 de la planification basée sur la topologie, Kueue peut ou non accepter les charges de travail en fonction de la topologie de l'infrastructure. Pour en savoir plus, consultez la section Quota de ressources très élevé.

  • Quota réaliste : Kueue n'accepte la charge de travail que si les ressources demandées par la charge de travail sont comprises dans ces limites de quota de ressources. En fonction des définitions de la planification basée sur la topologie, Kueue vérifie ensuite la topologie de l'infrastructure avant d'accepter la charge de travail. Pour en savoir plus, consultez la section Quota de ressources réaliste.

Toutes les références au quota de ressources dans les sections suivantes font référence au quota 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. Au lieu de cela, Kueue utilise 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. Incluez ensuite 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 de quota de ressources pour le système de mise en file d'attente des jobs Kueue :

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

Quota de ressources réaliste

L'exemple précédent n'a configuré 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, y compris le processeur, la mémoire et le GPU. Il s'agit de 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. Incluez ensuite 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 des jobs 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 la planification basée sur la topologie à l'aide de Kueue

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

  • Voici les types de requêtes de topologie disponibles (préférés ou requis) :

    • 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 accepte toujours une charge de travail qui ne correspond pas à ce niveau de topologie. Pour une charge de travail qui aurait pu s'adapter à 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'accepter cette charge de travail jusqu'à ce que l'ensemble de la charge de travail puisse s'adapter au niveau de topologie choisi.

  • Voici les niveaux de requêtes de topologie disponibles, qui vous permettent d'être plus ou moins spécifique concernant 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 exécutés en parallèle.

  • ANNOTATIONS_STRING : consultez le tableau suivant :

    Type et niveau de topologie demandés Description ANNOTATIONS_STRING
    Préféré pour s'exécuter 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 aux exigences de ressources de votre charge de travail, même si la capacité est fragmentée. Kueue planifiera vos pods de manière aussi compacte que possible. kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    Requis pour s'exécuter dans un hôte

    Cette configuration acceptera votre charge de travail uniquement s'il existe un hôte disposant de suffisamment de ressources pour répondre aux exigences de ressources 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 seul 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 aux exigences de ressources de votre charge de travail, 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"
    Requis pour s'exécuter dans un sous-bloc Cette configuration acceptera votre charge de travail uniquement s'il existe un sous-bloc disposant de 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 s'exécuter dans un sous-bloc Cette configuration acceptera votre charge de travail tant qu'il y aura suffisamment de ressources disponibles pour répondre aux exigences de ressources de votre charge de travail, 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"
    Requis pour s'exécuter dans un bloc Cette configuration acceptera votre charge de travail uniquement si les ressources disponibles dans un bloc répondent aux exigences de ressources de votre charge de travail. Si elle est acceptée, Kueue réduira 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 aux exigences de ressources de votre charge de travail, 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 la planification basée sur la topologie à l'aide de Kueue

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

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

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

Si le responsable est l'un des nœuds de calcul et 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 responsable 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 responsable n'est pas un nœud de calcul et ne nécessite pas de GPU

Si le responsable n'est pas l'un des nœuds de calcul en raison de la limitation de Kueue, il doit avoir le dernier index dans le PodGroup, en raison de la façon dont Kueue crée des PodSet. Si le responsable n'a 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 classement.

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"

Étape suivante