Orchestrer des charges de travail Multislice à l'aide de JobSet et de Kueue

Autopilot Standard

---

Ce tutoriel explique comment orchestrer plusieurs charges de travail Multislice sur Google Kubernetes Engine (GKE). Vous exécutez une charge de travail Jax à l'aide de TPU Multislice, JobSet et Kueue. Kueue met en œuvre la mise en file d'attente des jobs, en déterminant quand les Jobs doivent attendre et quand ils doivent démarrer, en fonction des quotas et d'une hiérarchie de partage équitable des ressources entre les équipes.

Nous vous recommandons de suivre ce tutoriel si vous utilisez des charges de travail qui nécessitent l'exécution simultanée de ressources TPU.

Avant d'utiliser des TPU dans GKE, nous vous recommandons de suivre le parcours de formation suivant :

1. Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
2. Apprenez-en plus sur les TPU Multislice dans GKE.

Objectifs

Ce tutoriel est destiné aux administrateurs GKE qui possèdent des clusters en mode GKE Standard et qui souhaitent exécuter des charges de travail Multislice pour la première fois.

Ce tutoriel couvre les étapes suivantes :

1. Préparez votre environnement avec un cluster GKE avec trois tranches de TPU v5e. Chaque tranche de TPU possède une topologie 2x4 avec huit puces. Par conséquent, 24 puces TPU v5e au total.
2. Créez les ressources Kueue pour vous assurer que les quotas sont partagés équitablement entre les charges de travail.
3. Exécutez votre charge de travail Multislice.

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 Google Cloud CLI pour cette tâche, installez puis initialisez gcloud CLI. Si vous avez déjà installé gcloud CLI, assurez-vous de disposer de la dernière version en exécutant la commande gcloud components update.
  Remarque : Pour les installations de gcloud CLI existantes, veillez à définir les propriétés compute/region et compute/zone. En définissant des emplacements par défaut, vous pouvez éviter les erreurs gcloud CLI telles que celle-ci : One of [--zone, --region] must be supplied: Please specify location.

• Installez JobSet v0.2.3 ou version ultérieure.

• Installez Kueue 0.4.1 ou une version ultérieure.

Préparer l'environnement

1. Dans la console Google Cloud, démarrez une instance Cloud Shell :
   Ouvrir Cloud Shell

2. Définissez les variables d'environnement par défaut :

   gcloud config set project PROJECT_ID
   gcloud config set compute/region COMPUTE_REGION
   

   Remplacez les valeurs suivantes :

   • PROJECT_ID : ID de votre projet Google Cloud.
   • COMPUTE_REGION : région Compute Engine.

Les clusters Autopilot qui exécutent la version 1.29.2-gke.1521000 ou une version ultérieure activent les TPU par défaut. Les TPU des clusters Autopilot sont configurés dans la spécification de la charge de travail. Pour en savoir plus, consultez la section Définir vos charges de travail Multislice avec des JobSets.

Créer un cluster GKE

Dans Cloud Shell, créez un cluster GKE :

Autopilot

gcloud container clusters create-auto multislice-cluster \
    --location=LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

Standard

gcloud container clusters create multislice-cluster \
    --location=LOCATION

Remplacez LOCATION par l'emplacement dans lequel vous souhaitez créer votre cluster. Vérifiez qu'il dispose de la capacité associée au type de machine ct5lp-hightpu-4t. La création du cluster peut prendre plusieurs minutes.

Si vous utilisez le mode GKE Autopilot, passez à la section Créer les ressources Kueue. Les clusters Autopilot qui exécutent la version 1.29.2-gke.1521000 ou une version ultérieure activent les TPU par défaut.

Créer trois pools de nœuds de tranche TPU en mode Standard

1. Créez le premier pool de nœuds nommé nodepool1 :

   gcloud beta container node-pools create nodepool1 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

   Remplacez NODE_LOCATION par une ou plusieurs zones de la région du cluster dans laquelle vous souhaitez créer les nœuds.

2. Créez le deuxième pool de nœuds nommé nodepool2 :

   gcloud beta container node-pools create nodepool2 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

3. Créez le troisième pool de nœuds nommé nodepool3 :

   gcloud beta container node-pools create nodepool3 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

GKE crée trois pools de nœuds. Chaque pool de nœuds est une tranche de TPU distincte.

Créer les ressources Kueue

1. Créez le fichier manifeste kueue.yaml suivant :

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ResourceFlavor
   metadata:
     name: "vlp-24"
   spec:
     nodeLabels:
       cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       cloud.google.com/gke-tpu-topology: 2x4
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ClusterQueue
   metadata:
     name: "cluster-queue"
   spec:
     namespaceSelector: {}
     queueingStrategy: BestEffortFIFO
     resourceGroups:
     - coveredResources: ["google.com/tpu"]
       flavors:
       - name: "vlp-24"
         resources:
         - name: "google.com/tpu"
           nominalQuota: 24
   
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: LocalQueue
   metadata:
     namespace: default
     name: multislice-queue
   spec:
     clusterQueue: cluster-queue
   

2. Appliquez le fichier manifeste kueue.yaml :

   kubectl apply -f kueue.yaml
   

   GKE crée les ressources Kueue suivantes :

• ResourceFlavor : abstraction des ressources d'un cluster. Dans cet exemple, GKE crée trois tranches de TPU avec une topologie 2x4. Chaque tranche de TPU possède une topologie 2x4 avec huit puces (24 puces TPU au total).
• ClusterQueue : file d'attente globale qui gère les charges de travail et les ressources du cluster.
• LocalQueue : groupes de charges de travail étroitement liées qui sont généralement exécutés par un seul locataire (utilisateur). Chaque LocalQueue pointe vers un ClusterQueue à partir duquel les ressources sont allouées pour exécuter ses charges de travail. Une charge de travail Kubeue est une abstraction représentant une charge de travail par lot. Dans ce cas, chaque charge de travail est un JobSet.

Définir vos charges de travail multilignes avec des JobSets

Dans cette section, vous allez créer trois objets JobSet. Ces JobSets exécutent une charge de travail Jax qui génère le nombre global de puces TPU dans la tranche, puis reste en veille pendant 60 secondes pour simuler la durée d'entraînement du modèle, puis se ferme.

1. Créez le fichier manifeste jobsets-multislice.yaml suivant :

   Autopilot

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-1slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                   resources:
                     limits:
                       google.com/tpu: 4
   
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-2slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 2
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-3slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 3
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   

   Standard

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-1slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                   resources:
                     limits:
                       google.com/tpu: 4
   
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-2slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 2
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-3slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 3
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   

2. Appliquez le fichier manifeste jobsets-multislice.yaml :

   kubectl apply -f jobsets-multislice.yaml
   

GKE crée les Jobs avec les demandes de ressources suivantes :

• La ressource JobSet multislice-1slice crée un Job qui nécessite une tranche de TPU au total.
• La ressource multislice-2slice JobSet crée deux Jobs qui nécessitent deux tranches de TPU au total.
• La ressource multislice-3slice SetSet crée trois Jobs qui nécessitent trois tranches de TPU au total.

Étant donné que le cluster ne comporte que trois tranches de TPU, les JobSets ne peuvent pas tous être exécutés en même temps. Lorsque Kueue met les trois JobSets multislice-3slice en file d'attente, ses Jobs s'exécutent seuls. Les multislice-1slice et multislice-2slice attendent et s'exécutent ensemble.

Vérifier que Kueue a admis les charges de travail

1. Vérifiez les charges de travail mises en file d'attente dans Kueue :

   kubectl get workloads
   

   Le résultat ressemble à ce qui suit :

   NAME                             QUEUE              ADMITTED BY     AGE
   jobset-multislice-1slice-2530a   multislice-queue                   3s
   jobset-multislice-2slice-ffb02   multislice-queue                   4s
   jobset-multislice-3slice-8c695   multislice-queue   cluster-queue   10s
   

Kueue met une ou plusieurs charges de travail en file d'attente, en fonction des ressources TPU dont elles ont besoin.

Surveiller les charges de travail

1. Surveillez les pods en cours d'exécution :

   kubectl get pods
   

   Le résultat ressemble à ce qui suit :

   NAME                                READY   STATUS      RESTARTS   AGE
   multislice-1slice-slice-0-0-pf2ll   1/1     Running     0          1s
   multislice-1slice-slice-0-1-55g62   1/1     Running     0          1s
   multislice-2slice-slice-0-0-f4hf7   1/1     Running     0          3s
   multislice-2slice-slice-0-1-c8kv7   1/1     Running     0          3s
   multislice-2slice-slice-1-0-7h46t   1/1     Running     0          3s
   multislice-2slice-slice-1-1-lj9hb   1/1     Running     0          3s
   multislice-3slice-slice-0-0-wzq9t   0/1     Completed   0          2m31s
   multislice-3slice-slice-0-1-zf4dp   0/1     Completed   0          2m30s
   multislice-3slice-slice-1-0-hbfn5   0/1     Completed   0          2m31s
   multislice-3slice-slice-1-1-45fgl   0/1     Completed   0          2m30s
   multislice-3slice-slice-2-0-wjbp4   0/1     Completed   0          2m30s
   multislice-3slice-slice-2-1-lwnvs   0/1     Completed   0          2m30s
   

   Vérifiez que GKE a d'abord planifié, créé et exécuté les pods pour multislice-3slice. GKE a ensuite exécuté les pods à partir des JobSets multislice-1slice et multislice-2slice.

Opération réussie : vous avez exécuté trois charges de travail Jax Multislice.

Activer la priorité et la préemption des charges de travail Kueue

Vous pouvez éventuellement attribuer des priorités aux charges de travail Kueue qui déterminent l'ordre dans lequel les charges de travail mises en file d'attente sont acceptées par Kueue.

Remarque : Selon la configuration de la règle de préemption de ClusterQueue, certains scénarios de préemption peuvent se produire. Par exemple, GKE peut préempter une charge de travail en cours d'exécution pour exécuter une charge de travail de priorité supérieure s'il n'y a pas de ressources suffisantes pour exécuter les deux charges de travail simultanément. GKE arrête et remet en file d'attente une charge de travail préemptée pour lui permettre de poursuivre le traitement si possible.

1. Mettez à jour votre ClusterQueue pour disposer d'une règle de préemption :

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ResourceFlavor
   metadata:
     name: "vlp-24"
   spec:
     nodeLabels:
       cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       cloud.google.com/gke-tpu-topology: 2x4
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ClusterQueue
   metadata:
     name: "cluster-queue"
   spec:
     namespaceSelector: {}
     resourceGroups:
     - coveredResources: ["google.com/tpu"]
       flavors:
       - name: "vlp-24"
         resources:
         - name: "google.com/tpu"
           nominalQuota: 24
    preemption:
       reclaimWithinCohort: Any
       withinClusterQueue: LowerPriority
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: LocalQueue
   metadata:
     namespace: default
     name: multislice-queue
   spec:
     clusterQueue: cluster-queue
   

2. Créez un PriorityClass pour chaque niveau de priorité distinct que vous souhaitez attribuer aux charges de travail :

   apiVersion: scheduling.k8s.io/v1
   kind: PriorityClass
   metadata:
     name: low-priority
   value: 100
   globalDefault: false
   description: "This low priority class should be used for some Pods only."
   

3. Attribuez le priorityClassName à votre JobSet :

   Autopilot

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: low-priority
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 priorityClassName: low-priority
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4 # Number of TPU chips per worker
   

   Standard

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: low-priority
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 priorityClassName: low-priority
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4 # Number of TPU chips per worker
     ```
   

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer le projet

Attention : La suppression d'un projet entraîne les effets décrits ci-dessous :
• Tout le contenu du projet est supprimé. Si vous avez utilisé un projet existant pour les tâches décrites dans ce document et que vous le supprimez, vous supprimerez également tout autre travail effectué dans le projet.
• Les ID de projets personnalisés sont perdus. Lorsque vous avez créé ce projet, vous avez peut-être créé un ID de projet personnalisé que vous souhaitez utiliser à l'avenir. Pour conserver les URL qui utilisent l'ID de projet, telle qu'une URL appspot.com, supprimez les ressources sélectionnées dans le projet au lieu de supprimer l'ensemble du projet.
1. In the Google Cloud console, go to the Manage resources page.

   Go to Manage resources

2. In the project list, select the project that you want to delete, and then click Delete.
3. In the dialog, type the project ID, and then click Shut down to delete the project.

Supprimer la ressource individuelle

1. Supprimez le système de quota Kueue :

   kubectl delete -n team-a localqueue
   kubectl delete -n team-b localqueue
   kubectl delete clusterqueue
   kubectl delete clusterqueue
   kubectl delete clusterqueue
   kubectl delete resourceflavor
   kubectl delete resourceflavor
   kubectl delete resourceflavor
   

2. Supprimez le fichier manifeste Kueue :

   VERSION=kueue.x-k8s.io/v1beta1
   kubectl delete -f \
       https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml
   

3. Supprimez le cluster à l'aide de la commande suivante :

   gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION
   

Étape suivante

• Apprenez-en plus sur Kueue.
• Apprenez à mettre en œuvre un système de mise en file d'attente de Job avec un partage de quotas entre espaces de noms sur GKE.