Optimiza el entrenamiento de IA en TPU con DWS y Kueue

En este documento, se explica cómo configurar un clúster de GKE para maximizar la disponibilidad del entrenamiento de IA en las unidades de procesamiento tensorial (TPU). Configurarás un sistema de respaldo automatizado con una herramienta de cola de trabajos de código abierto llamada Kueue y el Programador dinámico de cargas de trabajo (DWS) deGoogle Cloud.

La configuración que defines en este documento establece un grupo de nodos principal y un grupo de nodos de copia de seguridad:

  • A pedido (plan A): Este grupo de nodos es tu primera opción para los nodos. El sistema primero intenta programar trabajos en estas máquinas a pedido. El término a pedido significa que, una vez que estas máquinas comienzan a ejecutarse, tienes acceso confiable e ininterrumpido a ellas.
  • DWS flex-start (plan B): Este es tu grupo de nodos de copia de seguridad. Cuando las máquinas del plan A no están disponibles, el programa de programación de Kueue asigna automáticamente tu trabajo a este grupo del plan B. Luego, DWS busca el hardware del plan B, pero no garantiza el acceso inmediato porque es posible que ese hardware tampoco esté disponible. Sin embargo, DWS no se rinde: mantiene tu solicitud en fila durante un máximo de 7 días y proporciona automáticamente las máquinas en cuanto están disponibles.

Este enfoque minimiza el tiempo que los trabajos pasan esperando en la fila. Esto significa que no tienes que verificar manualmente los recursos disponibles ni volver a escribir tu secuencia de comandos para diferentes máquinas.

Descripción general de los pasos de configuración

Para configurar el sistema de resguardo automático, debes completar varios pasos de configuración. Es útil dividir esta configuración en dos categorías:

  • Tareas del administrador del clúster: Configuración de infraestructura única, como la creación del clúster de GKE, el aprovisionamiento de grupos de nodos y la instalación del controlador de programación de Kueue
  • Tareas de desarrollador de IA: Flujos de trabajo diarios y repetitivos, como definir los requisitos del trabajo de entrenamiento y enviar la carga de trabajo.

Incluso si realizas todos estos pasos por tu cuenta, tener en cuenta esta distinción ayuda a aclarar el proceso general.

Antes de configurar el sistema, revisa los pasos de configuración que realizarás.

Tema Tarea
Configura la infraestructura (administrador del clúster) 1. Crea un clúster de GKE.
2. Crea los grupos de nodos
3. Verifica el estado de flex-start en el grupo de nodos
Instala y configura Kueue (administrador del clúster) 1. Instala Kueue.
2. Define las reglas de configuración
Ejecuta un trabajo de entrenamiento (desarrollador de IA) 1. Crea el ConfigMap
2. Define el manifiesto de RayJob
3. Envía la carga de trabajo
4. Conéctate al RayJob
5. Verifica los registros

Conceptos clave

  • Grupo de nodos a pedido (plan A): Es el grupo de nodos principal y de alta prioridad. Tu trabajo siempre intenta usar este grupo primero.
  • Grupo de nodos de inicio flexible de DWS (plan B): Es el grupo de nodos de copia de seguridad. Si no hay máquinas disponibles en el grupo principal, el sistema usa automáticamente este grupo para buscar hardware disponible.
  • Kueue: Es un programa de programación que administra la cola de trabajos. Intercepta tu solicitud de trabajo y decide qué grupo de nodos usar (plan A o plan B).
  • Trabajo: Es la carga de trabajo de entrenamiento de IA que deseas ejecutar. En este documento, lo definirás con un manifiesto de RayJob.

Antes de comenzar

  1. En la consola de Google Cloud , en la página del selector de proyectos, selecciona o crea un proyecto de Google Cloud .

    Roles necesarios para seleccionar o crear un proyecto

    • Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
    • Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (roles/resourcemanager.projectCreator), que contiene el permiso resourcemanager.projects.create. Obtén más información para otorgar roles.

    Ir al selector de proyectos

  2. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

  3. Habilita las APIs de Google Kubernetes Engine y Cloud TPU.

    Roles necesarios para habilitar las APIs

    Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin), que contiene el permiso serviceusage.services.enable. Obtén más información para otorgar roles.

    Habilitar las API

  4. En la consola de Google Cloud , activa Cloud Shell.

    Activa Cloud Shell

  5. Asegúrate de tener suficiente cuota de recursos interrumpibles para usar las VMs de TPU con inicio flexible. Si la cuota predeterminada no es suficiente para tus necesidades, solicita una asignación más alta. Para obtener más detalles, consulta Cuotas de Cloud TPU y Configura el entorno de Cloud TPU.

Define las variables de entorno

Para simplificar los comandos que ejecutas en este documento, puedes configurar variables de entorno en Cloud Shell. Estas variables almacenan valores como el ID de tu proyecto de Google Cloud , los nombres de tus grupos de nodos y la ubicación de tu clúster de GKE.

Después de definir estas variables, puedes volver a usarlas en varios comandos haciendo referencia al nombre de la variable (por ejemplo, $CLUSTER_NAME) en lugar de volver a escribir o reemplazar los valores cada vez. Este enfoque facilita el proceso y reduce el riesgo de errores.

Para definir las siguientes variables de entorno útiles en Cloud Shell, ejecuta los siguientes comandos:

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"

A continuación, se explica cada una de estas variables de entorno:

  • PROJECT_ID: Es el ID de tu proyecto de Google Cloud .
  • PROJECT_NUMBER: Es el número de identificador único de tu proyecto (por ejemplo, 123456789012).
  • ZONE: Es la zona de procesamiento de tu clúster (por ejemplo, us-east5-b). Selecciona una zona con disponibilidad para el tipo de acelerador que elegiste. Consulta las cuotas de Cloud TPU o las cuotas de GPU para obtener información sobre la disponibilidad.
  • REGION: Es la región en la que creas los recursos del clúster (por ejemplo, us-east5).
  • CLUSTER_NAME: Es el nombre que eliges para tu clúster de GKE.
  • GKE_VERSION: Es la versión de GKE de tu clúster. Usa la versión 1.34 o una posterior.
  • ONDEMAND_NODEPOOL: Es el nombre de tu grupo de nodos estándar a pedido. (Este es tu grupo de nodos del plan A).
  • DWS_NODEPOOL: Es el nombre de tu grupo de nodos de inicio flexible de DWS. (Este es tu grupo de nodos del plan B).

Configura la infraestructura (administrador del clúster)

Como administrador del clúster, debes configurar el clúster de GKE y los grupos de nodos para que admitan el mecanismo de respaldo.

Crea un clúster de GKE

Primero, crea el clúster de GKE. Este clúster es el entorno en el que instalas el controlador de Kueue, configuras tus grupos de nodos y ejecutas tus trabajos de entrenamiento de IA. Para crear un clúster y conectarte a él, sigue estos pasos:

  1. Crea el clúster:

    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}

    Este comando usa las siguientes marcas clave:

    • --addons=RayOperator: Instala el operador de Ray en tu clúster. Necesitas este operador para administrar la carga de trabajo de RayJob que enviarás más adelante en este documento.
    • --enable-image-streaming: Permite que tu clúster extraiga imágenes de contenedores más rápido. Esta función reduce significativamente el tiempo que tardan en comenzar a ejecutarse las imágenes de contenedores grandes de IA.

  2. Recupera las credenciales del clúster para que la CLI de kubectl pueda conectarse a él. Este comando actualiza tu archivo de configuración de Kubernetes, que se almacena de forma predeterminada en el directorio ~/.kube/config:

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

Crea los grupos de nodos

Crea los grupos de nodos principal y de copia de seguridad para tu entorno: el grupo de nodos a pedido (Plan A) y el grupo de nodos de inicio flexible de DWS (Plan B):

  1. Crea el grupo de nodos a pedido: Este grupo sirve como recurso principal para los trabajos de entrenamiento:

    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

    En este ejemplo, las máquinas de tu primera opción son los aceleradores de TPU v6e. Especifica este hardware con la marca --machine-type=ct6e-standard-4t. Puedes cambiar este tipo de máquina para que coincida con el hardware, como las GPUs o las diferentes TPU, que desees para tu modelo de IA.

  2. Crea el grupo de nodos de inicio flexible de DWS: En este ejemplo, selecciona el mismo tipo de máquina (--machine-type=ct6e-standard-4t) que elegiste para el grupo principal a pedido. Tu grupo de nodos de Plan B no tiene que usar un tipo de máquina diferente. Como realmente quieres este hardware específico, elegirlo como tu plan B solo significa que cambiarás a otro método para adquirirlo si no está disponible de inmediato. Este método alternativo usa DWS, que busca continuamente el hardware disponible durante un máximo de 7 días:

    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

    Estos comandos usan las siguientes marcas clave:

    • --num-nodes=0, --min-nodes=0, --max-nodes=4 y --enable-autoscaling: Esta combinación permite que los grupos de nodos se ajusten de cero nodos cuando un trabajo los necesita y se reduzcan cuando están inactivos, lo que ayuda a ahorrar costos.
    • --tpu-topology: Define la disposición física de los chips de TPU. Especificas este diseño porque la disposición física de los chips afecta la velocidad con la que se ejecuta tu trabajo de entrenamiento distribuido.
    • --reservation-affinity=none: Ayuda a garantizar que el grupo de nodos no consuma el hardware que reservaste previamente. Google Cloud te permite reservar máquinas específicas para garantizar la disponibilidad. Si se configura este parámetro en none, se le indica al sistema que omita esas reservas y, en su lugar, solicite de forma dinámica máquinas no reservadas.
    • --enable-queued-provisioning y --flex-start: (Solo para el grupo del plan B) Estas marcas permiten que DWS aprovisione nodos para tu grupo del plan B a partir de la capacidad flexible cuando esté disponible.

Verifica el estado de flex-start en el grupo de nodos

Inspecciona el grupo de nodos de inicio flexible de DWS y verifica que el inicio flexible esté habilitado:

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

Si flex-start está habilitado, el resultado es True.

Instala y configura Kueue (administrador del clúster)

En esta sección, instalarás el controlador de Kueue en tu clúster. Recuerda que Kueue es un programa de planificación que administra la cola de trabajos. Intercepta tu solicitud de trabajo, decide qué grupo de nodos usar (a pedido o de inicio flexible de DWS) y, luego, asigna el trabajo.

Instala Kueue

Ejecuta el siguiente comando para instalar Kueue. Este comando descarga los manifiestos de instalación del repositorio oficial y los aplica a tu clúster:

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"

Define las reglas de configuración

Crea un manifiesto en YAML que defina las reglas de prioridad. Estas reglas le indican a Kueue que use primero el grupo a pedido y, luego, el grupo de inicio flexible de DWS:

  1. Crea un archivo llamado dws-tpu-queue.yaml con el siguiente contenido. Este archivo define dos variantes de recursos (a pedido y de inicio flexible de DWS) y una cola de clúster que los prioriza. Este archivo de configuración define la lógica que usa Kueue para controlar tus trabajos:

    • ResourceFlavor: Al comienzo de este documento, creaste dos grupos de nodos y les asignaste nombres con las variables de entorno ${ONDEMAND_NODEPOOL} y ${DWS_NODEPOOL}. Cuando se crearon esos grupos de nodos, GKE etiquetó automáticamente cada nodo de esos grupos con el nombre que elegiste para esas variables de entorno. La sección ResourceFlavor le indica a Kueue que busque nodos con esas etiquetas.
    • ClusterQueue: En esta sección del manifiesto, se define la regla de prioridad. Primero, se muestra la variante a pedido, por lo que Kueue intenta aprovisionar máquinas a pedido primero. Si Kueue no puede obtener esas máquinas, intenta aprovisionar máquinas de DWS de inicio flexible.
    • Quotas: El archivo establece una cuota, que es un límite para los recursos totales (como CPU, memoria y chips de TPU) que tus trabajos pueden usar en cualquier momento en el grupo de nodos bajo demanda. Cuando tus trabajos alcancen este límite, Kueue intentará automáticamente aprovisionar máquinas de inicio flexible de DWS (tus máquinas del plan B), que configuraste en dws-tpu-queue.yaml con un límite de cuota mucho más alto. 
      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. Aplica la configuración al clúster. El siguiente comando usa una herramienta de línea de comandos llamada envsubst para reemplazar las variables de marcador de posición que aparecen en el archivo dws-tpu-queue.yaml. envsubst reemplaza los marcadores de posición por los valores de las variables de entorno que definiste anteriormente:

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

Ejecuta un trabajo de entrenamiento (desarrollador de IA)

Como desarrollador de IA, defines y envías una carga de trabajo de entrenamiento creando un manifiesto de RayJob. En este manifiesto, especificas tus requisitos de recursos, y el sistema de respaldo automatizado, que el administrador del clúster configuró anteriormente con Kueue y DWS, controla los grupos de nodos subyacentes por ti.

En esta sección, realizarás los siguientes pasos:

  • Crea una secuencia de comandos de entrenamiento de Python.
  • Almacena esa secuencia de comandos en un ConfigMap de Kubernetes.
  • Implementa un RayJob que active el ConfigMap como un volumen para que la secuencia de comandos de entrenamiento se pueda ejecutar en los nodos.

Después de realizar estos pasos, Ray Train distribuye automáticamente la carga de trabajo de JAX en los nodos, y Kueue se encarga de obtener las máquinas que necesitas.

La secuencia de comandos de entrenamiento

Copia y pega la siguiente secuencia de comandos de Python en un archivo llamado 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()

La secuencia de comandos de entrenamiento usa JAX, una biblioteca de Python para la computación numérica de alto rendimiento, para entrenar un modelo de regresión lineal. Este es un ejemplo simplificado diseñado para demostrar cómo usar DWS y Kueue para la conmutación por recuperación automatizada, y no realiza paralelismo de datos ni paralelismo de modelos.

Ten en cuenta que la sección ScalingConfig de la secuencia de comandos de entrenamiento define los requisitos de hardware para el trabajo de entrenamiento. En esa sección, se solicita una topología de TPU de 4x4, que coincide con el diseño físico de los grupos de nodos que configuraste anteriormente.

Crea el ConfigMap

Sube el contenido de tu secuencia de comandos train.py a un objeto ConfigMap de Kubernetes. Esto permite que el clúster almacene la secuencia de comandos y la ponga a disposición de tu RayJob:

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

El RayJob que definas en la siguiente sección activará este ConfigMap como un volumen. Esto hace que el archivo de secuencia de comandos aparezca dentro de los contenedores de Ray para que el software de Ray pueda encontrarlo y ejecutarlo.

Aplica el manifiesto de RayJob

Crea un archivo llamado rayjob-tpu-v6e-dws.yaml con el siguiente contenido. Este manifiesto define tu trabajo de entrenamiento y le indica al sistema cómo enrutarlo:

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

Este manifiesto incluye tres configuraciones que hacen que funcione el sistema de resguardo:

  • Solicita hardware específico: La sección nodeSelector especifica el hardware que requiere tu secuencia de comandos (en este ejemplo, un tpu-v6e-slice con una topología de 4x4).
  • Selecciona la cola: La etiqueta kueue.x-k8s.io/queue-name enruta tu trabajo directamente a Kueue. Esto habilita la lógica de resguardo automatizada.
  • Tolerates DWS flex-start nodes: La sección tolerations permite que el trabajo se ejecute en tu grupo de nodos de Plan B. Debido a que GKE marca (taintea) especialmente los nodos de inicio flexible de DWS para que las cargas de trabajo normales no se ejecuten accidentalmente en ellos, tu trabajo debe tolerar de forma explícita el taint cloud.google.com/gke-queued.

Envía la carga de trabajo

Para demostrar que el sistema de resguardo funciona, debes enviar dos trabajos. El primer trabajo consume la capacidad on demand del plan A, lo que obliga al segundo trabajo a recurrir a la capacidad de inicio flexible del DWS del plan B.

Ejecuta el siguiente comando para enviar los dos trabajos. El comando usa un bucle for y envsubst para insertar un ID de trabajo único en el manifiesto para cada ejecución:

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

Después de que envías los trabajos, el sistema controla la carga de trabajo de la siguiente manera:

  1. Interceptación: Kueue detecta los trabajos con la etiqueta de la cola y los suspende temporalmente.
  2. Decisión: Kueue evalúa la disponibilidad de recursos según las reglas del administrador. Primero, verifica el grupo del plan A.
  3. Tarea:
    • Como los recursos del plan A están disponibles para el primer trabajo, Kueue le asigna el trabajo 1.
    • Dado que el trabajo 1 consume los recursos del plan A, Kueue asigna automáticamente el trabajo 2 al grupo del plan B (DWS de inicio flexible).
  4. Lanzamiento: Kueue reanuda los trabajos. Esta acción activa el escalador automático del clúster de GKE para aprovisionar los nodos y comenzar las secuencias de comandos de entrenamiento.

Conéctate al RayJob

Como paso de verificación final, puedes usar el comando kubectl port-forward para conectarte al panel de Ray y observar la ejecución de tus trabajos.

Para verificar el estado de tu primer trabajo, ejecuta el siguiente comando:

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

Después de ejecutar este comando, abre un navegador web y ve a http://localhost:8265. En el panel de Ray, puedes ver el estado del trabajo y las métricas registradas para verificar que ambos trabajos se completen correctamente en sus respectivos grupos de nodos.

También puedes ver los registros del primer trabajo ejecutando el siguiente comando:

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

El resultado truncado de la secuencia de comandos de entrenamiento debería parecerse al siguiente. Deberías ver los mensajes Training Complete! y Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded cerca del final del resultado:

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

Realiza una limpieza

Para evitar que se apliquen cargos a tu Google Cloud cuenta por los recursos usados en este documento, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra el proyecto

  • En la Google Cloud consola, ve a la página Administrar recursos.

    Ir a Administrar recursos

  • En la lista de proyectos, elige el proyecto que quieres borrar y haz clic en Borrar.
  • En el diálogo, escribe el ID del proyecto y, luego, haz clic en Cerrar para borrar el proyecto.

    • Borra los recursos individuales

    Si deseas conservar el proyecto de GGoogle Cloud que usaste en este documento, ejecuta el siguiente comando para borrar el clúster:

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

    Resumen

    En este documento, configuraste y probaste un entorno de entrenamiento de Ray. Este entorno usa un grupo de nodos principal y un grupo de DWS de copia de seguridad para maximizar la disponibilidad del hardware. Al recurrir automáticamente a DWS cuando las máquinas principales no están disponibles, minimizaste el tiempo que tus trabajos de entrenamiento pasan en la fila de espera.

    Para que esto funcione, realizaste los siguientes pasos:

    1. Se creó un clúster de GKE: Se estableció el entorno para alojar los grupos de nodos y las herramientas de programación.
    2. Configuraste los grupos de nodos: Creaste un grupo de nodos según demanda (plan A) y un grupo de nodos de DWS (plan B).
    3. Kueue instalado y configurado: Se implementó el controlador de Kueue y se aplicaron reglas de prioridad que indican al sistema que primero intente el plan A y, luego, recurra al plan B.
    4. Se creó un ConfigMap: Se implementó una secuencia de comandos de entrenamiento de JAX simplificada en el clúster para que sirva como carga de trabajo de prueba.
    5. Definiste un manifiesto de RayJob: Configuraste el trabajo para solicitar hardware específico, enrutar al controlador de Kueue y tolerar nodos de DWS.
    6. Se envió la carga de trabajo: Se enviaron dos trabajos para obligar a Kueue a enrutar automáticamente el segundo trabajo al plan B cuando se consumen los recursos del plan A.
    7. Verificamos los resultados: Usamos el reenvío de puertos para conectarnos al panel de Ray y confirmamos que ambos trabajos se ejecutaron correctamente.

    ¿Qué sigue?