Otimizar o treinamento de IA em TPUs com DWS e Kueue

Este documento ensina a configurar um cluster do GKE para maximizar a disponibilidade de treinamento de IA em Unidades de Processamento de Tensor (TPUs). Você configura um sistema de fallback automatizado usando uma ferramenta de enfileiramento de jobs de código aberto chamada Kueue e o Google Cloud Dynamic Workload Scheduler (DWS).

A configuração definida neste documento configura um pool de nós principal e um pool de nós de backup:

  • Sob demanda (plano A): esse pool de nós é sua primeira opção para nós. O sistema primeiro tenta programar jobs nessas máquinas sob demanda. O termo sob demanda significa que, depois que essas máquinas começam a ser executadas, você tem acesso confiável e ininterrupto a elas.
  • Início flexível do DWS (plano B): este é seu pool de nós de backup. Quando as máquinas do Plano A não estão disponíveis, o programa de programação do Kueue atribui automaticamente seu job a esse pool do Plano B. Em seguida, o DWS procura o hardware do plano B, mas não garante acesso imediato porque ele também pode estar indisponível. No entanto, o DWS não desiste: ele mantém sua solicitação na fila por até 7 dias e fornece automaticamente as máquinas assim que elas ficam disponíveis.

Essa abordagem minimiza o tempo que os jobs passam esperando na fila. Isso significa que você não precisa verificar manualmente os recursos disponíveis nem reescrever o script para diferentes máquinas.

Visão geral das etapas de configuração

Para configurar o sistema de failback automático, é preciso concluir várias etapas de configuração. Ajuda a dividir essa configuração em duas categorias:

  • Tarefas do administrador do cluster:configuração única da infraestrutura, como criar o cluster do GKE, provisionar pools de nós e instalar o controlador de programação do Kueue.
  • Tarefas de desenvolvedor de IA:fluxos de trabalho diários e repetitivos, como definir os requisitos do job de treinamento e enviar a carga de trabalho.

Mesmo que você execute todas essas etapas, manter essa distinção em mente ajuda a esclarecer o processo geral.

Antes de configurar o sistema, revise as etapas de configuração que você vai realizar.

Tópico Tarefa
Configurar a infraestrutura (administrador do cluster) 1. Crie um cluster do GKE
2. Crie os pools de nós
3. Verificar o status do início flexível no pool de nós
Instalar e configurar o Kueue (administrador do cluster) 1. Instale o Kueue
2. Definir as regras de configuração
Executar um job de treinamento (desenvolvedor de IA) 1. Crie o ConfigMap
2. Defina o manifesto do RayJob
3. Envie a carga de trabalho
4. Conecte-se ao RayJob
5. Verifique os registros

Principais conceitos

  • Pool de nós sob demanda (plano A):o pool de nós principal e de alta prioridade. Seu job sempre tenta usar esse pool primeiro.
  • Pool de nós de início flexível do DWS (plano B):o pool de nós de backup. Se as máquinas no pool principal não estiverem disponíveis, o sistema vai usar automaticamente esse pool para procurar hardware disponível.
  • Kueue:um programa de programação que gerencia a fila de jobs. Ele intercepta sua solicitação de job e decide qual pool de nós usar (Plano A ou Plano B).
  • Job:a carga de trabalho de treinamento de IA que você quer executar. Neste documento, você vai definir usando um manifesto do RayJob.

Antes de começar

  1. No console do Google Cloud , na página do seletor de projetos, selecione ou crie um projeto do Google Cloud .

    Funções necessárias para selecionar ou criar um projeto

    • Selecionar um projeto: não é necessário um papel específico do IAM para selecionar um projeto. Você pode escolher qualquer projeto em que tenha recebido um papel.
    • Criar um projeto: para criar um projeto, é necessário ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

    Acessar o seletor de projetos

  2. Verifique se o faturamento está ativado para o projeto do Google Cloud .

  3. Ative as APIs Google Kubernetes Engine e Cloud TPU.

    Funções necessárias para ativar APIs

    Para ativar as APIs, é necessário ter o papel do IAM de administrador de uso do serviço (roles/serviceusage.serviceUsageAdmin), que contém a permissão serviceusage.services.enable. Saiba como conceder papéis.

    Ativar as APIs

  4. No console do Google Cloud , ative o Cloud Shell.

    Ativar o Cloud Shell

  5. Verifique se você tem cota preemptiva suficiente para usar VMs de início flexível de TPU. Se a cota padrão não for suficiente para suas necessidades, solicite uma alocação maior. Para mais detalhes, consulte Cotas do Cloud TPU e Configurar o ambiente do Cloud TPU.

Definir variáveis de ambiente

Para simplificar os comandos executados neste documento, defina variáveis de ambiente no Cloud Shell. Essas variáveis armazenam valores como o ID do seu projeto Google Cloud , os nomes dos seus pools de nós e a localização do cluster do GKE.

Depois de definir essas variáveis, você pode reutilizá-las em vários comandos referenciando o nome da variável (por exemplo, $CLUSTER_NAME) em vez de digitar ou substituir valores a cada vez. Essa abordagem facilita o processo e reduz o risco de erros.

Para definir as seguintes variáveis de ambiente úteis no Cloud Shell, execute os comandos a seguir:

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"

Confira uma explicação dessas variáveis de ambiente:

  • PROJECT_ID: o ID do seu projeto do Google Cloud .
  • PROJECT_NUMBER: o número de identificador exclusivo do projeto (por exemplo, 123456789012).
  • ZONE: a zona do Compute do cluster (por exemplo, us-east5-b). Selecione uma zona com disponibilidade para o tipo de acelerador escolhido. Consulte as cotas de Cloud TPU ou de GPU para informações de disponibilidade.
  • REGION: a região em que você cria os recursos do cluster (por exemplo, us-east5).
  • CLUSTER_NAME: o nome escolhido para o cluster do GKE.
  • GKE_VERSION: a versão do GKE do cluster. Use a versão 1.34 ou posterior.
  • ONDEMAND_NODEPOOL: o nome do pool de nós padrão sob demanda. (Este é o pool de nós do seu Plano A).
  • DWS_NODEPOOL: o nome do pool de nós de início flexível do DWS. Esse é o pool de nós do plano B.

Configurar a infraestrutura (administrador do cluster)

Como administrador do cluster, você configura o cluster do GKE e os pools de nós para oferecer suporte ao mecanismo de substituição.

Criar um cluster do GKE

Primeiro, crie o cluster do GKE. Esse cluster é o ambiente em que você instala o controlador do Kueue, configura os pools de nós e executa os jobs de treinamento de IA. Para criar e se conectar a um cluster, siga estas etapas:

  1. Crie o 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}

    Esse comando usa as seguintes flags principais:

    • --addons=RayOperator: instala o operador do Ray no cluster. Você precisa desse operador para gerenciar a carga de trabalho do RayJob que será enviada mais tarde neste documento.
    • --enable-image-streaming: permite que o cluster extraia imagens de contêiner mais rápido. Esse recurso reduz significativamente o tempo necessário para que imagens de contêiner grandes de IA comecem a ser executadas.

  2. Recupere as credenciais do cluster para que a CLI kubectl possa se conectar a ele. Esse comando atualiza o arquivo de configuração do Kubernetes, que é armazenado por padrão no diretório ~/.kube/config:

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

Criar os pools de nós

Crie os pools de nós principal e de backup para seu ambiente: o pool de nós on demand (Plano A) e o pool de nós de início flexível do DWS (Plano B):

  1. Crie o pool de nós sob demanda: esse pool serve como o recurso principal para jobs de treinamento:

    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

    Neste exemplo, as primeiras máquinas escolhidas são aceleradores de TPU v6e. Especifique esse hardware usando a flag --machine-type=ct6e-standard-4t. Você pode mudar esse tipo de máquina para corresponder ao hardware, como GPUs ou TPUs diferentes, que você quer para seu modelo de IA.

  2. Crie o pool de nós de início flexível do DWS: neste exemplo, selecione o mesmo tipo de máquina (--machine-type=ct6e-standard-4t) escolhido para o pool principal sob demanda. O pool de nós do plano B não precisa usar um tipo de máquina diferente. Como você realmente quer esse hardware específico, o plano B é apenas uma mudança para um método diferente de aquisição se ele não estiver disponível imediatamente. Esse método alternativo usa o DWS, que pesquisa continuamente o hardware disponível por até 7 dias:

    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

    Esses comandos usam as seguintes flags principais:

    • --num-nodes=0, --min-nodes=0, --max-nodes=4 e --enable-autoscaling: essa combinação permite que os pools de nós sejam escalonados de zero nós quando um job precisa deles e sejam reduzidos quando estão ociosos, o que ajuda a economizar custos.
    • --tpu-topology: define a disposição física dos chips de TPU. Você especifica esse layout porque a disposição física dos chips afeta a velocidade de execução do job de treinamento distribuído.
    • --reservation-affinity=none: ajuda a garantir que o pool de nós não consuma o hardware pré-reservado. Google Cloud permite reservar máquinas específicas para garantir a disponibilidade. Definir essa flag como none instrui o sistema a ignorar essas reservas e solicitar dinamicamente máquinas não reservadas.
    • --enable-queued-provisioning e --flex-start: (somente pool do plano B) Essas flags permitem que o DWS provisione nós para o pool do plano B com capacidade flexível quando ela estiver disponível.

Verificar o status do início flexível no pool de nós

Inspecione o pool de nós de início flexível do DWS e verifique se o início flexível está ativado:

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

Se "flex-start" estiver ativado, a saída será True.

Instalar e configurar o Kueue (administrador do cluster)

Nesta seção, você vai instalar o controlador do Kueue no cluster. O Kueue é um programa de programação que gerencia a fila de jobs. Ele intercepta sua solicitação de job, decide qual pool de nós usar (sob demanda ou DWS início flexível) e atribui o job.

Instalar o Kueue

Execute o comando a seguir para instalar o Kueue. Esse comando baixa os manifestos de instalação do repositório oficial e os aplica ao seu 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"

Definir as regras de configuração

Crie um manifesto YAML que defina as regras de prioridade. Essas regras informam ao Kueue para usar primeiro o pool sob demanda e depois o pool de início flexível do DWS:

  1. Crie um arquivo chamado dws-tpu-queue.yaml com o conteúdo a seguir. Esse arquivo define dois variações de recursos (sob demanda e DWS início flexível) e uma fila de cluster que os prioriza. Esse arquivo de configuração define a lógica que o Kueue usa para processar seus jobs:

    • ResourceFlavor: no início deste documento, você criou dois pools de nós e atribuiu nomes a eles usando as variáveis de ambiente ${ONDEMAND_NODEPOOL} e ${DWS_NODEPOOL}. Ao criar esses pools de nós, o GKE rotulou automaticamente cada nó com o nome escolhido para as variáveis de ambiente. A seção ResourceFlavor informa ao Kueue para procurar nós com esses rótulos.
    • ClusterQueue: esta seção do manifesto define a regra de prioridade. Ele lista primeiro o tipo de instância sob demanda, para que o Kueue tente provisionar máquinas sob demanda primeiro. Se o Kueue não conseguir essas máquinas, ele tentará provisionar máquinas de início flexível do DWS.
    • Quotas: o arquivo define uma cota, que é um limite para o total de recursos (como CPU, memória e chips de TPU) que seus jobs podem usar a qualquer momento no pool de nós on demand. Quando seus jobs atingem esse limite, o Kueue tenta provisionar automaticamente máquinas de início flexível do DWS (suas máquinas do plano B), que você configurou em dws-tpu-queue.yaml com um limite de cota muito maior. 
      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. Aplique a configuração ao seu cluster. O comando a seguir usa uma ferramenta de linha de comando chamada envsubst para substituir variáveis de marcador de posição que aparecem no arquivo dws-tpu-queue.yaml. envsubst substitui os marcadores pelos valores das variáveis de ambiente que você definiu anteriormente:

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

Executar um job de treinamento (desenvolvedor de IA)

Como desenvolvedor de IA, você define e envia uma carga de trabalho de treinamento criando um manifesto RayJob. Você especifica os requisitos de recursos nesse manifesto, e o sistema de fallback automatizado, que o administrador do cluster configurou anteriormente com o Kueue e o DWS, processa os pools de nós subjacentes para você.

Nesta seção, você vai realizar as seguintes etapas:

  • Crie um script de treinamento em Python.
  • Armazene esse script em um ConfigMap do Kubernetes.
  • Implante um RayJob que monte o ConfigMap como um volume para que o script de treinamento possa ser executado nos nós.

Depois de realizar essas etapas, o Ray Train distribui automaticamente a carga de trabalho do JAX entre os nós, e o Kueue processa a obtenção das máquinas necessárias.

O script de treinamento

Copie e cole o seguinte script Python em um arquivo chamado 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()

O script de treinamento usa o JAX, uma biblioteca Python para computação numérica de alto desempenho, para treinar um modelo de regressão linear. Este script é um exemplo simplificado projetado para demonstrar como usar o DWS e o Kueue para fallback automatizado. Ele não realiza paralelismo de dados ou paralelismo de modelos.

A seção ScalingConfig do script de treinamento define os requisitos de hardware para o job de treinamento. Essa seção solicita uma topologia de TPU 4x4, que corresponde ao layout físico dos pools de nós configurados anteriormente.

Criar o ConfigMap

Faça upload do conteúdo do script train.py em um objeto ConfigMap do Kubernetes. Isso permite que o cluster armazene o script e o disponibilize para seu RayJob:

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

O RayJob definido na próxima seção vai montar esse ConfigMap como um volume. Isso faz com que o arquivo de script apareça dentro dos contêineres do Ray para que o software Ray possa encontrá-lo e executá-lo.

Aplicar o manifesto do RayJob

Crie um arquivo chamado rayjob-tpu-v6e-dws.yaml com o conteúdo a seguir. Esse manifesto define o job de treinamento e informa ao sistema como roteá-lo:

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

Esse manifesto inclui três configurações que fazem o sistema substituto funcionar:

  • Solicita hardware específico:a seção nodeSelector especifica o hardware necessário para o script. Neste exemplo, um tpu-v6e-slice com uma topologia 4x4.
  • Seleciona a fila:o rótulo kueue.x-k8s.io/queue-name encaminha seu job diretamente para o Kueue. Isso ativa a lógica de substituição automática.
  • Tolerância a nós de início flexível do DWS:a seção tolerations permite que o job seja executado no pool de nós do plano B. Como os nós de início flexível do DWS são marcados (tainted) pelo GKE para que as cargas de trabalho normais não sejam executadas neles por acidente, seu job precisa tolerar explicitamente o taint cloud.google.com/gke-queued.

Enviar a carga de trabalho

Para provar que o sistema de substituição funciona, envie dois jobs. O primeiro job consome a capacidade sob demanda do plano A, o que força o segundo job a voltar para a capacidade de início flexível do DWS do plano B.

Execute o comando a seguir para enviar os dois jobs. O comando usa um loop for e envsubst para injetar um ID de job exclusivo no manifesto de cada execução:

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

Depois de enviar os jobs, o sistema processa a carga de trabalho da seguinte maneira:

  1. Interceptação:o Kueue detecta os jobs usando o rótulo da fila e os suspende temporariamente.
  2. Decisão:o Kueue avalia a disponibilidade de recursos em relação às regras do administrador. Ela verifica primeiro o pool do Plano A.
  3. Atividade:
    • Como os recursos do Plano A estão disponíveis para o primeiro job, o Kueue atribui o Job 1 a ele.
    • Como o Job 1 consome os recursos do Plano A, o Kueue atribui automaticamente o Job 2 ao pool do Plano B (DWS início flexível).
  4. Início:o Kueue retoma os jobs. Essa ação aciona o escalonador automático de cluster do GKE para provisionar os nós e iniciar os scripts de treinamento.

Conectar-se ao RayJob

Como etapa final de verificação, use o comando kubectl port-forward para se conectar ao painel do Ray e acompanhar a execução dos jobs.

Para verificar o status do seu primeiro job, execute o seguinte comando:

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

Depois de executar esse comando, abra um navegador da Web e acesse http://localhost:8265. No painel do Ray, é possível conferir o status do job e as métricas informadas para verificar se os dois jobs foram concluídos com êxito nos respectivos pools de nós.

Você também pode conferir os registros do primeiro job executando o seguinte comando:

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

A saída truncada do script de treinamento será semelhante a esta. Você vai ver as mensagens Training Complete! e Job 'rayjob-tpu-v6e-dws-1-498t6' succeeded perto do final da saída:

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

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste documento, exclua o projeto que contém os recursos ou mantenha o projeto e exclua os recursos individuais.

Excluir o projeto

  • No console Google Cloud , acesse a página Gerenciar recursos.

    Acessar "Gerenciar recursos"

  • Na lista de projetos, selecione o projeto que você quer excluir e clique em Excluir .
  • Na caixa de diálogo, digite o ID do projeto e clique em Encerrar para excluí-lo.

    • Excluir recursos individuais

    Se você quiser manter o projeto do GGoogle Cloud usado neste documento, execute o comando a seguir para excluir o cluster:

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

    Resumo

    Neste documento, você configurou e testou um ambiente de treinamento do Ray. Esse ambiente usa um pool de nós principal e um pool de DWS de backup para maximizar a disponibilidade de hardware. Ao fazer fallback automático para o DWS quando as máquinas principais estão indisponíveis, você minimiza o tempo que os jobs de treinamento passam esperando na fila.

    Para fazer isso funcionar, você seguiu estas etapas:

    1. Criou um cluster do GKE:estabeleceu o ambiente para hospedar os pools de nós e as ferramentas de programação.
    2. Configurou os pools de nós:criou um pool de nós sob demanda (plano A) e um pool de nós do DWS (plano B).
    3. Instalou e configurou o Kueue:implantou o controlador do Kueue e aplicou regras de prioridade instruindo o sistema a tentar o Plano A primeiro e voltar ao Plano B.
    4. Criamos um ConfigMap:implantamos um script de treinamento do JAX simplificado no cluster para servir como carga de trabalho de teste.
    5. Definir um manifesto RayJob:configure o job para solicitar hardware específico, rotear para o controlador Kueue e tolerar nós do DWS.
    6. Envio da carga de trabalho:dois jobs foram enviados para forçar o Kueue a rotear automaticamente o segundo job para o Plano B quando os recursos do Plano A forem consumidos.
    7. Verificação dos resultados:usei o encaminhamento de porta para me conectar ao painel do Ray e confirmei que os dois jobs foram executados com sucesso.

    A seguir