Configura KubeRay con la TPU Trillium

Autopilot estándar

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

Este instructivo está dirigido a los administradores y operadores de plataformas y a los especialistas en IA y datos que desean aprender a configurar la inicialización de TPU Trillium con KubeRay para grupos de nodos de host único y varios hosts. En este instructivo, se muestra cómo ejecutar un script con Jax que verifica la inicialización correcta de la TPU. En este instructivo, no se implementa ningún modelo.

Antes de configurar KubeRay en GKE, asegúrate de estar familiarizado con las definiciones y la terminología de Ray en GKE.

Descripción general

En este instructivo, 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 haya realizado correctamente. JAX es una biblioteca de procesamiento numérico de alto rendimiento que admite cargas de trabajo de aprendizaje automático. KubeRay es un operador de Kubernetes que proporciona una forma unificada de implementar, administrar y supervisar aplicaciones de Ray en Kubernetes.

Las TPU Trillium (v6e) requieren variables de entorno y especificaciones de Pod específicas que difieren de las generaciones anteriores de TPU. En este instructivo, se proporcionan las configuraciones necesarias para implementar correctamente una carga de trabajo con KubeRay en las TPU de Trillium.

Antes de comenzar

Antes de comenzar, asegúrate de haber 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, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta gcloud components update para obtener la versión más reciente.
  Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar la propiedad compute/region. Si usas principalmente clústeres zonales, configura compute/zone en su lugar. Cuando configuras 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 debas especificar la ubicación en ciertos comandos si la ubicación de tu clúster difiere de la predeterminada que estableciste.

• Asegúrate de tener instalada la CLI de Ray (versión 2.37.0).

Activar Cloud Shell

Cloud Shell viene preinstalado con las herramientas de línea de comandos de gcloud, helm y kubectl que se usan en este instructivo.

1. Ve a la consola deGoogle Cloud .

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

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

   

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

Puedes configurar KubeRay en TPU en un clúster de GKE Autopilot o Standard. Te recomendamos que uses un clúster de Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Acerca de los modos de operación de GKE.

Autopilot

1. En Cloud Shell, ejecute el siguiente comando:

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

   Reemplaza lo siguiente:

   • CLUSTER_NAME es el nombre del clúster nuevo.
   • LOCATION: Es la región en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta la disponibilidad de TPU en GKE.

   GKE crea un clúster de Autopilot con el complemento Ray Operator habilitado. El complemento instala automáticamente el webhook de TPU de Ray en el plano de control del clúster.

2. Para comunicarte con tu clúster, configura kubectl :

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

Estándar

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

   Reemplaza lo siguiente:

   • CLUSTER_NAME es el nombre del clúster nuevo.
   • LOCATION: Es la región en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta Disponibilidad de TPU en GKE.

   La creación del clúster puede tomar varios minutos.

2. Para comunicarte con tu clúster, configura kubectl :

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

3. Puedes crear un grupo de nodos de porción de TPU de host único o de varios hosts:

Host único

En Cloud Shell, ejecute 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, ejecute 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

Ejecuta un recurso personalizado de RayJob

Cuando defines 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 instancias principales y de trabajadores) que deseas.
• Ejecuta un trabajo específico: El campo entrypoint dentro de RayJob especifica el comando o la secuencia de comandos que se ejecutará dentro del clúster de Ray creado. En este instructivo, 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 de RayJob, completa los siguientes pasos:

Host único

1. Crea el siguiente 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

2. Aplica el manifiesto

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

3. Verifica que el RayJob se haya creado y esté en ejecución:

   kubectl get rayjobs v6e-4-job
   

   El resultado es similar a este:

   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. Imprime el resultado de RayJob.

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

   El resultado es similar a este:

   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

1. Crea el siguiente 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

2. Aplica el manifiesto

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

3. Verifica que el RayJob v6e-16 se haya creado y esté en ejecución:

   kubectl get rayjobs v6e-16-job
   

   El resultado es similar a este:

   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. Imprime el resultado del RayJob v6e-16:

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

   El resultado es similar a este:

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

Visualiza el RayJob en el panel de Ray

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

Host único

1. Recupera el nombre del RayCluster generado para el RayJob:

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

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

3. Conéctate al panel de Ray a través de la redirección de puertos del servicio principal:

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

4. Abre un navegador web y, luego, ingresa la siguiente URL:

   http://localhost:8265/#/jobs
   

5. Visualiza el estado de RayJob y los registros pertinentes.

Varios hosts

1. Recupera el nombre del RayCluster generado para el RayJob:

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

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

3. Conéctate al panel de Ray a través de la redirección de puertos del servicio principal:

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

4. Abre un navegador web y, luego, ingresa la siguiente URL:

   http://localhost:8265/#/jobs
   

5. Visualiza el estado de RayJob y los registros pertinentes.

Ray establece un recurso TPU-{accelerator}-Head para identificar el nodo trabajador de Ray que corresponde al valor TPU_WORKER_ID=0. En el grupo de TPU de varios hosts, el nodo de Ray con TPU_WORKER_ID=0 tiene TPU-v6e-16-head: 1.0 establecido en sus recursos. Esta variable de entorno TPU_WORKER_ID se configura mediante un webhook de GKE de mutación para KubeRay.

Éxito: Ejecutaste correctamente un RayJob en un grupo de nodos de segmentación de TPU que ejecuta TPU Trillium. La finalización de RayJob indica que Jax se inicializó correctamente y que los nodos de TPU distribuidos están habilitados de forma correcta.

Limpia

Después de completar el instructivo, borra el RayJob para evitar cargos no deseados en tu cuenta:

Host único

kubectl delete rayjobs v6e-4-job

Varios hosts

kubectl delete rayjobs v6e-16-job

¿Qué sigue?

• Obtén más información sobre Ray en Kubernetes.
• Obtén información para entregar vLLM en GKE con TPUs.
• Obtén información para entregar SDXL en GKE con TPU.
• Obtén más información sobre las TPU en GKE.