notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.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": [
        "# Gemma deployment to GKE using TGI on GPU\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_gemma_deployment_on_gke.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_gemma_deployment_on_gke.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 downloading and deploying Gemma, open models from Google DeepMind using Text Generation Inference [TGI](https://github.com/), an efficient serving option to improve serving throughput. In this notebook we will deploy and serve TGI on GPUs. In this guide we specifically use L4 GPUs but this guide should also work for A100(40 GB), A100(80 GB), H100(80 GB) GPUs.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "Deploy and run inference for serving Gemma with TGI on GPUs.\n",
        "\n",
        "### GPUs\n",
        "\n",
        "GPUs let you accelerate specific workloads running on your nodes such as machine learning and data processing. GKE provides a range of machine type options for node configuration, including machine types with NVIDIA H100, L4, and A100 GPUs.\n",
        "\n",
        "Before you use GPUs in GKE, we recommend that you complete the following learning path:\n",
        "\n",
        "Learn about [current GPU version availability](https://cloud.google.com/compute/docs/gpus)\n",
        "\n",
        "Learn about [GPUs in GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/gpus)\n",
        "\n",
        "\n",
        "### TGI\n",
        "\n",
        "TGI is a highly optimized open-source LLM serving framework that can increase serving throughput on GPUs. TGI includes features such as:\n",
        "\n",
        "Optimized transformer implementation with PagedAttention\n",
        "Continuous batching to improve the overall serving throughput\n",
        "Tensor parallelism and distributed serving on multiple GPUs\n",
        "\n",
        "To learn more, refer to the [TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "855d6b96f291"
      },
      "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. Set Hugging Face access token in `HF_TOKEN` field. If you don't already have a \"read\" access token, follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create an access token with \"read\" permission. You can find your existing access tokens in the Hugging Face [Access Token](https://huggingface.co/settings/tokens) page.\n",
        "\n",
        "# @markdown 3. **[Optional]** Set `CLUSTER_NAME` if you want to use your own GKE cluster. If not set, this example will create a standard cluster with 2 NVIDIA L4 GPU accelerators.\n",
        "\n",
        "import datetime\n",
        "import os\n",
        "\n",
        "# The HuggingFace token used to download models.\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\"}\n",
        "assert HF_TOKEN, \"Set Hugging Face access token in `HF_TOKEN`.\"\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",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Set up gcloud.\n",
        "! gcloud config set project \"$PROJECT_ID\"\n",
        "! gcloud services enable container.googleapis.com\n",
        "\n",
        "# Add kubectl to the set of available tools.\n",
        "! mkdir -p /tools/google-cloud-sdk/.install\n",
        "! gcloud components install kubectl --quiet\n",
        "\n",
        "# The cluster name to create\n",
        "CLUSTER_NAME = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Use existing GKE cluster or create a new cluster.\n",
        "if CLUSTER_NAME:\n",
        "    ! gcloud container clusters get-credentials {CLUSTER_NAME} --location {REGION}\n",
        "else:\n",
        "    now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "    CLUSTER_NAME=f\"gke-gemma-cluster-{now}\"\n",
        "    ! gcloud container clusters create {CLUSTER_NAME} \\\n",
        "        --project={PROJECT_ID} \\\n",
        "        --region={REGION} \\\n",
        "        --workload-pool={PROJECT_ID}.svc.id.goog \\\n",
        "        --release-channel=rapid \\\n",
        "        --num-nodes=4\n",
        "    ! gcloud container node-pools create gpupool \\\n",
        "        --accelerator=type=nvidia-l4,count=2,gpu-driver-version=latest \\\n",
        "        --project={PROJECT_ID} \\\n",
        "        --location={REGION} \\\n",
        "        --node-locations={REGION}-a \\\n",
        "        --cluster={CLUSTER_NAME} \\\n",
        "        --machine-type=g2-standard-24 \\\n",
        "        --num-nodes=1\n",
        "\n",
        "# Create Kubernetes secret for Hugging Face credentials\n",
        "! kubectl create secret generic hf-secret \\\n",
        "    --from-literal=hf_api_token={HF_TOKEN} \\\n",
        "    --dry-run=client -o yaml > hf-secret.yaml\n",
        "\n",
        "! kubectl apply -f hf-secret.yaml"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6psJZY_zUDgj"
      },
      "outputs": [],
      "source": [
        "# @title Deploy TGI\n",
        "\n",
        "# @markdown This section deploys Gemma on TGI.\n",
        "\n",
        "# @markdown Select one of the following model version and size options:\n",
        "\n",
        "# The size of the model to launch\n",
        "MODEL_VERSION = \"1.1\"  # @param [\"1.0\", \"1.1\"]\n",
        "if MODEL_VERSION == \"1.1\":\n",
        "    version_string = \"-1.1\"\n",
        "else:\n",
        "    version_string = \"\"\n",
        "# The size of the model to launch\n",
        "MODEL_SIZE = \"2b\"  # @param [\"2b\", \"7b\"]\n",
        "\n",
        "# @markdown After the container is up, there will be another ~5 minutes to download the needed artifacts, the time depends on what runtime you are using to run your colab environment.\n",
        "\n",
        "# The number of GPUs to run: 1 for 2b, 2 for 7b\n",
        "GPU_COUNT = 1\n",
        "if MODEL_SIZE == \"7b\":\n",
        "    GPU_COUNT = 2\n",
        "\n",
        "# Ephemeral storage\n",
        "EPHEMERAL_STORAGE_SIZE = \"20Gi\"\n",
        "if MODEL_SIZE == \"7b\":\n",
        "    EPHEMERAL_STORAGE_SIZE = \"40Gi\"\n",
        "\n",
        "# Memory size\n",
        "MEMORY_SIZE = \"7Gi\"\n",
        "if MODEL_SIZE == \"7b\":\n",
        "    MEMORY_SIZE = \"25Gi\"\n",
        "\n",
        "GPU_SHARD = 1\n",
        "if MODEL_SIZE == \"7b\":\n",
        "    GPU_SHARD = 2\n",
        "\n",
        "CPU_LIMITS = 2\n",
        "if MODEL_SIZE == \"7b\":\n",
        "    CPU_LIMITS = 10\n",
        "\n",
        "K8S_YAML = f\"\"\"\n",
        "apiVersion: apps/v1\n",
        "kind: Deployment\n",
        "metadata:\n",
        "  name: tgi-gemma-deployment\n",
        "spec:\n",
        "  replicas: 1\n",
        "  selector:\n",
        "    matchLabels:\n",
        "      app: gemma-server\n",
        "  template:\n",
        "    metadata:\n",
        "      labels:\n",
        "        app: gemma-server\n",
        "        ai.gke.io/model: gemma{version_string}-{MODEL_SIZE}\n",
        "        ai.gke.io/inference-server: text-generation-inference\n",
        "        examples.ai.gke.io/source: user-guide\n",
        "    spec:\n",
        "      containers:\n",
        "      - name: inference-server\n",
        "        image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-hf-tgi-serve:20240220_0936_RC01\n",
        "        resources:\n",
        "          requests:\n",
        "            cpu: \"2\"\n",
        "            memory: {MEMORY_SIZE}\n",
        "            ephemeral-storage: {EPHEMERAL_STORAGE_SIZE}\n",
        "            nvidia.com/gpu: {GPU_COUNT}\n",
        "          limits:\n",
        "            cpu: {CPU_LIMITS}\n",
        "            memory: {MEMORY_SIZE}\n",
        "            ephemeral-storage: {EPHEMERAL_STORAGE_SIZE}\n",
        "            nvidia.com/gpu: {GPU_COUNT}\n",
        "        args:\n",
        "        - --model-id=$(MODEL_ID)\n",
        "        - --num-shard={GPU_SHARD}\n",
        "        env:\n",
        "        - name: MODEL_ID\n",
        "          value: google/gemma{version_string}-{MODEL_SIZE}-it\n",
        "        - name: PORT\n",
        "          value: \"8000\"\n",
        "        - name: HUGGING_FACE_HUB_TOKEN\n",
        "          valueFrom:\n",
        "            secretKeyRef:\n",
        "              name: hf-secret\n",
        "              key: hf_api_token\n",
        "        volumeMounts:\n",
        "        - mountPath: /dev/shm\n",
        "          name: dshm\n",
        "      volumes:\n",
        "      - name: dshm\n",
        "        emptyDir:\n",
        "          medium: Memory\n",
        "      nodeSelector:\n",
        "        cloud.google.com/gke-accelerator: nvidia-l4\n",
        "---\n",
        "apiVersion: v1\n",
        "kind: Service\n",
        "metadata:\n",
        "  name: llm-service\n",
        "spec:\n",
        "  selector:\n",
        "    app: gemma-server\n",
        "  type: ClusterIP\n",
        "  ports:\n",
        "  - protocol: TCP\n",
        "    port: 8000\n",
        "    targetPort: 8000\n",
        "\"\"\"\n",
        "\n",
        "with open(\"tgi.yaml\", \"w\") as f:\n",
        "    f.write(K8S_YAML)\n",
        "\n",
        "! kubectl apply -f tgi.yaml\n",
        "\n",
        "# Wait for container to be created.\n",
        "import time\n",
        "\n",
        "print(\"Waiting for container to be created...\\n\")\n",
        "while True:\n",
        "    shell_output = ! kubectl get pod\n",
        "    container_status = \"\\n\".join(shell_output)\n",
        "    if \"1/1\" in container_status:\n",
        "        break\n",
        "    time.sleep(5)\n",
        "\n",
        "print(container_status)\n",
        "\n",
        "# Wait for downloading artifacts.\n",
        "print(\"\\nDownloading artifacts...\")\n",
        "while True:\n",
        "    shell_output = ! kubectl logs -l app=gemma-server\n",
        "    logs = \"\\n\".join(shell_output)\n",
        "    if \"Connected\" in logs:\n",
        "        break\n",
        "    time.sleep(5)\n",
        "\n",
        "print(\"Server is up and running.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QDTasPgGW7EG"
      },
      "outputs": [],
      "source": [
        "# @title Prediction\n",
        "\n",
        "# @markdown Once the server is up and running, you may send prompts to local server for prediction.\n",
        "\n",
        "import json\n",
        "\n",
        "prompt = \"What are the top 5 most popular programming languages? Be brief.\"  # @param {type: \"string\"}\n",
        "temperature = 0.40  # @param {type: \"number\"}\n",
        "top_p = 0.1  # @param {type: \"number\"}\n",
        "max_tokens = 250  # @param {type: \"number\"}\n",
        "\n",
        "request = {\n",
        "    \"inputs\": prompt,\n",
        "    \"temperature\": temperature,\n",
        "    \"top_p\": top_p,\n",
        "    \"max_tokens\": max_tokens,\n",
        "}\n",
        "\n",
        "command = f\"\"\"kubectl exec -t $( kubectl get pod -l app=gemma-server -o jsonpath=\"{{.items[0].metadata.name}}\" ) -c inference-server -- curl -X POST http://localhost:8000/generate \\\n",
        "   -H \"Content-Type: application/json\" \\\n",
        "   -d '{json.dumps(request)}' \\\n",
        "   2> /dev/null\"\"\"\n",
        "\n",
        "output = !{command}\n",
        "print(\"Output:\")\n",
        "print(json.loads(output[0])[\"generated_text\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wbRmgoOZF6es"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "\n",
        "! kubectl delete deployments tgi-gemma-deployment\n",
        "! kubectl delete services llm-service\n",
        "! kubectl delete secrets hf-secret\n",
        "\n",
        "DELETE_CLUSTER = False # @param {type: \"boolean\"}\n",
        "\n",
        "if DELETE_CLUSTER:\n",
        "  ! gcloud container clusters delete {CLUSTER_NAME} \\\n",
        "    --region={REGION} \\\n",
        "    --quiet"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_deployment_on_gke.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}