diff --git a/notebooks/official/CODEOWNERS b/notebooks/official/CODEOWNERS index 6bd956ee2..6f29f7860 100644 --- a/notebooks/official/CODEOWNERS +++ b/notebooks/official/CODEOWNERS @@ -65,5 +65,8 @@ /generative_ai/nvidia_nim_vertexai.ipynb @sujituk /generative_ai/cambai_intro.ipynb @arnavmehta7 /generative_ai/qodo_intro.ipynb @talshef -/generative_ai/csm_intro.ipynb @elaidlaw - +/prediction/vertexai_serving_vllm/vertexai_serving_vllm_cpu_llama3_2_3B.ipynb @ravi-dalal +/prediction/vertexai_serving_vllm/vertexai_serving_vllm_gpu_llama3_2_3B.ipynb @ravi-dalal +/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_llama3_2_3B.ipynb @ravi-dalal +/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb @ravi-dalal +/generative_ai/csm_intro.ipynb @elaidlaw \ No newline at end of file diff --git a/notebooks/official/prediction/vertexai_serving_vllm/README.md b/notebooks/official/prediction/vertexai_serving_vllm/README.md new file mode 100644 index 000000000..31cbc55e0 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/README.md @@ -0,0 +1,30 @@ + + +# Serving models in Vertex AI using vLLM + +The notebooks in this directory demonstrate how Llama 3.2 3B open weight model can be served on Vertex AI using [vLLM](https://github.com/vllm-project/vllm.git). + +## Using TPU +This [colab notebook](vertexai_serving_vllm_tpu_llama3_2_3B.ipynb) shows how Llama 3.2 3B model can be deployed (downloaded from Hugging Face) to Vertex AI Endpoint using this repository on TPUs. + +This [colab notebook](vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb) shows how Llama 3.2 3B model can be deployed (downloaded from Google Cloud Storage) to Vertex AI Endpoint using this repository on TPUs. + +## Using GPU +This [colab notebook](vertexai_serving_vllm_gpu_llama3_2_3B.ipynb) shows how Llama 3.2 3B model can be deployed (downloaded from Hugging Face) to Vertex AI Endpoint using this repository on GPUs. + +## Using CPU +This [colab notebook](vertexai_serving_vllm_cpu_llama3_2_3B.ipynb) shows how Llama 3.2 3B model can be deployed (downloaded from Hugging Face) to Vertex AI Endpoint using this repository on CPUs. diff --git a/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/cloudbuild.yaml b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/cloudbuild.yaml new file mode 100644 index 000000000..42744937f --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/cloudbuild.yaml @@ -0,0 +1,37 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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 \ No newline at end of file diff --git a/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/dockerfile b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/dockerfile new file mode 100644 index 000000000..5521da1f6 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/dockerfile @@ -0,0 +1,33 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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"] \ No newline at end of file diff --git a/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/entrypoint.sh b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/entrypoint.sh new file mode 100755 index 000000000..0500664c2 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/cloud-build/entrypoint.sh @@ -0,0 +1,54 @@ +#!/bin/bash + +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -euo pipefail + +readonly LOCAL_MODEL_DIR=${LOCAL_MODEL_DIR:-"/tmp/model_dir"} + +download_model_from_gcs() { + gcs_uri=$1 + mkdir -p $LOCAL_MODEL_DIR + echo "Downloading model from $gcs_uri to local directory..." + if gcloud storage cp -r "$gcs_uri/*" "$LOCAL_MODEL_DIR"; then + echo "Model downloaded successfully to ${LOCAL_MODEL_DIR}." + else + echo "Failed to download model from Cloud Storage: $gcs_uri." >&2 + exit 1 + fi +} + + +updated_args=() +model_arg="--model=" +gcs_protocol="gs://" +for a in "$@"; do + if [[ $a == $model_arg* ]]; then + model_path=${a#*=} + echo $model_path + if [[ $model_path == $gcs_protocol* ]]; then + download_model_from_gcs $model_path + updated_args+=("--model=${LOCAL_MODEL_DIR}") + else + updated_args+=("--model=${model_path}") + fi + + else + updated_args+=("$a") + fi +done + +echo "Launch command: " "${updated_args[@]}" +exec "${updated_args[@]}" \ No newline at end of file diff --git a/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_cpu_llama3_2_3B.ipynb b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_cpu_llama3_2_3B.ipynb new file mode 100644 index 000000000..f9c4b3ab3 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_cpu_llama3_2_3B.ipynb @@ -0,0 +1,627 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f705f4be70e9" + }, + "outputs": [], + "source": [ + "# Copyright 2025 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8c9b0667d10c" + }, + "source": [ + "# Serving Open Models on Vertex AI using vLLM with CPU\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Vertex
Open in Vertex AI Workbench\n", + " \n", + " 		\n", + " \n", + " \"Google
Run in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "24d72ac84336" + }, + "source": [ + "## Overview\n", + "\n", + "There are multiple ways of serving open models (including open source and open weight) such as Llama 3.2 on Google Cloud Vertex AI. The Llama models are available in [Model Garden](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) and Model Garden allows a single click self-deployment of the models. This notebooks demonstrates how Llama 3.2 3B model can be served via Vertex AI Endpoint using a custom vLLM container image built for the CPU. This notebook does the following:\n", + "\n", + "- Builds a custom docker container image using vLLM source code\n", + "- Uploads the model to Model Registry using custom docker container image\n", + "- Creates a public Endpoint for Online Prediction\n", + "- Deploys model to the Endpoint\n", + "- Llama 3.2 3B model is downloaded from Hugging Face during deployment\n", + "- This custom container image can also be used for downloading model from Google Storage\n", + "\n", + "The code in this notebook can be used for serving other open models supported by vLLM. This notebook has been tested with Python 3.10 and `google-cloud-aiplatform` SDK Version `1.106.0`.\n", + "\n", + "To download the models from the Hugging Face, you need a Hugging Face token.\n", + " 1. Create a [Hugging Face account](https://huggingface.co/) if you don't have one.\n", + " 2. For **gated models** like Llama 3.2, ensure you have requested and been granted access on Hugging Face before proceeding.\n", + " 3. Generate an Access Token: Go to **Your Profile > Settings > Access Tokens**.\n", + " 4. Select **New Token**.\n", + " 5. Specify a Name and a Role of at least Read.\n", + " 6. Select **Generate a token**.\n", + " 7. Set the token in hf_token env below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "36df818a34bd" + }, + "source": [ + "## Get Started" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "26645caf62fe" + }, + "source": [ + "### Install Vertex AI SDK for Python and other required packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d4cf289f0d99" + }, + "outputs": [], + "source": [ + "!pip install --upgrade --quiet google-cloud-aiplatform" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "848322ec177e" + }, + "source": [ + "### Restart runtime (Colab only)\n", + "\n", + "To use the newly installed packages, you must restart the runtime on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b8d49bb74a53" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " import IPython\n", + "\n", + " app = IPython.Application.instance()\n", + " app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f332441fe51" + }, + "source": [ + "
\n", + "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "11947ae0fe5e" + }, + "source": [ + "### Authenticate your notebook environment (Colab only)\n", + "\n", + "Authenticate your environment on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "015bf6d5da75" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "722a10c66085" + }, + "source": [ + "### Set Google Cloud project and initialize Vertex AI SDK for Python\n", + "\n", + "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "66156945acb1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", + "DEVICE_TYPE = \"cpu\" # @param {type:\"string\"}\n", + "\n", + "import vertexai\n", + "\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "202ff51ba5de" + }, + "source": [ + "Set Project ID in active gcloud configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fc8b4506d7ea" + }, + "outputs": [], + "source": [ + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7cce406746d5" + }, + "source": [ + "## Create vLLM Customer Container Image for Vertex AI\n", + "\n", + "Vertex AI requires [requests](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#inference) and [responses](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#response_requirements) in specific formats. vLLM API server implements OpenAI API protocol and therefore, it does not support the Vertex AI request and response requirements. Therefore, the vLLM API server (vllm.entrypoints.openai.api_server.py) needs to be updated to support Vertex AI request and response formats." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f7ec1c7bdf8" + }, + "source": [ + "### Enable Artifact Registry API\n", + "Enable the Artifact Registry API service for the Google cloud project. This tutorial requires [gcloud CLI](https://cloud.google.com/sdk/docs/install) installed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6e93c586de15" + }, + "outputs": [], + "source": [ + "! gcloud components update --quiet && gcloud services enable artifactregistry.googleapis.com" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6c4e1d4d5e66" + }, + "source": [ + "### Create a private Docker repository\n", + "Create a Docker repository in [Artifact Registry](https://cloud.google.com/artifact-registry/docs/overview)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1f21a9e18e5c" + }, + "outputs": [], + "source": [ + "DOCKER_REPOSITORY = \"my-docker-repo\"\n", + "! gcloud artifacts repositories create {DOCKER_REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Vertex AI Docker repository\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "479fbec1d5e0" + }, + "source": [ + "### Build vLLM Custom Docker Container Image for CPU" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "bed9ed248a5e" + }, + "source": [ + "Clone vertex-ai-samples code reposistory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "13df52a033d3" + }, + "outputs": [], + "source": [ + "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b89cf1a29b80" + }, + "source": [ + "Build image using Cloud Build" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8c00680ff433" + }, + "outputs": [], + "source": [ + "! cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build \\\n", + " && 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=vllm-cpu-base" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "95c5ddf6f5c6" + }, + "source": [ + "## Deploy Model to Vertex AI Endpoint\n", + "\n", + "Following steps are required to serve model via a Vertex AI Prediction Endpoint:\n", + "- import model to model registry\n", + "- create a Online Prediction Endpoint\n", + "- Deploy the model to endpoint" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dd851913ae85" + }, + "source": [ + "### Define Variable" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5b6b731880ed" + }, + "outputs": [], + "source": [ + "hf_token = \"[your-hugging-face-auth-token]\" # @param {type:\"string\"}\n", + "model_name = \"cpu-llama3_2_3B-serve-vllm\" # @param {type:\"string\"}\n", + "model_id = \"meta-llama/Llama-3.2-3B\" # @param {type:\"string\"}\n", + "machine_type = \"c2-standard-16\" # @param {type:\"string\"}\n", + "DOCKER_URI = (\n", + " f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/vllm-{DEVICE_TYPE}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e0c1bb82c9e0" + }, + "source": [ + "### Import model to Model Registry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d36036e3f6f7" + }, + "outputs": [], + "source": [ + "from google.cloud import aiplatform\n", + "\n", + "\n", + "def upload_model(\n", + " model_name: str,\n", + " model_id: str,\n", + " hf_token: str,\n", + " docker_uri: str,\n", + ") -> aiplatform.Model:\n", + "\n", + " vllm_args = [\n", + " \"python3\",\n", + " \"-m\",\n", + " \"vllm.entrypoints.openai.api_server\",\n", + " \"--host=0.0.0.0\",\n", + " \"--port=8080\",\n", + " f\"--model={model_id}\",\n", + " \"--max-model-len=2048\",\n", + " ]\n", + "\n", + " env_vars = {\n", + " \"HF_TOKEN\": hf_token,\n", + " }\n", + "\n", + " model = aiplatform.Model.upload(\n", + " display_name=model_name,\n", + " serving_container_image_uri=docker_uri,\n", + " serving_container_args=vllm_args,\n", + " serving_container_ports=[8080],\n", + " serving_container_predict_route=\"/v1/completions\",\n", + " serving_container_health_route=\"/health\",\n", + " serving_container_environment_variables=env_vars,\n", + " serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n", + " serving_container_deployment_timeout=1800,\n", + " )\n", + " return model\n", + "\n", + "\n", + "vertexai_model = upload_model(\n", + " model_name=model_name, model_id=model_id, hf_token=hf_token, docker_uri=DOCKER_URI\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "124c8611bab3" + }, + "source": [ + "### Create Vertex AI Endpoint for Online Prediction" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "91591af558cf" + }, + "outputs": [], + "source": [ + "def create_model_endpoint(model_name: str) -> aiplatform.Endpoint:\n", + " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n", + " return endpoint\n", + "\n", + "\n", + "vertexai_endpoint = create_model_endpoint(model_name=model_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "15a1ba904c6e" + }, + "source": [ + "### Deploy Model to Endpoint\n", + "**NOTE**: The model deployment may take around 30 minutes to complete." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f77f1ede9afb" + }, + "outputs": [], + "source": [ + "def deploy_model(\n", + " model: aiplatform.Model,\n", + " endpoint: aiplatform.Endpoint,\n", + " model_name: str,\n", + " machine_type: str,\n", + "):\n", + " print(\n", + " f\"Deploying {model_name} to endpoint: {endpoint.resource_name} using machine type: {machine_type}\"\n", + " )\n", + " model.deploy(\n", + " endpoint=endpoint,\n", + " deployed_model_display_name=model_name,\n", + " machine_type=machine_type,\n", + " min_replica_count=1,\n", + " max_replica_count=4,\n", + " autoscaling_target_cpu_utilization=60,\n", + " traffic_percentage=100,\n", + " deploy_request_timeout=1800,\n", + " )\n", + "\n", + "\n", + "deploy_model(\n", + " model=vertexai_model,\n", + " endpoint=vertexai_endpoint,\n", + " model_name=model_name,\n", + " machine_type=machine_type,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "13f1059160b9" + }, + "source": [ + "## Test Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8ec42d0c4422" + }, + "outputs": [], + "source": [ + "import json\n", + "\n", + "PROMPT = \"Distance of moon from earth is \"\n", + "request_body = json.dumps(\n", + " {\n", + " \"prompt\": PROMPT,\n", + " \"temperature\": 0.0,\n", + " },\n", + ")\n", + "\n", + "raw_response = vertexai_endpoint.raw_predict(\n", + " body=request_body, headers={\"Content-Type\": \"application/json\"}\n", + ")\n", + "assert raw_response.status_code == 200\n", + "result = json.loads(raw_response.text)\n", + "\n", + "for choice in result[\"choices\"]:\n", + " print(choice)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fdc1973ace2f" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, delete the resources created in this tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8c8cc0c481e8" + }, + "source": [ + "### Delete Vertex AI Prediction Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6615b3d9888d" + }, + "outputs": [], + "source": [ + "vertexai_endpoint.delete(force=True, sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ebd767d8e848" + }, + "source": [ + "### Delete Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2089ddbdaf35" + }, + "outputs": [], + "source": [ + "vertexai_model.delete(sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9e35ebc98df3" + }, + "source": [ + "### Delete private docker repository" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "795a05a169bc" + }, + "outputs": [], + "source": [ + "! gcloud artifacts repositories delete {DOCKER_REPOSITORY} --location={LOCATION} --quiet" + ] + } + ], + "metadata": { + "colab": { + "name": "vertexai_serving_vllm_cpu_llama3_2_3B.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_gpu_llama3_2_3B.ipynb b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_gpu_llama3_2_3B.ipynb new file mode 100644 index 000000000..110e1b2a8 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_gpu_llama3_2_3B.ipynb @@ -0,0 +1,644 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f705f4be70e9" + }, + "outputs": [], + "source": [ + "# Copyright 2025 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0fc0c49e9ad2" + }, + "source": [ + "# Serving Open Models on Vertex AI using vLLM with GPU\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Vertex
Open in Vertex AI Workbench\n", + " \n", + " 		\n", + " \n", + " \"Google
Run in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "23aaa66bce19" + }, + "source": [ + "## Overview\n", + "\n", + "There are multiple ways of serving open models (including open source and open weight) such as Llama 3.2 on Google Cloud Vertex AI. The Llama models are available in [Model Garden](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) and Model Garden allows a single click self-deployment of the models. This notebooks demonstrates how Llama 3.2 3B model can be served via Vertex AI Endpoint using a custom vLLM container image built for the GPU. This notebook does the following:\n", + "\n", + "- Builds a custom docker container image using vLLM source code\n", + "- Uploads the model to Model Registry using custom docker container image\n", + "- Creates a public Endpoint for Online Prediction\n", + "- Deploys model to the Endpoint\n", + "- Llama 3.2 3B model is downloaded from Hugging Face during deployment\n", + "- This custom container image can also be used for downloading model from Google Storage\n", + "\n", + "The code in this notebook can be used for serving other open models supported by vLLM. This notebook has been tested with Python 3.10 and `google-cloud-aiplatform` SDK Version `1.106.0`.\n", + "\n", + "To download the models from the Hugging Face, you need a Hugging Face token.\n", + " 1. Create a [Hugging Face account](https://huggingface.co/) if you don't have one.\n", + " 2. For **gated models** like Llama 3.2, ensure you have requested and been granted access on Hugging Face before proceeding.\n", + " 3. Generate an Access Token: Go to **Your Profile > Settings > Access Tokens**.\n", + " 4. Select **New Token**.\n", + " 5. Specify a Name and a Role of at least Read.\n", + " 6. Select **Generate a token**.\n", + " 7. Set the token in hf_token env below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "36df818a34bd" + }, + "source": [ + "## Get Started" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "26645caf62fe" + }, + "source": [ + "### Install Vertex AI SDK for Python and other required packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d4cf289f0d99" + }, + "outputs": [], + "source": [ + "!pip install --upgrade --quiet google-cloud-aiplatform" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "848322ec177e" + }, + "source": [ + "### Restart runtime (Colab only)\n", + "\n", + "To use the newly installed packages, you must restart the runtime on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b8d49bb74a53" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " import IPython\n", + "\n", + " app = IPython.Application.instance()\n", + " app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f332441fe51" + }, + "source": [ + "
\n", + "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "11947ae0fe5e" + }, + "source": [ + "### Authenticate your notebook environment (Colab only)\n", + "\n", + "Authenticate your environment on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "015bf6d5da75" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "722a10c66085" + }, + "source": [ + "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n", + "\n", + "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "66156945acb1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", + "DEVICE_TYPE = \"gpu\" # @param {type:\"string\"}\n", + "\n", + "import vertexai\n", + "\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "202ff51ba5de" + }, + "source": [ + "Set Project ID in active gcloud configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fc8b4506d7ea" + }, + "outputs": [], + "source": [ + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7cce406746d5" + }, + "source": [ + "## Create vLLM Customer Container Image for Vertex AI\n", + "\n", + "Vertex AI requires [requests](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#inference) and [responses](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#response_requirements) in specific formats. vLLM API server implements OpenAI API protocol and therefore, it does not support the Vertex AI request and response requirements. Therefore, the vLLM API server (vllm.entrypoints.openai.api_server.py) needs to be updated to support Vertex AI request and response formats." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f7ec1c7bdf8" + }, + "source": [ + "### Enable Artifact Registry API\n", + "Enable the Artifact Registry API service for the Google cloud project. This tutorial requires [gcloud CLI](https://cloud.google.com/sdk/docs/install) installed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6e93c586de15" + }, + "outputs": [], + "source": [ + "! gcloud components update --quiet && gcloud services enable artifactregistry.googleapis.com" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6c4e1d4d5e66" + }, + "source": [ + "### Create a private Docker repository\n", + "Create a Docker repository in [Artifact Registry](https://cloud.google.com/artifact-registry/docs/overview)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1f21a9e18e5c" + }, + "outputs": [], + "source": [ + "DOCKER_REPOSITORY = \"my-docker-repo\"\n", + "! gcloud artifacts repositories create {DOCKER_REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Vertex AI Docker repository\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d8d9b176ce90" + }, + "source": [ + "### Build vLLM Custom Docker Container Image for GPU" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "84b6342dd2f0" + }, + "source": [ + "Clone vertex-ai-samples code reposistory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "13df52a033d3" + }, + "outputs": [], + "source": [ + "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b89cf1a29b80" + }, + "source": [ + "Build image using Cloud Build" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "a4ee2f849c7a" + }, + "outputs": [], + "source": [ + "! cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build \\\n", + " && 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=vllm/vllm-openai" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "95c5ddf6f5c6" + }, + "source": [ + "## Deploy Model to Vertex AI Endpoint\n", + "\n", + "Following steps are required to serve model via a Vertex AI Prediction Endpoint:\n", + "- import model to model registry\n", + "- create a Online Prediction Endpoint\n", + "- Deploy the model to endpoint" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dd851913ae85" + }, + "source": [ + "### Define Variable" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fb6a8f90dc7c" + }, + "outputs": [], + "source": [ + "hf_token = \"[your-hugging-face-auth-token]\" # @param {type:\"string\"}\n", + "model_name = \"gpu-llama3_2_3B-serve-vllm\" # @param {type:\"string\"}\n", + "model_id = \"meta-llama/Llama-3.2-3B\" # @param {type:\"string\"}\n", + "machine_type = \"g2-standard-8\" # @param {type:\"string\"}\n", + "accelerator_type = \"NVIDIA_L4\" # @param {type:\"string\"}\n", + "accelerator_count = 1 # @param {type:\"integer\"}\n", + "DOCKER_URI = (\n", + " f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/vllm-{DEVICE_TYPE}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e0c1bb82c9e0" + }, + "source": [ + "### Import model to Model Registry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "791245e8598e" + }, + "outputs": [], + "source": [ + "from google.cloud import aiplatform\n", + "\n", + "\n", + "def upload_model(\n", + " model_name: str,\n", + " model_id: str,\n", + " hf_token: str,\n", + " accelerator_count: int,\n", + " docker_uri: str,\n", + ") -> aiplatform.Model:\n", + "\n", + " vllm_args = [\n", + " \"python3\",\n", + " \"-m\",\n", + " \"vllm.entrypoints.openai.api_server\",\n", + " \"--host=0.0.0.0\",\n", + " \"--port=8080\",\n", + " f\"--model={model_id}\",\n", + " \"--max-model-len=2048\",\n", + " \"--gpu-memory-utilization=0.9\",\n", + " \"--enable-prefix-caching\",\n", + " f\"--tensor-parallel-size={accelerator_count}\",\n", + " ]\n", + "\n", + " env_vars = {\n", + " \"HF_TOKEN\": hf_token,\n", + " \"LD_LIBRARY_PATH\": \"$LD_LIBRARY_PATH:/usr/local/nvidia/lib64\",\n", + " }\n", + "\n", + " model = aiplatform.Model.upload(\n", + " display_name=model_name,\n", + " serving_container_image_uri=docker_uri,\n", + " serving_container_args=vllm_args,\n", + " serving_container_ports=[8080],\n", + " serving_container_predict_route=\"/v1/completions\",\n", + " serving_container_health_route=\"/health\",\n", + " serving_container_environment_variables=env_vars,\n", + " serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n", + " serving_container_deployment_timeout=1800,\n", + " )\n", + " return model\n", + "\n", + "\n", + "vertexai_model = upload_model(\n", + " model_name=model_name,\n", + " model_id=model_id,\n", + " hf_token=hf_token,\n", + " accelerator_count=int(accelerator_count),\n", + " docker_uri=DOCKER_URI,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "124c8611bab3" + }, + "source": [ + "### Create Vertex AI Endpoint for Online Prediction" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "91591af558cf" + }, + "outputs": [], + "source": [ + "def create_model_endpoint(model_name: str) -> aiplatform.Endpoint:\n", + " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n", + " return endpoint\n", + "\n", + "\n", + "vertexai_endpoint = create_model_endpoint(model_name=model_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "15a1ba904c6e" + }, + "source": [ + "### Deploy Model to Endpoint\n", + "**NOTE**: The model deployment may take around 30 minutes to complete." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3cc813e9a4c7" + }, + "outputs": [], + "source": [ + "def deploy_model(\n", + " model: aiplatform.Model,\n", + " endpoint: aiplatform.Endpoint,\n", + " model_name: str,\n", + " machine_type: str,\n", + " accelerator_type: str,\n", + " accelerator_count: int,\n", + "):\n", + " print(\n", + " f\"Deploying {model_name} to endpoint: {endpoint.resource_name} using machine type: {machine_type}\"\n", + " )\n", + " model.deploy(\n", + " endpoint=endpoint,\n", + " deployed_model_display_name=model_name,\n", + " machine_type=machine_type,\n", + " accelerator_type=accelerator_type,\n", + " accelerator_count=accelerator_count,\n", + " min_replica_count=1,\n", + " max_replica_count=4,\n", + " autoscaling_target_accelerator_duty_cycle=60,\n", + " traffic_percentage=100,\n", + " deploy_request_timeout=1800,\n", + " )\n", + "\n", + "\n", + "deploy_model(\n", + " model=vertexai_model,\n", + " endpoint=vertexai_endpoint,\n", + " model_name=model_name,\n", + " machine_type=machine_type,\n", + " accelerator_type=accelerator_type,\n", + " accelerator_count=int(accelerator_count),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "13f1059160b9" + }, + "source": [ + "## Test Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8ec42d0c4422" + }, + "outputs": [], + "source": [ + "import json\n", + "\n", + "PROMPT = \"Distance of moon from earth is \"\n", + "request_body = json.dumps(\n", + " {\n", + " \"prompt\": PROMPT,\n", + " \"temperature\": 0.0,\n", + " },\n", + ")\n", + "\n", + "raw_response = vertexai_endpoint.raw_predict(\n", + " body=request_body, headers={\"Content-Type\": \"application/json\"}\n", + ")\n", + "assert raw_response.status_code == 200\n", + "result = json.loads(raw_response.text)\n", + "\n", + "for choice in result[\"choices\"]:\n", + " print(choice)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fdc1973ace2f" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, delete the resources created in this tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f012708acb9e" + }, + "source": [ + "### Delete Vertex AI Prediction Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3e178bb42c2d" + }, + "outputs": [], + "source": [ + "vertexai_endpoint.delete(force=True, sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "56f628a19cd0" + }, + "source": [ + "### Delete Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bd6f96ed7026" + }, + "outputs": [], + "source": [ + "vertexai_model.delete(sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9e35ebc98df3" + }, + "source": [ + "### Delete private docker repository" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "795a05a169bc" + }, + "outputs": [], + "source": [ + "! gcloud artifacts repositories delete {DOCKER_REPOSITORY} --location={LOCATION} --quiet" + ] + } + ], + "metadata": { + "colab": { + "name": "vertexai_serving_vllm_gpu_llama3_2_3B.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb new file mode 100644 index 000000000..6f03cc511 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb @@ -0,0 +1,809 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f705f4be70e9" + }, + "outputs": [], + "source": [ + "# Copyright 2025 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0266150200d7" + }, + "source": [ + "# Serving Open Models on Vertex AI using vLLM with TPU\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Vertex
Open in Vertex AI Workbench\n", + " \n", + " 		\n", + " \n", + " \"Google
Run in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3b5f1be92f4c" + }, + "source": [ + "## Overview\n", + "\n", + "There are multiple ways of serving open models (including open source and open weight) such as Llama 3.2 on Google Cloud Vertex AI. The Llama models are available in [Model Garden](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) and Model Garden allows a single click self-deployment of the models. This notebooks demonstrates how Llama 3.2 3B model can be served via Vertex AI Endpoint using a custom vLLM container image built for the TPU. This notebook does the following:\n", + "\n", + "- Builds a custom docker container image using vLLM source code\n", + "- Uploads the model to Model Registry using custom docker container image\n", + "- Creates a public Endpoint for Online Prediction\n", + "- Deploys model to the Endpoint\n", + "- Llama 3.2 3B model is downloaded from Google Cloud Storage during deployment\n", + "- This custom container image can also be used for downloading model from Hugging Face\n", + "\n", + "The code in this notebook can be used for serving other open models supported by vLLM. This notebook has been tested with Python 3.10 and `google-cloud-aiplatform` SDK Version `1.106.0`.\n", + "\n", + "To download the models from the Hugging Face, you need a Hugging Face token.\n", + " 1. Create a [Hugging Face account](https://huggingface.co/) if you don't have one.\n", + " 2. For **gated models** like Llama 3.2, ensure you have requested and been granted access on Hugging Face before proceeding.\n", + " 3. Generate an Access Token: Go to **Your Profile > Settings > Access Tokens**.\n", + " 4. Select **New Token**.\n", + " 5. Specify a Name and a Role of at least Read.\n", + " 6. Select **Generate a token**.\n", + " 7. Set the token in hf_token env below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "36df818a34bd" + }, + "source": [ + "## Get Started" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "26645caf62fe" + }, + "source": [ + "### Install Vertex AI SDK for Python and other required packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d4cf289f0d99" + }, + "outputs": [], + "source": [ + "!pip install --upgrade --quiet google-cloud-aiplatform" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "848322ec177e" + }, + "source": [ + "### Restart runtime (Colab only)\n", + "\n", + "To use the newly installed packages, you must restart the runtime on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b8d49bb74a53" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " import IPython\n", + "\n", + " app = IPython.Application.instance()\n", + " app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f332441fe51" + }, + "source": [ + "
\n", + "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "11947ae0fe5e" + }, + "source": [ + "### Authenticate your notebook environment (Colab only)\n", + "\n", + "Authenticate your environment on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "015bf6d5da75" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "722a10c66085" + }, + "source": [ + "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n", + "\n", + "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "66156945acb1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", + "DEVICE_TYPE = \"tpu\" # @param {type:\"string\"}\n", + "\n", + "import vertexai\n", + "\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "202ff51ba5de" + }, + "source": [ + "Set Project ID in active gcloud configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fc8b4506d7ea" + }, + "outputs": [], + "source": [ + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "61583f53cbd1" + }, + "source": [ + "## Prerequisite\n", + "Upload Llama 3.2 3B model to Google Cloud Storage location `BUCKET_URI` before running this notebook. Model can be downloaded from [Hugging Face](https://huggingface.co/meta-llama/Llama-3.2-3B/tree/main) or other repositories and uploaded to Cloud Storage. The following commands are for downloading the model from Hugging Face." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "500fa6c83a27" + }, + "source": [ + "**NOTE:** Downloading model from Hugging Face requires manual input. Run the following command in a shell and when prompted for a password, use an access token with write permissions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4badd3aa3236" + }, + "outputs": [], + "source": [ + "! git lfs install\n", + "! git clone https://huggingface.co/meta-llama/Llama-3.2-3B" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "25ecc2f46bc3" + }, + "outputs": [], + "source": [ + "BUCKET_NAME = f\"{PROJECT_ID}-vertexai-models\" # @param {type:\"string\"}\n", + "BUCKET_URI = f\"gs://{BUCKET_NAME}/meta-llama\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a2f06fadd1b6" + }, + "source": [ + "**If your bucket doesn't already exist**: Run the following cell to create your Google Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2ca7035602da" + }, + "outputs": [], + "source": [ + "! gcloud storage buckets create \"gs://{BUCKET_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "09afcef54e86" + }, + "source": [ + "Upload downloaded model to Cloud Storage location." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "69151a368209" + }, + "outputs": [], + "source": [ + "! gcloud storage cp --recursive Llama-3.2-3B {BUCKET_URI}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "15d914484a1e" + }, + "source": [ + "### Create an IAM Service Account\n", + "When the model is deployed to a Vertex AI Endpoint, it needs to download the model from Cloud Storage bucket and therefore, create a user-managed service account with required permissions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "83dc6a31cc0f" + }, + "outputs": [], + "source": [ + "SERVICE_ACCOUNT_NAME = \"vertexai-endpoint-sa\"\n", + "SERVICE_ACCOUNT_DISPLAY_NAME = \"Vertex AI Endpoint Service Account\"\n", + "SERVICE_ACCOUNT_EMAIL = f\"{SERVICE_ACCOUNT_NAME}@{PROJECT_ID}.iam.gserviceaccount.com\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "601064146e08" + }, + "outputs": [], + "source": [ + "! gcloud iam service-accounts create {SERVICE_ACCOUNT_NAME} \\\n", + " --display-name=\"{SERVICE_ACCOUNT_DISPLAY_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fca977e82b40" + }, + "source": [ + "**NOTE:** You may have to wait for a few seconds before running the next command." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "06e05074ea0c" + }, + "outputs": [], + "source": [ + "! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n", + " --member=\"serviceAccount:{SERVICE_ACCOUNT_EMAIL}\" \\\n", + " --role=\"roles/storage.objectViewer\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fc25cd85ed0f" + }, + "source": [ + "## Create vLLM Customer Container Image for Vertex AI\n", + "\n", + "Vertex AI requires [requests](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#inference) and [responses](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#response_requirements) in specific formats. vLLM API server implements OpenAI API protocol and therefore, it does not support the Vertex AI request and response requirements. Therefore, the vLLM API server (vllm.entrypoints.openai.api_server.py) needs to be updated to support Vertex AI request and response formats." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1620c163c6ce" + }, + "source": [ + "### Enable Artifact Registry API\n", + "Enable the Artifact Registry API service for the Google cloud project. This tutorial requires [gcloud CLI](https://cloud.google.com/sdk/docs/install) installed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4e4296c4b4cd" + }, + "outputs": [], + "source": [ + "! gcloud components update --quiet && gcloud services enable artifactregistry.googleapis.com" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "8d82d43f89d4" + }, + "source": [ + "### Create a private Docker repository\n", + "Create a Docker repository in [Artifact Registry](https://cloud.google.com/artifact-registry/docs/overview)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8041aed2e208" + }, + "outputs": [], + "source": [ + "DOCKER_REPOSITORY = \"my-docker-repo\"\n", + "! gcloud artifacts repositories create {DOCKER_REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Vertex AI Docker repository\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fcbe6342b3f6" + }, + "source": [ + "### Build vLLM Custom Docker Container Image for TPU" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "324a4090cb19" + }, + "source": [ + "Clone vertex-ai-samples code reposistory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "e0943c9ccbca" + }, + "outputs": [], + "source": [ + "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "16a008a0ac1b" + }, + "source": [ + "Build image using Cloud Build" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "e3c86e4d310f" + }, + "outputs": [], + "source": [ + "! cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build \\\n", + " && 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=vllm/vllm-tpu:nightly" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "978d082be4b2" + }, + "source": [ + "## Deploy Model to Vertex AI Endpoint\n", + "\n", + "Following steps are required to serve model via a Vertex AI Prediction Endpoint:\n", + "- import model to model registry\n", + "- create a Online Prediction Endpoint\n", + "- Deploy the model to endpoint" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1fe80a1d509d" + }, + "source": [ + "### Define Variable" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9cb937a1b7b6" + }, + "outputs": [], + "source": [ + "model_name = \"tpu-gcs-llama3_2_3B-serve-vllm\" # @param {type:\"string\"}\n", + "model_id = f\"{BUCKET_URI}/Llama-3.2-3B\" # @param {type:\"string\"}\n", + "machine_type = \"ct5lp-hightpu-1t\" # @param {type:\"string\"}\n", + "tpu_count = 1 # @param {type:\"integer\"}\n", + "DOCKER_URI = (\n", + " f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/vllm-{DEVICE_TYPE}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "13c95aadf084" + }, + "source": [ + "### Import model to Model Registry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2eaff581c1bc" + }, + "outputs": [], + "source": [ + "from google.cloud import aiplatform\n", + "\n", + "\n", + "def upload_model(\n", + " model_name: str,\n", + " model_id: str,\n", + " tpu_count: int,\n", + " docker_uri: str,\n", + ") -> aiplatform.Model:\n", + "\n", + " vllm_args = [\n", + " \"python3\",\n", + " \"-m\",\n", + " \"vllm.entrypoints.openai.api_server\",\n", + " \"--host=0.0.0.0\",\n", + " \"--port=8080\",\n", + " f\"--model={model_id}\",\n", + " \"--max-model-len=2048\",\n", + " \"--enable-prefix-caching\",\n", + " f\"--tensor-parallel-size={tpu_count}\",\n", + " ]\n", + "\n", + " model = aiplatform.Model.upload(\n", + " display_name=model_name,\n", + " serving_container_image_uri=docker_uri,\n", + " serving_container_args=vllm_args,\n", + " serving_container_ports=[8080],\n", + " serving_container_predict_route=\"/v1/completions\",\n", + " serving_container_health_route=\"/health\",\n", + " serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n", + " serving_container_deployment_timeout=1800,\n", + " )\n", + " return model\n", + "\n", + "\n", + "vertexai_model = upload_model(\n", + " model_name=model_name,\n", + " model_id=model_id,\n", + " tpu_count=tpu_count,\n", + " docker_uri=DOCKER_URI,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f99daa0f012b" + }, + "source": [ + "### Create Vertex AI Endpoint for Online Prediction" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "efc6c5bcf498" + }, + "outputs": [], + "source": [ + "def create_model_endpoint(model_name: str) -> aiplatform.Endpoint:\n", + " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n", + " return endpoint\n", + "\n", + "\n", + "vertexai_endpoint = create_model_endpoint(model_name=model_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1d061256a3d4" + }, + "source": [ + "### Deploy Model to Endpoint\n", + "**NOTE**: The model deployment will take around 20-30 minutes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f7000c4082c2" + }, + "outputs": [], + "source": [ + "def deploy_model(\n", + " model: aiplatform.Model,\n", + " endpoint: aiplatform.Endpoint,\n", + " model_name: str,\n", + " machine_type: str,\n", + " service_account: str,\n", + "):\n", + " print(\n", + " f\"Deploying {model_name} to endpoint: {endpoint.resource_name} using machine type: {machine_type}\"\n", + " )\n", + " model.deploy(\n", + " endpoint=endpoint,\n", + " deployed_model_display_name=model_name,\n", + " machine_type=machine_type,\n", + " traffic_percentage=100,\n", + " deploy_request_timeout=1800,\n", + " service_account=service_account,\n", + " min_replica_count=1,\n", + " max_replica_count=1,\n", + " )\n", + "\n", + "\n", + "deploy_model(\n", + " model=vertexai_model,\n", + " endpoint=vertexai_endpoint,\n", + " model_name=model_name,\n", + " machine_type=machine_type,\n", + " service_account=SERVICE_ACCOUNT_EMAIL,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "07d227760d15" + }, + "source": [ + "## Test Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cc7e5ea24220" + }, + "outputs": [], + "source": [ + "import json\n", + "\n", + "PROMPT = \"Distance of moon from earth is \"\n", + "request_body = json.dumps(\n", + " {\n", + " \"prompt\": PROMPT,\n", + " \"temperature\": 0.0,\n", + " },\n", + ")\n", + "\n", + "raw_response = vertexai_endpoint.raw_predict(\n", + " body=request_body, headers={\"Content-Type\": \"application/json\"}\n", + ")\n", + "assert raw_response.status_code == 200\n", + "result = json.loads(raw_response.text)\n", + "\n", + "for choice in result[\"choices\"]:\n", + " print(choice)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "56878e4916ba" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, delete the resources created in this tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "96f9c8266e04" + }, + "source": [ + "### Delete Vertex AI Prediction Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "516e3e746298" + }, + "outputs": [], + "source": [ + "vertexai_endpoint.delete(force=True, sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fa9a5f688ec0" + }, + "source": [ + "### Delete Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "487f1f92beee" + }, + "outputs": [], + "source": [ + "vertexai_model.delete(sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dc44db149a4a" + }, + "source": [ + "### Delete private docker repository" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "4157c8318d20" + }, + "outputs": [], + "source": [ + "! gcloud artifacts repositories delete {DOCKER_REPOSITORY} --location={LOCATION} --quiet" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "cf55f6320692" + }, + "source": [ + "### Delete Cloud Storage Bucket" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6463e1fdc006" + }, + "outputs": [], + "source": [ + "! gcloud storage rm --recursive \"gs://{BUCKET_NAME}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "84cf93af001d" + }, + "source": [ + "### Delete Service Account" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "19252c94635e" + }, + "outputs": [], + "source": [ + "! gcloud iam service-accounts delete {SERVICE_ACCOUNT_NAME}" + ] + } + ], + "metadata": { + "colab": { + "name": "vertexai_serving_vllm_tpu_gcs_llama3_2_3B.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_llama3_2_3B.ipynb b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_llama3_2_3B.ipynb new file mode 100644 index 000000000..8a18da5d1 --- /dev/null +++ b/notebooks/official/prediction/vertexai_serving_vllm/vertexai_serving_vllm_tpu_llama3_2_3B.ipynb @@ -0,0 +1,633 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "f705f4be70e9" + }, + "outputs": [], + "source": [ + "# Copyright 2025 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "865bd455faa2" + }, + "source": [ + "# Serving Open Models on Vertex AI using vLLM with TPU\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Vertex
Open in Vertex AI Workbench\n", + " \n", + " 		\n", + " \n", + " \"Google
Run in Colab Enterprise\n", + " \n", + " 		\n", + " \n", + " \"GitHub
View on GitHub\n", + " \n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "206b4f5e4268" + }, + "source": [ + "## Overview\n", + "\n", + "There are multiple ways of serving open models (including open source and open weight) such as Llama 3.2 on Google Cloud Vertex AI. The Llama models are available in [Model Garden](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) and Model Garden allows a single click self-deployment of the models. This notebooks demonstrates how Llama 3.2 3B model can be served via Vertex AI Endpoint using a custom vLLM container image built for the TPU. This notebook does the following:\n", + "\n", + "- Builds a custom docker container image using vLLM source code\n", + "- Uploads the model to Model Registry using custom docker container image\n", + "- Creates a public Endpoint for Online Prediction\n", + "- Deploys model to the Endpoint\n", + "- Llama 3.2 3B model is downloaded from Hugging Face during deployment\n", + "- This custom container image can also be used for downloading model from Google Storage\n", + "\n", + "The code in this notebook can be used for serving other open models supported by vLLM. This notebook has been tested with Python 3.10 and `google-cloud-aiplatform` SDK Version `1.106.0`.\n", + "\n", + "To download the models from the Hugging Face, you need a Hugging Face token.\n", + " 1. Create a [Hugging Face account](https://huggingface.co/) if you don't have one.\n", + " 2. For **gated models** like Llama 3.2, ensure you have requested and been granted access on Hugging Face before proceeding.\n", + " 3. Generate an Access Token: Go to **Your Profile > Settings > Access Tokens**.\n", + " 4. Select **New Token**.\n", + " 5. Specify a Name and a Role of at least Read.\n", + " 6. Select **Generate a token**.\n", + " 7. Set the token in hf_token env below." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "36df818a34bd" + }, + "source": [ + "## Get Started" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "26645caf62fe" + }, + "source": [ + "### Install Vertex AI SDK for Python and other required packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d4cf289f0d99" + }, + "outputs": [], + "source": [ + "!pip install --upgrade --quiet google-cloud-aiplatform" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "848322ec177e" + }, + "source": [ + "### Restart runtime (Colab only)\n", + "\n", + "To use the newly installed packages, you must restart the runtime on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "b8d49bb74a53" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " import IPython\n", + "\n", + " app = IPython.Application.instance()\n", + " app.kernel.do_shutdown(True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f332441fe51" + }, + "source": [ + "
\n", + "⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "11947ae0fe5e" + }, + "source": [ + "### Authenticate your notebook environment (Colab only)\n", + "\n", + "Authenticate your environment on Google Colab." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "015bf6d5da75" + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if \"google.colab\" in sys.modules:\n", + "\n", + " from google.colab import auth\n", + "\n", + " auth.authenticate_user()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "722a10c66085" + }, + "source": [ + "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n", + "\n", + "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "66156945acb1" + }, + "outputs": [], + "source": [ + "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n", + "LOCATION = \"us-central1\" # @param {type:\"string\"}\n", + "DEVICE_TYPE = \"tpu\" # @param {type:\"string\"}\n", + "\n", + "import vertexai\n", + "\n", + "vertexai.init(project=PROJECT_ID, location=LOCATION)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "202ff51ba5de" + }, + "source": [ + "Set Project ID in active gcloud configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fc8b4506d7ea" + }, + "outputs": [], + "source": [ + "! gcloud config set project {PROJECT_ID}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "7cce406746d5" + }, + "source": [ + "## Create vLLM Customer Container Image for Vertex AI\n", + "\n", + "Vertex AI requires [requests](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#inference) and [responses](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#response_requirements) in specific formats. vLLM API server implements OpenAI API protocol and therefore, it does not support the Vertex AI request and response requirements. Therefore, the vLLM API server (vllm.entrypoints.openai.api_server.py) needs to be updated to support Vertex AI request and response formats." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "2f7ec1c7bdf8" + }, + "source": [ + "### Enable Artifact Registry API\n", + "Enable the Artifact Registry API service for the Google cloud project. This tutorial requires [gcloud CLI](https://cloud.google.com/sdk/docs/install) installed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6e93c586de15" + }, + "outputs": [], + "source": [ + "! gcloud components update --quiet && gcloud services enable artifactregistry.googleapis.com" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6c4e1d4d5e66" + }, + "source": [ + "### Create a private Docker repository\n", + "Create a Docker repository in [Artifact Registry](https://cloud.google.com/artifact-registry/docs/overview)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1f21a9e18e5c" + }, + "outputs": [], + "source": [ + "DOCKER_REPOSITORY = \"my-docker-repo\"\n", + "! gcloud artifacts repositories create {DOCKER_REPOSITORY} --repository-format=docker --location={LOCATION} --description=\"Vertex AI Docker repository\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "abf2da8199cc" + }, + "source": [ + "### Build vLLM Custom Docker Container Image for TPU" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "84b6342dd2f0" + }, + "source": [ + "Clone vertex-ai-samples code reposistory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "13df52a033d3" + }, + "outputs": [], + "source": [ + "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "b89cf1a29b80" + }, + "source": [ + "Build image using Cloud Build" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8c00680ff433" + }, + "outputs": [], + "source": [ + "! cd vertex-ai-samples/notebooks/official/prediction/vertexai_serving_vllm/cloud-build \\\n", + " && 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=vllm/vllm-tpu:nightly" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "95c5ddf6f5c6" + }, + "source": [ + "## Deploy Model to Vertex AI Endpoint\n", + "\n", + "Following steps are required to serve model via a Vertex AI Prediction Endpoint:\n", + "- import model to model registry\n", + "- create a Online Prediction Endpoint\n", + "- Deploy the model to endpoint" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "dd851913ae85" + }, + "source": [ + "### Define Variable" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "743ef67b7b22" + }, + "outputs": [], + "source": [ + "hf_token = \"[your-hugging-face-auth-token]\" # @param {type:\"string\"}\n", + "model_name = \"tpu-llama3_2_3B-serve-vllm\" # @param {type:\"string\"}\n", + "model_id = \"meta-llama/Llama-3.2-3B\" # @param {type:\"string\"}\n", + "machine_type = \"ct5lp-hightpu-1t\" # @param {type:\"string\"}\n", + "tpu_count = 1 # @param {type:\"integer\"}\n", + "DOCKER_URI = (\n", + " f\"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/vllm-{DEVICE_TYPE}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "e0c1bb82c9e0" + }, + "source": [ + "### Import model to Model Registry" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "74adebc02ac2" + }, + "outputs": [], + "source": [ + "from google.cloud import aiplatform\n", + "\n", + "\n", + "def upload_model(\n", + " model_name: str,\n", + " model_id: str,\n", + " hf_token: str,\n", + " tpu_count: int,\n", + " docker_uri: str,\n", + ") -> aiplatform.Model:\n", + "\n", + " vllm_args = [\n", + " \"python3\",\n", + " \"-m\",\n", + " \"vllm.entrypoints.openai.api_server\",\n", + " \"--host=0.0.0.0\",\n", + " \"--port=8080\",\n", + " f\"--model={model_id}\",\n", + " \"--max-model-len=2048\",\n", + " \"--enable-prefix-caching\",\n", + " f\"--tensor-parallel-size={tpu_count}\",\n", + " ]\n", + "\n", + " env_vars = {\"HF_TOKEN\": hf_token}\n", + "\n", + " model = aiplatform.Model.upload(\n", + " display_name=model_name,\n", + " serving_container_image_uri=docker_uri,\n", + " serving_container_args=vllm_args,\n", + " serving_container_ports=[8080],\n", + " serving_container_predict_route=\"/v1/completions\",\n", + " serving_container_health_route=\"/health\",\n", + " serving_container_environment_variables=env_vars,\n", + " serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n", + " serving_container_deployment_timeout=1800,\n", + " )\n", + " return model\n", + "\n", + "\n", + "vertexai_model = upload_model(\n", + " model_name=model_name,\n", + " model_id=model_id,\n", + " hf_token=hf_token,\n", + " tpu_count=tpu_count,\n", + " docker_uri=DOCKER_URI,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "124c8611bab3" + }, + "source": [ + "### Create Vertex AI Endpoint for Online Prediction" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "91591af558cf" + }, + "outputs": [], + "source": [ + "def create_model_endpoint(model_name: str) -> aiplatform.Endpoint:\n", + " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n", + " return endpoint\n", + "\n", + "\n", + "vertexai_endpoint = create_model_endpoint(model_name=model_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a4d9c500a13e" + }, + "source": [ + "### Deploy Model to Endpoint\n", + "**NOTE**: The model deployment will take around 20-30 minutes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "309ae44ffdfb" + }, + "outputs": [], + "source": [ + "def deploy_model(\n", + " model: aiplatform.Model,\n", + " endpoint: aiplatform.Endpoint,\n", + " model_name: str,\n", + " machine_type: str,\n", + "):\n", + " print(\n", + " f\"Deploying {model_name} to endpoint: {endpoint.resource_name} using machine type: {machine_type}\"\n", + " )\n", + " model.deploy(\n", + " endpoint=endpoint,\n", + " deployed_model_display_name=model_name,\n", + " machine_type=machine_type,\n", + " min_replica_count=1,\n", + " max_replica_count=4,\n", + " autoscaling_target_request_count_per_minute=60,\n", + " traffic_percentage=100,\n", + " deploy_request_timeout=1800,\n", + " )\n", + "\n", + "\n", + "deploy_model(\n", + " model=vertexai_model,\n", + " endpoint=vertexai_endpoint,\n", + " model_name=model_name,\n", + " machine_type=machine_type,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "13f1059160b9" + }, + "source": [ + "## Test Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "8ec42d0c4422" + }, + "outputs": [], + "source": [ + "import json\n", + "\n", + "PROMPT = \"Distance of moon from earth is \"\n", + "request_body = json.dumps(\n", + " {\n", + " \"prompt\": PROMPT,\n", + " \"temperature\": 0.0,\n", + " },\n", + ")\n", + "\n", + "raw_response = vertexai_endpoint.raw_predict(\n", + " body=request_body, headers={\"Content-Type\": \"application/json\"}\n", + ")\n", + "assert raw_response.status_code == 200\n", + "result = json.loads(raw_response.text)\n", + "\n", + "for choice in result[\"choices\"]:\n", + " print(choice)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fdc1973ace2f" + }, + "source": [ + "## Cleaning up\n", + "\n", + "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n", + "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n", + "\n", + "Otherwise, delete the resources created in this tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f012708acb9e" + }, + "source": [ + "### Delete Vertex AI Prediction Endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3e178bb42c2d" + }, + "outputs": [], + "source": [ + "vertexai_endpoint.delete(force=True, sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "56f628a19cd0" + }, + "source": [ + "### Delete Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "bd6f96ed7026" + }, + "outputs": [], + "source": [ + "vertexai_model.delete(sync=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9e35ebc98df3" + }, + "source": [ + "### Delete private docker repository" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "795a05a169bc" + }, + "outputs": [], + "source": [ + "! gcloud artifacts repositories delete {DOCKER_REPOSITORY} --location={LOCATION} --quiet" + ] + } + ], + "metadata": { + "colab": { + "name": "vertexai_serving_vllm_tpu_llama3_2_3B.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +}