Se usó la API de Cloud Translation para traducir esta página.

Implementa una aplicación de Ray Serve con un modelo de difusión estable en Google Kubernetes Engine (GKE) con TPU

En esta guía, se muestra cómo implementar y entregar un modelo de difusión estable en Google Kubernetes Engine (GKE) con TPU, Ray Serve y el complemento Ray Operator.

Esta guía está dirigida a clientes de IA generativa, usuarios nuevos o existentes de GKE, ingenieros de AA, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en el uso de las capacidades de organización de contenedores de Kubernetes para entregar modelos con Ray.

Información acerca de Ray y Ray Serve

Ray es un framework de procesamiento escalable de código abierto para aplicaciones de IA/AA. Ray Serve es una biblioteca de entrega de modelos para Ray que se usa para el escalamiento y la entrega de modelos en un entorno distribuido. Para obtener más información, consulta Ray Serve en la documentación de Ray.

Acerca de las TPU

Las unidades de procesamiento tensorial (TPU) son aceleradores de hardware especializados diseñados para acelerar significativamente el entrenamiento y la inferencia de los modelos de aprendizaje automático a gran escala. El uso de Ray con TPU te permite escalar sin problemas aplicaciones de AA de alto rendimiento. Para obtener más información sobre las TPU, consulta Introducción a Cloud TPU en la documentación de Cloud TPU.

Acerca del webhook de inicialización de TPU de KubeRay

Como parte del complemento de operador de Ray, GKE proporciona validación y mutaciónwebhooks que manejan la programación de Pods de TPU y ciertas variables de entorno de TPU que requieren los frameworks comoJAX para la inicialización del contenedor. El webhook de TPU de KubeRay muta los Pods con la etiqueta app.kubernetes.io/name: kuberay que solicita TPUs con las siguientes propiedades:

TPU_WORKER_ID: Es un número entero único para cada Pod de trabajo en la porción de TPU.
TPU_WORKER_HOSTNAMES: Es una lista de nombres de host de DNS para todos los trabajadores de TPU que necesitan comunicarse entre sí dentro de la porción. Esta variable solo se inserta para los grupos de TPU en un grupo de varios hosts.
replicaIndex: Es una etiqueta de Pod que contiene un identificador único para la réplica del grupo de trabajadores a la que pertenece el Pod. Esto es útil para grupos de trabajadores de varios hosts, en los que varios Pods de trabajadores pueden pertenecer a la misma réplica, y Ray lo usa para habilitar el ajuste de escala automático de varios hosts.
TPU_NAME: Es una cadena que representa el Podslice de TPU de GKE al que pertenece este Pod, configurado con el mismo valor que la etiqueta replicaIndex.
podAffinity: Garantiza que GKE programe Pods de TPU con etiquetas replicaIndex coincidentes en el mismo grupo de nodos. Esto permite que GKE escala las TPU de varios hosts de forma atómica por grupos de nodos, en lugar de nodos individuales.

Prepara el entorno

Para preparar tu entorno, sigue estos pasos:

Para iniciar una sesión de Cloud Shell desde la Google Cloud consola, haz clic en Activar Cloud Shell en la Google Cloud consola. Esto inicia una sesión en el panel inferior de la consola de Google Cloud .
Establece 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
```
Reemplaza lo siguiente:
- PROJECT_ID: Tu Google Cloud ID del proyecto.
- CLUSTER_VERSION: la versión de GKE que se usará. Debe ser 1.30.1 o una versión 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

Crea un clúster con un grupo de nodos TPU

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

Crea un clúster en modo Standard 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 TPU de host único:

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

Para usar TPUs con el modo Standard, debes seleccionar lo siguiente:

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

Configura un recurso de RayCluster con TPU

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

Configura TPU `nodeSelector`

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

Actualiza el manifiesto ray-cluster.yaml para programar tu Pod en un Podslice 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

Configura un recurso de contenedor de TPU

Para usar un acelerador de TPU, debes especificar la cantidad de chips TPU que GKE debe asignar a cada Pod utilizando la configuración delgoogle.com/tpu recurso limits y requests en el campo del contenedor de TPU de tu manifiesto de RayCluster workerGroupSpecs.

Actualiza el manifiesto ray-cluster.yaml con solicitudes y límites 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"

Configura el grupo de trabajo `numOfHosts`

KubeRay v1.1.0 agrega un campo numOfHosts al recurso personalizado de RayCluster, que especifica la cantidad de hosts de TPU que se crearán por réplica de grupo de trabajadores. En el caso de los grupos de trabajadores de varios hosts, las réplicas se tratan como Podslices en lugar de trabajadores individuales, y se crean numOfHosts nodos trabajador por réplica.

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

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

Crea un recurso personalizado de RayService

Crea un recurso personalizado RayService:

Revisa el siguiente 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

En este manifiesto, se describe un recurso personalizado de 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 trabajador tendrá 4 chips TPU v4.

El nodo TPU pertenece a un solo Podslice de TPU v4 con una topología 2x2x1. Para crear un grupo de trabajo de varios hosts, reemplaza los valores de gke-tpu nodeSelector, los límites y las solicitudes de contenedores de google.com/tpu y los valores de numOfHosts con tu configuración de varios hosts. Para obtener más información sobre las topologías de varios hosts de TPU, consulta Arquitectura del sistema en la documentación de Cloud TPU.

Aplica el manifiesto al clúster:
```
kubectl apply -f ray-service-tpu.yaml
```
Verifica que el recurso de RayService esté en ejecución:
```
kubectl get rayservices
```
El resultado es similar a este:
```
NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2
```
En este resultado, Running en la columna SERVICE STATUS indica que el recurso de RayCluster está listo.

(Opcional) Visualiza el panel de Ray

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

Establece una sesión de redirección de puertos en el panel 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 Entregar.

Envía mensajes al servidor de modelos

Establece una sesión de redirección de puertos al extremo 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 un mensaje de texto a imagen al servidor del modelo de difusión estable:
```
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.

Implementa una aplicación de Ray Serve con un modelo de difusión estable en Google Kubernetes Engine (GKE) con TPU Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.