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

Autopilot Standard

Ce tutoriel explique comment exécuter une charge de travail Jax à l'aide de TPU Multislice dans Google Kubernetes Engine (GKE) 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.

Ce tutoriel explique comment orchestrer plusieurs charges de travail Multislice nécessitant 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 :

Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
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 :

Préparez votre environnement avec un cluster GKE avec trois tranches de TPU v5e. Chaque tranche de TPU possède une topologie 2x4 et quatre puces par hôte. Par conséquent, 24 puces TPU v5e au total.
Créez les ressources Kueue pour vous assurer que les quotas sont partagés équitablement entre les charges de travail.
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

Dans la console Google Cloud, démarrez une instance Cloud Shell :
Ouvrir Cloud Shell
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 sur les 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 TPU en mode Standard

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.

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

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

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

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, trois tranches de TPU ont une topologie 2x4 et quatre puces par hôte, soit 24 puces TPU.
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.

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

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

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

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.

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.

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

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."

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

In the Google Cloud console, go to the Manage resources page.
Go to Manage resources
In the project list, select the project that you want to delete, and then click Delete.
In the dialog, type the project ID, and then click Shut down to delete the project.

Supprimer la ressource individuelle

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

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

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

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

Étapes suivantes

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.