Configure o KubeRay com o TPU Trillium

Autopilot Padrão

Este tutorial mostra como configurar o KubeRay com a TPU Trillium no Google Kubernetes Engine (GKE). Saiba como configurar configurações de TPU de anfitrião único e vários anfitriões, incluindo as variáveis de ambiente necessárias e as especificações de pods para o TPU Trillium.

Este tutorial destina-se a administradores e operadores da plataforma, bem como a especialistas em dados e IA que querem saber como configurar a inicialização do TPU Trillium com o KubeRay para pools de nós de anfitrião único e vários anfitriões. Este tutorial demonstra como executar um script com o Jax que verifica a inicialização bem-sucedida da TPU. Este tutorial não implementa um modelo.

Antes de configurar o KubeRay no GKE, certifique-se de que conhece as definições e a terminologia do Ray no GKE.

Vista geral

Este tutorial mostra como executar um script Python com o Jax que verifica se a inicialização do TPU Trillium com o KubeRay foi bem-sucedida. O JAX é uma biblioteca de computação numérica de elevado desempenho que suporta cargas de trabalho de aprendizagem automática. O KubeRay é um operador do Kubernetes que oferece uma forma unificada de implementar, gerir e monitorizar aplicações Ray no Kubernetes.

As TPUs Trillium (v6e) requerem variáveis de ambiente e especificações de Pod específicas que diferem das gerações de TPUs anteriores. Este tutorial fornece as configurações necessárias para implementar com êxito uma carga de trabalho com o KubeRay em TPUs Trillium.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

• Ative a API Google Kubernetes Engine.
Ative a API Google Kubernetes Engine
• Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute gcloud components update para obter a versão mais recente.
  Nota: para instalações existentes da CLI gcloud, certifique-se de que define a compute/region propriedade. Se usar principalmente clusters zonais, defina o compute/zone. Ao definir uma localização predefinida, pode evitar erros na CLI gcloud, como os seguintes: One of [--zone, --region] must be supplied: Please specify location. Pode ter de especificar a localização em determinados comandos se a localização do seu cluster for diferente da predefinição que definiu.

• Certifique-se de que tem a CLI Ray (versão 2.37.0) instalada.

Ativar Cloud Shell

A Cloud Shell vem pré-instalada com as ferramentas de linha de comandos gcloud, helm e kubectl que são usadas neste tutorial.

1. Aceda à Google Cloud consola.

2. Na parte superior da Google Cloud janela da consola, clique no botão Ativar Cloud Shell .

   É aberta uma sessão do Cloud Shell num novo frame na Google Cloud consola e é apresentado um comando.

   

Crie um cluster e um node pool do GKE

Pode configurar o KubeRay em TPUs num cluster do GKE Autopilot ou Standard. Recomendamos que use um cluster do Autopilot para uma experiência do Kubernetes totalmente gerida. Para escolher o modo de funcionamento do GKE mais adequado às suas cargas de trabalho, consulte o artigo Acerca dos modos de funcionamento do GKE.

Piloto automático

1. No Cloud Shell, execute o seguinte comando:

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

   Substitua o seguinte:

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

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

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

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

Standard

1. No Cloud Shell, crie um cluster Standard que ative o suplemento 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 o seguinte:

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

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

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

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

3. Pode criar um conjunto de nós de segmentação do TPU de host único ou de vários hosts:

Anfitrião único

No Cloud Shell, execute o seguinte 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 anfitriões

No Cloud Shell, execute o seguinte 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

Execute um recurso personalizado RayJob

Ao definir um manifesto RayJob, indica ao KubeRay que faça o seguinte:

• Criar um RayCluster: a especificação RayJob inclui um rayClusterSpec que define a configuração do cluster Ray (grupos principais e de trabalho) que quer.
• Executar uma tarefa específica: o campo entrypoint no RayJob especifica o comando ou o script a executar no cluster do Ray criado. Neste tutorial, o entrypoint é um script Python (tpu_list_devices.py) concebido para validar a inicialização da TPU Trillium.

Para criar um recurso personalizado RayJob, conclua os seguintes passos:

Anfitrião ú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 é semelhante ao 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 é semelhante ao 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 anfitriões

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 é semelhante ao 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 o resultado do RayJob v6e-16:

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

   O resultado é semelhante ao 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 -- ---------------------------------
   

Veja o RayJob no painel de controlo do Ray

Verifique se o GKE criou o serviço RayCluster e ligue-se também à instância do RayCluster.

Anfitrião ú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. Estabeleça ligação ao painel de controlo do Ray encaminhando a porta do serviço principal:

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

4. Abra um navegador de Internet e introduza o seguinte URL:

   http://localhost:8265/#/jobs
   

5. Veja o estado do RayJob e os registos relevantes.

Vários anfitriões

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. Estabeleça ligação ao painel de controlo do Ray encaminhando a porta do serviço principal:

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

4. Abra um navegador de Internet e introduza o seguinte URL:

   http://localhost:8265/#/jobs
   

5. Veja o estado do RayJob e os registos relevantes.

O Ray define um recurso TPU-{accelerator}-Head para identificar o nó de trabalho do Ray que corresponde ao valor TPU_WORKER_ID=0. No grupo de TPUs com vários anfitriões, o nó do Ray com TPU_WORKER_ID=0 tem TPU-v6e-16-head: 1.0 definido nos respetivos recursos. Esta variável de ambiente TPU_WORKER_ID é definida por um webhook de mutação do GKE para o KubeRay.

Êxito: executou com êxito um RayJob num conjunto de nós de fatia de TPU que executa o TPU Trillium. A conclusão do RayJob indica que o Jax foi inicializado com êxito e que os nós da TPU distribuída estão corretamente ativados.

Limpar

Depois de concluir o tutorial, para evitar a ocorrência de cobranças indesejadas na sua conta, elimine o RayJob:

Anfitrião único

kubectl delete rayjobs v6e-4-job

Vários anfitriões

kubectl delete rayjobs v6e-16-job

O que se segue?

• Saiba mais sobre o Ray no Kubernetes.
• Saiba como servir vLLM no GKE com TPUs.
• Saiba como publicar o SDXL no GKE com TPUs.
• Saiba mais acerca das TPUs no GKE.