Configurar o KubeRay com a TPU Trillium

Autopilot Standard

Neste tutorial, mostramos como configurar o KubeRay com a TPU Trillium no Google Kubernetes Engine (GKE). Aprenda a configurar TPUs de host único e vários hosts, incluindo variáveis de ambiente e especificações de pod necessárias para TPU Trillium.

Este tutorial é destinado a administradores e operadores da plataforma e especialistas em dados e IA que querem aprender a configurar a inicialização da TPU Trillium com o KubeRay para pools de nós de host único e vários hosts. Este tutorial demonstra como executar um script com o Jax que verifica a inicialização da TPU. Este tutorial não implanta um modelo.

Antes de configurar o KubeRay no GKE, familiarize-se com as definições e terminologia do Ray no GKE.

Visão geral

Neste tutorial, mostramos como executar um script Python com Jax que verifica se a inicialização do TPU Trillium com KubeRay foi bem-sucedida. O Jax é uma biblioteca de computação numérica de alto desempenho que oferece suporte a cargas de trabalho de machine learning. O KubeRay é um operador do Kubernetes que oferece uma maneira unificada de implantar, gerenciar e monitorar aplicativos do Ray no Kubernetes.

As TPUs Trillium (v6e) exigem variáveis de ambiente e especificações de pod específicas que são diferentes das gerações anteriores de TPU. Este tutorial fornece as configurações necessárias para implantar uma carga de trabalho com KubeRay em TPUs Trillium.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

• Ativar a API Google Kubernetes Engine.
Ativar a API Google Kubernetes Engine
• Se você quiser usar a CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a gcloud CLI anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da gcloud CLI não sejam compatíveis com a execução dos comandos neste documento.
  Observação: em instalações prévias da gcloud CLI, defina a propriedade compute/region. Se você usa principalmente clusters zonais, defina compute/zone. Ao definir um local padrão, é possível evitar erros na gcloud CLI como: One of [--zone, --region] must be supplied: Please specify location. Talvez seja necessário especificar o local em determinados comandos se o local do cluster for diferente do padrão definido.

• Verifique se você instalou a CLI do Ray (versão 2.37.0).

Ativar o Cloud Shell

O Cloud Shell vem pré-instalado com gcloud, helm e Ferramentas de linha de comando kubectl usadas neste tutorial.

1. Acesse o console doGoogle Cloud .

2. Na parte de cima da janela do console do Google Cloud , clique no botão Ativar o Cloud Shell .

   Uma sessão do Cloud Shell é aberta em um novo frame no console Google Cloud e exibe um prompt de linha de comando.

   

Criar um cluster do GKE e um pool de nós

É possível configurar o KubeRay em TPUs em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Sobre os modos de operação do GKE.

Piloto automático

1. No Cloud Shell, execute este comando:

   gcloud container clusters create-auto CLUSTER_NAME \
       --enable-ray-operator \
       --release-channel=rapid \
       --location=LOCATION
   

   Substitua:

   • CLUSTER_NAME: o nome do novo cluster;
   • LOCATION: a região em que sua capacidade de TPU Trillium está disponível. Para mais informações, consulte Disponibilidade da TPU no GKE.

   O GKE cria um cluster do Autopilot com o complemento do operador do Ray ativado. O complemento instala automaticamente o webhook do Ray TPU no plano de controle do cluster.

2. Para se comunicar com o cluster, configure kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

Padrão

1. No Cloud Shell, crie um cluster Standard que ative o complemento do operador Ray executando o seguinte comando :

   gcloud container clusters create CLUSTER_NAME \
     --location LOCATION \
     --addons=RayOperator \
     --cluster-version=1.33 \
     --machine-type=n1-standard-16
   

   Substitua:

   • CLUSTER_NAME: o nome do novo cluster;
   • LOCATION: a região em que sua capacidade de TPU Trillium está disponível. Para mais informações, consulte Disponibilidade da TPU no GKE.

   A criação do cluster pode levar vários minutos.

2. Para se comunicar com o cluster, configure kubectl :

   gcloud container clusters get-credentials CLUSTER_NAME --location=LOCATION
   

3. É possível criar um pool de nós de fração de TPU de host único ou vários hosts:

Host único

No Cloud Shell, execute este comando:

gcloud container node-pools create v6e-4 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=1 \
    --threads-per-core=1 \
    --tpu-topology=2x2

Vários hosts

No Cloud Shell, execute este comando:

gcloud container node-pools create v6e-16 \
    --location=us-central2-b \
    --cluster=CLUSTER_NAME \
    --machine-type=ct6e-standard-4t \
    --num-nodes=4 \
    --threads-per-core=1 \
    --tpu-topology=4x4

Executar um recurso personalizado RayJob

Ao definir um manifesto do RayJob, você instrui o KubeRay a fazer o seguinte:

• Criar um RayCluster:a especificação RayJob inclui um rayClusterSpec que define a configuração do cluster do Ray (grupos principais e de worker) que você quer.
• Executar um job específico:o campo entrypoint em RayJob especifica o comando ou script a ser executado no cluster do Ray criado. Neste tutorial, o entrypoint é um script Python (tpu_list_devices.py) projetado para verificar a inicialização da TPU Trillium.

Para criar um recurso personalizado RayJob, siga estas etapas:

Host único

1. Crie o seguinte manifesto ray-job.tpu-v6e-singlehost.yaml:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-4-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
       -   replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 1
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
               -   name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 2x2

2. Aplique o manifesto:

   kubectl apply -f ray-job.tpu-v6e-singlehost.yaml
   

3. Verifique se o RayJob foi criado e está em execução:

   kubectl get rayjobs v6e-4-job
   

   O resultado será o seguinte:

   NAME      JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME       START TIME  END TIME   AGE
   v6e-4-job PENDING      Running             v6e-4-job-raycluster   2024-10-15T23:15:22Z  20s
   

4. Imprima a saída do RayJob.

   kubectl logs -l=job-name=v6e-4-job
   

   O resultado será o seguinte:

   2024-10-15 16:15:40,222 INFO cli.py:300 -- ray job stop v6e-4-job-hzq5q
   2024-10-15 16:15:40,246 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-15 16:15:40,112 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-15 16:15:50,181 INFO worker.py:1461 -- Using address 10.84.1.25:6379 set in the environment variable RAY_ADDRESS
   2024-10-15 16:15:50,181 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.25:6379...
   2024-10-15 16:15:50,186 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.25:8265
   ['TPU cores:4']
   2024-10-15 16:16:12,349 SUCC cli.py:63 -- -------------------------------------
   2024-10-15 16:16:12,349 SUCC cli.py:64 -- Job 'v6e-4-job-hzq5q' succeeded
   2024-10-15 16:16:12,349 SUCC cli.py:65 -- -------------------------------------
   

Vários hosts

1. Crie o seguinte manifesto ray-job.tpu-v6e-multihost.yaml:

   apiVersion: ray.io/v1
   kind: RayJob
   metadata:
     name: v6e-16-job
   spec:
     entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
     runtimeEnvYAML: |
       working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
       pip:
         - jax[tpu]==0.4.33
         - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
     rayClusterSpec:
       rayVersion: '2.43.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             -   name: ray-head
                 image: rayproject/ray:2.43.0-py310
                 ports:
                   - containerPort: 6379
                     name: gcs-server
                   - containerPort: 8265
                     name: dashboard
                   - containerPort: 10001
                     name: client
                 resources:
                   limits:
                     cpu: "8"
                     memory: 40G
                   requests:
                     cpu: "8"
                     memory: 40G
       workerGroupSpecs:
         - replicas: 1
           minReplicas: 1
           maxReplicas: 1
           numOfHosts: 4
           groupName: tpu-group
           rayStartParams: {}
           template:
             spec:
               containers:
                 - name: ray-worker
                   image: rayproject/ray:2.43.0-py310
                   resources:
                     limits:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                     requests:
                       cpu: "24"
                       google.com/tpu: "4"
                       memory: 200G
                   env:
                   - name: NODE_IP
                     valueFrom:
                       fieldRef:
                         fieldPath: status.hostIP
                   - name: VBAR_CONTROL_SERVICE_URL
                     value: $(NODE_IP):8353
                   - name: JAX_PLATFORMS
                     value: tpu,cpu
                   - name: ENABLE_PJRT_COMPATIBILITY
                     value: "true"
                   ports:
                   - containerPort: 8081
                     name: mxla
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: 4x4

2. Aplique o manifesto:

   kubectl apply -f ray-job.tpu-v6e-multihost.yaml
   

3. Verifique se o RayJob v6e-16 foi criado e está em execução:

   kubectl get rayjobs v6e-16-job
   

   O resultado será o seguinte:

   NAME         JOB STATUS   DEPLOYMENT STATUS   RAY CLUSTER NAME              START TIME             END TIME   AGE
   v6e-16-job                Running             v6e-16-job-raycluster-qr6vk   2024-10-16T19:28:19Z              66s
   

4. Imprima a saída do RayJob v6e-16:

   kubectl logs -l=job-name=v6e-16-job
   

   O resultado será o seguinte:

   2024-10-16 12:21:33,986 INFO cli.py:300 -- ray job stop v6e-16-job-z44s7
   2024-10-16 12:21:34,011 INFO cli.py:307 -- Tailing logs until the job exits (disable with --no-wait):
   2024-10-16 12:21:33,826 INFO job_manager.py:528 -- Runtime env is setting up.
   2024-10-16 12:21:46,327 INFO worker.py:1461 -- Using address 10.84.1.61:6379 set in the environment variable RAY_ADDRESS
   2024-10-16 12:21:46,327 INFO worker.py:1601 -- Connecting to existing Ray cluster at address: 10.84.1.61:6379...
   2024-10-16 12:21:46,333 INFO worker.py:1777 -- Connected to Ray cluster. View the dashboard at 10.84.1.61:8265
   ['TPU cores:16', 'TPU cores:16', 'TPU cores:16', 'TPU cores:16']
   2024-10-16 12:22:12,156 SUCC cli.py:63 -- ---------------------------------
   2024-10-16 12:22:12,156 SUCC cli.py:64 -- Job 'v6e-16-job-z44s7' succeeded
   2024-10-16 12:22:12,156 SUCC cli.py:65 -- ---------------------------------
   

Ver o RayJob no painel do Ray

Verifique se o GKE criou o serviço RayCluster e conecte-se à instância do RayCluster.

Host único

1. Recupere o nome do RayCluster gerado para o RayJob:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-4-job -o jsonpath='{.status.rayClusterName}')
   

2. Recupere o nome do serviço principal do RayCluster:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Conecte-se ao painel do Ray encaminhando o serviço principal:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Abra um navegador da Web e digite o seguinte URL:

   http://localhost:8265/#/jobs
   

5. Confira o status do RayJob e os registros relevantes.

Vários hosts

1. Recupere o nome do RayCluster gerado para o RayJob:

   export RAYCLUSTER_NAME=$(kubectl get rayjob v6e-16-job -o jsonpath='{.status.rayClusterName}')
   

2. Recupere o nome do serviço principal do RayCluster:

   export HEAD_SVC=$(kubectl get svc -l ray.io/cluster=$RAYCLUSTER_NAME,ray.io/node-type=head -o jsonpath='{.items[0].metadata.name}')
   

3. Conecte-se ao painel do Ray encaminhando o serviço principal:

   kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
   

4. Abra um navegador da Web e digite o seguinte URL:

   http://localhost:8265/#/jobs
   

5. Confira o status do RayJob e os registros relevantes.

O Ray define um recurso TPU-{accelerator}-Head para identificar o nó de worker do Ray que corresponde ao valor TPU_WORKER_ID=0. No grupo de TPU de vários hosts, o nó do Ray com TPU_WORKER_ID=0 tem TPU-v6e-16-head: 1.0 definido nos recursos. Essa variável de ambiente TPU_WORKER_ID é definida por um webhook mutante do GKE para o KubeRay.

Sucesso: você executou um RayJob em um pool de nós de fatia de TPU que executa TPU Trillium. A conclusão do RayJob indica que o Jax foi inicializado com sucesso e que os nodes da TPU distribuídos estão ativados corretamente.

Limpar

Depois de concluir o tutorial, exclua o RayJob para evitar cobranças indesejadas na sua conta:

Host único

kubectl delete rayjobs v6e-4-job

Vários hosts

kubectl delete rayjobs v6e-16-job

A seguir

• Saiba mais sobre o Ray no Kubernetes.
• Saiba como disponibilizar o vLLM no GKE com TPUs.
• Saiba como disponibilizar o SDXL no GKE com TPUs.
• Saiba mais sobre TPUs no GKE.