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

En esta guía se muestra un ejemplo de cómo desplegar y servir un modelo de Stable Diffusion en Google Kubernetes Engine (GKE) con Ray Serve y el complemento Ray Operator como implementación de ejemplo.

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.

Puedes usar un recurso RayCluster o RayService para desplegar tus aplicaciones Ray Serve. Deberías usar un recurso RayService en producción por los siguientes motivos:

  • Actualizaciones in situ de aplicaciones RayService
  • Actualización sin tiempo de inactividad de los recursos de RayCluster
  • Aplicaciones Ray Serve de alta disponibilidad

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=rayserve-cluster
    export COMPUTE_REGION=us-central1
    export COMPUTE_ZONE=us-central1-c
    export CLUSTER_VERSION=CLUSTER_VERSION
    export TUTORIAL_HOME=`pwd`
    

    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
    
  5. Crea un entorno virtual de Python:

    venv

    python -m venv myenv && \
    source myenv/bin/activate
    

    Conda

    1. Instala Conda.

    2. Ejecuta estos comandos:

      conda create -c conda-forge python=3.9.19 -n myenv && \
      conda activate myenv
      

    Cuando despliegas una aplicación de Serve con serve run, Ray espera que la versión de Python del cliente local coincida con la versión utilizada en el clúster de Ray. La imagen rayproject/ray:2.37.0 usa Python 3.9. Si usas otra versión del cliente, selecciona la imagen de Ray correspondiente.

  6. Instala las dependencias necesarias para ejecutar la aplicación Serve:

    pip install ray[serve]==2.37.0
    pip install torch
    pip install requests
    

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

Crea un clúster de GKE Autopilot o Estándar con un grupo de nodos de GPU:

Autopilot

Crea un clúster de Autopilot:

gcloud container clusters create-auto ${CLUSTER_NAME}  \
    --enable-ray-operator \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

Estándar

  1. Crea un clúster estándar:

    gcloud container clusters create ${CLUSTER_NAME} \
        --addons=RayOperator \
        --cluster-version=${CLUSTER_VERSION}  \
        --machine-type=c3d-standard-8 \
        --location=${COMPUTE_ZONE} \
        --num-nodes=1
    
  2. Crea un grupo de nodos de GPU:

    gcloud container node-pools create gpu-pool \
        --cluster=${CLUSTER_NAME} \
        --machine-type=g2-standard-8 \
        --location=${COMPUTE_ZONE} \
        --num-nodes=1 \
        --accelerator type=nvidia-l4,count=1,gpu-driver-version=latest
    

Desplegar un recurso RayCluster

Para desplegar un recurso RayCluster, sigue estos pasos:

  1. Revisa el siguiente archivo de manifiesto:

    apiVersion: ray.io/v1
    kind: RayCluster
    metadata:
      name: stable-diffusion-cluster
    spec:
      rayVersion: '2.37.0'
      headGroupSpec:
        rayStartParams:
          dashboard-host: '0.0.0.0'
        template:
          metadata:
          spec:
            containers:
            - name: ray-head
              image: rayproject/ray:2.37.0
              ports:
              - containerPort: 6379
                name: gcs
              - containerPort: 8265
                name: dashboard
              - containerPort: 10001
                name: client
              - containerPort: 8000
                name: serve
              resources:
                limits:
                  cpu: "2"
                  ephemeral-storage: "15Gi"
                  memory: "8Gi"
                requests:
                  cpu: "2"
                  ephemeral-storage: "15Gi"
                  memory: "8Gi"
            nodeSelector:
              cloud.google.com/machine-family: c3d
      workerGroupSpecs:
      - replicas: 1
        minReplicas: 1
        maxReplicas: 4
        groupName: gpu-group
        rayStartParams: {}
        template:
          spec:
            containers:
            - name: ray-worker
              image: rayproject/ray:2.37.0-gpu
              resources:
                limits:
                  cpu: 4
                  memory: "16Gi"
                  nvidia.com/gpu: 1
                requests:
                  cpu: 3
                  memory: "16Gi"
                  nvidia.com/gpu: 1
            nodeSelector:
              cloud.google.com/gke-accelerator: nvidia-l4

    Este manifiesto describe un recurso RayCluster.

  2. Aplica el manifiesto a tu clúster:

    kubectl apply -f ray-cluster.yaml
    
  3. Verifica que el recurso RayCluster esté listo:

    kubectl get raycluster
    

    El resultado debería ser similar al siguiente:

    NAME                       DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY   GPUS   STATUS   AGE
    stable-diffusion-cluster   2                 2                   6      20Gi     0      ready    33s
    

    En este resultado, ready en la columna STATUS indica que el recurso RayCluster está listo.

Conectarse al recurso RayCluster

Para conectarte al recurso RayCluster, haz lo siguiente:

  1. Verifica que GKE haya creado el servicio RayCluster:

    kubectl get svc stable-diffusion-cluster-head-svc
    

    El resultado debería ser similar al siguiente:

    NAME                             TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)                                AGE
    pytorch-mnist-cluster-head-svc   ClusterIP   34.118.238.247   <none>        10001/TCP,8265/TCP,6379/TCP,8080/TCP   109s
    
  2. Establece sesiones de reenvío de puertos al nodo principal de Ray:

    kubectl port-forward svc/stable-diffusion-cluster-head-svc 8265:8265 2>&1 >/dev/null &
    kubectl port-forward svc/stable-diffusion-cluster-head-svc 10001:10001 2>&1 >/dev/null &
    
  3. Verifica que el cliente de Ray puede conectarse al clúster de Ray mediante localhost:

    ray list nodes --address http://localhost:8265
    

    El resultado debería ser similar al siguiente:

    ======== List: 2024-06-19 15:15:15.707336 ========
    Stats:
    ------------------------------
    Total: 3
    
    Table:
    ------------------------------
        NODE_ID                                                   NODE_IP     IS_HEAD_NODE    STATE    NODE_NAME    RESOURCES_TOTAL                 LABELS
    0  1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2  10.28.1.21  False           ALIVE    10.28.1.21   CPU: 2.0                        ray.io/node_id: 1d07447d7d124db641052a3443ed882f913510dbe866719ac36667d2
    # Several lines of output omitted
    

Ejecutar una aplicación de Ray Serve

Para ejecutar una aplicación de Ray Serve, sigue estos pasos:

  1. Ejecuta la aplicación Stable Diffusion Ray Serve:

    serve run stable_diffusion:entrypoint --working-dir=. --runtime-env-json='{"pip": ["torch", "torchvision", "diffusers==0.12.1", "huggingface_hub==0.25.2", "transformers", "fastapi==0.113.0"], "excludes": ["myenv"]}' --address ray://localhost:10001
    
    

    El resultado debería ser similar al siguiente:

    2024-06-19 18:20:58,444 INFO scripts.py:499 -- Running import path: 'stable_diffusion:entrypoint'.
    2024-06-19 18:20:59,730 INFO packaging.py:530 -- Creating a file package for local directory '.'.
    2024-06-19 18:21:04,833 INFO handle.py:126 -- Created DeploymentHandle 'hyil6u9f' for Deployment(name='StableDiffusionV2', app='default').
    2024-06-19 18:21:04,834 INFO handle.py:126 -- Created DeploymentHandle 'xo25rl4k' for Deployment(name='StableDiffusionV2', app='default').
    2024-06-19 18:21:04,836 INFO handle.py:126 -- Created DeploymentHandle '57x9u4fp' for Deployment(name='APIIngress', app='default').
    2024-06-19 18:21:04,836 INFO handle.py:126 -- Created DeploymentHandle 'xr6kt85t' for Deployment(name='StableDiffusionV2', app='default').
    2024-06-19 18:21:04,836 INFO handle.py:126 -- Created DeploymentHandle 'g54qagbz' for Deployment(name='APIIngress', app='default').
    2024-06-19 18:21:19,139 INFO handle.py:126 -- Created DeploymentHandle 'iwuz00mv' for Deployment(name='APIIngress', app='default').
    2024-06-19 18:21:19,139 INFO api.py:583 -- Deployed app 'default' successfully.
    
  2. Establece una sesión de redirección de puertos al puerto de Ray Serve (8000):

    kubectl port-forward svc/stable-diffusion-cluster-head-svc 8000:8000 2>&1 >/dev/null &
    
  3. Ejecuta la secuencia de comandos de Python:

    python generate_image.py
    

    La secuencia de comandos genera una imagen en un archivo llamado output.png. La imagen es similar a la siguiente:

    Una playa al atardecer. Imagen generada por Stable Diffusion.

Desplegar un RayService

El recurso personalizado RayService gestiona el ciclo de vida de un recurso RayCluster y de una aplicación Ray Serve.

Para obtener más información sobre RayService, consulta los artículos Deploy Ray Serve Applications (Desplegar aplicaciones Ray Serve) y Production Guide (Guía de producción) en la documentación de Ray.

Para implementar un recurso RayService, sigue estos pasos:

  1. Revisa el siguiente archivo de manifiesto:

    apiVersion: ray.io/v1
    kind: RayService
    metadata:
      name: stable-diffusion
    spec:
      serveConfigV2: |
        applications:
          - name: stable_diffusion
            import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion:entrypoint
            runtime_env:
              working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
              pip: ["diffusers==0.12.1", "torch", "torchvision", "huggingface_hub==0.25.2", "transformers"]
      rayClusterConfig:
        rayVersion: '2.37.0'
        headGroupSpec:
          rayStartParams:
            dashboard-host: '0.0.0.0'
          template:
            spec:
              containers:
              - name: ray-head
                image:  rayproject/ray:2.37.0
                ports:
                - containerPort: 6379
                  name: gcs
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
                - containerPort: 8000
                  name: serve
                resources:
                  limits:
                    cpu: "2"
                    ephemeral-storage: "15Gi"
                    memory: "8Gi"
                  requests:
                    cpu: "2"
                    ephemeral-storage: "15Gi"
                    memory: "8Gi"
              nodeSelector:
                cloud.google.com/machine-family: c3d
        workerGroupSpecs:
        - replicas: 1
          minReplicas: 1
          maxReplicas: 4
          groupName: gpu-group
          rayStartParams: {}
          template:
            spec:
              containers:
              - name: ray-worker
                image: rayproject/ray:2.37.0-gpu
                resources:
                  limits:
                    cpu: 4
                    memory: "16Gi"
                    nvidia.com/gpu: 1
                  requests:
                    cpu: 3
                    memory: "16Gi"
                    nvidia.com/gpu: 1
              nodeSelector:
                cloud.google.com/gke-accelerator: nvidia-l4

    Este manifiesto describe un recurso personalizado de RayService.

  2. Aplica el manifiesto a tu clúster:

    kubectl apply -f ray-service.yaml
    
  3. Verifica que el servicio esté listo:

    kubectl get svc stable-diffusion-serve-svc
    

    El resultado debería ser similar al siguiente:

    NAME                         TYPE        CLUSTER-IP     EXTERNAL-IP   PORT(S)    AGE
    
    stable-diffusion-serve-svc   ClusterIP   34.118.236.0   <none>        8000/TCP   31m
    
  4. Configura la redirección de puertos al servicio Ray Serve:

    kubectl port-forward svc/stable-diffusion-serve-svc 8000:8000 2>&1 >/dev/null &
    
  5. Ejecuta la secuencia de comandos de Python de la sección anterior:

    python generate_image.py
    

    La secuencia de comandos genera una imagen similar a la imagen generada en la sección anterior.