Implantar modelos abertos com um contêiner vLLM personalizado

Embora as várias opções de disponibilização de modelos da Vertex AI sejam suficientes para muitos casos de uso, talvez seja necessário usar suas próprias imagens de contêiner para disponibilizar modelos na Vertex AI. Este documento descreve como usar uma imagem de contêiner personalizada do vLLM para veicular modelos na Vertex AI em CPUs, GPUs ou TPUs. Para mais informações sobre os modelos compatíveis com vLLM, consulte a documentação do vLLM.

O servidor da API vLLM implementa o protocolo da API OpenAI, mas não oferece suporte aos requisitos de solicitação e resposta da Vertex AI. Portanto, use uma solicitação de inferência bruta da Vertex AI para receber inferências de modelos implantados na Vertex AI usando um endpoint de previsão. Para mais informações sobre o método Raw Prediction no SDK da Vertex AI para Python, consulte a documentação do SDK para Python.

É possível usar modelos do Hugging Face e do Cloud Storage. Essa abordagem oferece flexibilidade, permitindo que você aproveite o hub de modelos orientado pela comunidade (Hugging Face) e os recursos otimizados de transferência e segurança de dados do Cloud Storage para gerenciamento interno de modelos ou versões refinadas.

O vLLM faz o download dos modelos do Hugging Face se um token de acesso do Hugging Face for fornecido. Caso contrário, o vLLM pressupõe que o modelo está disponível no disco local. A imagem de contêiner personalizada permite que a Vertex AI faça o download do modelo deGoogle Cloud além do Hugging Face.

Antes de começar

  1. No seu projeto Google Cloud , ative as APIs Vertex AI e Artifact Registry.

    gcloud services enable aiplatform.googleapis.com \
    artifactregistry.googleapis.com

  2. Configure a Google Cloud CLI com o ID do projeto e inicialize o SDK da Vertex AI.

    PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
import vertexai
vertexai.init(project=PROJECT_ID, location=LOCATION)

    gcloud config set project {PROJECT_ID}

  3. Criar um repositório do Docker no Artifact Registry.

    gcloud artifacts repositories create DOCKER_REPOSITORY \
    --repository-format=docker \
    --location=LOCATION \
    --description="Vertex AI Docker repository"

  4. Opcional: se você estiver baixando modelos do Hugging Face, obtenha um token do Hugging Face.

    1. Crie uma conta do Hugging Face se você não tiver uma.
    2. Para modelos restritos, como o Llama 3.2, solicite e receba acesso no Hugging Face antes de continuar.
    3. Gere um token de acesso: acesse Seu perfil > Configurações > Tokens de acesso.
    4. Selecione Novo token.
    5. Especifique um nome e uma função de pelo menos Leitura.
    6. Selecione Gerar um token.
    7. Salve esse token para as etapas de implantação.

Preparar arquivos de build de contêiner

O Dockerfile a seguir cria a imagem de contêiner personalizada do vLLM para GPUs, TPUs e CPUs. Esse contêiner personalizado faz o download de modelos do Hugging Face ou do 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"]

Crie a imagem de contêiner personalizada usando o Cloud Build. O arquivo de configuração cloudbuild.yaml a seguir mostra como criar a imagem para várias plataformas usando o mesmo 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

Os arquivos estão disponíveis no repositório googlecloudplatform/vertex-ai-samples do GitHub. Clone o repositório para usá-los:

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

Crie e envie a imagem do contêiner

Crie a imagem de contêiner personalizada usando o Cloud Build enviando o arquivo cloudbuild.yaml. Use substituições para especificar o tipo de dispositivo de destino (GPU, TPU ou CPU) e a imagem de base correspondente.

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

Depois que o build for concluído, configure o Docker para autenticar com o Artifact Registry:

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

Fazer upload do modelo para o Model Registry e implantar

Faça upload do modelo para o Vertex AI Model Registry, crie um endpoint e implante o modelo seguindo estas etapas. Este exemplo usa o Llama 3.2 3B, mas você pode adaptá-lo para outros modelos.

  1. Defina variáveis de modelo e implantação. Defina a variável DOCKER_URI como a imagem criada na etapa anterior (por exemplo, para GPU):

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

    Defina variáveis para o token do Hugging Face e as propriedades do modelo. Por exemplo, para implantação 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. Faça upload do modelo para o Model Registry. A função upload_model varia um pouco dependendo do tipo de dispositivo devido a diferentes argumentos de vLLM e variáveis de ambiente.

    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. Crie um endpoint

    endpoint = aiplatform.Endpoint.create(display_name=f"model_name-endpoint")

  4. Implante o modelo no endpoint. A implantação do modelo pode levar de 20 a 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,
)

    Para TPUs, omita os parâmetros accelerator_type e accelerator_count e use autoscaling_target_request_count_per_minute=60. Para CPUs, omita os parâmetros accelerator_type e accelerator_count e use autoscaling_target_cpu_utilization=60.

Carregar modelos do Cloud Storage

O contêiner personalizado faz o download do modelo de um local do Cloud Storage em vez de fazer o download do Hugging Face. Ao usar o Cloud Storage:

  • Defina o parâmetro model_id na função upload_model como um URI do Cloud Storage, por exemplo, gs://<var>my-bucket</var>/<var>my-models</var>/<var>llama_3_2_3B</var>.
  • Omita a variável HF_TOKEN de env_vars ao chamar upload_model.
  • Ao chamar model.deploy, especifique um service_account que tenha permissões para ler do bucket do Cloud Storage.

Criar uma conta de serviço do IAM para acesso ao Cloud Storage

Se o modelo estiver no Cloud Storage, crie uma conta de serviço que os endpoints Prediction da Vertex possam usar para acessar os artefatos do 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"

Ao implantar, transmita o e-mail da conta de serviço para o método deploy: service_account=<var>SERVICE_ACCOUNT_EMAIL</var>.

Receber previsões usando o endpoint

Depois de implantar o modelo no endpoint, verifique a resposta dele usando 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)

Exemplo de saída:

{
  "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
}

A seguir