Esta página foi traduzida pela API Cloud Translation.

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 queiram 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 o comando gcloud components update para obter a versão mais recente. As versões anteriores da CLI gcloud podem não suportar a execução dos comandos neste documento.
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 é pré-instalada com as ferramentas de linha de comandos gcloud, helm e kubectl usadas neste tutorial.

Aceda à Google Cloud consola.
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

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.

Para comunicar com o cluster, configure o seguinte: kubectl

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

Standard

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.

Para comunicar com o cluster, configure o seguinte: kubectl

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

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

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

Aplique o manifesto:

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

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

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

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

Aplique o manifesto:

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

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

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

Recupere o nome do RayCluster gerado para o RayJob:

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

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}')

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 &
```
Abra um navegador de Internet e introduza o seguinte URL:
```
http://localhost:8265/#/jobs
```
Veja o estado do RayJob e os registos relevantes.

Vários anfitriões

Recupere o nome do RayCluster gerado para o RayJob:

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

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}')

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 &
```
Abra um navegador de Internet e introduza o seguinte URL:
```
http://localhost:8265/#/jobs
```
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.

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.