カスタム vLLM コンテナを使用してオープンモデルをデプロイする

多くのユースケースでは、さまざまな Vertex AI モデル サービング オプションで十分ですが、Vertex AI でモデルをサービングするために独自のコンテナ イメージを使用する必要がある場合があります。このドキュメントでは、vLLM カスタム コンテナ イメージを使用して、CPU、GPU、TPU で Vertex AI のモデルをサービングする方法について説明します。vLLM でサポートされているモデルの詳細については、vLLM のドキュメントをご覧ください。

vLLM API サーバーは OpenAI API プロトコルを実装していますが、Vertex AI のリクエストとレスポンスの要件はサポートしていません。したがって、予測エンドポイントを使用して Vertex AI にデプロイされたモデルから推論を取得するには、Vertex AI の未加工の推論リクエストを使用する必要があります。Vertex AI Python SDK の Raw Prediction メソッドの詳細については、Python SDK のドキュメントをご覧ください。

Hugging Face と Cloud Storage の両方からモデルを取得できます。このアプローチは柔軟性が高く、コミュニティ主導のモデルハブ（Hugging Face）と、内部モデル管理またはファインチューニング バージョンの Cloud Storage の最適化されたデータ転送機能とセキュリティ機能を活用できます。

Hugging Face アクセス トークンが指定されている場合、vLLM は Hugging Face からモデルをダウンロードします。それ以外の場合、vLLM はモデルがローカル ディスクで使用可能であると想定します。カスタム コンテナ イメージを使用すると、Vertex AI は Hugging Face に加えてGoogle Cloud からモデルをダウンロードできます。

始める前に

1. Google Cloud プロジェクトで、Vertex AI API と 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. アクセス トークンを生成する: [Your Profile] > [Settings] > [Access Tokens] の順に移動します。
   4. [New Token] を選択します。
   5. 名前と、少なくとも Read ロールを指定します。
   6. [Generate a token] を選択します。
   7. このトークンはデプロイ手順で使用します。

コンテナ ビルドファイルを準備する

次の 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

コンテナ イメージをビルドして push する

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

ビルドが完了したら、Artifact Registry で認証するように Docker を構成します。

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 からモデルを読み込む

カスタム コンテナは、Hugging Face からダウンロードするのではなく、Cloud Storage の場所からモデルをダウンロードします。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
}

次のステップ

• オープンモデルのサービング オプションを選択する
• Model as a Service（MaaS）を使用してオープンモデルを使用する
• Model Garden からオープンモデルをデプロイする
• ビルド済みコンテナを使用してオープンモデルをデプロイする