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

Estándar

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.

Objetivos

Crear un clúster de GKE con un grupo de nodos TPU
Implementar un clúster de Ray con TPU
Implementa un recurso personalizado de RayService.
Interactuar con el servidor de modelos de difusión estable

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

Para obtener una estimación de costos en función del uso previsto, usa la calculadora de precios.

Es posible que los usuarios de Google Cloud nuevos cumplan con los requisitos para acceder a una prueba gratuita.

Cuando completes las tareas que se describen en este documento, podrás borrar los recursos que creaste para evitar que se te siga facturando. Para obtener más información, consulta Realiza una limpieza.

Antes de comenzar

Cloud Shell tiene preinstalado el software que necesitas para este instructivo, incluidos kubectl y gcloud CLI. Si no usas Cloud Shell, instala gcloud CLI.

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

Nota: Si instalaste gcloud CLI anteriormente, asegúrate de tener la versión más reciente ejecutando gcloud components update.

Si usas un proveedor de identidad externo (IdP), primero debes Acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Create a Google Cloud project:
```
gcloud projects create PROJECT_ID
```
Replace PROJECT_ID with a name for the Google Cloud project you are creating.
Select the Google Cloud project that you created:
```
gcloud config set project PROJECT_ID
```
Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

Nota: Si instalaste gcloud CLI anteriormente, asegúrate de tener la versión más reciente ejecutando gcloud components update.

Si usas un proveedor de identidad externo (IdP), primero debes Acceder a la gcloud CLI con tu identidad federada.

Para inicializar gcloud CLI, ejecuta el siguiente comando:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Create a Google Cloud project:
```
gcloud projects create PROJECT_ID
```
Replace PROJECT_ID with a name for the Google Cloud project you are creating.
Select the Google Cloud project that you created:
```
gcloud config set project PROJECT_ID
```
Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.clusterAdmin, roles/container.admin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
```
Replace the following:
- PROJECT_ID: Your project ID.
- USER_IDENTIFIER: The identifier for your user account. For example, myemail@example.com.
- ROLE: The IAM role that you grant to your user account.

Asegúrate de que la cuota sea suficiente

Asegúrate de que tu proyecto Google Cloud tenga suficiente cuota de TPU en tu región o zona de Compute Engine. Para obtener más información, consulta Asegúrate de tener cuotas suficientes de TPU y GKE en la documentación de Cloud TPU. Es posible que también debas aumentar tus cuotas para los siguientes elementos:

Persistent Disk SSD (GB)
Direcciones IP en uso

Prepare 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 Google Cloud consola.
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: Es el ID de tu proyecto. Google Cloud
- 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.

Limpia

Borra el proyecto

Precaución: Borrar un proyecto tiene las siguientes consecuencias:

Se borra todo su contenido. Si usaste un proyecto existente para las tareas de este documento, cuando lo borres, también se borrará cualquier otro trabajo que hayas realizado en el proyecto.
Se pierden los IDs personalizados de los proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID personalizado del proyecto que desees utilizar en el futuro. Para conservar las URLs que usan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto, en lugar de todo el proyecto.

Si planeas explorar varios instructivos, arquitecturas o guías de inicio rápido, reutilizar proyectos puede ser útil para no exceder los límites de cuota de los proyectos.

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

Borra los recursos individuales

Para borrar el clúster, escribe lo siguiente:

gcloud container clusters delete ${CLUSTER_NAME}

¿Qué sigue?

Obtén más información sobre Ray en Kubernetes.
Explora la documentación de KubeRay.
Explora arquitecturas de referencia, diagramas y prácticas recomendadas sobre Google Cloud. Consulta nuestro Cloud Architecture Center.

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.

Información acerca de Ray y Ray Serve

Acerca de las TPU

Acerca del webhook de inicialización de TPU de KubeRay

Objetivos

Costos

Antes de comenzar

Asegúrate de que la cuota sea suficiente

Prepare el entorno

Crea un clúster con un grupo de nodos TPU

Configura un recurso de RayCluster con TPU

Configura TPU nodeSelector

Configura un recurso de contenedor de TPU

Configura el grupo de trabajo numOfHosts

Crea un recurso personalizado de RayService

(Opcional) Visualiza el panel de Ray

Envía mensajes al servidor de modelos

Limpia

Borra el proyecto

Borra los recursos individuales

¿Qué sigue?

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

Configura TPU `nodeSelector`

Configura el grupo de trabajo `numOfHosts`