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.

Objetivos

  • Crea un clúster de GKE con un grupo de nodos de TPU.
  • Despliega un clúster de Ray con TPUs.
  • Despliega un recurso personalizado RayService.
  • Interactúa con el servidor del modelo Stable Diffusion.

Costes

En este documento, se utilizan los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costes basada en el uso previsto, utiliza la calculadora de precios.

Los usuarios nuevos Google Cloud pueden disfrutar de una prueba gratuita.

Cuando termines las tareas que se describen en este documento, puedes evitar que se te siga facturando eliminando los recursos que has creado. Para obtener más información, consulta la sección Limpiar.

Antes de empezar

Cloud Shell tiene preinstalado el software que necesitas para este tutorial, incluido kubectl y la CLI de gcloud. Si no usas Cloud Shell, instala la CLI de gcloud.

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

  2. Install the Google Cloud CLI.

  3. Si utilizas un proveedor de identidades (IdP) externo, primero debes iniciar sesión en la CLI de gcloud con tu identidad federada.

  4. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  5. 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 (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.

  6. Verify that billing is enabled for your Google Cloud project.

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

  8. Install the Google Cloud CLI.

  9. Si utilizas un proveedor de identidades (IdP) externo, primero debes iniciar sesión en la CLI de gcloud con tu identidad federada.

  10. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  11. 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 (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.

  12. Verify that billing is enabled for your Google Cloud project.

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

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

    15. Asegúrate de que haya suficiente cuota

    Comprueba que tu proyecto Google Cloud tiene suficiente cuota de TPU en tu región o zona de Compute Engine. Para obtener más información, consulta Asegurarse de que se dispone de cuotas suficientes de TPU y GKE en la documentación de Cloud TPU. También es posible que tengas que aumentar tus cuotas de lo siguiente:

    • Disco persistente SSD (GB)
    • Direcciones IP en uso

    Prepara tu entorno

    Para preparar tu entorno, sigue estos pasos:

    1. Inicia una sesión de Cloud Shell desde la Google Cloud consolaIcono de activación de Cloud Shell 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 .

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

    3. Clona el repositorio de GitHub:

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

    4. 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:

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

    2. 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 una 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 una porción 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 configurando los recursos google.com/tpulimits 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:

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

    2. Aplica el manifiesto a tu clúster:

      kubectl apply -f ray-service-tpu.yaml

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

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

    2. En un navegador web, ve a http://localhost:8265/.

    3. Haz clic en la pestaña Servir.

    Enviar peticiones al servidor del modelo

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

    2. Abre una nueva sesión de Cloud Shell.

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

      Imagen generada por Stable Diffusion con 8 secciones: una silla verde, un hombre de pie fuera de una casa, un robot en la calle, una familia sentada en una mesa, un médico caminando por un parque, un dragón volando, un retrato de osos de estilo japonés y una cascada.

    Limpieza

    Eliminar el proyecto

      Delete a Google Cloud project:

      gcloud projects delete PROJECT_ID

    Eliminar recursos concretos

    Para eliminar el clúster, escribe lo siguiente:

    gcloud container clusters delete ${CLUSTER_NAME}

    Siguientes pasos