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 configurações de TPU 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 mostra como executar um script com o Jax que verifica a inicialização bem-sucedida 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 Trillium TPU 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 o 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 gcloud components update.
  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á assim:

   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á assim:

   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á assim:

   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á assim:

   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 de 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.