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

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 à l'aide de la commande gcloud config set:

   gcloud config set project PROJECT_ID
   

   Remplacez PROJECT_ID par votre Google Cloud ID de projet.

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 :

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

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

Dans cette section, vous allez créer des pools de nœuds TPU à l'aide de la commande gcloud beta container node-pools create.

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 \
       --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 \
       --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 \
       --project=PROJECT_ID
   

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

Dans les étapes précédentes, vous avez utilisé la commande gcloud beta container node-pools create pour créer les pools de nœuds. Ces commandes utilisent les options suivantes:

• --node-locations: liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée les pools de nœuds.
• --machine-type : type de machine à utiliser pour les nœuds. Dans ce cas, vous avez utilisé ct5lp-hightpu-4t. Pour en savoir plus sur les types de machines compatibles avec les TPU, utilisez le tableau de la section Choisir la version de TPU.
• --tpu-topology: topologie TPU à utiliser pour le pool de nœuds. Dans ce cas, vous avez utilisé 2x4. Pour en savoir plus sur les topologies TPU, consultez la section Choisir la topologie TPU.

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. Un jobset est une API de charge de travail qui vous permet de gérer un groupe de tâches Kubernetes en tant qu'unité. Le cas d'utilisation le plus courant d'un JobSet est l'entraînement distribué, mais vous pouvez également l'utiliser pour exécuter des charges de travail par lot.

Les JobSets suivants 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. Installez l'API JobSet dans votre cluster:

   VERSION=v0.8.1
   kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml
   

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

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

Les métriques et tableaux de bord d'observabilité des jobsets et des pools de nœuds dans la consoleGoogle Cloud sont disponibles en disponibilité générale.

Tableaux de bord

Pour afficher l'état de vos pools de nœuds TPU multi-hôtes sur GKE, accédez au tableau de bord État des pools de nœuds TPU GKE fourni par Cloud Monitoring:

Accéder à l'état du pool de nœuds TPU GKE

Pour en savoir plus, consultez Surveiller les métriques de santé des nœuds et des pools de nœuds TPU.

Sur la page Kubernetes Engine AI/ML de la consoleGoogle Cloud , l'onglet Déploiement de l'IA > Jobs affiche un tableau de bord de surveillance des JobSets avec des informations complètes sur l'état et les performances des JobSets et de leur infrastructure sous-jacente, comme l'état du JobSet, la disponibilité des réplicas et l'état des réplicas. Le tableau de bord inclut également des métriques d'infrastructure, y compris des métriques de processeur, de GPU, de TPU, de mémoire et de stockage. Pour en savoir plus, consultez Surveiller l'état d'un JobSet à l'aide de métriques.

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

Surveiller l'état des JobSets à l'aide de métriques

Pour savoir si un JobSet s'exécute comme prévu ou pour déduire s'il a été interrompu, vous pouvez utiliser les métriques Prometheus du package de métriques JobSet, comme kube_jobset_succeeded_replicas.

Notez que les métriques d'état des Jobsets ne sont compatibles qu'avec la version 1.32.1-gke.135700 ou ultérieure de GKE. Les métriques de santé JobSet sont activées par défaut dans les nouveaux clusters créés avec des versions compatibles. Pour les clusters existants qui sont mis à niveau vers des versions compatibles, les clients doivent activer manuellement le package de métriques JobSet. Pour en savoir plus, consultez la documentation.

Pour ce tutoriel, vérifiez l'achèvement du JobSet à l'aide de cette requête PromQL:

kube_jobset_succeeded_replicas{
  cluster="multislice-cluster",
  jobset_name=~"mulitslice-.*"}

Surveiller le temps d'activité, les temps de restauration (TTR) et les temps entre interruptions (TBI) du JobSet

Les métriques suivantes sont utiles pour surveiller la disponibilité d'un JobSet:

• kubernetes.io/jobset/uptime: durée totale pendant laquelle le JobSet a été disponible.
• kubernetes.io/jobset/times_to_recover: répartition de la période de récupération pour un JobSet. Chaque exemple indique un seul événement de récupération à partir d'une période d'indisponibilité pour le JobSet.
• kubernetes.io/jobset/times_between_interruptions: distribution de l'intervalle entre la fin de l'interruption précédente et le début de l'interruption actuelle pour un JobSet. Chaque exemple indique une seule durée entre l'interruption précédente et l'interruption actuelle.

Ces métriques s'appliquent aux JobSets qui comportent exactement un seul job GPU ou TPU répliqué. Le calcul des métriques ne repose que sur la disponibilité de cette seule tâche répliquée. Les métriques sont compatibles avec toutes les versions de GKE.

Pour afficher le temps de disponibilité des JobSets que vous avez utilisés dans ce tutoriel, exécutez la requête PromQL suivante:

avg_over_time(
  kubernetes_io:jobset_uptime{
    monitored_resource="k8s_entity", entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}])

Pour afficher les distributions de TBI pour les JobSets de ce tutoriel, exécutez la requête PromQL suivante:

histogram_quantile(0.50,
  sum_over_time(
    kubernetes_io:jobset_times_between_interruptions_bucket{
      monitored_resource="k8s_entity",entity_type="jobset",
      entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Vous pouvez étendre l'intervalle de la requête à un horizon temporel plus long, par exemple sept jours, et calculer le temps moyen entre interruptions (MTBI) sur cette période:

sum(sum_over_time(
  kubernetes_io:jobset_times_between_interruptions_sum{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
  kubernetes_io:jobset_times_between_interruptions_count{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Pour afficher les distributions de TTR, vous pouvez exécuter les requêtes PromQL suivantes:

histogram_quantile(0.50,
  sum_over_time(
    kubernetes_io:jobset_times_to_recover_bucket{
      monitored_resource="k8s_entity",entity_type="jobset",
      entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Après avoir augmenté l'intervalle de requêtes sur une période plus longue, par exemple sept jours, vous pouvez calculer le temps moyen de récupération (MTTR) sur cette période:

sum(sum_over_time(
  kubernetes_io:jobset_times_to_recover_sum{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
  kubernetes_io:jobset_times_to_recover_count{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

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.

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
   

GKE inclut une stratégie de préemption, qui définit la manière dont Kueue attribue les ressources disponibles. La règle spécifie qu'une charge de travail peut être préemptée si une charge de travail de priorité supérieure a besoin des ressources. Les charges de travail ayant une valeur de priorité inférieure sont plus susceptibles d'être préemptées par des charges de travail de priorité supérieure.