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

Desplegar una aplicación de Ray Serve con un modelo de Stable Diffusion en Google Kubernetes Engine (GKE) con TPUs

En esta guía se muestra cómo desplegar y servir un modelo de Stable Diffusion en Google Kubernetes Engine (GKE) mediante TPUs, Ray Serve y el complemento Ray Operator.

Esta guía está dirigida a clientes de IA generativa, usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en usar las funciones de orquestación de contenedores de Kubernetes para servir modelos con Ray.

Acerca de Ray y Ray Serve

Ray es un framework de computación escalable de código abierto para aplicaciones de IA y aprendizaje automático. Ray Serve es una biblioteca de servicio de modelos para Ray que se usa para escalar y servir modelos en un entorno distribuido. Para obtener más información, consulta Ray Serve en la documentación de Ray.

Acerca de las TPUs

Las unidades de procesamiento de tensor (TPUs) son aceleradores por hardware especializados diseñados para agilizar significativamente el entrenamiento y la inferencia de modelos de aprendizaje automático a gran escala. Usar Ray con TPUs te permite escalar sin problemas aplicaciones de aprendizaje automático de alto rendimiento. Para obtener más información sobre las TPUs, consulta el artículo Introducción a las TPUs de Cloud en la documentación de las TPUs de Cloud.

Acerca del webhook de inicialización de TPU de KubeRay

Como parte del complemento Ray Operator, GKE proporciona webhooks de validación y mutación que gestionan la programación de pods de TPU y determinadas variables de entorno de TPU que necesitan frameworks como JAX para la inicialización de contenedores. El webhook de KubeRay TPU muta los pods con la etiqueta app.kubernetes.io/name: kuberay que solicitan TPUs con las siguientes propiedades:

TPU_WORKER_ID: un número entero único para cada pod de trabajador del segmento de TPU.
TPU_WORKER_HOSTNAMES: lista de nombres de host de DNS de todos los trabajadores de TPU que necesitan comunicarse entre sí en el segmento. Esta variable solo se inserta en los pods de TPU de un grupo de varios hosts.
replicaIndex: etiqueta de pod que contiene un identificador único de la réplica del grupo de trabajo al que pertenece el pod. Esto resulta útil para los grupos de trabajadores de varios hosts, en los que varios pods de trabajadores pueden pertenecer a la misma réplica. Ray lo usa para habilitar el escalado automático de varios hosts.
TPU_NAME: cadena que representa el PodSlice de TPU de GKE al que pertenece este pod. Tiene el mismo valor que la etiqueta replicaIndex.
podAffinity: asegura que GKE programe pods de TPU con etiquetas replicaIndexcoincidentes en el mismo grupo de nodos. De esta forma, GKE puede escalar las TPUs de varios hosts de forma atómica por grupos de nodos, en lugar de por nodos individuales.

Prepara tu entorno

Para preparar tu entorno, sigue estos pasos:

Inicia una sesión de Cloud Shell desde la Google Cloud consola haciendo clic en Activar Cloud Shell en la Google Cloud consola. Se iniciará una sesión en el panel inferior de la consola Google Cloud .
Define las variables de entorno:
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=ray-cluster
export COMPUTE_REGION=us-central2-b
export CLUSTER_VERSION=CLUSTER_VERSION
```
Haz los cambios siguientes:
- PROJECT_ID: tu Google Cloud ID de proyecto.
- CLUSTER_VERSION: la versión de GKE que se va a usar. Debe ser 1.30.1 o posterior.

Clona el repositorio de GitHub:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples

Cambia al directorio de trabajo:

cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion

Crear un clúster con un grupo de nodos de TPU

Crea un clúster de GKE Standard con un grupo de nodos de TPU:

Crea un clúster en modo Estándar con el operador de Ray habilitado:

gcloud container clusters create ${CLUSTER_NAME} \
    --addons=RayOperator \
    --machine-type=n1-standard-8 \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

Crea un grupo de nodos de TPU de un solo host:

gcloud container node-pools create tpu-pool \
    --location=${COMPUTE_REGION} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct4p-hightpu-4t \
    --num-nodes=1

Para usar las TPUs con el modo Estándar, debes seleccionar lo siguiente:

Una ubicación de Compute Engine con capacidad para aceleradores de TPU
Un tipo de máquina compatible con la TPU
La topología física del TPU PodSlice

Configurar un recurso de RayCluster con TPUs

Configura el manifiesto de RayCluster para preparar tu carga de trabajo de TPU:

Configurar la TPU `nodeSelector`

GKE usa nodeSelectors de Kubernetes para asegurarse de que las cargas de trabajo de TPU se programan en la topología y el acelerador de TPU adecuados. Para obtener más información sobre cómo seleccionar nodeSelectors de TPU, consulta Desplegar cargas de trabajo de TPU en GKE Standard.

Actualiza el manifiesto de ray-cluster.yaml para programar tu pod en un segmento de pod de TPU v4 con una topología 2x2x1:

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

Configurar un recurso de contenedor de TPU

Para usar un acelerador de TPU, debes especificar el número de chips de TPU que GKE debe asignar a cada pod. Para ello, configura los campos google.com/tpuresource limits y requests en el campo del contenedor de TPU de tu manifiesto de RayCluster workerGroupSpecs.

Actualiza el manifiesto de ray-cluster.yaml con los límites y las solicitudes de recursos:

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

Configurar el grupo de trabajadores `numOfHosts`

KubeRay 1.1.0 añade un campo numOfHosts al recurso personalizado RayCluster, que especifica el número de hosts de TPU que se van a crear por réplica de grupo de trabajadores. En los grupos de trabajadores multihost, las réplicas se tratan como PodSlices en lugar de como trabajadores individuales, y se crean numOfHosts nodos de trabajador por réplica.

Actualiza el archivo de manifiesto ray-cluster.yaml con lo siguiente:

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

Crear un recurso personalizado RayService

Crea un recurso personalizado RayService:

Revisa el siguiente archivo de manifiesto:

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: stable-diffusion-tpu
spec:
  serveConfigV2: |
    applications:
      - name: stable_diffusion
        import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
        runtime_env:
          working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
          pip:
            - diffusers==0.7.2
            - flax
            - jax[tpu]==0.4.11
            - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            - fastapi
  rayClusterConfig:
    rayVersion: '2.9.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-head
            image: rayproject/ray-ml:2.9.0-py310
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "2"
                memory: "8G"
              requests:
                cpu: "2"
                memory: "8G"
    workerGroupSpecs:
    - replicas: 1
      minReplicas: 1
      maxReplicas: 10
      numOfHosts: 1
      groupName: tpu-group
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-worker
            image: rayproject/ray-ml:2.9.0-py310
            resources:
              limits:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
              requests:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x1

Este manifiesto describe un recurso personalizado RayService que crea un recurso RayCluster con 1 nodo principal y un grupo de trabajadores de TPU con una topología de 2x2x1, lo que significa que cada nodo de trabajador tendrá 4 chips de TPU v4.

El nodo de TPU pertenece a un único segmento de pod de TPU v4 con una topología de 2x2x1. Para crear un grupo de trabajadores multihost, sustituye los valores de gke-tpu nodeSelector, los límites y las solicitudes de google.com/tpu del contenedor y los valores de numOfHosts por tu configuración multihost. Para obtener más información sobre las topologías multihost de TPU, consulta Arquitectura del sistema en la documentación de TPU de Cloud.

Aplica el manifiesto a tu clúster:
```
kubectl apply -f ray-service-tpu.yaml
```
Comprueba que el recurso RayService se está ejecutando:
```
kubectl get rayservices
```
El resultado debería ser similar al siguiente:
```
NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2
```
En este resultado, Running en la columna SERVICE STATUS indica que el recurso RayService está listo.

(Opcional) Ver el panel de control de Ray

Puedes ver tu implementación de Ray Serve y los registros pertinentes en el panel de control de Ray.

Establece una sesión de reenvío de puertos al panel de control de Ray desde el servicio principal de Ray:
```
kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
```
En un navegador web, ve a http://localhost:8265/.
Haz clic en la pestaña Servir.

Enviar peticiones al servidor del modelo

Establece una sesión de reenvío de puertos al endpoint de Serve desde el servicio principal de Ray:
```
kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
```
Abre una nueva sesión de Cloud Shell.
Envía una petición de texto a imagen al servidor del modelo Stable Diffusion:
```
python stable_diffusion_tpu_req.py  --save_pictures
```
Los resultados de la inferencia de difusión estable se guardan en un archivo llamado diffusion_results.png.

Desplegar una aplicación de Ray Serve con un modelo de Stable Diffusion en Google Kubernetes Engine (GKE) con TPUs Organízate con las colecciones Guarda y clasifica el contenido según tus preferencias.