notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - CodeGemma Model (Deployment)\n",
        "\n",
        "<table><tbody><tr>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_codegemma_deployment_on_vertex.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying CodeGemma models\n",
        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
        " * on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy CodeGemma with Hex-LLM on TPU\n",
        "- Deploy CodeGemma with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nLsuvskfhOv4"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
        "\n",
        "# @markdown 3. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 4. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Import the necessary packages\n",
        "# Upgrade Vertex AI SDK.\n",
        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import os\n",
        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "models, endpoints = {}, {}\n",
        "# Dedicated endpoint not supported yet\n",
        "use_dedicated_endpoint = False\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"codegemma\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "45c8c5438737"
      },
      "outputs": [],
      "source": [
        "# @title Access CodeGemma Models\n",
        "\n",
        "# @markdown If you already obtained access to CodeGemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
        "# @markdown Alternatively, you can also load the original CodeGemma models for serving from Vertex AI after accepting the agreement.\n",
        "# @markdown **Select and fill one of the two following sections.**\n",
        "LOAD_MODEL_FROM = \"Google Cloud\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
        "\n",
        "# @markdown #### Access CodeGemma models on Vertex AI\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [CodeGemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/364) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
        "# @markdown 3. After accepting the agreement of CodeGemma, a `https://` link containing CodeGemma pretrained and finetuned models will be shared.\n",
        "# @markdown 4. Paste the link in the `VERTEX_MODEL_GARDEN_CODEGEMMA` field below.\n",
        "# @markdown **Note:** This will unzip and copy the CodeGemma model artifacts to your Cloud Storage bucket, which will take around 30 minutes.\n",
        "\n",
        "VERTEX_MODEL_GARDEN_CODEGEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "VERTEX_MODEL_GARDEN_CODEGEMMA = VERTEX_MODEL_GARDEN_CODEGEMMA.replace(\n",
        "    \"gs://\", \"https://storage.googleapis.com/\", 1\n",
        ")\n",
        "\n",
        "# @markdown *--- Or ---*\n",
        "\n",
        "# @markdown ##Enable the Vertex AI API## Access CodeGemma models on HuggingFace\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the CodeGemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
        "    assert (\n",
        "        HF_TOKEN\n",
        "    ), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
        "\n",
        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
        "    assert (\n",
        "        VERTEX_MODEL_GARDEN_CODEGEMMA\n",
        "    ), \"Click the agreement of CodeGemma in Vertex AI Model Garden, and get the URL to CodeGemma model artifacts.\"\n",
        "\n",
        "    # Only use the last part in case a full command is pasted.\n",
        "    signed_url = VERTEX_MODEL_GARDEN_CODEGEMMA.split(\" \")[-1].strip('\"')\n",
        "\n",
        "    ! mkdir -p ./codegemma\n",
        "    ! curl -X GET \"{signed_url}\" | tar -xzvf - -C ./codegemma/\n",
        "    ! gsutil -m cp -R ./codegemma/* {BUCKET_URI}\n",
        "\n",
        "    model_path_prefix = BUCKET_URI.strip(\"/\") + \"/codegemma\"\n",
        "else:\n",
        "    model_path_prefix = \"google/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8neJc8CnDDpu"
      },
      "source": [
        "## Deploy CodeGemma models with Hex-LLM on TPU\n",
        "\n",
        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E8OiHHNNE_wj"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "\n",
        "# @markdown Set the model to deploy.\n",
        "\n",
        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "hf_model_id = \"google/\" + MODEL_ID\n",
        "\n",
        "# The pre-built serving docker image for Hex-LLM.\n",
        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
        "\n",
        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
        "if \"2b\" in model_id:\n",
        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy CodeGemma 2B models.\n",
        "    machine_type = \"ct5lp-hightpu-1t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 1\n",
        "else:\n",
        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy CodeGemma 7B models.\n",
        "    machine_type = \"ct5lp-hightpu-4t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 4\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "# Server parameters.\n",
        "tensor_parallel_size = accelerator_count\n",
        "hbm_utilization_factor = 0.6  # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
        "max_running_seqs = 256\n",
        "\n",
        "# Endpoint configurations.\n",
        "min_replica_count = 1\n",
        "max_replica_count = 1\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    tensor_parallel_size: int = 1,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
        "    tpu_topology: str = \"1x1\",\n",
        "    hbm_utilization_factor: float = 0.6,\n",
        "    max_running_seqs: int = 256,\n",
        "    max_model_len: int = 4096,\n",
        "    endpoint_id: str = \"\",\n",
        "    min_replica_count: int = 1,\n",
        "    max_replica_count: int = 1,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
        "    if endpoint_id:\n",
        "        aip_endpoint_name = (\n",
        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
        "        )\n",
        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "    else:\n",
        "        endpoint = aiplatform.Endpoint.create(\n",
        "            display_name=f\"{model_name}-endpoint\",\n",
        "            location=TPU_DEPLOYMENT_REGION,\n",
        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "        )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    if not tensor_parallel_size:\n",
        "        tensor_parallel_size = int(machine_type[-2])\n",
        "\n",
        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
        "\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
        "        f\"--num_hosts={num_hosts}\",\n",
        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
        "        f\"--max_running_seqs={max_running_seqs}\",\n",
        "        f\"--max_model_len={max_model_len}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
        "    except:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
        "        serving_container_args=hexllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "        location=TPU_DEPLOYMENT_REGION,\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        min_replica_count=min_replica_count,\n",
        "        max_replica_count=max_replica_count,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
        "    model_id=model_id,\n",
        "    base_model_id=hf_model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    tensor_parallel_size=tensor_parallel_size,\n",
        "    hbm_utilization_factor=hbm_utilization_factor,\n",
        "    max_running_seqs=max_running_seqs,\n",
        "    min_replica_count=min_replica_count,\n",
        "    max_replica_count=max_replica_count,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nkUaMxIus6Pv"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
        "#   name of the endpoint `endpoint` created in the cell above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint:\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoints[\"hexllm_tpu\"] = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
        "max_tokens = 500  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
        "top_k = 1  # @param {type:\"integer\"}\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "    },\n",
        "]\n",
        "response = endpoints[\"hexllm_tpu\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "# \"<|file_separator|>\" is the end of the file token.\n",
        "for prediction in response.predictions:\n",
        "    print(prediction.split(\"<|file_separator|>\")[0])\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YKZ4CBJ2kYaW"
      },
      "source": [
        "## Deploy CodeGemma models with vLLM on GPU\n",
        "\n",
        "[vLLM](https://github.com/vllm-project/vllm) is a high-throughput GPU Large Language Model (LLM) serving library which implements a number of optimizations including paged attention and continuous batching.\n",
        "\n",
        "Note that V100 GPUs generally offer better throughput and latency performance than L4 GPUs, while L4 GPUs are generally more cost efficient than V100 GPUs. The serving efficiency of L4, V100 and T4 GPUs is inferior to that of A100 GPUs, but L4, V100 and T4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
        "\n",
        "CodeGemma model weights are stored in bfloat16 precision. L4 and A100 GPUs are needed for vLLM serving at bfloat16 precision. V100 and T4 GPUs can support vLLM serving at float32 and float16 precision, and they are also meaningful deployment configurations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "tQIEisUajS6t"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "\n",
        "# @markdown This section uploads prebuilt CodeGemma model to Model Registry and deploys with [vLLM](https://github.com/vllm-project/vllm) to a Vertex AI Endpoint. It takes 15 to 30 minutes to finish depending on the model and the accelerator.\n",
        "\n",
        "# @markdown Set the model to deploy.\n",
        "\n",
        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
        "# The pre-built serving docker image for vLLM.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
        "\n",
        "# Find Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "# @markdown L4 GPUs are good serving solutions and are more cost effective than A100s.\n",
        "\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
        "\n",
        "if accelerator_type == \"NVIDIA_L4\":\n",
        "    machine_type = \"g2-standard-12\"\n",
        "    accelerator_count = 1\n",
        "elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "    machine_type = \"a2-highgpu-1g\"\n",
        "    accelerator_count = 1\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "# Larger setting of `max-model-len` can lead to higher requirements on\n",
        "# `gpu-memory-utilization` and GPU configuration. Larger setting of\n",
        "# `gpu-memory-utilization` increases the risk of running out of GPU memory with\n",
        "# long prompts.\n",
        "max_model_len = 2048\n",
        "gpu_memory_utilization = 0.9\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        "    max_model_len: int = 4096,\n",
        "    dtype: str = \"auto\",\n",
        "    enable_trust_remote_code: bool = False,\n",
        "    enforce_eager: bool = False,\n",
        "    enable_lora: bool = False,\n",
        "    max_loras: int = 1,\n",
        "    max_cpu_loras: int = 8,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        "    max_num_seqs: int = 256,\n",
        "    model_type: str = None,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
        "    vllm_args = [\n",
        "        \"python\",\n",
        "        \"-m\",\n",
        "        \"vllm.entrypoints.api_server\",\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=8080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        f\"--max-loras={max_loras}\",\n",
        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
        "        f\"--max-num-seqs={max_num_seqs}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    if enable_trust_remote_code:\n",
        "        vllm_args.append(\"--trust-remote-code\")\n",
        "\n",
        "    if enforce_eager:\n",
        "        vllm_args.append(\"--enforce-eager\")\n",
        "\n",
        "    if enable_lora:\n",
        "        vllm_args.append(\"--enable-lora\")\n",
        "\n",
        "    if model_type:\n",
        "        vllm_args.append(f\"--model-type={model_type}\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=\"codegemma-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_model_len=max_model_len,\n",
        "    gpu_memory_utilization=gpu_memory_utilization,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RRR11SWykYaX"
      },
      "source": [
        "Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64). Setting `raw_response` to `True` allows you to obtain raw outputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3f5a1e1de60d"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoints[\"vllm_gpu\"] = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
        "max_tokens = 500  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
        "top_k = 1  # @param {type:\"integer\"}\n",
        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
        "raw_response = True  # @param {type:\"boolean\"}\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
        "response = endpoints[\"vllm_gpu\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "# \"<|file_separator|>\" is the end of the file token.\n",
        "for prediction in response.predictions:\n",
        "    print(prediction.split(\"<|file_separator|>\")[0])\n",
        "\n",
        "# @markdown Click \"Show code\" to see more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kMYK62ZOk-qA"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @title Delete the model and endpoint\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_codegemma_deployment_on_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}