使用自定义 vLLM 容器部署开放模型

虽然各种 Vertex AI 模型服务选项足以满足许多使用场景，但您可能需要使用自己的容器映像在 Vertex AI 上提供模型服务。本文档介绍了如何使用 vLLM 自定义容器映像在 Vertex AI 的 CPU、GPU 或 TPU 上提供模型服务。如需详细了解 vLLM 支持的模型，请参阅 vLLM 文档。

vLLM API 服务器实现了 OpenAI API 协议，但它不支持 Vertex AI 请求和响应要求。因此，您必须使用 Vertex AI 原始推理请求，才能通过预测端点从部署到 Vertex AI 的模型获取推理结果。如需详细了解 Vertex AI Python SDK 中的原始预测方法，请参阅 Python SDK 文档。

您可以从 Hugging Face 和 Cloud Storage 获取模型。这种方法非常灵活，可让您利用社区驱动的模型中心 (Hugging Face) 以及 Cloud Storage 针对内部模型管理或微调版本进行优化的数据传输和安全功能。

如果提供了 Hugging Face 访问令牌，vLLM 会从 Hugging Face 下载模型。否则，vLLM 会假定模型在本地磁盘上可用。除了 Hugging Face 之外，该自定义容器映像还允许 Vertex AI 从Google Cloud 下载模型。

准备工作

1. 在您的 Google Cloud 项目中，启用 Vertex AI 和 Artifact Registry API。

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

2. 使用您的项目 ID 配置 Google Cloud CLI 并初始化 Vertex AI SDK。

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

   gcloud config set project {PROJECT_ID}
   

3. 在 Artifact Registry 中创建 Docker 代码库。

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

4. 可选：如果是从 Hugging Face 下载模型，请获取 Hugging Face 令牌。

   1. 如果您还没有 Hugging Face 账号，请创建一个。
   2. 对于像 Llama 3.2 这样的门控模型，请先在 Hugging Face 上申请并获得访问权限，然后再继续操作。
   3. 生成访问令牌：前往您的个人资料 > 设置 > 访问令牌。
   4. 选择新建令牌 (New Token)。
   5. 指定名称和至少拥有读取权限的角色。
   6. 选择生成令牌。
   7. 保存此令牌以供部署步骤使用。

准备容器 build 文件

以下 Dockerfile 会构建用于 GPU、TPU 和 CPU 的 vLLM 自定义容器映像。此自定义容器可从 Hugging Face 或 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"]

使用 Cloud Build 构建自定义容器映像。以下 cloudbuild.yaml 配置文件展示了如何使用同一 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

这些文件可在 googlecloudplatform/vertex-ai-samples GitHub 代码库中找到。克隆该代码库以使用这些文件：

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

构建并推送容器映像

通过提交 cloudbuild.yaml 文件，使用 Cloud Build 构建自定义容器映像。使用替代变量指定目标设备类型（可以是 GPU、TPU 或 CPU）以及相应的基础映像。

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

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

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

构建完成后，配置 Docker 以向 Artifact Registry 进行身份验证：

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

将模型上传到 Model Registry 并部署模型

完成以下步骤，将模型上传到 Vertex AI Model Registry，创建端点并部署模型。此示例使用 Llama 3.2 3B，但您也可以针对其他模型进行调整。

1. 定义模型和部署变量。将 DOCKER_URI 变量设置为您在上一步中构建的映像（例如，针对 GPU）：

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

   定义 Hugging Face 令牌和模型属性的变量。例如，对于 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. 将模型上传到 Model Registry。upload_model 函数会因设备类型而略有不同，因为不同设备类型的 vLLM 参数和环境变量也会不同。

   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. 创建端点。

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

4. 将模型部署到端点。模型部署可能需要 20 到 30 分钟。

   # 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,
   )
   

   对于 TPU，请省略 accelerator_type 和 accelerator_count 参数，并使用 autoscaling_target_request_count_per_minute=60。对于 CPU，请省略 accelerator_type 和 accelerator_count 参数，并使用 autoscaling_target_cpu_utilization=60。

从 Cloud Storage 加载模型

自定义容器会从 Cloud Storage 位置下载模型，而不是从 Hugging Face 下载。使用 Cloud Storage 时：

• 将 upload_model 函数中的 model_id 参数设置为 Cloud Storage URI，例如 gs://<var>my-bucket</var>/<var>my-models</var>/<var>llama_3_2_3B</var>。
• 调用 upload_model 时，从 env_vars 中省略 HF_TOKEN 变量。
• 调用 model.deploy 时，指定一个有权从 Cloud Storage 存储桶读取数据的 service_account。

创建用于访问 Cloud Storage 的 IAM 服务账号

如果您的模型位于 Cloud Storage 中，请创建一个服务账号，供 Vertex Prediction 端点用来访问模型制品。

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"

部署时，将服务账号邮箱传递给 deploy 方法：service_account=<var>SERVICE_ACCOUNT_EMAIL</var>。

使用端点获取预测结果

成功将模型部署到端点后，使用 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)

输出示例：

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

后续步骤

• 选择开放模型服务选项
• 使用以模型即服务 (MaaS) 形式提供的开放模型
• 从 Model Garden 部署开放模型
• 使用预构建容器部署开放模型