Optimiser l'entraînement de l'IA sur les TPU avec DWS et Kueue

Ce document explique comment configurer un cluster GKE pour maximiser la disponibilité de l'entraînement de l'IA sur les Tensor Processing Units (TPU). Vous configurez un système de secours automatisé à l'aide d'un outil de mise en file d'attente de jobs Open Source appelé Kueue et du planificateur de charges de travail dynamique (DWS) deGoogle Cloud.

La configuration que vous définissez dans ce document permet de configurer un pool de nœuds principal et un pool de nœuds de sauvegarde :

  • À la demande (plan A) : ce pool de nœuds est votre premier choix pour les nœuds. Le système tente d'abord de planifier les jobs sur ces machines à la demande. Le terme à la demande signifie qu'une fois ces machines en cours d'exécution, vous y avez accès de manière fiable et ininterrompue.
  • Démarrage flexible DWS (plan B) : il s'agit de votre pool de nœuds de sauvegarde. Lorsque les machines du plan A ne sont pas disponibles, le programme de planification Kueue attribue automatiquement votre job à ce pool du plan B. Le DWS recherche ensuite le matériel du plan B, mais il ne garantit pas un accès immédiat, car ce matériel peut également être indisponible. Toutefois, DWS ne renonce pas : il met votre demande en file d'attente pendant sept jours maximum et fournit automatiquement les machines dès qu'elles sont disponibles.

Cette approche minimise le temps d'attente des jobs. Cela signifie que vous n'avez pas à vérifier manuellement les ressources disponibles ni à réécrire votre script pour différentes machines.

Présentation des étapes de configuration

Pour configurer le système de secours automatique, vous devez effectuer plusieurs étapes de configuration. Il est utile de diviser cette configuration en deux catégories :

  • Tâches de l'administrateur de cluster : configuration ponctuelle de l'infrastructure, comme la création du cluster GKE, le provisionnement des pools de nœuds et l'installation du contrôleur de planification Kueue.
  • Tâches des développeurs d'IA : workflows quotidiens répétitifs, comme la définition des exigences des jobs d'entraînement et l'envoi de la charge de travail.

Même si vous effectuez toutes ces étapes vous-même, garder cette distinction à l'esprit permet de clarifier le processus global.

Avant de configurer le système, examinez les étapes de configuration que vous allez effectuer.

Sujet Tâche
Configurer l'infrastructure (administrateur de cluster) 1. Créez un cluster GKE
2. Créez les pools de nœuds.
3. Vérifier l'état du démarrage flexible dans le pool de nœuds
Installer et configurer Kueue (administrateur de cluster) 1. Installez Kueue
2. Définir les règles de configuration
Exécuter un job d'entraînement (développeur IA) 1. Créez le ConfigMap
. Définissez le fichier manifeste RayJob.
3. Envoyez la charge de travail
4. Se connecter au RayJob
5 Vérifier les journaux

Concepts clés

  • Pool de nœuds à la demande (plan A) : pool de nœuds principal à haute priorité. Votre job essaie toujours d'utiliser ce pool en premier.
  • Pool de nœuds DWS à démarrage flexible (plan B) : pool de nœuds de sauvegarde. Si les machines du pool principal ne sont pas disponibles, le système utilise automatiquement ce pool pour rechercher du matériel disponible.
  • Kueue : programme de planification qui gère la file d'attente des jobs. Il intercepte votre demande de job et décide quel pool de nœuds utiliser (plan A ou plan B).
  • Job : charge de travail d'entraînement de l'IA que vous souhaitez exécuter. Dans ce document, vous le définissez à l'aide d'un fichier manifeste RayJob.

Avant de commencer

  1. Dans la console Google Cloud , sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud .

    Rôles requis pour sélectionner ou créer un projet

    • Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
    • Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

    Accéder au sélecteur de projet

  2. Vérifiez que la facturation est activée pour votre projet Google Cloud .

  3. Activez les API Google Kubernetes Engine et Cloud TPU.

    Rôles requis pour activer les API

    Pour activer les API, vous avez besoin du rôle IAM Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin), qui contient l'autorisation serviceusage.services.enable. Découvrez comment attribuer des rôles.

    Activer les API

  4. Dans la console Google Cloud , activez Cloud Shell.

    Activer Cloud Shell

  5. Assurez-vous de disposer d'un quota préemptif suffisant pour utiliser les VM TPU à démarrage flexible. Si le quota par défaut ne répond pas à vos besoins, demandez une allocation plus élevée. Pour en savoir plus, consultez Quotas Cloud TPU et Configurer l'environnement Cloud TPU.

Définir des variables d'environnement

Pour simplifier les commandes que vous exécutez dans ce document, vous pouvez définir des variables d'environnement dans Cloud Shell. Ces variables stockent des valeurs telles que l'ID de votre projet Google Cloud , les noms de vos pools de nœuds et l'emplacement de votre cluster GKE.

Une fois ces variables définies, vous pouvez les réutiliser dans plusieurs commandes en faisant référence au nom de la variable (par exemple, $CLUSTER_NAME) au lieu de retaper ou de remplacer les valeurs à chaque fois. Cette approche facilite le processus et réduit le risque d'erreurs.

Pour définir les variables d'environnement utiles suivantes dans Cloud Shell, exécutez les commandes suivantes :

export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="us-east5-b"
export REGION="us-east5"
export CLUSTER_NAME="tpu-cluster"
export GKE_VERSION="1.34"
export ONDEMAND_NODEPOOL="on-demand-pool"
export DWS_NODEPOOL="dws-pool"

Voici une explication de ces variables d'environnement :

  • PROJECT_ID : ID de votre projet Google Cloud .
  • PROJECT_NUMBER : numéro d'identification unique de votre projet (par exemple, 123456789012).
  • ZONE : zone de calcul de votre cluster (par exemple, us-east5-b). Sélectionnez une zone où le type d'accélérateur choisi est disponible. Pour en savoir plus sur la disponibilité, consultez les quotas Cloud TPU ou les quotas de GPU.
  • REGION : région dans laquelle vous créez les ressources du cluster (par exemple, us-east5).
  • CLUSTER_NAME : nom que vous choisissez pour votre cluster GKE.
  • GKE_VERSION : version GKE de votre cluster. Utilisez la version 1.34 ou ultérieure.
  • ONDEMAND_NODEPOOL : nom de votre pool de nœuds standards à la demande. (Il s'agit de votre pool de nœuds du plan A.)
  • DWS_NODEPOOL : nom de votre pool de nœuds DWS à démarrage flexible. (Il s'agit de votre pool de nœuds du plan B.)

Configurer l'infrastructure (administrateur de cluster)

En tant qu'administrateur de cluster, vous configurez le cluster GKE et les pools de nœuds pour qu'ils soient compatibles avec le mécanisme de secours.

Créer un cluster GKE

Commencez par créer le cluster GKE. Ce cluster est l'environnement dans lequel vous installez le contrôleur Kueue, configurez vos pools de nœuds et exécutez vos tâches d'entraînement d'IA. Pour créer un cluster et vous y connecter, procédez comme suit :

  1. Créez le cluster :

    gcloud container clusters create ${CLUSTER_NAME} \
  --cluster-version=${GKE_VERSION} \
  --machine-type=n2-standard-16 \
  --location=${ZONE} \
  --enable-image-streaming \
  --addons=RayOperator \
  --project=${PROJECT_ID}

    Cette commande utilise les options clés suivantes :

    • --addons=RayOperator : installe l'opérateur Ray sur votre cluster. Vous avez besoin de cet opérateur pour gérer la charge de travail RayJob que vous enverrez plus tard dans ce document.
    • --enable-image-streaming : permet à votre cluster d'extraire les images de conteneurs plus rapidement. Cette fonctionnalité réduit considérablement le temps nécessaire au démarrage des grandes images de conteneur d'IA.

  2. Récupérez les identifiants du cluster afin que l'CLI kubectl puisse s'y connecter. Cette commande met à jour votre fichier de configuration Kubernetes, qui est stocké par défaut dans le répertoire ~/.kube/config :

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --location=${ZONE} \
  --project=${PROJECT_ID}

Créer les pools de nœuds

Créez les pools de nœuds principal et de sauvegarde pour votre environnement : le pool de nœuds à la demande (Plan A) et le pool de nœuds DWS flex-start (Plan B) :

  1. Créez le pool de nœuds à la demande : ce pool sert de ressource principale pour les jobs d'entraînement :

    gcloud container node-pools create ${ONDEMAND_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Dans cet exemple, vos machines de premier choix sont des accélérateurs TPU v6e. Vous spécifiez ce matériel à l'aide de l'option --machine-type=ct6e-standard-4t. Vous pouvez modifier ce type de machine pour qu'il corresponde au matériel souhaité pour votre modèle d'IA, comme des GPU ou différents TPU.

  2. Créez le pool de nœuds à démarrage flexible DWS : dans cet exemple, vous sélectionnez le même type de machine (--machine-type=ct6e-standard-4t) que celui que vous avez choisi pour le pool à la demande principal. Votre pool de nœuds du plan B n'a pas besoin d'utiliser un type de machine différent. Si vous souhaitez vraiment ce matériel spécifique, le fait de le choisir comme solution de secours signifie simplement que vous passerez à une autre méthode pour l'obtenir s'il n'est pas disponible immédiatement. Cette autre méthode utilise DWS, qui recherche en continu le matériel disponible pendant sept jours maximum :

    gcloud container node-pools create ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --reservation-affinity=none \
  --enable-autoscaling \
  --enable-queued-provisioning \
  --flex-start \
  --num-nodes=0 \
  --min-nodes=0 \
  --max-nodes=4

    Ces commandes utilisent les options clés suivantes :

    • --num-nodes=0, --min-nodes=0, --max-nodes=4 et --enable-autoscaling : cette combinaison permet aux pools de nœuds de passer de zéro nœud à un nombre plus élevé lorsque des jobs en ont besoin, puis de revenir à zéro lorsqu'ils sont inactifs, ce qui permet de réduire les coûts.
    • --tpu-topology : définit la disposition physique des puces TPU. Vous spécifiez cette mise en page, car la disposition physique des puces affecte la vitesse d'exécution de votre tâche d'entraînement distribuée.
    • --reservation-affinity=none : permet de s'assurer que le pool de nœuds ne consomme pas votre matériel pré-réservé. Google Cloud vous permet de réserver des machines spécifiques pour garantir leur disponibilité. Définir cet indicateur sur none indique au système d'ignorer ces réservations et de demander à la place des machines non réservées de manière dynamique.
    • --enable-queued-provisioning et --flex-start : (pool du forfait B uniquement) Ces indicateurs permettent à DWS de provisionner des nœuds pour votre pool du forfait B à partir de la capacité flexible lorsqu'elle devient disponible.

Vérifier l'état du démarrage flexible dans le pool de nœuds

Inspectez le pool de nœuds à démarrage flexible DWS et vérifiez que le démarrage flexible est activé :

gcloud container node-pools describe ${DWS_NODEPOOL} \
  --cluster=${CLUSTER_NAME} \
  --location=${ZONE} \
  --format="get(config.flexStart)"

Si flex-start est activé, le résultat est True.

Installer et configurer Kueue (administrateur de cluster)

Dans cette section, vous allez installer le contrôleur Kueue sur votre cluster. Rappelons que Kueue est un programme de planification qui gère la file d'attente des jobs. Il intercepte votre demande de job, décide du pool de nœuds à utiliser (à la demande ou DWS flex-start), puis attribue le job.

Installer Kueue

Exécutez la commande suivante pour installer Kueue. Cette commande télécharge les fichiers manifestes d'installation depuis le dépôt officiel et les applique à votre cluster :

helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
  --namespace kueue-system \
  --create-namespace \
  --set "controllerManager.featureGates[0].name=ElasticJobsViaWorkloadSlices" \
  --set "controllerManager.featureGates[0].enabled=true"

Définir les règles de configuration

Créez un fichier manifeste YAML qui définit les règles de priorité. Ces règles indiquent à Kueue d'utiliser d'abord le pool à la demande, puis le pool à démarrage flexible DWS :

  1. Créez un fichier nommé dws-tpu-queue.yaml avec le contenu suivant : Ce fichier définit deux saveurs de ressources (à la demande et DWS à démarrage flexible) et une file d'attente de cluster qui les hiérarchise. Ce fichier de configuration définit la logique utilisée par Kueue pour gérer vos jobs :

    • ResourceFlavor : au début de ce document, vous avez créé deux pools de nœuds et leur avez attribué des noms à l'aide des variables d'environnement ${ONDEMAND_NODEPOOL} et ${DWS_NODEPOOL}. Lorsqu'il a créé ces pools de nœuds, GKE a automatiquement attribué à chaque nœud de ces pools le nom que vous avez choisi pour ces variables d'environnement. La section ResourceFlavor indique à Kueue de rechercher les nœuds portant ces libellés.
    • ClusterQueue : cette section du fichier manifeste définit la règle de priorité. Il liste d'abord le type à la demande, de sorte que Kueue tente d'abord de provisionner des machines à la demande. Si Kueue ne peut pas obtenir ces machines, il tente de provisionner des machines DWS à démarrage flexible à la place.
    • Quotas : le fichier définit un quota, qui est une limite sur le total des ressources (telles que les processeurs, la mémoire et les puces TPU) que vos jobs peuvent utiliser à tout moment dans le pool de nœuds à la demande. Lorsque vos jobs atteignent cette limite, Kueue tente automatiquement de provisionner des machines DWS à démarrage flexible (vos machines de plan B), que vous avez configurées dans dws-tpu-queue.yaml avec une limite de quota beaucoup plus élevée. 
      apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-cpu"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: default-pool
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "on-demand"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${ONDEMAND_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "dws"
spec:
  nodeLabels:
    cloud.google.com/gke-nodepool: ${DWS_NODEPOOL}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {}
  resourceGroups:
    - coveredResources: ["cpu", "memory", "google.com/tpu"]
      flavors:
        - name: "default-cpu" # Used for Ray Head Pod.
          resources:
            - name: "cpu"
              nominalQuota: 10
            - name: "memory"
              nominalQuota: 20Gi
            - name: "google.com/tpu"
              nominalQuota: 0
        - name: "on-demand" # First choice: on-demand node-pool.
          resources:
            - name: "cpu"
              nominalQuota: 40
            - name: "memory"
              nominalQuota: 75Gi
            - name: "google.com/tpu"
              nominalQuota: 16
        - name: "dws" # If on-demand is unavailable, fallback to DWS.
          resources:
            - name: "cpu"
              nominalQuota: 1000000000
            - name: "memory"
              nominalQuota: 1000000000Gi
            - name: "google.com/tpu"
              nominalQuota: 1000000000 # "Infinite" quota
  admissionChecksStrategy:
    admissionChecks:
      - name: "dws-prov"
        onFlavors: [dws]
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: dws-prov
spec:
  controllerName: kueue.x-k8s.io/provisioning-request
  parameters:
    apiGroup: kueue.x-k8s.io
    kind: ProvisioningRequestConfig
    name: dws-config
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ProvisioningRequestConfig
metadata:
  name: dws-config
spec:
  provisioningClassName: queued-provisioning.gke.io
  managedResources:
    - google.com/tpu

  2. Appliquez la configuration à votre cluster. La commande suivante utilise un outil de ligne de commande appelé envsubst pour remplacer les variables d'espace réservé qui apparaissent dans le fichier dws-tpu-queue.yaml. envsubst remplace les espaces réservés par les valeurs des variables d'environnement que vous avez définies précédemment :

    envsubst < dws-tpu-queue.yaml | kubectl apply -f -

Exécuter un job d'entraînement (développeur d'IA)

En tant que développeur d'IA, vous définissez et envoyez une charge de travail d'entraînement en créant un fichier manifeste RayJob. Vous spécifiez vos besoins en ressources dans ce fichier manifeste. Le système de secours automatisé, que l'administrateur de cluster a configuré précédemment avec Kueue et DWS, gère les pools de nœuds sous-jacents pour vous.

Dans cette section, vous allez effectuer les étapes suivantes :

  • Créez un script d'entraînement Python.
  • Stockez ce script dans un ConfigMap Kubernetes.
  • Déployez un RayJob qui installe le ConfigMap en tant que volume afin que le script d'entraînement puisse être exécuté sur les nœuds.

Après avoir effectué ces étapes, Ray Train distribue automatiquement la charge de travail JAX sur les nœuds, et Kueue gère l'obtention des machines dont vous avez besoin.

Script d'entraînement

Copiez et collez le script Python suivant dans un fichier nommé train.py :

import time
import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.v2.jax import JaxTrainer
import jax
import jax.numpy as jnp

def train_func():
    # JaxTrainer handles JAX distributed setup.
    print(f"Local Devices: {jax.local_devices()}")

    # Simple Linear Regression Training Loop.
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (1000, 10))
    w_true = jax.random.normal(key, (10, 1))
    y = jnp.dot(x, w_true)

    # Initialize weights
    w = jnp.zeros((10, 1))
    learning_rate = 0.1

    @jax.jit
    def update(w, x, y):
        y_pred = jnp.dot(x, w)
        loss = jnp.mean((y_pred - y) ** 2)
        grad = jax.grad(lambda w: jnp.mean((jnp.dot(x, w) - y) ** 2))(w)
        return w - learning_rate * grad, loss

    # Training loop
    print("Starting training...")
    for epoch in range(50):
        w, loss = update(w, x, y)
        if epoch % 10 == 0:
            train.report({"loss": loss.item(), "epoch": epoch})
            print(f"Epoch {epoch}: Loss {loss:.4f}")

    print("Training Complete!")
    # Allow metrics to sync before closing
    time.sleep(3)

def main():
    scaling_config = ScalingConfig(
        num_workers=4,
        resources_per_worker={"TPU": 4},
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E"
    )

    trainer = JaxTrainer(
        train_loop_per_worker=train_func,
        scaling_config=scaling_config
    )

    result = trainer.fit()
    print(f"Run Result: {result.metrics}")

if __name__ == "__main__":
    main()

Le script d'entraînement utilise JAX, une bibliothèque Python pour le calcul numérique hautes performances, afin d'entraîner un modèle de régression linéaire. Ce script est un exemple simplifié conçu pour montrer comment utiliser DWS et Kueue pour le basculement automatique. Il n'effectue pas de parallélisme des données ni de parallélisme des modèles.

Notez que la section ScalingConfig du script d'entraînement définit les exigences matérielles pour la tâche d'entraînement. Cette section demande une topologie TPU 4x4, qui correspond à la disposition physique des pools de nœuds que vous avez configurés précédemment.

Créer le fichier ConfigMap

Importez le contenu de votre script train.py dans un objet Kubernetes ConfigMap. Cela permet au cluster de stocker le script et de le mettre à la disposition de votre RayJob :

kubectl create configmap jax-train-script --from-file=train.py

Le RayJob que vous définirez dans la section suivante montera ce ConfigMap en tant que volume. Le fichier de script apparaît alors dans les conteneurs Ray, ce qui permet au logiciel Ray de le trouver et de l'exécuter.

Appliquer le fichier manifeste RayJob

Créez un fichier nommé rayjob-tpu-v6e-dws.yaml avec le contenu suivant : Ce fichier manifeste définit votre tâche d'entraînement et indique au système comment l'acheminer :

apiVersion: ray.io/v1
kind: RayJob
metadata:
  name: rayjob-tpu-v6e-dws-${JOB_ID}
  labels:
    kueue.x-k8s.io/queue-name: user-queue
  annotations:
    kueue.x-k8s.io/elastic-job: "true"
spec:
  shutdownAfterJobFinishes: true
  entrypoint: python /app/train.py
  runtimeEnvYAML: |
    pip:
      - jax[tpu]==0.8.2
      - pandas==2.3.3
  rayClusterSpec:
    enableInTreeAutoscaling: true
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-nodepool: default-pool
          containers:
            - name: ray-head
              image: rayproject/ray:2.53.0-py311
              ports:
                - containerPort: 6379
                  name: gcs-server
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
              resources:
                limits:
                  cpu: "2"
                  memory: "4Gi"
                requests:
                  cpu: "2"
                  memory: "4Gi"
              volumeMounts:
                - mountPath: /app
                  name: train-script-volume
          volumes:
            - name: train-script-volume
              configMap:
                name: jax-train-script
    workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 2
        numOfHosts: 4
        groupName: tpu-group
        rayStartParams: {}
        template:
          spec:
            tolerations:
              - key: "google.com/tpu"
                operator: "Exists"
                effect: "NoSchedule"
              - key: "cloud.google.com/gke-queued"
                operator: "Exists"
                effect: "NoSchedule"
            containers:
              - name: ray-worker
                image: rayproject/ray:2.53.0-py311
                resources:
                  limits:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                  requests:
                    cpu: "8"
                    google.com/tpu: "4"
                    memory: "16Gi"
                volumeMounts:
                  - mountPath: /app
                    name: train-script-volume
            volumes:
              - name: train-script-volume
                configMap:
                  name: jax-train-script
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: 4x4

Ce fichier manifeste inclut trois configurations qui permettent au système de secours de fonctionner :

  • Demande du matériel spécifique : la section nodeSelector spécifie le matériel requis par votre script (dans cet exemple, un tpu-v6e-slice avec une topologie 4x4).
  • Sélectionnez la file d'attente : le libellé kueue.x-k8s.io/queue-name achemine votre job directement vers Kueue. Cela active la logique de secours automatique.
  • Tolère les nœuds à démarrage flexible DWS : la section tolerations permet au job de s'exécuter sur votre pool de nœuds du plan B. Étant donné que les nœuds à démarrage flexible DWS sont spécialement marqués (rejetés) par GKE afin que les charges de travail normales ne s'exécutent pas accidentellement dessus, votre job doit tolérer explicitement le rejet cloud.google.com/gke-queued.

Envoyer la charge de travail

Pour prouver que le système de secours fonctionne, vous devez envoyer deux tâches. La première tâche consomme la capacité à la demande du plan A, ce qui oblige la deuxième tâche à revenir à la capacité à démarrage flexible du plan B du programmeur de charge de travail dynamique.

Exécutez la commande suivante pour envoyer les deux tâches. La commande utilise une boucle for et envsubst pour injecter un ID de job unique dans le fichier manifeste pour chaque exécution :

for i in 1 2; do
  export JOB_ID=$i
  envsubst < rayjob-tpu-v6e-dws.yaml | kubectl apply -f -
  echo "Submitted Job $i"
  sleep 2
done

Une fois les jobs envoyés, le système gère la charge de travail comme suit :

  1. Interception : Kueue détecte les jobs à l'aide du libellé de file d'attente et les suspend temporairement.
  2. Décision : Kueue évalue la disponibilité des ressources par rapport aux règles de l'administrateur. Elle vérifie d'abord le pool A.
  3. Devoir :
    • Étant donné que les ressources du plan A sont disponibles pour le premier job, Kueue y attribue le job 1.
    • Étant donné que le job 1 consomme les ressources du plan A, Kueue attribue automatiquement le job 2 au pool du plan B (DWS flex-start).
  4. Lancement : Kueue réactive les jobs. Cette action déclenche l'autoscaler de cluster GKE pour provisionner les nœuds et démarrer les scripts d'entraînement.

Se connecter à RayJob

Pour terminer la vérification, vous pouvez utiliser la commande kubectl port-forward pour vous connecter au tableau de bord Ray et observer l'exécution de vos jobs.

Pour vérifier l'état de votre premier job, exécutez la commande suivante :

kubectl port-forward service/rayjob-tpu-v6e-dws-1-head-svc 8265:8265 &

Après avoir exécuté cette commande, ouvrez un navigateur Web et accédez à http://localhost:8265. Dans le tableau de bord Ray, vous pouvez afficher l'état des jobs et les métriques signalées pour vérifier que les deux jobs se terminent correctement sur leurs pools de nœuds respectifs.

Vous pouvez également afficher les journaux du premier job en exécutant la commande suivante :

kubectl logs job/rayjob-tpu-v6e-dws-1

La sortie tronquée du script d'entraînement doit ressembler à ce qui suit. Les messages Training Complete! et Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded devraient s'afficher vers la fin du résultat :

(pid=, ip=10.68.3.4) 5] XLA::TPU program HBM usage: 52.5K / 31.25G
(pid=, ip=10.68.9.4) :2152] XLA::TPU program VMEM usage: 141.0K / 128.00M [repeated 5x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722540     855 deepsea_compiler_backend.cc:2163] Total hbm usage >= 260.14M: [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167]     reserved           204B [repeated 19x across cluster]
(pid=, ip=10.68.9.4) I0320 03:59:34.722542     855 deepsea_compiler_backend.cc:2163]     program           70.0K  [repeated 5x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     arguments            0B  [repeated 12x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777627     888 deepsea_compiler_backend.cc:2163] Output size 0B; shares 0B with arguments. [repeated 14x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777625     888 deepsea_compiler_backend.cc:2163] Total host usage >= 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777626     888 deepsea_compiler_backend.cc:2163]     program         unknown size  [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777634     888 deepsea_compiler_backend.cc:2167] Program sflag requirement 224B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167]     scoped              40B [repeated 21x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167] Program vmem requirement 141.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program smem requirement 40B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program host requirement 0B: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777637     888 deepsea_compiler_backend.cc:2167] Program hbm requirement 70.0K: [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2167]     overlays          70.0K [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777638     888 deepsea_compiler_backend.cc:2175] XLA::TPU program SMEM usage: 1.9K / 1.00M (3 parameters) [repeated 7x across cluster]
(pid=, ip=10.68.6.4) I0320 03:59:34.777636     888 deepsea_compiler_backend.cc:2167]     HLO temp          76.0K (0.0% utilization: Unpadded (0B) Padded (0B), 100.0% fragmentation (76.0K)) [repeated 14x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Training Complete! [repeated 3x across cluster]
(RayTrainWorker pid=542, ip=10.68.6.4) Epoch 40: Loss 0.0000 [repeated 3x across cluster]
2026-03-20 03:59:51,008 SUCC cli.py:65 -- ------------------------------------------
2026-03-20 03:59:51,008 SUCC cli.py:66 -- Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded
2026-03-20 03:59:51,008 SUCC cli.py:67 -- ------------------------------------------

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce document ne 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

  • Dans la console Google Cloud , accédez à la page Gérer les ressources.

    Accéder à la page "Gérer les ressources"

  • Dans la liste des projets, sélectionnez le projet que vous souhaitez supprimer, puis cliquez sur Supprimer.
  • Dans la boîte de dialogue, saisissez l'ID du projet, puis cliquez sur Arrêter pour supprimer le projet.

    • Supprimer les ressources individuelles

    Si vous souhaitez conserver le projet GGoogle Cloud que vous avez utilisé dans ce document, exécutez la commande suivante pour supprimer le cluster :

    gcloud container clusters delete ${CLUSTER_NAME} \
    --location=${ZONE} \
    --project=${PROJECT_ID} \
    --quiet

    Résumé

    Dans ce document, vous avez configuré et testé un environnement d'entraînement Ray. Cet environnement utilise un pool de nœuds principal et un pool DWS de sauvegarde pour maximiser la disponibilité du matériel. En basculant automatiquement vers DWS lorsque les machines principales ne sont pas disponibles, vous avez réduit le temps d'attente de vos jobs d'entraînement.

    Pour que cela fonctionne, vous avez effectué les étapes suivantes :

    1. Création d'un cluster GKE : l'environnement permettant d'héberger les pools de nœuds et les outils de planification a été établi.
    2. Configuration des pools de nœuds : création d'un pool de nœuds à la demande (forfait A) et d'un pool de nœuds DWS (forfait B).
    3. Kueue installé et configuré : le contrôleur Kueue a été déployé et des règles de priorité ont été appliquées pour indiquer au système d'essayer d'abord le plan A et de passer au plan B en cas d'échec.
    4. Création d'un ConfigMap : déploiement d'un script d'entraînement JAX simplifié sur le cluster pour servir de charge de travail de test.
    5. Définition d'un fichier manifeste RayJob : configuration de la tâche pour demander du matériel spécifique, acheminer vers le contrôleur Kueue et tolérer les nœuds DWS.
    6. Envoyez la charge de travail : envoyez deux jobs pour forcer Kueue à router automatiquement le deuxième job vers le plan B lorsque les ressources du plan A sont consommées.
    7. Vérification des résultats : utilisé le transfert de port pour se connecter au tableau de bord Ray et confirmé que les deux jobs s'étaient exécutés correctement.

    Étapes suivantes