{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "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": "99c1c3fc2ca5"
      },
      "source": [
        "# Deploying Llama 3 on Google Kubernetes Engine with Cloud Functions and vLLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXqCSRBUX_6Q"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Fuse-cases%2Fdeploy_llama3_vllm_gke_cloud_function.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/use-cases/deploy_llama3_vllm_gke_cloud_function.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPhuLTRXZAy_"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [KC Ayyagari](https://github.com/krishchyt) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "This notebook demonstrates how to deploy and serve Llama 3 models on Google Kubernetes Engine (GKE) using GPUs, and how to integrate this deployment with a Cloud Function to create an accessible API endpoint. This uses Virtual Large Language Model [vLLM](https://developers.googleblog.com/en/inference-with-gemma-using-dataflow-and-vllm/#:~:text=model%20frameworks%20simple.-,What%20is%20vLLM%3F,-vLLM%20is%20an) inference server.\n",
        "\n",
        "## Objective\n",
        "\n",
        "*   Deploy and run inference for serving Llama 3 on GPUs.\n",
        "*   Create a Cloud Function.\n",
        "*   Call Llama from the Cloud Function.\n",
        "\n",
        "![Workflow diagram](https://github.com/krishchyt/kc-demoassets/blob/main/gke-cf-workflow.png?raw=true)\n",
        "\n",
        "## Key Steps\n",
        "\n",
        "### Setup\n",
        "\n",
        "*   Authenticates with Google Cloud.\n",
        "*   Sets up the Google Cloud project, region, network, and subnet configurations.\n",
        "*   Enables the necessary Google Cloud services (Container Registry, VPC Access).\n",
        "*   Installs `kubectl`.\n",
        "\n",
        "### GKE Cluster Creation\n",
        "\n",
        "*   Creates an Autopilot GKE cluster with specified configurations (region, network, private nodes, etc.).\n",
        "\n",
        "### Llama 3 Deployment on GKE\n",
        "\n",
        "*   Defines Kubernetes deployment and service configurations for deploying a selected Llama 3 model on GKE using vLLM.\n",
        "*   Allows the user to select from a range of Llama 3 models (Llama-3-2-1B-Instruct, Llama-3-2-11B-Vision, etc.).\n",
        "*   Creates a Kubernetes YAML file based on the selected model and applies it to the cluster.\n",
        "*   Waits for the container to be created and the server to be up and running.\n",
        "\n",
        "### Endpoint Testing\n",
        "\n",
        "*   Retrieves the Internal Load Balancer IP address for the Llama service.\n",
        "*   Sends a test request to the Llama 3 endpoint and prints the response.\n",
        "\n",
        "### Cloud Function Integration\n",
        "\n",
        "*   Creates a service account with the necessary permissions.\n",
        "*   Prepares the code for a Cloud Function that will act as an API endpoint for the Llama 3 model.\n",
        "*   Creates a VPC connector to allow the Cloud Function to access the GKE cluster.\n",
        "*   Deploys the Cloud Function, configuring environment variables (including the Llama endpoint IP) and other settings.\n",
        "\n",
        "### Cleanup\n",
        "\n",
        "*   Provides options to delete the deployment, cluster, Cloud Function, and VPC connector to avoid unnecessary charges.\n",
        "\n",
        "## Use Cases\n",
        "\n",
        "This setup enables various use cases, including:\n",
        "\n",
        "*   **Chatbots/Conversational AI:** Building a chatbot that can answer user questions, provide information, or engage in conversations. The Cloud Function acts as the API endpoint, receiving user input and passing it to the Llama 3 model running on GKE for generating responses.\n",
        "*   **Content Generation:** Generating different kinds of creative text formats, like poems, code, scripts, musical pieces, email, letters, etc. The Cloud Function receives a prompt and parameters (e.g., length, style) and uses Llama 3 to generate the content.\n",
        "*   **Question Answering:** Providing answers to specific questions based on a knowledge base. The Cloud Function receives the question, and Llama 3 extracts the relevant information and formulates an answer.\n",
        "*   **Code Generation/Completion:** Assisting developers by generating code snippets or completing code based on context. The Cloud Function receives the code context and uses Llama 3 to suggest code completions or generate entire functions.\n",
        "*   **Image Captioning and Visual Question Answering (if using a vision-enabled model):** If you're using a Llama 3 model with vision capabilities, you can use this setup for image captioning (generating descriptions of images) or visual question answering (answering questions about the content of an image).\n",
        "*   **Data Analysis and Insights:** Use Llama 3 to analyze text data and extract insights. The Cloud Function can receive text data, pass it to Llama 3 for analysis, and then return the insights.\n",
        "*   **Workflow Automation:** Integrate Llama 3 into automated workflows to perform tasks such as summarizing documents, extracting key information, or translating text.\n",
        "\n",
        "\n",
        "## Benefits\n",
        "\n",
        "*   **Scalability:** GKE allows you to scale the Llama 3 deployment based on demand.\n",
        "*   **Cost-Effectiveness:** Cloud Functions are serverless and only charged when used.\n",
        "*   **Flexibility:** You can easily update the Llama 3 model or the Cloud Function code.\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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vyf-fGQ3Znts"
      },
      "outputs": [],
      "source": [
        "%pip install requests>=2.26.0 functions-framework>=3.0.0 google-cloud-logging>=3.0.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PZAiVWhA-efF"
      },
      "outputs": [],
      "source": [
        "! gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJMLYKTIUG-c"
      },
      "source": [
        "## Setup Google Cloud project\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "855d6b96f291"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Updated property [core/project].\n",
            "\n",
            "All components are up to date.\n"
          ]
        }
      ],
      "source": [
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 3. **[Optional]** Set `CLUSTER_NAME` if you want to use your own GKE cluster. If not set, this example will create an auto-pilot cluster in the specified project.\n",
        "import datetime\n",
        "\n",
        "# Get the default region and project for launching jobs.\n",
        "\n",
        "REGION = \"us-central1\"  # @param [\"us-central1\", \"us-west1\", \"us-east4\"]\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "assert PROJECT_ID\n",
        "\n",
        "# Enter sure network is in same region\n",
        "NETWORK_NAME = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Make sure the subnet is in same region\n",
        "SUBNET_NAME = \"\"  # @param {type:\"string\"}\n",
        "NETWORK_URI = f\"projects/{PROJECT_ID}/global/networks/{NETWORK_NAME}\"\n",
        "SUBNET_URI = f\"projects/{PROJECT_ID}/regions/{REGION}/subnetworks/{SUBNET_NAME}\"\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",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "\n",
        "CLUSTER_NAME = f\"{CLUSTER_NAME}-{now}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oOtW4baPURcZ"
      },
      "source": [
        "## Set up GKE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4vjRPrOTf2t"
      },
      "source": [
        "### Create auto pilot cluster"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "70W0Fa-qTei0"
      },
      "outputs": [],
      "source": [
        "# create auto-pilot cluster\n",
        "!gcloud beta container --project {PROJECT_ID} clusters create-auto {CLUSTER_NAME} --region {REGION} --release-channel \"regular\" --tier \"standard\" --enable-private-nodes --enable-ip-access --no-enable-google-cloud-access --network {NETWORK_URI} --subnetwork {SUBNET_URI} --cluster-ipv4-cidr \"/17\" --binauthz-evaluation-mode=DISABLED"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3ixcamqUTqu8"
      },
      "source": [
        "### Deploy Llama on GKE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6psJZY_zUDgj"
      },
      "outputs": [],
      "source": [
        "# @markdown This section deploys llama 3.2 on GKE.\n",
        "\n",
        "# @markdown The model deployment takes about 5 to 15 minutes to complete. Larger models may take longer.\n",
        "\n",
        "# @markdown Select the model to deploy:\n",
        "# fmt: off\n",
        "MODEL_NAME = \"Llama-3-2-1B-Instruct\"  # @param [ 'Llama-3-2-1B-Instruct', 'Llama-3-2-11B-Vision', 'Llama-3-2-11B-Vision-Instruct', 'Llama-3-2-3B', 'Llama-3-2-3B-Instruct', 'Llama-3-2-1B']\n",
        "# fmt: on\n",
        "ARGS_TEMPLATE = \"\"\"args:\n",
        "        - python\n",
        "        - -m\n",
        "        - vllm.entrypoints.api_server\n",
        "        - --host 0.0.0.0\n",
        "        - --port 7080\n",
        "        - --model=gs://vertex-model-garden-public-us/llama3.2/{}\n",
        "        - --tensor-parallel-size {}\n",
        "        - --swap-space 16\n",
        "        - --gpu-memory-utilization 0.95\n",
        "        {}\n",
        "        - --max-num-seqs {}\n",
        "        {}\n",
        "        - --enable-auto-tool-choice\n",
        "        {}\n",
        "        - --disable-log-stats\n",
        "        {}\"\"\"\n",
        "# Model_name, tensor_size, - --model-type=llama3.1, max_seqs, - --enforce-eager, - --limit_mm_per_prompt='image=1', - --max-model-len 8192\n",
        "\n",
        "\n",
        "def generate_args(missing_args):\n",
        "    args = ARGS_TEMPLATE.format(\n",
        "        missing_args[0],\n",
        "        missing_args[1],\n",
        "        missing_args[2],\n",
        "        missing_args[3],\n",
        "        missing_args[4],\n",
        "        missing_args[5],\n",
        "        missing_args[6],\n",
        "    )\n",
        "    lines = args.splitlines()\n",
        "    non_empty_lines = [line for line in lines if line.strip()]\n",
        "    return \"\\n\".join(non_empty_lines)\n",
        "\n",
        "\n",
        "attr = {\n",
        "    \"Llama-3-2-11B-Vision\": [\n",
        "        [\n",
        "            \"Llama-3.2-11B-Vision\",\n",
        "            \"2\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "            \"12\",\n",
        "            \"- --enforce-eager\",\n",
        "            \"- --limit_mm_per_prompt='image=1'\",\n",
        "            \"- --max-model-len 8192\",\n",
        "        ],\n",
        "        15,\n",
        "        \"58Gi\",\n",
        "        \"120Gi\",\n",
        "        2,\n",
        "    ],\n",
        "    \"Llama-3-2-11B-Vision-Instruct\": [\n",
        "        [\n",
        "            \"Llama-3.2-11B-Vision-Instruct\",\n",
        "            \"2\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "            \"12\",\n",
        "            \"- --enforce-eager\",\n",
        "            \"- --limit_mm_per_prompt='image=1'\",\n",
        "            \"- --max-model-len 8192\",\n",
        "        ],\n",
        "        15,\n",
        "        \"58Gi\",\n",
        "        \"120Gi\",\n",
        "        2,\n",
        "    ],\n",
        "    \"Llama-3-2-3B\": [\n",
        "        [\n",
        "            \"Llama-3.2-3B\",\n",
        "            \"1\",\n",
        "            \"\",\n",
        "            \"64\",\n",
        "            \"\",\n",
        "            \"\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "        ],\n",
        "        10,\n",
        "        \"39Gi\",\n",
        "        \"100Gi\",\n",
        "        1,\n",
        "    ],\n",
        "    \"Llama-3-2-3B-Instruct\": [\n",
        "        [\n",
        "            \"Llama-3.2-3B-Instruct\",\n",
        "            \"1\",\n",
        "            \"\",\n",
        "            \"64\",\n",
        "            \"\",\n",
        "            \"\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "        ],\n",
        "        10,\n",
        "        \"39Gi\",\n",
        "        \"100Gi\",\n",
        "        1,\n",
        "    ],\n",
        "    \"Llama-3-2-1B\": [\n",
        "        [\n",
        "            \"Llama-3.2-1B\",\n",
        "            \"1\",\n",
        "            \"\",\n",
        "            \"64\",\n",
        "            \"\",\n",
        "            \"\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "        ],\n",
        "        8,\n",
        "        \"29Gi\",\n",
        "        \"80Gi\",\n",
        "        1,\n",
        "    ],\n",
        "    \"Llama-3-2-1B-Instruct\": [\n",
        "        [\n",
        "            \"Llama-3.2-1B-Instruct\",\n",
        "            \"1\",\n",
        "            \"\",\n",
        "            \"64\",\n",
        "            \"\",\n",
        "            \"\",\n",
        "            \"- --tool-call-parser=vertex-llama-3\",\n",
        "        ],\n",
        "        8,\n",
        "        \"29Gi\",\n",
        "        \"80Gi\",\n",
        "        1,\n",
        "    ],\n",
        "}\n",
        "\n",
        "model_attr = attr[MODEL_NAME]\n",
        "ARGS = generate_args(model_attr[0])\n",
        "CPU_LIMITS = model_attr[1]\n",
        "MEMORY_SIZE = model_attr[2]\n",
        "EPHEMERAL_STORAGE_SIZE = model_attr[3]\n",
        "GPU_COUNT = model_attr[4]\n",
        "\n",
        "\n",
        "K8S_YAML = f\"\"\"\n",
        "apiVersion: apps/v1\n",
        "kind: Deployment\n",
        "metadata:\n",
        "  name: llama-deployment\n",
        "spec:\n",
        "  replicas: 1\n",
        "  selector:\n",
        "    matchLabels:\n",
        "      app: llama-server\n",
        "  template:\n",
        "    metadata:\n",
        "      labels:\n",
        "        app: llama-server\n",
        "        ai.gke.io/model: {MODEL_NAME}\n",
        "        ai.gke.io/inference-server: vllm\n",
        "        examples.ai.gke.io/source: model-garden\n",
        "    spec:\n",
        "      containers:\n",
        "      - name: inference-server\n",
        "        image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241007_2233_RC00\n",
        "        resources:\n",
        "          requests:\n",
        "            cpu: {CPU_LIMITS}\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",
        "        env:\n",
        "        - name: MODEL_ID\n",
        "          value: 'meta-llama/{MODEL_NAME}'\n",
        "        - name: DEPLOY_SOURCE\n",
        "          value: 'UI_NATIVE_MODEL'\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: llama-service\n",
        "  annotations:\n",
        "    cloud.google.com/load-balancer-type: \"Internal\"\n",
        "spec:\n",
        "  selector:\n",
        "    app: llama-server # Should match the labels on your Pods\n",
        "  type: LoadBalancer\n",
        "  ports:\n",
        "  - protocol: TCP\n",
        "    port: 8000\n",
        "    targetPort: 7080\n",
        "\"\"\"\n",
        "\n",
        "with open(\"llama_32.yaml\", \"w\") as f:\n",
        "    f.write(K8S_YAML)\n",
        "\n",
        "! kubectl apply -f llama_32.yaml"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JgMxqkSdT45A"
      },
      "source": [
        "### Wait for container to be created.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bvEZBNGKD2Hp"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "MAX_WAIT_TIME = 600  # 10 minutes in seconds\n",
        "start_time = time.time()\n",
        "end_time = time.time() + MAX_WAIT_TIME\n",
        "\n",
        "print(\"Waiting for container to be created...\\n\")\n",
        "while start_time < end_time:\n",
        "    shell_output = ! kubectl get pod -l app=llama-server\n",
        "    container_status = \"\\n\".join(shell_output)\n",
        "    if \"1/1\" in container_status:\n",
        "        break\n",
        "    time.sleep(15)\n",
        "    start_time += 15\n",
        "\n",
        "if start_time > end_time:\n",
        "    print(\"Deployment took longer than expected\")\n",
        "\n",
        "print(container_status)\n",
        "\n",
        "# Wait for downloading artifacts.\n",
        "start_time = time.time()\n",
        "end_time = time.time() + MAX_WAIT_TIME\n",
        "print(\"\\nDownloading artifacts...\")\n",
        "while start_time < end_time:\n",
        "    shell_output = ! kubectl logs -l app=llama-server\n",
        "    logs = \"\\n\".join(shell_output)\n",
        "    if \"Connected\" in logs or \"Uvicorn running\" in logs:\n",
        "        break\n",
        "    time.sleep(15)\n",
        "    start_time += 15\n",
        "\n",
        "if start_time > end_time:\n",
        "    print(\"Deployment took longer than expected\")\n",
        "\n",
        "\n",
        "print(\"\\nServer is up and running!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qhkc2pVpcoLq"
      },
      "source": [
        "### Testing endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JxWMiOWgNfM1"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import subprocess\n",
        "import sys\n",
        "import time\n",
        "\n",
        "import requests  # Required: pip install requests\n",
        "\n",
        "# --- Configuration ---\n",
        "SERVICE_NAME = \"llama-service\"\n",
        "NAMESPACE = \"default\"  # Adjust if your service is in a different namespace\n",
        "SERVICE_PORT = 8000  # The 'port' specified in your Service YAML\n",
        "API_PATH = \"/generate\"  # The specific API path for vLLM generation\n",
        "\n",
        "# --- Request Payload ---\n",
        "user_message = \"What is AI?\"  # @param {type: \"string\"}\n",
        "max_tokens = 50  # @param {type:\"integer\"}\n",
        "temperature = 0.9  # @param {type:\"number\"}\n",
        "\n",
        "request_payload = {\n",
        "    \"prompt\": user_message,\n",
        "    \"max_tokens\": max_tokens,\n",
        "    \"temperature\": temperature,\n",
        "}\n",
        "\n",
        "\n",
        "# --- Function to Get Internal Load Balancer IP ---\n",
        "def get_internal_lb_ip(service_name, namespace, retries=5, delay=10):\n",
        "    \"\"\"Fetches the Internal Load Balancer IP for a Kubernetes Service.\"\"\"\n",
        "    command = [\n",
        "        \"kubectl\",\n",
        "        \"get\",\n",
        "        \"service\",\n",
        "        service_name,\n",
        "        \"-n\",\n",
        "        namespace,\n",
        "        \"-o\",\n",
        "        \"jsonpath={.status.loadBalancer.ingress[0].ip}\",\n",
        "    ]\n",
        "    print(\n",
        "        f\"Attempting to get IP for service '{service_name}' in namespace '{namespace}'...\"\n",
        "    )\n",
        "    for attempt in range(retries):\n",
        "        try:\n",
        "            result = subprocess.run(\n",
        "                command, capture_output=True, text=True, check=True, timeout=30\n",
        "            )\n",
        "            ip_address = result.stdout.strip()\n",
        "            if ip_address:\n",
        "                print(f\"Successfully found Internal LB IP: {ip_address}\")\n",
        "                return ip_address\n",
        "            print(\n",
        "                f\"Attempt {attempt + 1}/{retries}: IP address not assigned yet. Waiting {delay} seconds...\"\n",
        "            )\n",
        "        except subprocess.CalledProcessError as e:\n",
        "            # Handle cases where the service might not exist or JSON path fails temporarily\n",
        "            print(\n",
        "                f\"Attempt {attempt + 1}/{retries}: Error getting service IP: {e}. stderr: {e.stderr}. Waiting {delay} seconds...\"\n",
        "            )\n",
        "        except subprocess.TimeoutExpired:\n",
        "            print(\n",
        "                f\"Attempt {attempt + 1}/{retries}: 'kubectl' command timed out. Waiting {delay} seconds...\"\n",
        "            )\n",
        "        except FileNotFoundError:\n",
        "            print(\n",
        "                \"Error: 'kubectl' command not found. Please ensure kubectl is installed and in your PATH.\"\n",
        "            )\n",
        "            return None\n",
        "\n",
        "        if attempt < retries - 1:\n",
        "            time.sleep(delay)\n",
        "\n",
        "    print(\n",
        "        f\"Error: Could not retrieve Internal LB IP for service '{service_name}' after {retries} attempts.\"\n",
        "    )\n",
        "    return None\n",
        "\n",
        "\n",
        "# --- Main Execution ---\n",
        "internal_ip = get_internal_lb_ip(SERVICE_NAME, NAMESPACE)\n",
        "\n",
        "if internal_ip:\n",
        "    # Construct the full endpoint URL\n",
        "    endpoint_url = f\"http://{internal_ip}:{SERVICE_PORT}{API_PATH}\"\n",
        "    print(f\"Calling endpoint: {endpoint_url}\")\n",
        "\n",
        "    headers = {\"Content-Type\": \"application/json\"}\n",
        "\n",
        "    try:\n",
        "        # Make the POST request\n",
        "        response = requests.post(\n",
        "            endpoint_url,\n",
        "            headers=headers,\n",
        "            json=request_payload,  # requests library handles json serialization\n",
        "            timeout=120,  # Set a reasonable timeout (in seconds)\n",
        "        )\n",
        "\n",
        "        # Raise an exception for bad status codes (4xx or 5xx)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        # Parse the JSON response\n",
        "        response_data = response.json()\n",
        "\n",
        "        # Extract and print the prediction (adjust path if vLLM format differs)\n",
        "        # Assuming the structure is similar to the output you got via kubectl exec\n",
        "        if (\n",
        "            \"predictions\" in response_data\n",
        "            and isinstance(response_data[\"predictions\"], list)\n",
        "            and len(response_data[\"predictions\"]) > 0\n",
        "        ):\n",
        "            print(\"\\n--- Response ---\")\n",
        "            print(response_data[\"predictions\"][0])\n",
        "        elif (\n",
        "            \"text\" in response_data\n",
        "            and isinstance(response_data[\"text\"], list)\n",
        "            and len(response_data[\"text\"]) > 0\n",
        "        ):\n",
        "            # Handle potential alternative vLLM output format like {'text': ['...']}\n",
        "            print(\"\\n--- Response ---\")\n",
        "            print(response_data[\"text\"][0])\n",
        "        else:\n",
        "            print(\n",
        "                \"\\n--- Full Response (Could not find 'predictions' or 'text' array) ---\"\n",
        "            )\n",
        "            print(json.dumps(response_data, indent=2))\n",
        "\n",
        "    except requests.exceptions.ConnectionError as e:\n",
        "        print(\n",
        "            f\"\\nError: Could not connect to the endpoint {endpoint_url}.\",\n",
        "            file=sys.stderr,\n",
        "        )\n",
        "        print(\n",
        "            \"Check network connectivity and if the service/pods are running.\",\n",
        "            file=sys.stderr,\n",
        "        )\n",
        "        print(f\"   Details: {e}\", file=sys.stderr)\n",
        "    except requests.exceptions.Timeout:\n",
        "        print(\n",
        "            f\"\\nError: Request timed out while calling {endpoint_url}.\", file=sys.stderr\n",
        "        )\n",
        "    except requests.exceptions.RequestException as e:\n",
        "        print(\n",
        "            f\"\\nError: An error occurred during the request to {endpoint_url}: {e}\",\n",
        "            file=sys.stderr,\n",
        "        )\n",
        "        # Print response body if available, might contain useful error info from the server\n",
        "        if e.response is not None:\n",
        "            print(\"--- Server Response ---\", file=sys.stderr)\n",
        "            try:\n",
        "                print(json.dumps(e.response.json(), indent=2), file=sys.stderr)\n",
        "            except json.JSONDecodeError:\n",
        "                print(e.response.text, file=sys.stderr)\n",
        "    except json.JSONDecodeError:\n",
        "        print(\n",
        "            \"\\nError: Could not decode the JSON response from the server.\",\n",
        "            file=sys.stderr,\n",
        "        )\n",
        "        print(\"--- Raw Response ---\", file=sys.stderr)\n",
        "        print(response.text, file=sys.stderr)\n",
        "else:\n",
        "    print(\n",
        "        \"\\nExecution failed: Could not determine the endpoint IP address.\",\n",
        "        file=sys.stderr,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lc1FaXLKcrmW"
      },
      "source": [
        "## Cloud Function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ICtkVx7ShVC7"
      },
      "source": [
        "### Create service account and assign right permissions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4JDnEAqZcqsx"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT_NAME = \"\"  # @param {type: \"string\"}\n",
        "\n",
        "!gcloud iam service-accounts create {SERVICE_ACCOUNT_NAME} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  --description=\"llama service account\" \\\n",
        "  --display-name=\"llama testing service account\"\n",
        "\n",
        "SERVICE_ACCOUNT = f\"{SERVICE_ACCOUNT_NAME}@{PROJECT_ID}.iam.gserviceaccount.com\"\n",
        "\n",
        "for role in ['aiplatform.user', 'storage.objectAdmin', 'artifactregistry.reader', 'run.developer', 'run.invoker']:\n",
        "\n",
        "    ! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
        "      --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
        "      --role=roles/{role} --condition=None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBCL4vkCyDtV"
      },
      "source": [
        "### Prepare code for Cloud function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f7CjHee3dFqT"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path as path\n",
        "\n",
        "ROOT_PATH = path.cwd()\n",
        "TUTORIAL_PATH = ROOT_PATH / \"tutorial\"\n",
        "BUILD_PATH = TUTORIAL_PATH / \"build\"\n",
        "\n",
        "TUTORIAL_PATH.mkdir(parents=True, exist_ok=True)\n",
        "BUILD_PATH.mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q0Ud8hWedfMe"
      },
      "outputs": [],
      "source": [
        "requirements = \"\"\"\n",
        "requests>=2.26.0\n",
        "functions-framework>=3.0.0\n",
        "google-cloud-logging>=3.0.0\n",
        "\"\"\"\n",
        "\n",
        "with open(BUILD_PATH / \"requirements.txt\", \"w\") as f:\n",
        "    f.write(requirements)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0I079nDYdqMo"
      },
      "outputs": [],
      "source": [
        "cloud_function_code = '''\n",
        "import os\n",
        "import json\n",
        "import requests\n",
        "import logging\n",
        "import functions_framework # Required by GCF\n",
        "\n",
        "# Configure logging\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "\n",
        "# --- Configuration (Partially from Environment) ---\n",
        "# Fetch required info from environment variables set during deployment\n",
        "LLAMA_ILB_IP = os.environ.get(\"LLAMA_ENDPOINT_IP\") # Expecting the IP here\n",
        "SERVICE_PORT = os.environ.get(\"LLAMA_ENDPOINT_PORT\", \"8000\") # Default to 8000 if not set\n",
        "API_PATH = \"/generate\" # The specific API path for vLLM generation\n",
        "REQUEST_TIMEOUT = 120 # Timeout for the request to the Llama service\n",
        "\n",
        "@functions_framework.http # Defines this as an HTTP-triggered function\n",
        "def call_llama_service(request):\n",
        "    \"\"\"\n",
        "    Google Cloud Function entry point.\n",
        "    Expects a POST request with JSON body containing:\n",
        "    {\n",
        "        \"prompt\": \"Your question here\",\n",
        "        \"max_tokens\": 50,      // optional\n",
        "        \"temperature\": 0.9     // optional\n",
        "    }\n",
        "    \"\"\"\n",
        "    if not LLAMA_ILB_IP:\n",
        "        logging.error(\"Environment variable LLAMA_ENDPOINT_IP is not set.\")\n",
        "        return (\"Internal Server Error: Service endpoint IP not configured.\", 500)\n",
        "\n",
        "    if request.method != 'POST':\n",
        "        logging.warning(f\"Received non-POST request method: {request.method}\")\n",
        "        return ('Method Not Allowed', 405)\n",
        "\n",
        "    try:\n",
        "        request_json = request.get_json(silent=True)\n",
        "        if not request_json or 'prompt' not in request_json:\n",
        "            logging.error(\"Invalid request: Missing JSON payload or 'prompt' key.\")\n",
        "            return (\"Invalid request: Missing JSON payload or 'prompt' key.\", 400)\n",
        "\n",
        "        # --- Get Parameters from Request ---\n",
        "        prompt = request_json['prompt']\n",
        "        max_tokens = request_json.get('max_tokens', 50) # Default if not provided\n",
        "        temperature = request_json.get('temperature', 0.9) # Default if not provided\n",
        "\n",
        "        request_payload = {\n",
        "            \"prompt\": prompt,\n",
        "            \"max_tokens\": max_tokens,\n",
        "            \"temperature\": temperature,\n",
        "        }\n",
        "\n",
        "        # Construct the full endpoint URL\n",
        "        endpoint_url = f\"http://{LLAMA_ILB_IP}:{SERVICE_PORT}{API_PATH}\"\n",
        "        logging.info(f\"Calling Llama endpoint: {endpoint_url}\")\n",
        "        # Log payload without potentially sensitive prompt details in production if needed\n",
        "        logging.info(f\"Payload (excluding prompt): max_tokens={max_tokens}, temperature={temperature}\")\n",
        "\n",
        "        headers = {\"Content-Type\": \"application/json\"}\n",
        "\n",
        "        # --- Make the POST request ---\n",
        "        response = requests.post(\n",
        "            endpoint_url,\n",
        "            headers=headers,\n",
        "            json=request_payload,\n",
        "            timeout=REQUEST_TIMEOUT\n",
        "        )\n",
        "\n",
        "        # Raise an exception for bad status codes (4xx or 5xx)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        # Parse the JSON response from Llama service\n",
        "        response_data = response.json()\n",
        "        logging.info(\"Successfully received response from Llama service.\")\n",
        "\n",
        "        # --- Process and Return Response ---\n",
        "        # Extract prediction (adjust keys if needed based on actual vLLM response)\n",
        "        prediction = None\n",
        "        if \"predictions\" in response_data and isinstance(response_data[\"predictions\"], list) and len(response_data[\"predictions\"]) > 0:\n",
        "             prediction = response_data[\"predictions\"][0]\n",
        "        elif \"text\" in response_data and isinstance(response_data[\"text\"], list) and len(response_data[\"text\"]) > 0:\n",
        "             prediction = response_data[\"text\"][0]\n",
        "        else:\n",
        "             logging.warning(\"Could not find 'predictions' or 'text' in response. Returning full response.\")\n",
        "             # Return the full response if the expected key isn't found\n",
        "             return (response_data, 200) # Return raw JSON with 200 OK\n",
        "\n",
        "        if prediction is not None:\n",
        "            # Return only the prediction text/data\n",
        "            return ({ \"prediction\": prediction }, 200) # Return prediction in a structured way\n",
        "        else:\n",
        "            logging.error(\"Prediction key found but content was empty/invalid.\")\n",
        "            return (\"Error processing model response\", 500)\n",
        "\n",
        "\n",
        "    except requests.exceptions.ConnectionError as e:\n",
        "        logging.error(f\"Connection Error calling {endpoint_url}: {e}\")\n",
        "        return (f\"Could not connect to the backend service: {e}\", 503) # 503 Service Unavailable\n",
        "    except requests.exceptions.Timeout:\n",
        "        logging.error(f\"Request timed out calling {endpoint_url}\")\n",
        "        return (\"Backend service timed out.\", 504) # 504 Gateway Timeout\n",
        "    except requests.exceptions.RequestException as e:\n",
        "        logging.error(f\"RequestException calling {endpoint_url}: {e}\")\n",
        "        # Log response body if available\n",
        "        error_details = str(e)\n",
        "        if e.response is not None:\n",
        "             logging.error(f\"Backend service response status: {e.response.status_code}\")\n",
        "             try:\n",
        "                  error_details = json.dumps(e.response.json())\n",
        "                  logging.error(f\"Backend service response body: {error_details}\")\n",
        "             except json.JSONDecodeError:\n",
        "                  error_details = e.response.text\n",
        "                  logging.error(f\"Backend service response body (non-JSON): {error_details}\")\n",
        "        return (f\"Error calling backend service: {error_details}\", 502) # 502 Bad Gateway\n",
        "    except json.JSONDecodeError:\n",
        "        logging.error(f\"Could not decode the JSON response from the Llama service. Raw response: {response.text}\")\n",
        "        return (\"Invalid response format from backend service.\", 502)\n",
        "    except Exception as e:\n",
        "        logging.exception(f\"An unexpected error occurred: {e}\") # Log full traceback\n",
        "        return (\"An internal server error occurred.\", 500)\n",
        "'''\n",
        "\n",
        "\n",
        "with open(BUILD_PATH / \"main.py\", \"w\") as f:\n",
        "    f.write(cloud_function_code)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1o01XKBmx5gs"
      },
      "source": [
        "### Create VPC connector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wZdxnxIvj6_w"
      },
      "outputs": [],
      "source": [
        "! gcloud config set project {PROJECT_ID}\n",
        "! gcloud services enable vpcaccess.googleapis.com"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jme6o84CorG0"
      },
      "outputs": [],
      "source": [
        "VPC_CONNECTOR_NAME = \"\" # @param {type: \"string\"}\n",
        "RANGE = \"10.6.0.0/28\" # @param {type: \"string\"}\n",
        "\n",
        "!gcloud compute networks vpc-access connectors create {VPC_CONNECTOR_NAME} \\\n",
        " --region {REGION} \\\n",
        " --network {NETWORK_NAME} \\\n",
        " --range {RANGE} \\\n",
        " --min-instances 2 \\\n",
        " --max-instances 5 \\\n",
        " --machine-type \"e2-micro\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXV5eikfx8no"
      },
      "source": [
        "### Create function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtI5uQQpfsh4"
      },
      "outputs": [],
      "source": [
        "# --- Set these variables ---\n",
        "FUNCTION_NAME=\"test-llama-gke-1\" # @param {type: \"string\"}\n",
        "ILB_IP= internal_ip # The IP address you found above\n",
        "\n",
        "# --- Deployment Command ---\n",
        "!gcloud functions deploy {FUNCTION_NAME} \\\n",
        "    --gen2 \\\n",
        "    --runtime=\"python311\" \\\n",
        "    --source={str(BUILD_PATH)} \\\n",
        "    --region={REGION} \\\n",
        "    --entry-point=call_llama_service \\\n",
        "    --trigger-http \\\n",
        "    --no-allow-unauthenticated \\\n",
        "    --set-env-vars=LLAMA_ENDPOINT_IP={ILB_IP} \\\n",
        "    --quiet \\\n",
        "    --service-account={SERVICE_ACCOUNT} \\\n",
        "    --timeout=600 \\\n",
        "    --memory=2Gb \\\n",
        "    --concurrency=2 \\\n",
        "    --min-instances=2 \\\n",
        "    --project {PROJECT_ID} \\\n",
        "    --vpc-connector {VPC_CONNECTOR_NAME}\n",
        "\n",
        "\n",
        "print(f\"Function: {FUNCTION_NAME} deployed successfully.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wbRmgoOZF6es"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "DELETE_DEPLOYMENT = False  # @param {type: \"boolean\"}\n",
        "DELETE_CLUSTER = False  # @param {type: \"boolean\"}\n",
        "\n",
        "if DELETE_DEPLOYMENT:\n",
        "    ! kubectl delete deployments llama-deployment\n",
        "    ! kubectl delete services llama-service\n",
        "\n",
        "if DELETE_CLUSTER:\n",
        "    ! gcloud container clusters delete {CLUSTER_NAME} --region={REGION} --quiet\n",
        "\n",
        "! gcloud functions delete {FUNCTION_NAME} --region={REGION} --quiet\n",
        "\n",
        "! gcloud compute networks vpc-access connectors delete {VPC_CONNECTOR_NAME} --region {REGION} --quiet"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "deploy_llama3_vllm_gke_cloud_function.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
