Desplegar modelos abiertos con un contenedor vLLM personalizado

Aunque las distintas opciones de servicio de modelos de Vertex AI son suficientes para muchos casos prácticos, es posible que tengas que usar tus propias imágenes de contenedor para ofrecer modelos en Vertex AI. En este documento se describe cómo usar una imagen de contenedor personalizada de vLLM para ofrecer modelos en Vertex AI en CPUs, GPUs o TPUs. Para obtener más información sobre los modelos compatibles con vLLM, consulta la documentación de vLLM.

El servidor de la API vLLM implementa el protocolo de la API de OpenAI, pero no admite los requisitos de solicitud y respuesta de Vertex AI. Por lo tanto, debes usar una Raw Inference Request de Vertex AI para obtener inferencias de modelos desplegados en Vertex AI mediante un Prediction Endpoint. Para obtener más información sobre el método Raw Prediction del SDK de Vertex AI Python, consulta la documentación del SDK de Python.

Puedes obtener modelos de Hugging Face y de Cloud Storage. Este enfoque ofrece flexibilidad, lo que te permite aprovechar el centro de modelos basado en la comunidad (Hugging Face) y las funciones optimizadas de transferencia de datos y seguridad de Cloud Storage para la gestión interna de modelos o versiones ajustadas.

vLLM descarga los modelos de Hugging Face si se proporciona un token de acceso de Hugging Face. De lo contrario, vLLM asume que el modelo está disponible en el disco local. La imagen de contenedor personalizada permite que Vertex AI descargue el modelo deGoogle Cloud además de Hugging Face.

Antes de empezar

  1. En tu Google Cloud proyecto, habilita las APIs Vertex AI y Artifact Registry.

    gcloud services enable aiplatform.googleapis.com \
        artifactregistry.googleapis.com
    
  2. Configura la CLI de Google Cloud con el ID de tu proyecto e inicializa el SDK de Vertex AI.

    PROJECT_ID = "PROJECT_ID"
    LOCATION = "LOCATION"
    import vertexai
    vertexai.init(project=PROJECT_ID, location=LOCATION)
    
    gcloud config set project {PROJECT_ID}
    
  3. Crea un repositorio de Docker en Artifact Registry.

    gcloud artifacts repositories create DOCKER_REPOSITORY \
        --repository-format=docker \
        --location=LOCATION \
        --description="Vertex AI Docker repository"
    
  4. Opcional: Si vas a descargar modelos de Hugging Face, obtén un token de Hugging Face.

    1. Crea una cuenta de Hugging Face si no tienes una.
    2. En el caso de los modelos protegidos, como Llama 3.2, solicita y recibe acceso en Hugging Face antes de continuar.
    3. Genera un token de acceso: ve a Tu perfil > Configuración > Tokens de acceso.
    4. Selecciona New Token (Nuevo token).
    5. Especifica un nombre y un rol de al menos Lectura.
    6. Selecciona Generar un token.
    7. Guarda este token para los pasos de implementación.

Preparar los archivos de compilación del contenedor

El siguiente Dockerfile crea la imagen de contenedor personalizada de vLLM para GPUs, TPUs y CPUs. Este contenedor personalizado descarga modelos de Hugging Face o Cloud Storage.

ARG BASE_IMAGE
FROM ${BASE_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive
# Install gcloud SDK
RUN apt-get update && \
    apt-get install -y apt-utils git apt-transport-https gnupg ca-certificates curl \
    && echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list \
    && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg \
    && apt-get update -y && apt-get install google-cloud-cli -y \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /workspace/vllm

# Copy entrypoint.sh to the container
COPY ./entrypoint.sh /workspace/vllm/vertexai/entrypoint.sh
RUN chmod +x /workspace/vllm/vertexai/entrypoint.sh

ENTRYPOINT ["/workspace/vllm/vertexai/entrypoint.sh"]

Crea la imagen de contenedor personalizada con Cloud Build. El siguiente archivo de configuración cloudbuild.yaml muestra cómo compilar la imagen para varias plataformas con el mismo Dockerfile.

steps:
-   name: 'gcr.io/cloud-builders/docker'
  automapSubstitutions: true
  script: |
      #!/usr/bin/env bash
      set -euo pipefail
      device_type_param=${_DEVICE_TYPE}
      device_type=${device_type_param,,}
      base_image=${_BASE_IMAGE}
      image_name="vllm-${_DEVICE_TYPE}"
      if [[ $device_type == "cpu" ]]; then
        echo "Quietly building open source vLLM CPU container image"
        git clone https://github.com/vllm-project/vllm.git
        cd vllm && DOCKER_BUILDKIT=1 docker build -t $base_image -f docker/Dockerfile.cpu . -q
        cd ..
      fi
      echo "Quietly building container image for: $device_type"
      docker build -t $LOCATION-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/$image_name --build-arg BASE_IMAGE=$base_image . -q
      docker push $LOCATION-docker.pkg.dev/$PROJECT_ID/${_REPOSITORY}/$image_name
substitutions:
    _DEVICE_TYPE: gpu
    _BASE_IMAGE: vllm/vllm-openai
    _REPOSITORY: my-docker-repo

Los archivos están disponibles en el repositorio de GitHub googlecloudplatform/vertex-ai-samples. Clona el repositorio para usarlos:

git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git

Crea y envía la imagen de contenedor

Crea la imagen de contenedor personalizada con Cloud Build enviando el archivo cloudbuild.yaml. Usa sustituciones para especificar el tipo de dispositivo de destino, que puede ser GPU, TPU o CPU, y la imagen base correspondiente.

GPU

DEVICE_TYPE="gpu"
BASE_IMAGE="vllm/vllm-openai"
cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build && \
gcloud builds submit \
    --config=cloudbuild.yaml \
    --region=LOCATION \
    --timeout="2h" \
    --machine-type=e2-highcpu-32 \
    --substitutions=_REPOSITORY=DOCKER_REPOSITORY,_DEVICE_TYPE=$DEVICE_TYPE,_BASE_IMAGE=$BASE_IMAGE

TPU

DEVICE_TYPE="tpu"
BASE_IMAGE="vllm/vllm-tpu:nightly"
cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build && \
gcloud builds submit \
    --config=cloudbuild.yaml \
    --region=LOCATION \
    --timeout="2h" \
    --machine-type=e2-highcpu-32 \
    --substitutions=_REPOSITORY=DOCKER_REPOSITORY,_DEVICE_TYPE=$DEVICE_TYPE,_BASE_IMAGE=$BASE_IMAGE

CPU

DEVICE_TYPE="cpu"
BASE_IMAGE="vllm-cpu-base"
cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build && \
gcloud builds submit \
    --config=cloudbuild.yaml \
    --region=LOCATION \
    --timeout="2h" \
    --machine-type=e2-highcpu-32 \
    --substitutions=_REPOSITORY=DOCKER_REPOSITORY,_DEVICE_TYPE=$DEVICE_TYPE,_BASE_IMAGE=$BASE_IMAGE

Una vez que se haya completado la compilación, configura Docker para que se autentique con Artifact Registry:

gcloud auth configure-docker LOCATION-docker.pkg.dev --quiet

Subir un modelo al registro de modelos y desplegarlo

Sube tu modelo al registro de modelos de Vertex AI, crea un endpoint y despliega el modelo siguiendo estos pasos. En este ejemplo se usa Llama 3.2 3B, pero puedes adaptarlo a otros modelos.

  1. Define las variables de modelo y de implementación. Asigna a la variable DOCKER_URI la imagen que has creado en el paso anterior (por ejemplo, para la GPU):

    DOCKER_URI = f"LOCATION-docker.pkg.dev/PROJECT_ID/DOCKER_REPOSITORY/vllm-gpu"
    

    Define variables para el token de Hugging Face y las propiedades del modelo. Por ejemplo, en el caso de las implementaciones de GPU:

    hf_token = "your-hugging-face-auth-token"
    model_name = "gpu-llama3_2_3B-serve-vllm"
    model_id = "meta-llama/Llama-3.2-3B"
    machine_type = "g2-standard-8"
    accelerator_type = "NVIDIA_L4"
    accelerator_count = 1
    
  2. Sube el modelo al registro de modelos. La función upload_model varía ligeramente en función del tipo de dispositivo debido a los diferentes argumentos de vLLM y variables de entorno.

    from google.cloud import aiplatform
    
    def upload_model_gpu(model_name, model_id, hf_token, accelerator_count, docker_uri):
        vllm_args = [
            "python3", "-m", "vllm.entrypoints.openai.api_server",
            "--host=0.0.0.0", "--port=8080", f"--model={model_id}",
            "--max-model-len=2048", "--gpu-memory-utilization=0.9",
            "--enable-prefix-caching", f"--tensor-parallel-size={accelerator_count}",
        ]
        env_vars = {
            "HF_TOKEN": hf_token,
            "LD_LIBRARY_PATH": "$LD_LIBRARY_PATH:/usr/local/nvidia/lib64",
        }
        model = aiplatform.Model.upload(
            display_name=model_name,
            serving_container_image_uri=docker_uri,
            serving_container_args=vllm_args,
            serving_container_ports=[8080],
            serving_container_predict_route="/v1/completions",
            serving_container_health_route="/health",
            serving_container_environment_variables=env_vars,
            serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
            serving_container_deployment_timeout=1800,
        )
        return model
    
    def upload_model_tpu(model_name, model_id, hf_token, tpu_count, docker_uri):
        vllm_args = [
            "python3", "-m", "vllm.entrypoints.openai.api_server",
            "--host=0.0.0.0", "--port=8080", f"--model={model_id}",
            "--max-model-len=2048", "--enable-prefix-caching",
            f"--tensor-parallel-size={tpu_count}",
        ]
        env_vars = {"HF_TOKEN": hf_token}
        model = aiplatform.Model.upload(
            display_name=model_name,
            serving_container_image_uri=docker_uri,
            serving_container_args=vllm_args,
            serving_container_ports=[8080],
            serving_container_predict_route="/v1/completions",
            serving_container_health_route="/health",
            serving_container_environment_variables=env_vars,
            serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
            serving_container_deployment_timeout=1800,
        )
        return model
    
    def upload_model_cpu(model_name, model_id, hf_token, docker_uri):
        vllm_args = [
            "python3", "-m", "vllm.entrypoints.openai.api_server",
            "--host=0.0.0.0", "--port=8080", f"--model={model_id}",
            "--max-model-len=2048",
        ]
        env_vars = {"HF_TOKEN": hf_token}
        model = aiplatform.Model.upload(
            display_name=model_name,
            serving_container_image_uri=docker_uri,
            serving_container_args=vllm_args,
            serving_container_ports=[8080],
            serving_container_predict_route="/v1/completions",
            serving_container_health_route="/health",
            serving_container_environment_variables=env_vars,
            serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
            serving_container_deployment_timeout=1800,
        )
        return model
    
    # Example for GPU:
    vertexai_model = upload_model_gpu(model_name, model_id, hf_token, accelerator_count, DOCKER_URI)
    
  3. Crea un endpoint.

    endpoint = aiplatform.Endpoint.create(display_name=f"model_name-endpoint")
    
  4. Despliega el modelo en el endpoint. El despliegue del modelo puede tardar entre 20 y 30 minutos.

    # Example for GPU:
    vertexai_model.deploy(
        endpoint=endpoint,
        deployed_model_display_name=model_name,
        machine_type=machine_type,
        accelerator_type=accelerator_type,
        accelerator_count=accelerator_count,
        traffic_percentage=100,
        deploy_request_timeout=1800,
        min_replica_count=1,
        max_replica_count=4,
        autoscaling_target_accelerator_duty_cycle=60,
    )
    

    En el caso de las TPUs, omite los parámetros accelerator_type y accelerator_count, y usa autoscaling_target_request_count_per_minute=60. En el caso de las CPUs, omita los parámetros accelerator_type y accelerator_count, y use autoscaling_target_cpu_utilization=60.

Cargar modelos de Cloud Storage

El contenedor personalizado descarga el modelo de una ubicación de Cloud Storage en lugar de descargarlo de Hugging Face. Cuando usas Cloud Storage:

  • Asigna el parámetro model_id de la función upload_model a un URI de Cloud Storage, por ejemplo, gs://<var>my-bucket</var>/<var>my-models</var>/<var>llama_3_2_3B</var>.
  • Omite la variable HF_TOKEN de env_vars cuando llames a upload_model.
  • Cuando llames a model.deploy, especifica un service_account que tenga permisos de lectura en el segmento de Cloud Storage.

Crear una cuenta de servicio de IAM para acceder a Cloud Storage

Si tu modelo está en Cloud Storage, crea una cuenta de servicio que los endpoints de Vertex Prediction puedan usar para acceder a los artefactos del modelo.

SERVICE_ACCOUNT_NAME = "vertexai-endpoint-sa"
SERVICE_ACCOUNT_EMAIL = f"SERVICE_ACCOUNT_NAME@PROJECT_ID.iam.gserviceaccount.com"
gcloud iam service-accounts create SERVICE_ACCOUNT_NAME \
    --display-name="Vertex AI Endpoint Service Account"

# Grant storage read permission
gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:SERVICE_ACCOUNT_EMAIL" \
    --role="roles/storage.objectViewer"

Cuando realices la implementación, pasa el correo de la cuenta de servicio al método deploy: service_account=<var>SERVICE_ACCOUNT_EMAIL</var>.

Obtener predicciones mediante un endpoint

Una vez que hayas desplegado correctamente el modelo en el endpoint, verifica la respuesta del modelo con raw_predict.

import json

PROMPT = "Distance of moon from earth is "
request_body = json.dumps(
    {
        "prompt": PROMPT,
        "temperature": 0.0,
    },
)

raw_response = endpoint.raw_predict(
    body=request_body, headers={"Content-Type": "application/json"}
)
assert raw_response.status_code == 200
result = json.loads(raw_response.text)

for choice in result["choices"]:
    print(choice)

Ejemplo:

{
  "index": 0,
  "text": "384,400 km. The moon is 1/4 of the earth's",
  "logprobs": null,
  "finish_reason": "length",
  "stop_reason": null,
  "prompt_logprobs": null
}

Siguientes pasos