Esta página se ha traducido con Cloud Translation API.

Configurar KubeRay con TPU Trillium

Autopilot Standard

En este tutorial se muestra cómo configurar KubeRay con TPU Trillium en Google Kubernetes Engine (GKE). Aprende a configurar las TPUs de un solo host y de varios hosts, incluidas las variables de entorno y las especificaciones de Pod necesarias para TPU Trillium.

Este tutorial está dirigido a administradores y operadores de la plataforma, así como a especialistas en datos e IA que quieran aprender a configurar la inicialización de TPU Trillium con KubeRay para grupos de nodos de un solo host y de varios hosts. En este tutorial se muestra cómo ejecutar una secuencia de comandos con Jax que verifica que la inicialización de la TPU se ha realizado correctamente. En este tutorial no se despliega ningún modelo.

Antes de configurar KubeRay en GKE, asegúrate de que conoces las definiciones y la terminología de Ray en GKE.

Información general

En este tutorial se muestra cómo ejecutar una secuencia de comandos de Python con Jax que verifica que la inicialización de TPU Trillium con KubeRay se ha realizado correctamente. Jax es una biblioteca de computación numérica de alto rendimiento que admite cargas de trabajo de aprendizaje automático. KubeRay es un operador de Kubernetes que proporciona una forma unificada de desplegar, gestionar y monitorizar aplicaciones de Ray en Kubernetes.

Las TPUs Trillium (v6e) requieren variables de entorno y especificaciones de Pod concretas que difieren de las de las generaciones anteriores de TPUs. En este tutorial se proporcionan las configuraciones necesarias para desplegar correctamente una carga de trabajo con KubeRay en las TPUs de Trillium.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

Habilita la API de Google Kubernetes Engine.

Habilitar la API de Google Kubernetes Engine

Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la gcloud CLI, obtén la versión más reciente ejecutando gcloud components update.
Nota: Si ya tienes instalada la CLI gcloud, asegúrate de definir la compute/region propiedad. Si usas principalmente clústeres zonales, define compute/zone en su lugar. Si defines una ubicación predeterminada, puedes evitar errores en gcloud CLI como el siguiente: One of [--zone, --region] must be supplied: Please specify location. Es posible que tengas que especificar la ubicación en determinados comandos si la ubicación de tu clúster es diferente de la predeterminada que hayas definido.

Asegúrate de que tienes instalada la CLI de Ray (versión 2.37.0).

Activar Cloud Shell

Cloud Shell tiene preinstaladas las herramientas de línea de comandos gcloud, helm y kubectl que se usan en este tutorial.

Ve a la Google Cloud consola.
En la parte superior de la ventana de la consola Google Cloud , haz clic en el botón Activar Cloud Shell .

Se abrirá una sesión de Cloud Shell dentro de un nuevo marco en la consola y se mostrará en ella un mensaje de la línea de comandos. Google Cloud

Crear un clúster y un grupo de nodos de GKE

Puedes configurar KubeRay en TPUs en un clúster Autopilot o Estándar de GKE. Te recomendamos que uses un clúster de Autopilot para disfrutar de una experiencia de Kubernetes totalmente gestionada. Para elegir el modo de funcionamiento de GKE que mejor se adapte a tus cargas de trabajo, consulta Acerca de los modos de funcionamiento de GKE.

Autopilot

En Cloud Shell, ejecuta el siguiente comando:
```
gcloud container clusters create-auto CLUSTER_NAME \
    --enable-ray-operator \
    --release-channel=rapid \
    --location=LOCATION
```
Haz los cambios siguientes:
- CLUSTER_NAME: el nombre del nuevo clúster.
- LOCATION: la región en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta Disponibilidad de las TPU en GKE.
GKE crea un clúster de Autopilot con el complemento del operador de Ray habilitado. El complemento instala automáticamente el webhook de Ray TPU en el plano de control del clúster.

Para comunicarte con tu clúster, configura kubectl :

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

Estándar

En Cloud Shell, crea un clúster estándar que habilite el complemento del operador de Ray ejecutando el siguiente comando :
```
gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --addons=RayOperator \
  --cluster-version=1.33 \
  --machine-type=n1-standard-16
```
Haz los cambios siguientes:
- CLUSTER_NAME: el nombre del nuevo clúster.
- LOCATION: la región en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta Disponibilidad de las TPU en GKE.
La creación del clúster puede tardar varios minutos.

Para comunicarte con tu clúster, configura kubectl :

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

Puedes crear un grupo de nodos de TPU de un solo host o de varios hosts:

Un solo host

En Cloud Shell, ejecuta el siguiente 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

Varios hosts

En Cloud Shell, ejecuta el siguiente 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

Ejecutar un recurso personalizado RayJob

Al definir un manifiesto de RayJob, le indicas a KubeRay que haga lo siguiente:

Crea un RayCluster: la especificación de RayJob incluye un rayClusterSpec que define la configuración del clúster de Ray (grupos de encabezado y de trabajador) que quieras.
Ejecutar un trabajo específico: el campo entrypoint de RayJob especifica el comando o la secuencia de comandos que se va a ejecutar en el clúster de Ray creado. En este tutorial, entrypoint es una secuencia de comandos de Python (tpu_list_devices.py) diseñada para verificar la inicialización de la TPU Trillium.

Para crear un recurso personalizado RayJob, sigue estos pasos:

Un solo host

Crea el siguiente archivo de manifiesto 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

Aplica el archivo de manifiesto:

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

Verifica que se ha creado y se está ejecutando el RayJob:

kubectl get rayjobs v6e-4-job

El resultado debería ser similar al siguiente:

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

Imprime el resultado de RayJob.

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

El resultado debería ser similar al siguiente:

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

Varios hosts

Crea el siguiente archivo de manifiesto 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

Aplica el archivo de manifiesto:

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

Verifica que se haya creado y se esté ejecutando el RayJob v6e-16:

kubectl get rayjobs v6e-16-job

El resultado debería ser similar al siguiente:

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

Imprime el resultado de RayJob v6e-16:

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

El resultado debería ser similar al siguiente:

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 el RayJob en el panel de control de Ray

Verifica que GKE haya creado el servicio RayCluster y conéctate a la instancia de RayCluster.

Un solo host

Recupera el nombre del RayCluster generado para el RayJob:

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

Recupera el nombre del servicio principal de RayCluster:

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

Conéctate al panel de control de Ray reenviando el puerto del servicio principal:
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
Abre un navegador web e introduce la siguiente URL:
```
http://localhost:8265/#/jobs
```
Consulta el estado de RayJob y los registros pertinentes.

Varios hosts

Recupera el nombre del RayCluster generado para el RayJob:

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

Recupera el nombre del servicio principal de RayCluster:

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

Conéctate al panel de control de Ray reenviando el puerto del servicio principal:
```
kubectl port-forward svc/$HEAD_SVC 8265:8265 2>&1 >/dev/null &
```
Abre un navegador web e introduce la siguiente URL:
```
http://localhost:8265/#/jobs
```
Consulta el estado de RayJob y los registros pertinentes.

Ray define un recurso TPU-{accelerator}-Head para identificar el nodo de trabajador de Ray que corresponde al valor TPU_WORKER_ID=0. En el grupo de TPUs multihost, el nodo de Ray con TPU_WORKER_ID=0 tiene TPU-v6e-16-head: 1.0 en sus recursos. Esta variable de TPU_WORKER_IDentorno la define un webhook de GKE mutante para KubeRay.

Limpieza

Una vez que hayas completado el tutorial, elimina el RayJob para evitar que se te cobren cargos no deseados en tu cuenta:

Un solo host

kubectl delete rayjobs v6e-4-job

Varios hosts

kubectl delete rayjobs v6e-16-job

Siguientes pasos

Consulta información sobre Ray en Kubernetes.
Consulta cómo servir vLLM en GKE con TPUs.
Consulta cómo servir SDXL en GKE con TPUs.
Consulta más información sobre las TPUs en GKE.

Configurar KubeRay con TPU Trillium Organízate con las colecciones Guarda y clasifica el contenido según tus preferencias.

Información general

Antes de empezar

Activar Cloud Shell

Crear un clúster y un grupo de nodos de GKE

Autopilot

Estándar

Un solo host

Varios hosts

Ejecutar un recurso personalizado RayJob

Un solo host

Varios hosts

Ver el RayJob en el panel de control de Ray

Un solo host

Varios hosts

Limpieza

Un solo host

Varios hosts

Siguientes pasos

Configurar KubeRay con TPU Trillium