Ottimizzare l'addestramento dell'AI sulle TPU con DWS e Kueue

Questo documento spiega come configurare un cluster GKE per massimizzare la disponibilità dell'addestramento AI sulle Tensor Processing Unit (TPU). Configura un sistema di failover automatico utilizzando uno strumento di accodamento dei job open source chiamato Kueue e Dynamic Workload Scheduler (DWS) diGoogle Cloud.

La configurazione definita in questo documento configura un pool di nodi primario e un pool di nodi di backup:

  • On demand (piano A): questo pool di nodi è la tua prima scelta per i nodi. Il sistema tenta innanzitutto di pianificare i job su queste macchine on demand. Il termine on demand indica che, una volta avviate, queste macchine ti garantiscono un accesso affidabile e ininterrotto.
  • DWS flex-start (Piano B): Questo è il tuo pool di nodi di backup. Quando le macchine del piano A non sono disponibili, il programma di pianificazione Kueue assegna automaticamente il tuo job a questo pool del piano B. DWS cerca quindi l'hardware del piano B, ma non garantisce l'accesso immediato perché anche questo hardware potrebbe non essere disponibile. DWS non si arrende: mantiene la tua richiesta in coda per un massimo di 7 giorni e fornisce automaticamente le macchine non appena sono disponibili.

Questo approccio riduce al minimo il tempo di attesa dei job in coda. Ciò significa che non devi controllare manualmente le risorse disponibili o riscrivere lo script per macchine diverse.

Panoramica dei passaggi di configurazione

Per configurare il sistema di failover automatico, devi completare diversi passaggi di configurazione. È utile dividere questa configurazione in due categorie:

  • Attività dell'amministratore del cluster:configurazione dell'infrastruttura una tantum, ad esempio la creazione del cluster GKE, il provisioning dei pool di nodi e l'installazione del controller di pianificazione Kueue.
  • Attività dello sviluppatore di AI: flussi di lavoro quotidiani ripetitivi, come la definizione dei requisiti del job di addestramento e l'invio del workload.

Anche se esegui tutti questi passaggi personalmente, tenere presente questa distinzione aiuta a chiarire il processo complessivo.

Prima di configurare il sistema, esamina i passaggi di configurazione che eseguirai.

Argomento Attività
Configura l'infrastruttura (amministratore cluster) 1. Crea un cluster GKE
2. Crea i node pool
3. Verificare lo stato dell'avvio flessibile nel pool di nodi
Installare e configurare Kueue (amministratore cluster) 1. Installa Kueue
2. Definisci le regole di configurazione
Esegui un job di addestramento (AI Developer) 1. Crea il ConfigMap
2. Definisci il manifest RayJob
3. Invia il carico di lavoro
4. Connettiti a RayJob
5. Controllare i log

Concetti fondamentali

  • Pool di nodi on demand (piano A):il pool di nodi principale ad alta priorità. Il tuo job tenta sempre di utilizzare prima questo pool.
  • Node pool di nodi avvio flessibile (piano B):il pool di nodi di backup. Se le macchine nel pool principale non sono disponibili, il sistema utilizza automaticamente questo pool per cercare hardware disponibile.
  • Kueue:un programma di pianificazione che gestisce la coda dei job. Intercetta la richiesta di job e decide qualepool di nodil utilizzare (piano A o piano B).
  • Job:il workload di addestramento AI che vuoi eseguire. In questo documento, lo definiamo utilizzando un manifest RayJob.

Prima di iniziare

  1. Nella console Google Cloud , nella pagina di selezione del progetto, seleziona o crea un progetto Google Cloud .

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

    Vai al selettore di progetti

  2. Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  3. Attiva le API Google Kubernetes Engine e Cloud TPU.

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli.

    Abilita le API

  4. Nella console Google Cloud , attiva Cloud Shell.

    Attiva Cloud Shell

  5. Assicurati di avere una quota prerilasciabile sufficiente per utilizzare le VM TPU con avvio flessibile. Se la quota predefinita non è sufficiente per le tue esigenze, richiedi un'allocazione più elevata. Per i dettagli, consulta Quote di Cloud TPU e Configura l'ambiente Cloud TPU.

Definisci le variabili di ambiente

Per semplificare i comandi eseguiti in questo documento, puoi impostare le variabili di ambiente in Cloud Shell. Queste variabili memorizzano valori come l'ID del tuo progetto Google Cloud , i nomi dei tuoi pool di nodi e la posizione del tuo cluster GKE.

Dopo aver definito queste variabili, puoi riutilizzarle in più comandi facendo riferimento al nome della variabile (ad esempio, $CLUSTER_NAME) anziché digitare o sostituire i valori ogni volta. Questo approccio semplifica la procedura e riduce il rischio di errori.

Per definire le seguenti variabili di ambiente utili in Cloud Shell, esegui i seguenti comandi:

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"

Ecco una spiegazione di queste variabili di ambiente:

  • PROJECT_ID: l'ID del tuo Google Cloud progetto.
  • PROJECT_NUMBER: il numero identificatore univoco del tuo progetto (ad esempio, 123456789012).
  • ZONE: la zona di computing del cluster (ad esempio, us-east5-b). Seleziona una zona con disponibilità per il tipo di acceleratore scelto. Per informazioni sulla disponibilità, consulta Quote di Cloud TPU o Quote di GPU.
  • REGION: la regione in cui creare le risorse del cluster (ad esempio, us-east5).
  • CLUSTER_NAME: il nome che scegli per il tuo cluster GKE.
  • GKE_VERSION: la versione di GKE del tuo cluster. Utilizza la versione 1.34 o successive.
  • ONDEMAND_NODEPOOL: il nome del tuo pool di nodi standard on demand. (questo è pool di nodil del piano A).
  • DWS_NODEPOOL: il nome del tuo pool di nodi con avvio flessibile DWS. (Questo è il tuo pool di nodil di backup).

Configura l'infrastruttura (amministratore cluster)

In qualità di amministratore del cluster, configuri il cluster GKE e i node pool per supportare il meccanismo di fallback.

Crea un cluster GKE

Innanzitutto, crea il cluster GKE. Questo cluster è l'ambiente in cui installi il controller Kueue, configuri i node pool ed esegui i job di addestramento AI. Per creare un cluster e connetterti a esso:

  1. Crea il 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}

    Questo comando utilizza i seguenti flag chiave:

    • --addons=RayOperator: installa l'operatore Ray sul cluster. Questo operatore ti serve per gestire il workload RayJob che invierai più avanti in questo documento.
    • --enable-image-streaming: consente al cluster di estrarre più rapidamente le immagini container. Questa funzionalità riduce notevolmente il tempo necessario per l'avvio delle immagini dei container AI di grandi dimensioni.

  2. Recupera le credenziali del cluster in modo che la CLI kubectl possa connettersi al cluster. Questo comando aggiorna il file di configurazione di Kubernetes, archiviato per impostazione predefinita nella directory ~/.kube/config:

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

Crea i node pool

Crea i pool di nodi primario e di backup per il tuo ambiente: il pool di nodi on demand (piano A) e il pool di nodi con avvio flessibile DWS (piano B):

  1. Crea il node pool on demand: questo pool funge da risorsa principale per i job di addestramento:

    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

    In questo esempio, le macchine di prima scelta sono gli acceleratori TPU v6e. Specifica questo hardware utilizzando il flag --machine-type=ct6e-standard-4t. Puoi modificare questo tipo di macchina in modo che corrisponda all'hardware, ad esempio GPU o TPU diverse, che vuoi per il tuo modello di AI.

  2. Crea il pool di nodi con avvio flessibile DWS: in questo esempio, seleziona lo stesso tipo di macchina (--machine-type=ct6e-standard-4t) scelto per il pool di nodi on demand principale. Il pool di nodi del piano B non deve utilizzare un tipo di macchina diverso. Poiché vuoi davvero questo hardware specifico, sceglierlo come opzione di riserva significa solo che passerai a un altro metodo di acquisto se non è disponibile immediatamente. Questo metodo alternativo utilizza DWS, che cerca continuamente l'hardware disponibile fino a 7 giorni:

    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

    Questi comandi utilizzano i seguenti flag delle chiavi:

    • --num-nodes=0, --min-nodes=0, --max-nodes=4 e --enable-autoscaling: questa combinazione consente ai pool di nodi di scalare da zero nodi quando un job ne ha bisogno e di ridimensionarsi quando sono inattivi, il che contribuisce a ridurre i costi.
    • --tpu-topology: definisce la disposizione fisica dei chip TPU. Specifichi questo layout perché la disposizione fisica dei chip influisce sulla velocità di esecuzione del job di addestramento distribuito.
    • --reservation-affinity=none: contribuisce a garantire che il pool di nodi non consumi l'hardware prenotato in anticipo. Google Cloud consente di prenotare macchine specifiche per garantire la disponibilità. Se imposti questo flag su none, il sistema ignora queste prenotazioni e richiede dinamicamente macchine non riservate.
    • --enable-queued-provisioning e --flex-start: (solo pool del piano B) Questi flag consentono a DWS di eseguire il provisioning dei nodi per il pool del piano B dalla capacità flessibile quando diventa disponibile.

Verifica lo stato dell'avvio flessibile nel pool di nodi

Ispeziona il pool di nodi DWS con avvio flessibile e verifica che l'avvio flessibile sia abilitato:

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

Se l'avvio flessibile è abilitato, l'output è True.

Installare e configurare Kueue (amministratore del cluster)

In questa sezione, installerai il controller Kueue sul tuo cluster. Ricorda che Kueue è un programma di pianificazione che gestisce la coda dei job. Intercetta la richiesta di job, decide quale pool di nodi utilizzare (on demand o DWS avvio flessibile) e poi assegna il job.

Installa Kueue

Esegui questo comando per installare Kueue. Questo comando scarica i manifest di installazione dal repository ufficiale e li applica al tuo 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"

Definisci le regole di configurazione

Crea un manifest YAML che definisca le regole di priorità. Queste regole indicano a Kueue di utilizzare prima il pool on demand e poi il pool DWS ad avvio flessibile:

  1. Crea un file denominato dws-tpu-queue.yaml con il seguente contenuto. Questo file definisce due varianti di risorse (on demand e DWS ad avvio flessibile) e una coda del cluster che assegna loro la priorità. Questo file di configurazione definisce la logica utilizzata da Kueue per gestire i job:

    • ResourceFlavor: all'inizio di questo documento, hai creato due pool di nodi e hai assegnato loro dei nomi utilizzando le variabili di ambiente ${ONDEMAND_NODEPOOL} e ${DWS_NODEPOOL}. Quando ha creato questi pool di nodi, GKE ha etichettato automaticamente ogni nodo di questi pool con il nome scelto per queste variabili di ambiente. La sezione ResourceFlavor indica a Kueue di cercare nodi con queste etichette.
    • ClusterQueue: questa sezione del manifest definisce la regola di priorità. Elenca prima il tipo on demand, quindi Kueue tenta di eseguire il provisioning delle macchine on demand. Se Kueue non riesce a ottenere queste macchine, prova a eseguire il provisioning di macchine DWS ad avvio flessibile.
    • Quotas: il file imposta una quota, ovvero un limite alle risorse totali (come CPU, memoria e chip TPU) che i tuoi job possono utilizzare in qualsiasi momento nel pool di nodi on demand. Quando i tuoi job raggiungono questo limite, Kueue tenta automaticamente di eseguire il provisioning delle macchine DWS con avvio flessibile (le macchine di riserva), che hai configurato in dws-tpu-queue.yaml con un limite di quota molto più elevato. 
      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. Applica la configurazione al cluster. Il seguente comando utilizza uno strumento a riga di comando chiamato envsubst per sostituire le variabili segnaposto che appaiono nel file dws-tpu-queue.yaml. envsubst sostituisce i segnaposto con i valori delle variabili di ambiente che hai definito in precedenza:

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

Esegui un job di addestramento (sviluppatore di AI)

In qualità di sviluppatore di AI, definisci e invii un carico di lavoro di addestramento creando un manifest RayJob. Specifichi i requisiti delle risorse in questo manifest e il sistema di fallback automatico, configurato in precedenza dall'amministratore del cluster con Kueue e DWS, gestisce i pool di nodi sottostanti.

In questa sezione, esegui i seguenti passaggi:

  • Crea uno script di addestramento Python.
  • Archivia lo script in un ConfigMap Kubernetes.
  • Esegui il deployment di un RayJob che monta ConfigMap come volume in modo che lo script di addestramento possa essere eseguito sui nodi.

Dopo aver eseguito questi passaggi, Ray Train distribuisce automaticamente il carico di lavoro JAX tra i nodi e Kueue gestisce l'ottenimento delle macchine necessarie.

Lo script di addestramento

Copia e incolla il seguente script Python in un file chiamato 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()

Lo script di addestramento utilizza JAX, una libreria Python per il calcolo numerico ad alte prestazioni, per addestrare un modello di regressione lineare. Questo script è un esempio semplificato progettato per mostrare come utilizzare DWS e Kueue per il fallback automatico e non esegue il parallelismo dei dati o il parallelismo del modello.

Tieni presente che la sezione ScalingConfig dello script di addestramento definisce i requisiti hardware per il job di addestramento. Questa sezione richiede una topologia TPU 4x4, che corrisponde al layout fisico dei node pool configurati in precedenza.

Crea ConfigMap

Carica i contenuti dello script train.py in un oggetto Kubernetes ConfigMap. In questo modo, il cluster può archiviare lo script e renderlo disponibile per RayJob:

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

Il RayJob che definisci nella sezione successiva monta questo ConfigMap come volume. In questo modo, il file di script viene visualizzato all'interno dei container Ray in modo che il software Ray possa trovarlo ed eseguirlo.

Applica il manifest RayJob

Crea un file denominato rayjob-tpu-v6e-dws.yaml con il seguente contenuto. Questo manifest definisce il job di addestramento e indica al sistema come instradarlo:

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

Questo manifest include tre configurazioni che consentono il funzionamento del sistema di fallback:

  • Richiede hardware specifico:la sezione nodeSelector specifica l'hardware richiesto dallo script (in questo esempio, una tpu-v6e-slice con una topologia 4x4).
  • Seleziona la coda:l'etichetta kueue.x-k8s.io/queue-name indirizza il job direttamente a Kueue. In questo modo viene attivata la logica di fallback automatica.
  • Tollera i nodi di avvio flessibile DWS:la sezione tolerations consente al job di essere eseguito sul tuo pool di nodi di piano B. Poiché i nodi con inizio flessibile DWS sono contrassegnati in modo speciale (tainted) da GKE in modo che i carichi di lavoro normali non vengano eseguiti accidentalmente, il tuo job deve tollerare esplicitamente il taint cloud.google.com/gke-queued.

Invia il workload

Per dimostrare che il sistema di fallback funziona, devi inviare due job. Il primo job utilizza la capacità on demand del piano A, il che costringe il secondo job a utilizzare la capacità di avvio flessibile DWS del piano B.

Esegui questo comando per inviare i due job. Il comando utilizza un ciclo for e envsubst per inserire un ID job univoco nel manifest per ogni esecuzione:

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

Dopo l'invio dei job, il sistema gestisce il workload nel seguente modo:

  1. Intercettazione:Kueue rileva i job utilizzando l'etichetta della coda e li sospende temporaneamente.
  2. Decisione: Kueue valuta la disponibilità delle risorse in base alle regole dell'amministratore. Controlla prima il pool del piano A.
  3. Compito:
    • Poiché le risorse del piano A sono disponibili per il primo job, Kueue assegna il job 1 a questo piano.
    • Poiché il job 1 utilizza le risorse del piano A, Kueue assegna automaticamente il job 2 al pool del piano B (avvio flessibile DWS).
  4. Lancio:Kueue riattiva i job. Questa azione attiva il gestore della scalabilità automatica del cluster GKE per eseguire il provisioning dei nodi e avviare gli script di addestramento.

Connettiti a RayJob

Come passaggio di verifica finale, puoi utilizzare il comando kubectl port-forward per connetterti alla dashboard Ray e osservare l'esecuzione dei job.

Per controllare lo stato del primo job, esegui questo comando:

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

Dopo aver eseguito questo comando, apri un browser web e vai a http://localhost:8265. Nella dashboard Ray, puoi visualizzare lo stato del job e le metriche riportate per verificare che entrambi i job vengano completati correttamente nei rispettivi pool di nodi.

Puoi anche visualizzare i log del primo job eseguendo questo comando:

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

L'output troncato dello script di addestramento dovrebbe essere simile al seguente. Dovresti vedere i messaggi Training Complete! e Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded verso la fine dell'output:

(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 -- ------------------------------------------

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo documento, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina il progetto

  • Nella console Google Cloud , vai alla pagina Gestisci risorse.

    Vai a Gestisci risorse

  • Nell'elenco dei progetti, seleziona quello che vuoi eliminare, quindi fai clic su Elimina.
  • Nella finestra di dialogo, digita l'ID del progetto e fai clic su Chiudi per eliminare il progetto.

    • Elimina le singole risorse

    Se vuoi conservare il progetto GGoogle Cloud che hai utilizzato in questo documento, esegui il seguente comando per eliminare il cluster:

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

    Riepilogo

    In questo documento hai configurato e testato un ambiente di addestramento Ray. Questo ambiente utilizza unpool di nodil principale e un pool DWS di backup per massimizzare la disponibilità dell'hardware. Eseguendo automaticamente il failover su DWS quando le macchine principali non sono disponibili, hai ridotto al minimo il tempo di attesa dei job di addestramento.

    Per far funzionare questa configurazione, hai eseguito i seguenti passaggi:

    1. È stato creato un cluster GKE:è stato creato l'ambiente per ospitare i node pool e gli strumenti di pianificazione.
    2. Configurato i pool di nodi:è stato creato un pool di nodi on demand (piano A) e un pool di nodi DWS (piano B).
    3. Kueue installato e configurato:è stato eseguito il deployment del controller Kueue e sono state applicate regole di priorità che indicano al sistema di provare prima il piano A e di ripiegare sul piano B.
    4. Creazione di un ConfigMap:deployment di uno script di addestramento JAX semplificato nel cluster per fungere da workload di test.
    5. Definito un manifest RayJob:configurato il job per richiedere hardware specifico, indirizzare al controller Kueue e tollerare i nodi DWS.
    6. Invio del workload:invio di due job per forzare Kueue a instradare automaticamente il secondo job al piano B quando le risorse del piano A sono utilizzate.
    7. Verifica dei risultati:ha utilizzato l'inoltro delle porte per connettersi alla dashboard Ray e ha confermato che entrambi i job sono stati eseguiti correttamente.

    Passaggi successivi