Implemente modelos abertos com um contentor vLLM personalizado

Embora as várias opções de apresentação de modelos do Vertex AI sejam suficientes para muitos exemplos de utilização, pode ter de usar as suas próprias imagens de contentores para apresentar modelos no Vertex AI. Este documento descreve como usar uma imagem de contentor personalizada de vLLM para publicar modelos no Vertex AI em CPUs, GPUs ou TPUs. Para mais informações sobre os modelos suportados pelo vLLM, consulte a documentação do vLLM.

O servidor da API vLLM implementa o protocolo da API OpenAI, mas não suporta os requisitos de pedido e resposta da Vertex AI. Por conseguinte, tem de usar um pedido de inferência bruta do Vertex AI para obter inferências de modelos implementados no Vertex AI através de um ponto final de previsão. Para mais informações sobre o método Raw Prediction no Vertex AI Python SDK, consulte a documentação do Python SDK.

Pode obter modelos do Hugging Face e do Cloud Storage. Esta abordagem oferece flexibilidade, o que lhe permite tirar partido do hub de modelos orientado pela comunidade (Hugging Face) e das capacidades de transferência de dados e segurança otimizadas do Cloud Storage para a gestão de modelos internos ou versões ajustadas.

O vLLM transfere os modelos do Hugging Face se for fornecido um token de acesso do Hugging Face. Caso contrário, o vLLM assume que o modelo está disponível no disco local. A imagem do contentor personalizado permite que o Vertex AI transfira o modelo de Google Cloud além do Hugging Face.

Antes de começar

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

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

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

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

    gcloud config set project {PROJECT_ID}

  3. Crie um repositório Docker no Artifact Registry.

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

  4. Opcional: se transferir modelos do Hugging Face, obtenha um token do Hugging Face.

    1. Crie uma conta do Hugging Face se não tiver uma.
    2. Para modelos restritos, como o Llama 3.2, peça e receba acesso no HuggingFace antes de continuar.
    3. Gere um token de acesso: aceda a O seu perfil > Definiçõ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. Guarde este token para os passos de implementação.

Prepare os ficheiros de compilação do contentor

O comando Dockerfile cria a imagem de contentor personalizada do vLLM para GPUs, TPUs e CPUs. Este contentor personalizado transfere 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 do contentor personalizada com o Cloud Build. O ficheiro de configuração cloudbuild.yaml seguinte 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 ficheiros estão disponíveis no repositório do googlecloudplatform/vertex-ai-samples GitHub. Clone o repositório para os usar:

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

Crie e envie a imagem de contentor

Crie a imagem do contentor personalizado com o Cloud Build enviando o ficheiro cloudbuild.yaml. Use substituições para especificar o tipo de dispositivo de destino, que pode ser GPU, TPU ou CPU, e a imagem 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

Após a compilação terminar, configure o Docker para autenticar com o Artifact Registry:

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

Carregue o modelo para o Registo de modelos e implemente-o

Carregue o seu modelo para o Registo de modelos Vertex AI, crie um ponto final e implemente o modelo através destes passos. Este exemplo usa o Llama 3.2 3B, mas pode adaptá-lo para outros modelos.

  1. Defina as variáveis do modelo e da implementação. Defina a variável DOCKER_URI para a imagem que criou no passo anterior (por exemplo, para a 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 a implementação de GPUs:

    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. Carregue o modelo para o Registo de modelos. A função upload_model varia ligeiramente consoante o tipo de dispositivo devido a diferentes argumentos 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 ponto final.

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

  4. Implemente o modelo no ponto final. A implementação do modelo pode demorar entre 20 e 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 as 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.

Carregue modelos a partir do Cloud Storage

O contentor personalizado transfere o modelo a partir de uma localização do Cloud Storage em vez de o transferir a partir do Hugging Face. Quando usa o Cloud Storage:

  • Defina o parâmetro model_id na função upload_model para 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 quando chamar upload_model.
  • Quando chamar model.deploy, especifique um service_account que tenha autorizações para ler a partir do contentor do Cloud Storage.

Crie uma conta de serviço da IAM para acesso ao Cloud Storage

Se o seu modelo estiver no Cloud Storage, crie uma conta de serviço que os pontos finais da Vertex AI Prediction possam usar para aceder aos artefactos 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"

Quando implementar, transmita o email da conta de serviço para o método deploy: service_account=<var>SERVICE_ACCOUNT_EMAIL</var>.

Receba previsões através do ponto final

Depois de implementar o modelo com êxito no ponto final, valide a resposta do modelo com 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
}

O que se segue?