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.

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 generar una estimación de costos en función del uso previsto, usa la calculadora de precios. Es posible que los usuarios nuevos de Google Cloud califiquen para obtener una prueba gratuita.

Cuando finalices las tareas que se describen en este documento, puedes borrar los recursos que creaste para evitar que continúe la facturación. Para obtener más información, consulta Cómo realizar 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.

  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. To initialize the gcloud CLI, run the following command: 

    gcloud init

  4. Create or select a Google Cloud project.

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

  5. Make sure that billing is enabled for your Google Cloud project.

  6. Enable the GKE API: 

    gcloud services enable container.googleapis.com
  7. Install the Google Cloud CLI.

  8. To initialize the gcloud CLI, run the following command: 

    gcloud init

  9. Create or select a Google Cloud project.

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

  10. Make sure that billing is enabled for your Google Cloud project.

  11. Enable the GKE API: 

    gcloud services enable container.googleapis.com

  12. 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_IDENTIFIER" --role=ROLE
    • Replace PROJECT_ID with your project ID.

    • Replace USER_IDENTIFIER with the identifier for your user account. For example, user:myemail@example.com.

    • Replace ROLE with each individual role.

Asegúrate de que la cuota sea suficiente

Asegúrate de que tu proyecto de 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:

  1. Para iniciar una sesión de Cloud Shell desde la consola de Google Cloud, haz clic en Ícono de activación de Cloud ShellActivar Cloud Shell en la consola de Google Cloud. Esto inicia una sesión en el panel inferior de la consola de Google Cloud.

  2. 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: El ID del proyecto de Google Cloud.
    • CLUSTER_VERSION: la versión de GKE que se usará. Debe ser 1.30.1 o una versión 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

Crea un clúster con un grupo de nodos TPU

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

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

  2. 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 \
    --tpu-topology=2x2x1

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:

  1. 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:
            - pydantic<2
            - google-api-python-client
            - pillow
            - diffusers==0.7.2
            - transformers==4.24.0
            - flax
            - ml_dtypes==0.2.0
            - 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: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: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.

  2. Aplica el manifiesto al clúster:

    kubectl apply -f ray-service-tpu.yaml

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

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

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

  3. Haz clic en la pestaña Entregar.

Envía mensajes al servidor de modelos

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

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

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

    Imagen generada a través de difusión estable con 8 secciones: una silla verde, un hombre de pie frente a una casa, un robot en la calle, una familia sentada frente a una mesa, un médico caminando en un parque, un dragón volador, un retrato de osos al estilo japonés y una cascada.

Limpia

Borra el proyecto

    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?