{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AjmN8CQB6xv3"
      },
      "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": "X222YYLg6ugv"
      },
      "source": [
        "# Accelerate LLM Inference with EAGLE Speculative Decoding on Vertex AI\n",
        "\n",
        "<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/benchmarking_eagle_on_vertex_ai.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/benchmarking_eagle_on_vertex_ai.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/benchmarking_eagle_on_vertex_ai.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",
        "<p>\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/benchmarking_eagle_on_vertex_ai.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/benchmarking_eagle_on_vertex_ai.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/benchmarking_eagle_on_vertex_ai.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/benchmarking_eagle_on_vertex_ai.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/benchmarking_eagle_on_vertex_ai.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>\n",
        "</p>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w1m8j-zV6ugv"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [Ivan Nardini](https://github.com/inardini) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gOTlpB4u6ugw"
      },
      "source": [
        "## Overview\n",
        "\n",
        "### Why EAGLE Matters\n",
        "\n",
        "Large Language Models (LLMs) generate text one token at a time, which creates a fundamental bottleneck: each token requires a full forward pass through the model. For production applications serving thousands of users, this sequential generation limits throughput and increases costs.\n",
        "\n",
        "**EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)** solves this by using speculative decoding:\n",
        "- A small, fast **draft model** predicts multiple tokens in parallel\n",
        "- The main model **verifies** these predictions in a single forward pass\n",
        "- Correct predictions are accepted (saving time), incorrect ones are discarded and regenerated\n",
        "- **Result**: 1.5-2x faster inference with identical output quality\n",
        "\n",
        "### What You'll Learn\n",
        "\n",
        "This tutorial demonstrates how to benchmark EAGLE's performance improvement on Vertex AI using real production workloads:\n",
        "\n",
        "1. **Deploy two Llama 4 Scout endpoints**: baseline (standard) and EAGLE-enabled\n",
        "2. **Run controlled benchmarks** using vLLM's industry-standard tooling and ShareGPT dataset (real user conversations)\n",
        "3. **Measure key metrics** across varying concurrency levels:\n",
        "   - **TTFT (Time to First Token)**: User-perceived latency - how quickly responses start\n",
        "   - **TPOT (Time Per Output Token)**: Generation speed - affects streaming smoothness\n",
        "   - **Throughput**: System capacity - tokens and requests per second\n",
        "   - **Scalability**: How performance changes under concurrent load\n",
        "4. **Visualize results** to quantify EAGLE's speedup and identify optimal configurations\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "Before starting, ensure you have:\n",
        "\n",
        "- **Google Cloud Project** with billing enabled\n",
        "- **Vertex AI API** enabled ([enable here](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com))\n",
        "- **GPU Quota**: 8x NVIDIA H100 80GB GPUs in your selected region\n",
        "  - Check quota: `gcloud compute regions describe <region> | grep 'NVIDIA_H100_80GB'`\n",
        "  - Request increase: [Quota page](https://console.cloud.google.com/iam-admin/quotas)\n",
        "- **Python 3.10+** (automatically available in Colab/Workbench)\n",
        "- **Hugging Face account** with access to Llama 4 model (requires Meta license acceptance)\n",
        "\n",
        "**Estimated Time:** 90-120 minutes (mostly deployment wait time)  \n",
        "**Estimated Cost:** ~$250 (mostly machine and GPU hours during benchmarking)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7L-7bvPR6ugw"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6VsC0vXq6ugw"
      },
      "source": [
        "### Install Required Packages\n",
        "\n",
        "**Note:** After running this cell, **you will need to restart the runtime**. This is expected behavior when installing new Python packages.\n",
        "\n",
        "- In **Colab**: Click the \"Restart Runtime\" button that appears\n",
        "- In **Vertex AI Workbench**: Kernel → Restart Kernel\n",
        "\n",
        "After restarting, **continue from the next cell** (do not re-run this installation cell)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mLwAQe656ugw"
      },
      "outputs": [],
      "source": [
        "# Install packages with pinned versions for reproducibility\n",
        "%pip install --upgrade --quiet \\\n",
        "    'google-cloud-aiplatform>=1.70.0' \\\n",
        "    'transformers>=4.45.0' \\\n",
        "    'huggingface-hub>=0.26.0' \\\n",
        "    'hf-transfer>=0.1.8' \\\n",
        "    'vllm==0.11.0' \\\n",
        "    'pandas>=2.0.0' \\\n",
        "    'matplotlib>=3.7.0' \\\n",
        "    'seaborn>=0.13.0'\n",
        "\n",
        "print(\"\\n\" + \"=\"*80)\n",
        "print(\"✅ Installation completed successfully!\")\n",
        "print(\"=\"*80)\n",
        "print(\"⚠️  NEXT STEP: Please restart your runtime now.\")\n",
        "print(\"   - Colab: Click 'Runtime' → 'Restart session' button above\")\n",
        "print(\"   - Workbench: 'Kernel' → 'Restart Kernel'\")\n",
        "print(\"   - Then continue from the cell below (skip this installation cell)\")\n",
        "print(\"=\"*80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1-WRJi3m6ugw"
      },
      "source": [
        "### Authenticate Your Environment\n",
        "\n",
        "**Colab users only**: Run this cell to authenticate your Google Cloud account. This allows the notebook to access Vertex AI services.\n",
        "\n",
        "**Vertex AI Workbench users**: Skip this cell - you're already authenticated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HOI2APLk6ugw"
      },
      "outputs": [],
      "source": [
        "# import sys\n",
        "\n",
        "# # Only authenticate in Colab environment\n",
        "# if \"google.colab\" in sys.modules:\n",
        "#     from google.colab import auth\n",
        "#     auth.authenticate_user()\n",
        "#     print(\"✅ Authentication successful!\")\n",
        "# else:\n",
        "#     print(\"ℹ️  Running in Vertex AI Workbench - already authenticated\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gap_mtsv6ugw"
      },
      "source": [
        "### Set Google Cloud Project Information\n",
        "\n",
        "Configure your Google Cloud project ID and region. The project must have:\n",
        "- Vertex AI API enabled\n",
        "- Sufficient GPU quota (8x H100 80GB recommended)\n",
        "\n",
        "**Recommended regions for H100 availability**:\n",
        "- `us-central1` (Iowa)\n",
        "- `us-east4` (Northern Virginia)\n",
        "- `europe-west4` (Netherlands)\n",
        "- `asia-southeast1` (Singapore)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qlKdeh146ugw"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# Configure these values for your environment\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "LOCATION = \"asia-southeast1\"  # @param {type: \"string\", placeholder: \"us-central1\", isTemplate: true}\n",
        "# fmt: on\n",
        "\n",
        "# Auto-detect project ID if not provided\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "    if not PROJECT_ID:\n",
        "        raise ValueError(\n",
        "            \"❌ PROJECT_ID not set. Please set it in the cell above or \"\n",
        "            \"set GOOGLE_CLOUD_PROJECT environment variable\"\n",
        "        )\n",
        "\n",
        "# Initialize Vertex AI SDK\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"✅ Vertex AI initialized successfully!\")\n",
        "print(\"=\" * 80)\n",
        "print(f\"   Project ID: {PROJECT_ID}\")\n",
        "print(f\"   Location:   {LOCATION}\")\n",
        "print(\"=\" * 80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0elIiVNF6ugw"
      },
      "source": [
        "### Import Required Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ceW7gfM6ugx"
      },
      "outputs": [],
      "source": [
        "# Standard libraries\n",
        "import json\n",
        "import subprocess\n",
        "import urllib.request\n",
        "from pathlib import Path\n",
        "\n",
        "import google.auth\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Data processing and visualization\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "from google.auth.transport.requests import Request\n",
        "\n",
        "# Hugging Face libraries\n",
        "from huggingface_hub import login, snapshot_download\n",
        "from vertexai import model_garden\n",
        "\n",
        "print(\"✅ Libraries imported successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3MBjFkt6ugx"
      },
      "source": [
        "## Model Deployment\n",
        "\n",
        "We'll deploy two versions of Llama 4 Scout (17B parameters, 16 experts):\n",
        "1. **Baseline**: Standard configuration (no EAGLE)\n",
        "2. **EAGLE-enabled**: With speculative decoding enabled\n",
        "\n",
        "Both use identical hardware (8x H100 80GB GPUs with tensor parallelism) to ensure fair comparison."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kPcP39K_6ugx"
      },
      "source": [
        "### Deploy Baseline Model (Without EAGLE)\n",
        "\n",
        "This deployment creates a baseline for comparison. We'll measure its performance, then compare against EAGLE.\n",
        "\n",
        "**Key Configuration Parameters:**\n",
        "\n",
        "| Parameter | Value | Purpose |\n",
        "|-----------|-------|----------|\n",
        "| `machine_type` | `a3-highgpu-8g` | VM with 8x H100 80GB GPUs |\n",
        "| `accelerator_type` | `NVIDIA_H100_80GB` | Latest generation GPU (9.5x faster than A100) |\n",
        "| `accelerator_count` | `8` | Number of GPUs for tensor parallelism |\n",
        "| `--tp` | `8` | Tensor parallelism degree (splits model across GPUs) |\n",
        "| `--attention-backend` | `fa3` | FlashAttention 3 (optimized attention computation) |\n",
        "| `--context-length` | `131072` | Maximum sequence length (128K tokens) |\n",
        "\n",
        "**Expected deployment time**: 10-15 minutes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qFHTU45E6ugx"
      },
      "outputs": [],
      "source": [
        "# Model configuration\n",
        "MODEL_NAME = \"meta/llama4@llama-4-scout-17b-16e-instruct\"\n",
        "MODEL_GCS_PATH = (\n",
        "    \"gs://vertex-model-garden-restricted-us/llama4/Llama-4-Scout-17B-16E-Instruct\"\n",
        ")\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"🚀 DEPLOYING BASELINE MODEL\")\n",
        "print(\"=\" * 80)\n",
        "print(f\"Model: {MODEL_NAME}\")\n",
        "print(\"Configuration: 8x H100 GPUs, No EAGLE\")\n",
        "print(\"\\n⏳ Starting deployment... (this will take 10-15 minutes)\")\n",
        "print(\"=\" * 80)\n",
        "\n",
        "# Baseline deployment arguments (no speculative decoding)\n",
        "baseline_args = [\n",
        "    f\"--model={MODEL_GCS_PATH}\",\n",
        "    \"--attention-backend=fa3\",  # FlashAttention 3 for optimal performance\n",
        "    \"--context-length=131072\",  # 128K context window\n",
        "    \"--chat-template=Llama-4\",\n",
        "    \"--tp=8\",  # Tensor parallelism across 8 GPUs\n",
        "    \"--enable-multimodal\",\n",
        "    \"--tool-call-parser=pythonic\",\n",
        "    \"--chat-template=sglang/examples/chat_template/tool_chat_template_llama4_pythonic.jinja\",\n",
        "]\n",
        "\n",
        "try:\n",
        "    # Initialize model from Model Garden\n",
        "    baseline_model = model_garden.OpenModel(MODEL_NAME)\n",
        "\n",
        "    # Deploy to dedicated endpoint\n",
        "    baseline_endpoint = baseline_model.deploy(\n",
        "        model_display_name=\"baseline-llama4-scout-17b-16e-instruct\",\n",
        "        endpoint_display_name=\"baseline-llama4-scout-17b-16e-instruct\",\n",
        "        serving_container_image_uri=\"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/sglang-serve.cu124.0-4.ubuntu2204.py310:model-garden.sglang-0-4-release_20250831.00_p0\",\n",
        "        machine_type=\"a3-highgpu-8g\",\n",
        "        accelerator_type=\"NVIDIA_H100_80GB\",\n",
        "        accelerator_count=8,\n",
        "        use_dedicated_endpoint=True,\n",
        "        accept_eula=True,\n",
        "        serving_container_args=baseline_args,\n",
        "        serving_container_environment_variables={\n",
        "            \"MODEL_ID\": \"meta-llama/Llama-4-Scout-17B-16E-Instruct\",\n",
        "            \"DEPLOY_SOURCE\": \"UI_NATIVE_MODEL\",\n",
        "        },\n",
        "        serving_container_ports=[30000],\n",
        "        serving_container_health_route=\"/health\",\n",
        "        serving_container_predict_route=\"/vertex_generate\",\n",
        "    )\n",
        "\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"✅ BASELINE ENDPOINT DEPLOYED SUCCESSFULLY!\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"   Endpoint ID: {baseline_endpoint.name}\")\n",
        "    print(f\"   Resource Name: {baseline_endpoint.resource_name}\")\n",
        "    print(\"   Status: READY\")\n",
        "    print(\"=\" * 80)\n",
        "\n",
        "except Exception as e:\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"❌ DEPLOYMENT FAILED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"Error: {e!s}\")\n",
        "    print(\"\\nCommon issues:\")\n",
        "    print(\"  - Insufficient GPU quota (need 8x H100 80GB)\")\n",
        "    print(\"  - Region doesn't have H100s available\")\n",
        "    print(\"  - Billing not enabled on project\")\n",
        "    print(\"=\" * 80)\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-XiFYo-06ugx"
      },
      "source": [
        "### Deploy EAGLE-Enabled Model\n",
        "\n",
        "This deployment adds EAGLE speculative decoding on top of the same base model and hardware.\n",
        "\n",
        "**EAGLE-Specific Parameters:**\n",
        "\n",
        "| Parameter | Value | Purpose |\n",
        "|-----------|-------|----------|\n",
        "| `--speculative-algo` | `EAGLE3` | Activates EAGLE version 3 (latest) |\n",
        "| `--speculative-draft-model-path` | `gs://...EAGLE3...` | Pre-trained draft model for Llama 4 Scout |\n",
        "| `--speculative-num-steps` | `3` | Speculation depth (how many tokens to predict ahead) |\n",
        "| `--speculative-num-draft-tokens` | `8` | Tokens generated per speculation step |\n",
        "| `--speculative-eagle-topk` | `4` | Top-K sampling for draft model (balances speed/quality) |\n",
        "\n",
        "**How these parameters affect performance:**\n",
        "- **Higher `num-steps`**: More speculation → better speedup but higher overhead (sweet spot: 2-5)\n",
        "- **Higher `num-draft-tokens`**: More tokens per step → better for long outputs (sweet spot: 4-12)\n",
        "- **Higher `topk`**: More diverse predictions → better acceptance rate but slower draft (sweet spot: 3-5)\n",
        "\n",
        "**Expected deployment time**: 10-15 minutes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SaV2bxoY6ugx"
      },
      "outputs": [],
      "source": [
        "print(\"=\" * 80)\n",
        "print(\"🚀 DEPLOYING EAGLE-ENABLED MODEL\")\n",
        "print(\"=\" * 80)\n",
        "print(f\"Model: {MODEL_NAME}\")\n",
        "print(\"Configuration: 8x H100 GPUs, EAGLE Speculative Decoding\")\n",
        "print(\"\\n⏳ Starting deployment... (this will take 10-15 minutes)\")\n",
        "print(\"=\" * 80)\n",
        "\n",
        "# EAGLE deployment arguments (adds speculative decoding parameters)\n",
        "eagle_args = [\n",
        "    f\"--model={MODEL_GCS_PATH}\",\n",
        "    \"--attention-backend=fa3\",\n",
        "    \"--context-length=131072\",\n",
        "    \"--chat-template=Llama-4\",\n",
        "    \"--tp=8\",\n",
        "    \"--enable-multimodal\",\n",
        "    \"--tool-call-parser=pythonic\",\n",
        "    \"--chat-template=sglang/examples/chat_template/tool_chat_template_llama4_pythonic.jinja\",\n",
        "    # EAGLE-specific configuration\n",
        "    \"--speculative-algo=EAGLE3\",\n",
        "    \"--speculative-draft-model-path=gs://vertex-model-garden-restricted-us/llama4/Llama-4-Scout-17B-16E-Instruct-EAGLE3-20250829/\",\n",
        "    \"--speculative-num-steps=3\",\n",
        "    \"--speculative-eagle-topk=4\",\n",
        "    \"--speculative-num-draft-tokens=8\",\n",
        "]\n",
        "\n",
        "try:\n",
        "    # Initialize model from Model Garden\n",
        "    eagle_model = model_garden.OpenModel(MODEL_NAME)\n",
        "\n",
        "    # Deploy to dedicated endpoint\n",
        "    eagle_endpoint = eagle_model.deploy(\n",
        "        model_display_name=\"eagle-llama4-scout-17b-16e-instruct\",\n",
        "        endpoint_display_name=\"eagle-llama4-scout-17b-16e-instruct\",\n",
        "        serving_container_image_uri=\"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/sglang-serve.cu124.0-4.ubuntu2204.py310:model-garden.sglang-0-4-release_20250831.00_p0\",\n",
        "        machine_type=\"a3-highgpu-8g\",\n",
        "        accelerator_type=\"NVIDIA_H100_80GB\",\n",
        "        accelerator_count=8,\n",
        "        use_dedicated_endpoint=True,\n",
        "        accept_eula=True,\n",
        "        serving_container_args=eagle_args,\n",
        "        serving_container_environment_variables={\n",
        "            \"MODEL_ID\": \"meta-llama/Llama-4-Scout-17B-16E-Instruct\",\n",
        "            \"DEPLOY_SOURCE\": \"UI_NATIVE_MODEL\",\n",
        "        },\n",
        "        serving_container_ports=[30000],\n",
        "        serving_container_health_route=\"/health\",\n",
        "        serving_container_predict_route=\"/vertex_generate\",\n",
        "    )\n",
        "\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"✅ EAGLE ENDPOINT DEPLOYED SUCCESSFULLY!\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"   Endpoint ID: {eagle_endpoint.name}\")\n",
        "    print(f\"   Resource Name: {eagle_endpoint.resource_name}\")\n",
        "    print(\"   Status: READY\")\n",
        "    print(\"=\" * 80)\n",
        "    print(\"\\n🎯 Both endpoints are now ready for benchmarking!\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"❌ DEPLOYMENT FAILED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"Error: {e!s}\")\n",
        "    print(\"\\nCommon issues:\")\n",
        "    print(\"  - Insufficient GPU quota (need 8x H100 80GB)\")\n",
        "    print(\"  - Region doesn't have H100s available\")\n",
        "    print(\"  - Billing not enabled on project\")\n",
        "    print(\"=\" * 80)\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LYNg084l6ugx"
      },
      "source": [
        "## Prepare Benchmark Dataset\n",
        "\n",
        "We'll use the **ShareGPT dataset**, which contains real user-assistant conversations from production systems.\n",
        "\n",
        "**Why ShareGPT?**\n",
        "- **Realistic workload**: Real conversations, not synthetic prompts\n",
        "- **Variable lengths**: Tests model performance across different input/output sizes\n",
        "- **Industry standard**: Used by major LLM serving frameworks for benchmarking"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H3ONN7dG6ugx"
      },
      "source": [
        "### Download ShareGPT Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RheEgLao6ugx"
      },
      "outputs": [],
      "source": [
        "DATASET_URL = \"https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json\"\n",
        "DATASET_PATH = \"/tmp/ShareGPT_V3_unfiltered_cleaned_split.json\"\n",
        "\n",
        "print(\"📥 Downloading ShareGPT dataset...\")\n",
        "\n",
        "try:\n",
        "    if not os.path.exists(DATASET_PATH):\n",
        "        urllib.request.urlretrieve(DATASET_URL, DATASET_PATH)\n",
        "        print(f\"✅ Dataset downloaded to {DATASET_PATH}\")\n",
        "    else:\n",
        "        print(f\"ℹ️  Dataset already exists at {DATASET_PATH}\")\n",
        "\n",
        "    # Preview dataset structure\n",
        "    with open(DATASET_PATH) as f:\n",
        "        data = json.load(f)\n",
        "\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"📊 DATASET INFORMATION\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"   Total conversations: {len(data):,}\")\n",
        "    print(f\"   Sample conversation keys: {list(data[0].keys())}\")\n",
        "    print(\"\\n   Sample conversation structure:\")\n",
        "    print(f\"   - ID: {data[0]['id']}\")\n",
        "    print(f\"   - Turns: {len(data[0]['conversations'])} messages\")\n",
        "    print(\"=\" * 80)\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"\\n❌ Failed to download dataset: {e}\")\n",
        "    print(\"Please check your internet connection and try again.\")\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-CKzYabL6ugx"
      },
      "source": [
        "### Download Model Artifacts for Tokenization\n",
        "\n",
        "We need the model's tokenizer to accurately measure prompt/response lengths during benchmarking.\n",
        "\n",
        "**Why download locally?**\n",
        "- vLLM's benchmark tool needs the tokenizer to count tokens accurately\n",
        "- We only download configuration files (less than 10MB), not the full model weights\n",
        "- Using `hf_transfer` library for 2-5x faster downloads\n",
        "\n",
        "**Note:** This requires a Hugging Face account with access to Llama 4 (requires accepting Meta's license)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0brX1tYx6ugx"
      },
      "outputs": [],
      "source": [
        "# Authenticate to Hugging Face\n",
        "# You'll be prompted to enter your HF token (get one at https://huggingface.co/settings/tokens)\n",
        "print(\"🔐 Authenticating to Hugging Face...\")\n",
        "print(\"   You'll need a token with access to meta-llama/Llama-4-Scout-17B-16E-Instruct\")\n",
        "print(\"   Get your token at: https://huggingface.co/settings/tokens\\n\")\n",
        "\n",
        "try:\n",
        "    login()\n",
        "    print(\"✅ Hugging Face authentication successful!\")\n",
        "except Exception as e:\n",
        "    print(f\"❌ Authentication failed: {e}\")\n",
        "    raise"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O2pAecen6ugx"
      },
      "outputs": [],
      "source": [
        "# Configure download settings\n",
        "HF_HOME = \"hf_cache\"\n",
        "MODEL_ID = \"meta-llama/Llama-4-Scout-17B-16E-Instruct\"\n",
        "LOCAL_DIR = Path(f\"{HF_HOME}/{MODEL_ID}\")\n",
        "\n",
        "# Create directory\n",
        "LOCAL_DIR.parent.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "# Enable fast transfers (2-5x faster using Rust-based hf_transfer)\n",
        "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"\n",
        "os.environ[\"HF_HOME\"] = HF_HOME\n",
        "\n",
        "# Only download configuration files (not model weights)\n",
        "allow_patterns = [\"*.json\", \"tokenizer.model\", \"*.txt\"]\n",
        "\n",
        "print(\"📥 Downloading model artifacts from Hugging Face...\")\n",
        "print(f\"   Model: {MODEL_ID}\")\n",
        "print(\"   Using fast transfer (hf_transfer enabled)\")\n",
        "print(f\"   Downloading to: {LOCAL_DIR}\\n\")\n",
        "\n",
        "try:\n",
        "    snapshot_download(\n",
        "        repo_id=MODEL_ID,\n",
        "        local_dir=str(LOCAL_DIR),\n",
        "        allow_patterns=allow_patterns,\n",
        "        resume_download=True,\n",
        "    )\n",
        "\n",
        "    print(\"\\n✅ Model artifacts downloaded successfully!\")\n",
        "    print(f\"   Location: {LOCAL_DIR}\")\n",
        "\n",
        "    # Store path for benchmarking\n",
        "    MODEL_PATH = str(LOCAL_DIR)\n",
        "    print(f\"\\n   MODEL_PATH set to: {MODEL_PATH}\")\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"\\n❌ Download failed: {e}\")\n",
        "    print(\"\\nCommon issues:\")\n",
        "    print(\"  - No access to Llama 4 model (need to accept Meta license)\")\n",
        "    print(\"  - Invalid Hugging Face token\")\n",
        "    print(\"  - Network connectivity issues\")\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rO3XNv9n6ugy"
      },
      "source": [
        "## Quick Smoke Test (100 Prompts)\n",
        "\n",
        "Before running the full benchmark, let's verify both endpoints work correctly with a quick test.\n",
        "\n",
        "**What this tests:**\n",
        "- Both endpoints are responding correctly\n",
        "- Authentication is working\n",
        "- Basic performance sanity check\n",
        "\n",
        "**This is NOT the main benchmark** - just a validation step. The comprehensive benchmark comes next."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lxwb1mZo6ugy"
      },
      "source": [
        "### Apply Vertex AI Compatibility Patch\n",
        "\n",
        "**Why is this patch needed?**\n",
        "\n",
        "vLLM's benchmark tool was originally designed for OpenAI's API, but Vertex AI has slightly different requirements:\n",
        "\n",
        "1. **Vertex AI doesn't support `stream_options`** parameter (OpenAI-specific)\n",
        "2. **Vertex AI uses `max_tokens`** instead of `max_completion_tokens`\n",
        "3. **Vertex AI has longer timeouts** for large models (6 hours vs 5 minutes)\n",
        "\n",
        "This patch modifies vLLM's request function to be compatible with Vertex AI's API format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8hm_y_EV6ugy"
      },
      "outputs": [],
      "source": [
        "import vllm.benchmarks.lib.endpoint_request_func as endpoint_func\n",
        "\n",
        "# Store the original function\n",
        "_original_async_request_openai_chat_completions = (\n",
        "    endpoint_func.async_request_openai_chat_completions\n",
        ")\n",
        "\n",
        "\n",
        "async def patched_async_request_openai_chat_completions(\n",
        "    request_func_input: endpoint_func.RequestFuncInput,\n",
        "    pbar=None,\n",
        ") -> endpoint_func.RequestFuncOutput:\n",
        "    \"\"\"Patched version compatible with Vertex AI's chat completions endpoint.\"\"\"\n",
        "    # Import necessary modules\n",
        "    import json\n",
        "    import os\n",
        "    import sys\n",
        "    import time\n",
        "    import traceback\n",
        "\n",
        "    import aiohttp\n",
        "\n",
        "    api_url = request_func_input.api_url\n",
        "    assert api_url.endswith((\"chat/completions\", \"profile\"))\n",
        "\n",
        "    # Set longer timeout for Vertex AI (6 hours for large model inference)\n",
        "    async with aiohttp.ClientSession(\n",
        "        trust_env=True, timeout=aiohttp.ClientTimeout(total=6 * 60 * 60)\n",
        "    ) as session:\n",
        "        # Build request content (text + optional multimodal)\n",
        "        content = [{\"type\": \"text\", \"text\": request_func_input.prompt}]\n",
        "        if request_func_input.multi_modal_content:\n",
        "            mm_content = request_func_input.multi_modal_content\n",
        "            if isinstance(mm_content, list):\n",
        "                content.extend(mm_content)\n",
        "            elif isinstance(mm_content, dict):\n",
        "                content.append(mm_content)\n",
        "\n",
        "        # Build payload with Vertex AI-compatible parameters\n",
        "        payload = {\n",
        "            \"model\": request_func_input.model_name\n",
        "            if request_func_input.model_name\n",
        "            else request_func_input.model,\n",
        "            \"messages\": [{\"role\": \"user\", \"content\": content}],\n",
        "            \"temperature\": 0.0,\n",
        "            \"max_tokens\": request_func_input.output_len,  # Vertex AI uses max_tokens, not max_completion_tokens\n",
        "            \"stream\": True,\n",
        "            # Removed stream_options - not supported by Vertex AI\n",
        "        }\n",
        "        if request_func_input.ignore_eos:\n",
        "            payload[\"ignore_eos\"] = request_func_input.ignore_eos\n",
        "        if request_func_input.extra_body:\n",
        "            payload.update(request_func_input.extra_body)\n",
        "\n",
        "        # Set authentication header\n",
        "        headers = {\n",
        "            \"Content-Type\": \"application/json\",\n",
        "            \"Authorization\": f\"Bearer {os.environ.get('OPENAI_API_KEY')}\",\n",
        "        }\n",
        "        if request_func_input.request_id:\n",
        "            headers[\"x-request-id\"] = request_func_input.request_id\n",
        "\n",
        "        # Initialize output metrics\n",
        "        output = endpoint_func.RequestFuncOutput()\n",
        "        output.prompt_len = request_func_input.prompt_len\n",
        "\n",
        "        generated_text = \"\"\n",
        "        ttft = 0.0  # Time to first token\n",
        "        st = time.perf_counter()\n",
        "        most_recent_timestamp = st\n",
        "\n",
        "        try:\n",
        "            async with session.post(\n",
        "                url=api_url, json=payload, headers=headers\n",
        "            ) as response:\n",
        "                if response.status == 200:\n",
        "                    # Parse streaming response\n",
        "                    async for chunk_bytes in response.content:\n",
        "                        chunk_bytes = chunk_bytes.strip()\n",
        "                        if not chunk_bytes:\n",
        "                            continue\n",
        "                        chunk_bytes = chunk_bytes.decode(\"utf-8\")\n",
        "                        if chunk_bytes.startswith(\":\"):\n",
        "                            continue\n",
        "\n",
        "                        chunk = chunk_bytes.removeprefix(\"data: \")\n",
        "                        if chunk != \"[DONE]\":\n",
        "                            timestamp = time.perf_counter()\n",
        "                            data = json.loads(chunk)\n",
        "\n",
        "                            if choices := data.get(\"choices\"):\n",
        "                                content = choices[0][\"delta\"].get(\"content\")\n",
        "                                if ttft == 0.0:\n",
        "                                    ttft = timestamp - st\n",
        "                                    output.ttft = ttft\n",
        "                                else:\n",
        "                                    # Record inter-token latency\n",
        "                                    output.itl.append(timestamp - most_recent_timestamp)\n",
        "                                generated_text += content or \"\"\n",
        "                            elif usage := data.get(\"usage\"):\n",
        "                                output.output_tokens = usage.get(\"completion_tokens\")\n",
        "\n",
        "                            most_recent_timestamp = timestamp\n",
        "\n",
        "                    output.generated_text = generated_text\n",
        "                    output.success = True\n",
        "                    output.latency = most_recent_timestamp - st\n",
        "                else:\n",
        "                    output.error = response.reason or \"\"\n",
        "                    output.success = False\n",
        "        except Exception:\n",
        "            output.success = False\n",
        "            exc_info = sys.exc_info()\n",
        "            output.error = \"\".join(traceback.format_exception(*exc_info))\n",
        "\n",
        "        if pbar:\n",
        "            pbar.update(1)\n",
        "        return output\n",
        "\n",
        "\n",
        "# Apply the monkey patch\n",
        "endpoint_func.async_request_openai_chat_completions = (\n",
        "    patched_async_request_openai_chat_completions\n",
        ")\n",
        "endpoint_func.ASYNC_REQUEST_FUNCS[\"openai-chat\"] = (\n",
        "    patched_async_request_openai_chat_completions\n",
        ")\n",
        "\n",
        "print(\"✅ Vertex AI compatibility patch applied successfully!\")\n",
        "print(\"   - Removed unsupported stream_options parameter\")\n",
        "print(\"   - Changed max_completion_tokens → max_tokens\")\n",
        "print(\"   - Increased timeout to 6 hours for large model inference\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AKCzSmJ26ugy"
      },
      "source": [
        "### Run Smoke Test on Baseline Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Fd1MoHL6ugy"
      },
      "outputs": [],
      "source": [
        "# Get authentication token\n",
        "creds, project = google.auth.default()\n",
        "auth_req = Request()\n",
        "creds.refresh(auth_req)\n",
        "\n",
        "# Set environment variables for vLLM benchmark\n",
        "os.environ[\"OPENAI_API_KEY\"] = creds.token\n",
        "os.environ[\"HF_HUB_OFFLINE\"] = \"1\"  # Use cached model artifacts\n",
        "\n",
        "# Construct Vertex AI endpoint URL\n",
        "baseline_dns = baseline_endpoint.gca_resource.dedicated_endpoint_dns\n",
        "baseline_url = f\"https://{baseline_dns}/v1beta1/{baseline_endpoint.resource_name}\"\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"🧪 SMOKE TEST: BASELINE ENDPOINT\")\n",
        "print(\"=\" * 80)\n",
        "print(f\"Endpoint: {baseline_endpoint.display_name}\")\n",
        "print(\"Test size: 100 prompts from ShareGPT\")\n",
        "print(\"Purpose: Verify endpoint is working correctly\")\n",
        "print(\"\\n⏳ Running test...\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "try:\n",
        "    result = subprocess.run(\n",
        "        [\n",
        "            \"vllm\",\n",
        "            \"bench\",\n",
        "            \"serve\",\n",
        "            \"--backend\",\n",
        "            \"openai-chat\",\n",
        "            \"--base-url\",\n",
        "            baseline_url,\n",
        "            \"--endpoint\",\n",
        "            \"/chat/completions\",\n",
        "            \"--model\",\n",
        "            \"\",\n",
        "            \"--tokenizer\",\n",
        "            MODEL_PATH,\n",
        "            \"--dataset-name\",\n",
        "            \"sharegpt\",\n",
        "            \"--dataset-path\",\n",
        "            DATASET_PATH,\n",
        "            \"--num-prompts\",\n",
        "            \"100\",\n",
        "        ],\n",
        "        check=True,\n",
        "        capture_output=False,\n",
        "    )\n",
        "\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"✅ BASELINE SMOKE TEST PASSED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(\"   Endpoint is responding correctly\")\n",
        "    print(\"   Ready for full benchmark\")\n",
        "    print(\"=\" * 80)\n",
        "\n",
        "except subprocess.CalledProcessError as e:\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"❌ SMOKE TEST FAILED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"Error: {e}\")\n",
        "    print(\"\\nPlease check:\")\n",
        "    print(\"  - Endpoint is fully deployed and healthy\")\n",
        "    print(\"  - Authentication token is valid (may need refresh)\")\n",
        "    print(\"  - Network connectivity to Vertex AI\")\n",
        "    print(\"=\" * 80)\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kw4AEIFe6ugy"
      },
      "source": [
        "### Run Smoke Test on EAGLE Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Scx2toKY6ugy"
      },
      "outputs": [],
      "source": [
        "# Refresh token (smoke tests may take several minutes)\n",
        "creds.refresh(auth_req)\n",
        "os.environ[\"OPENAI_API_KEY\"] = creds.token\n",
        "\n",
        "# Construct EAGLE endpoint URL\n",
        "eagle_dns = eagle_endpoint.gca_resource.dedicated_endpoint_dns\n",
        "eagle_url = f\"https://{eagle_dns}/v1beta1/{eagle_endpoint.resource_name}\"\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"🧪 SMOKE TEST: EAGLE ENDPOINT\")\n",
        "print(\"=\" * 80)\n",
        "print(f\"Endpoint: {eagle_endpoint.display_name}\")\n",
        "print(\"Test size: 100 prompts from ShareGPT\")\n",
        "print(\"Purpose: Verify EAGLE endpoint is working correctly\")\n",
        "print(\"\\n⏳ Running test...\")\n",
        "print(\"=\" * 80 + \"\\n\")\n",
        "\n",
        "try:\n",
        "    result = subprocess.run(\n",
        "        [\n",
        "            \"vllm\",\n",
        "            \"bench\",\n",
        "            \"serve\",\n",
        "            \"--backend\",\n",
        "            \"openai-chat\",\n",
        "            \"--base-url\",\n",
        "            eagle_url,\n",
        "            \"--endpoint\",\n",
        "            \"/chat/completions\",\n",
        "            \"--model\",\n",
        "            \"\",\n",
        "            \"--tokenizer\",\n",
        "            MODEL_PATH,\n",
        "            \"--dataset-name\",\n",
        "            \"sharegpt\",\n",
        "            \"--dataset-path\",\n",
        "            DATASET_PATH,\n",
        "            \"--num-prompts\",\n",
        "            \"100\",\n",
        "        ],\n",
        "        check=True,\n",
        "        capture_output=False,\n",
        "    )\n",
        "\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"✅ EAGLE SMOKE TEST PASSED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(\"   Endpoint is responding correctly\")\n",
        "    print(\"   Ready for full benchmark\")\n",
        "    print(\"=\" * 80)\n",
        "    print(\"\\n🎯 Both endpoints validated! Ready for comprehensive benchmarking.\")\n",
        "\n",
        "except subprocess.CalledProcessError as e:\n",
        "    print(\"\\n\" + \"=\" * 80)\n",
        "    print(\"❌ SMOKE TEST FAILED\")\n",
        "    print(\"=\" * 80)\n",
        "    print(f\"Error: {e}\")\n",
        "    print(\"\\nPlease check:\")\n",
        "    print(\"  - Endpoint is fully deployed and healthy\")\n",
        "    print(\"  - EAGLE draft model loaded correctly\")\n",
        "    print(\"  - Authentication token is valid (may need refresh)\")\n",
        "    print(\"  - Network connectivity to Vertex AI\")\n",
        "    print(\"=\" * 80)\n",
        "    raise"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AKJr1HOk6ugy"
      },
      "source": [
        "## Main Benchmark: Concurrency Sweep\n",
        "\n",
        "Now we'll run the comprehensive benchmark that tests both endpoints across multiple concurrency levels.\n",
        "\n",
        "**What is concurrency testing?**\n",
        "- Simulates multiple users sending requests simultaneously\n",
        "- Tests how the system scales under load\n",
        "- Reveals bottlenecks and optimal configuration\n",
        "\n",
        "\n",
        "**Benchmark setup**:\n",
        "\n",
        "- **Concurrency levels tested**: 1, 2, 4, 6, 8, 10 concurrent requests\n",
        "- **Prompts per level**: 1,000 (statistical significance)\n",
        "- **Total requests**: 12,000 (6,000 per endpoint)\n",
        "- **Expected duration**: 30-45 minutes per endpoint\n",
        "\n",
        "**What we'll measure:**\n",
        "\n",
        "| Metric | Formula | Why It Matters |\n",
        "|--------|---------|----------------|\n",
        "| **TTFT** (Time to First Token) | Time from request sent to first token received | User-perceived latency - how quickly responses start appearing |\n",
        "| **TPOT** (Time Per Output Token) | `(Total generation time - TTFT) / (num tokens - 1)` | Streaming smoothness - affects how fast text appears to stream |\n",
        "| **ITL** (Inter-Token Latency) | Time between consecutive tokens | Latency variation - consistent ITL = smooth streaming |\n",
        "| **Throughput** | `Total output tokens / Total time` | System capacity - how many tokens/sec the system can handle |\n",
        "| **Request Throughput** | `Total requests / Total time` | Request capacity - how many requests/sec the system can handle |\n",
        "\n",
        "**Why use median instead of mean?**\n",
        "- Medians are robust to outliers (e.g., one slow request doesn't skew results)\n",
        "- Better represents \"typical\" user experience\n",
        "- Industry standard for latency reporting (along with P99 for tail latency)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mexl7M4w6ugy"
      },
      "source": [
        "### Benchmark Baseline Across Concurrency Levels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y2h14t5f6ugy"
      },
      "outputs": [],
      "source": [
        "%%writefile run_baseline_concurrency.py\n",
        "\"\"\"Benchmark baseline endpoint across multiple concurrency levels.\"\"\"\n",
        "import subprocess\n",
        "import sys\n",
        "import os\n",
        "import google.auth\n",
        "from google.auth.transport.requests import Request\n",
        "\n",
        "def get_fresh_token():\n",
        "    \"\"\"Get a fresh authentication token from Google Cloud.\"\"\"\n",
        "    try:\n",
        "        creds, project = google.auth.default()\n",
        "        auth_req = Request()\n",
        "        creds.refresh(auth_req)\n",
        "        return creds.token\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed to refresh token: {e}\")\n",
        "        raise\n",
        "\n",
        "# Get configuration from environment\n",
        "baseline_url = os.environ[\"BASELINE_URL\"]\n",
        "model_path = os.environ[\"MODEL_PATH\"]\n",
        "dataset_path = os.environ[\"DATASET_PATH\"]\n",
        "\n",
        "# Concurrency levels to test\n",
        "CONCURRENCY_LEVELS = [1, 2, 4, 6, 8, 10]\n",
        "\n",
        "print(\"=\"*80)\n",
        "print(\"📊 COMPREHENSIVE BENCHMARK: BASELINE MODEL\")\n",
        "print(\"=\"*80)\n",
        "print(f\"Endpoint: {baseline_url}\")\n",
        "print(f\"Concurrency levels: {CONCURRENCY_LEVELS}\")\n",
        "print(f\"Prompts per level: 1,000\")\n",
        "print(f\"Total requests: {len(CONCURRENCY_LEVELS) * 1000:,}\")\n",
        "print(\"=\"*80)\n",
        "\n",
        "# Create output directory\n",
        "os.makedirs(\"benchmarks/baseline_concurrency\", exist_ok=True)\n",
        "\n",
        "# Run benchmark for each concurrency level\n",
        "for i, concurrency in enumerate(CONCURRENCY_LEVELS, 1):\n",
        "    print(f\"\\n{'='*80}\")\n",
        "    print(f\"🔄 BASELINE - Concurrency {concurrency} ({i}/{len(CONCURRENCY_LEVELS)})\")\n",
        "    print(f\"{'='*80}\")\n",
        "    print(f\"⏳ Running 1,000 requests with max concurrency={concurrency}...\\n\")\n",
        "\n",
        "    # Refresh token before each run (benchmarks can be long)\n",
        "    token = get_fresh_token()\n",
        "    os.environ[\"OPENAI_API_KEY\"] = token\n",
        "\n",
        "    try:\n",
        "        subprocess.run(\n",
        "            [\n",
        "                \"vllm\", \"bench\", \"serve\",\n",
        "                \"--backend\", \"openai-chat\",\n",
        "                \"--base-url\", baseline_url,\n",
        "                \"--endpoint\", \"/chat/completions\",\n",
        "                \"--model\", \"\",\n",
        "                \"--tokenizer\", model_path,\n",
        "                \"--dataset-name\", \"sharegpt\",\n",
        "                \"--dataset-path\", dataset_path,\n",
        "                \"--num-prompts\", \"1000\",\n",
        "                \"--max-concurrency\", str(concurrency),\n",
        "                \"--save-result\",\n",
        "                \"--result-dir\", \"benchmarks/baseline_concurrency\",\n",
        "                \"--result-filename\", f\"baseline_c{concurrency}.json\"\n",
        "            ],\n",
        "            check=True,\n",
        "        )\n",
        "\n",
        "        print(f\"\\n✅ Concurrency {concurrency} completed successfully\")\n",
        "        print(f\"   Results saved to: benchmarks/baseline_concurrency/baseline_c{concurrency}.json\")\n",
        "\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        print(f\"\\n❌ Benchmark failed at concurrency {concurrency}\")\n",
        "        print(f\"   Error: {e}\")\n",
        "        print(f\"   Continuing with next concurrency level...\")\n",
        "        continue\n",
        "\n",
        "print(\"\\n\" + \"=\"*80)\n",
        "print(\"✅ BASELINE CONCURRENCY SWEEP COMPLETED\")\n",
        "print(\"=\"*80)\n",
        "print(f\"   Results saved to: benchmarks/baseline_concurrency/\")\n",
        "print(\"=\"*80)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WEHM18zU6ugy"
      },
      "outputs": [],
      "source": [
        "# Set environment variables for the script\n",
        "os.environ[\"BASELINE_URL\"] = baseline_url\n",
        "os.environ[\"MODEL_PATH\"] = MODEL_PATH\n",
        "os.environ[\"DATASET_PATH\"] = DATASET_PATH\n",
        "\n",
        "print(\"🚀 Starting baseline concurrency sweep...\")\n",
        "print(\"   This will take approximately 30-45 minutes\\n\")\n",
        "\n",
        "# Run the benchmark script\n",
        "!python run_baseline_concurrency.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Eok6JBP6ugy"
      },
      "source": [
        "### Benchmark EAGLE Across Concurrency Levels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d-2QDmfk6ugy"
      },
      "outputs": [],
      "source": [
        "%%writefile run_eagle_concurrency.py\n",
        "\"\"\"Benchmark EAGLE endpoint across multiple concurrency levels.\"\"\"\n",
        "import subprocess\n",
        "import sys\n",
        "import os\n",
        "import google.auth\n",
        "from google.auth.transport.requests import Request\n",
        "\n",
        "def get_fresh_token():\n",
        "    \"\"\"Get a fresh authentication token from Google Cloud.\"\"\"\n",
        "    try:\n",
        "        creds, project = google.auth.default()\n",
        "        auth_req = Request()\n",
        "        creds.refresh(auth_req)\n",
        "        return creds.token\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed to refresh token: {e}\")\n",
        "        raise\n",
        "\n",
        "# Get configuration from environment\n",
        "eagle_url = os.environ[\"EAGLE_URL\"]\n",
        "model_path = os.environ[\"MODEL_PATH\"]\n",
        "dataset_path = os.environ[\"DATASET_PATH\"]\n",
        "\n",
        "# Concurrency levels to test (same as baseline for fair comparison)\n",
        "CONCURRENCY_LEVELS = [1, 2, 4, 6, 8, 10]\n",
        "\n",
        "print(\"=\"*80)\n",
        "print(\"📊 COMPREHENSIVE BENCHMARK: EAGLE MODEL\")\n",
        "print(\"=\"*80)\n",
        "print(f\"Endpoint: {eagle_url}\")\n",
        "print(f\"Concurrency levels: {CONCURRENCY_LEVELS}\")\n",
        "print(f\"Prompts per level: 1,000\")\n",
        "print(f\"Total requests: {len(CONCURRENCY_LEVELS) * 1000:,}\")\n",
        "print(\"=\"*80)\n",
        "\n",
        "# Create output directory\n",
        "os.makedirs(\"benchmarks/eagle_concurrency\", exist_ok=True)\n",
        "\n",
        "# Run benchmark for each concurrency level\n",
        "for i, concurrency in enumerate(CONCURRENCY_LEVELS, 1):\n",
        "    print(f\"\\n{'='*80}\")\n",
        "    print(f\"🔄 EAGLE - Concurrency {concurrency} ({i}/{len(CONCURRENCY_LEVELS)})\")\n",
        "    print(f\"{'='*80}\")\n",
        "    print(f\"⏳ Running 1,000 requests with max concurrency={concurrency}...\\n\")\n",
        "\n",
        "    # Refresh token before each run (benchmarks can be long)\n",
        "    token = get_fresh_token()\n",
        "    os.environ[\"OPENAI_API_KEY\"] = token\n",
        "\n",
        "    try:\n",
        "        subprocess.run(\n",
        "            [\n",
        "                \"vllm\", \"bench\", \"serve\",\n",
        "                \"--backend\", \"openai-chat\",\n",
        "                \"--base-url\", eagle_url,\n",
        "                \"--endpoint\", \"/chat/completions\",\n",
        "                \"--model\", \"\",\n",
        "                \"--tokenizer\", model_path,\n",
        "                \"--dataset-name\", \"sharegpt\",\n",
        "                \"--dataset-path\", dataset_path,\n",
        "                \"--num-prompts\", \"1000\",\n",
        "                \"--max-concurrency\", str(concurrency),\n",
        "                \"--save-result\",\n",
        "                \"--result-dir\", \"benchmarks/eagle_concurrency\",\n",
        "                \"--result-filename\", f\"eagle_c{concurrency}.json\"\n",
        "            ],\n",
        "            check=True,\n",
        "        )\n",
        "\n",
        "        print(f\"\\n✅ Concurrency {concurrency} completed successfully\")\n",
        "        print(f\"   Results saved to: benchmarks/eagle_concurrency/eagle_c{concurrency}.json\")\n",
        "\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        print(f\"\\n❌ Benchmark failed at concurrency {concurrency}\")\n",
        "        print(f\"   Error: {e}\")\n",
        "        print(f\"   Continuing with next concurrency level...\")\n",
        "        continue\n",
        "\n",
        "print(\"\\n\" + \"=\"*80)\n",
        "print(\"✅ EAGLE CONCURRENCY SWEEP COMPLETED\")\n",
        "print(\"=\"*80)\n",
        "print(f\"   Results saved to: benchmarks/eagle_concurrency/\")\n",
        "print(\"=\"*80)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0zo30uNT6ugy"
      },
      "outputs": [],
      "source": [
        "# Set environment variables for the script\n",
        "os.environ[\"EAGLE_URL\"] = eagle_url\n",
        "os.environ[\"MODEL_PATH\"] = MODEL_PATH\n",
        "os.environ[\"DATASET_PATH\"] = DATASET_PATH\n",
        "\n",
        "print(\"🚀 Starting EAGLE concurrency sweep...\")\n",
        "print(\"   This will take approximately 30-45 minutes\\n\")\n",
        "\n",
        "# Run the benchmark script\n",
        "!python run_eagle_concurrency.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aqq2Ebm26ugy"
      },
      "source": [
        "## Analysis and Visualization\n",
        "\n",
        "Now let's analyze the benchmark results to quantify EAGLE's performance improvement."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_vbkh4e6ug1"
      },
      "source": [
        "### Load and Parse Benchmark Results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RtxRG1eh6ug2"
      },
      "outputs": [],
      "source": [
        "# Load results for all concurrency levels\n",
        "baseline_results = {}\n",
        "eagle_results = {}\n",
        "concurrency_levels = [1, 2, 4, 6, 8, 10]\n",
        "\n",
        "print(\"📂 Loading benchmark results...\\n\")\n",
        "\n",
        "try:\n",
        "    for concurrency in concurrency_levels:\n",
        "        # Load baseline results\n",
        "        baseline_file = f\"benchmarks/baseline_concurrency/baseline_c{concurrency}.json\"\n",
        "        with open(baseline_file) as f:\n",
        "            baseline_results[concurrency] = json.load(f)\n",
        "        print(f\"✅ Loaded baseline concurrency {concurrency}\")\n",
        "\n",
        "        # Load EAGLE results\n",
        "        eagle_file = f\"benchmarks/eagle_concurrency/eagle_c{concurrency}.json\"\n",
        "        with open(eagle_file) as f:\n",
        "            eagle_results[concurrency] = json.load(f)\n",
        "        print(f\"✅ Loaded EAGLE concurrency {concurrency}\")\n",
        "\n",
        "    print(\"\\n✅ All results loaded successfully!\")\n",
        "\n",
        "except FileNotFoundError as e:\n",
        "    print(f\"\\n❌ Failed to load results: {e}\")\n",
        "    print(\"   Make sure all benchmarks completed successfully\")\n",
        "    raise"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lUyrLXBb6ug2"
      },
      "outputs": [],
      "source": [
        "# Extract key metrics for each concurrency level\n",
        "baseline_metrics = []\n",
        "eagle_metrics = []\n",
        "\n",
        "for concurrency in concurrency_levels:\n",
        "    baseline_data = baseline_results[concurrency]\n",
        "    eagle_data = eagle_results[concurrency]\n",
        "\n",
        "    baseline_metrics.append(\n",
        "        {\n",
        "            \"Concurrency\": concurrency,\n",
        "            \"TTFT (ms)\": baseline_data[\"median_ttft_ms\"],\n",
        "            \"TPOT (ms)\": baseline_data[\"median_tpot_ms\"],\n",
        "            \"Throughput (tok/s)\": baseline_data[\"output_throughput\"],\n",
        "            \"Request Throughput (req/s)\": baseline_data[\"request_throughput\"],\n",
        "        }\n",
        "    )\n",
        "\n",
        "    eagle_metrics.append(\n",
        "        {\n",
        "            \"Concurrency\": concurrency,\n",
        "            \"TTFT (ms)\": eagle_data[\"median_ttft_ms\"],\n",
        "            \"TPOT (ms)\": eagle_data[\"median_tpot_ms\"],\n",
        "            \"Throughput (tok/s)\": eagle_data[\"output_throughput\"],\n",
        "            \"Request Throughput (req/s)\": eagle_data[\"request_throughput\"],\n",
        "        }\n",
        "    )\n",
        "\n",
        "# Create DataFrames for easy comparison\n",
        "baseline_df = pd.DataFrame(baseline_metrics)\n",
        "eagle_df = pd.DataFrame(eagle_metrics)\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"📊 BASELINE PERFORMANCE BY CONCURRENCY\")\n",
        "print(\"=\" * 80)\n",
        "print(baseline_df.to_string(index=False))\n",
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"📊 EAGLE PERFORMANCE BY CONCURRENCY\")\n",
        "print(\"=\" * 80)\n",
        "print(eagle_df.to_string(index=False))\n",
        "print(\"=\" * 80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVbWxkza6ug2"
      },
      "source": [
        "### Calculate Performance Improvements"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cc07GRcv6ug2"
      },
      "outputs": [],
      "source": [
        "# Calculate percentage improvements at each concurrency level\n",
        "improvements = []\n",
        "\n",
        "for i, concurrency in enumerate(concurrency_levels):\n",
        "    baseline = baseline_metrics[i]\n",
        "    eagle = eagle_metrics[i]\n",
        "\n",
        "    # Calculate improvements (negative = worse, positive = better)\n",
        "    ttft_improvement = (\n",
        "        (baseline[\"TTFT (ms)\"] - eagle[\"TTFT (ms)\"]) / baseline[\"TTFT (ms)\"]\n",
        "    ) * 100\n",
        "    tpot_improvement = (\n",
        "        (baseline[\"TPOT (ms)\"] - eagle[\"TPOT (ms)\"]) / baseline[\"TPOT (ms)\"]\n",
        "    ) * 100\n",
        "    throughput_improvement = (\n",
        "        (eagle[\"Throughput (tok/s)\"] - baseline[\"Throughput (tok/s)\"])\n",
        "        / baseline[\"Throughput (tok/s)\"]\n",
        "    ) * 100\n",
        "    req_throughput_improvement = (\n",
        "        (eagle[\"Request Throughput (req/s)\"] - baseline[\"Request Throughput (req/s)\"])\n",
        "        / baseline[\"Request Throughput (req/s)\"]\n",
        "    ) * 100\n",
        "\n",
        "    improvements.append(\n",
        "        {\n",
        "            \"Concurrency\": concurrency,\n",
        "            \"TTFT Improvement (%)\": ttft_improvement,\n",
        "            \"TPOT Improvement (%)\": tpot_improvement,\n",
        "            \"Throughput Speedup (%)\": throughput_improvement,\n",
        "            \"Req Throughput Speedup (%)\": req_throughput_improvement,\n",
        "        }\n",
        "    )\n",
        "\n",
        "improvements_df = pd.DataFrame(improvements)\n",
        "\n",
        "print(\"=\" * 80)\n",
        "print(\"📈 EAGLE PERFORMANCE IMPROVEMENTS OVER BASELINE\")\n",
        "print(\"=\" * 80)\n",
        "print(\"   (Positive values = EAGLE is better)\")\n",
        "print(\"=\" * 80)\n",
        "print(improvements_df.to_string(index=False))\n",
        "print(\"=\" * 80)\n",
        "\n",
        "# Calculate average improvements\n",
        "avg_throughput_speedup = improvements_df[\"Throughput Speedup (%)\"].mean()\n",
        "avg_req_speedup = improvements_df[\"Req Throughput Speedup (%)\"].mean()\n",
        "avg_ttft_improvement = improvements_df[\"TTFT Improvement (%)\"].mean()\n",
        "avg_tpot_improvement = improvements_df[\"TPOT Improvement (%)\"].mean()\n",
        "\n",
        "print(\"\\n\" + \"=\" * 80)\n",
        "print(\"🎯 KEY TAKEAWAYS (Averaged Across All Concurrency Levels)\")\n",
        "print(\"=\" * 80)\n",
        "print(\n",
        "    f\"   Token Throughput: {avg_throughput_speedup:+.1f}% {'faster' if avg_throughput_speedup > 0 else 'slower'} with EAGLE\"\n",
        ")\n",
        "print(\n",
        "    f\"   Request Throughput: {avg_req_speedup:+.1f}% {'faster' if avg_req_speedup > 0 else 'slower'} with EAGLE\"\n",
        ")\n",
        "print(\n",
        "    f\"   Time to First Token: {avg_ttft_improvement:+.1f}% {'faster' if avg_ttft_improvement > 0 else 'slower'} with EAGLE\"\n",
        ")\n",
        "print(\n",
        "    f\"   Time Per Output Token: {avg_tpot_improvement:+.1f}% {'faster' if avg_tpot_improvement > 0 else 'slower'} with EAGLE\"\n",
        ")\n",
        "print(\"=\" * 80)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wUTzj-iu6ug2"
      },
      "source": [
        "### Visualize Performance Comparison\n",
        "\n",
        "**How to interpret these charts:**\n",
        "\n",
        "1. **TTFT Chart (Top Left)**: Lower is better - shows how quickly responses start\n",
        "   - EAGLE may have slightly higher TTFT due to draft model overhead\n",
        "   - This is expected and acceptable if overall throughput improves\n",
        "\n",
        "2. **TPOT Chart (Top Right)**: Lower is better - shows per-token generation speed\n",
        "   - EAGLE should show lower TPOT (faster per-token generation)\n",
        "   - This is where EAGLE's speedup comes from\n",
        "\n",
        "3. **Token Throughput Chart (Bottom Left)**: Higher is better - shows system capacity\n",
        "   - EAGLE should show higher throughput (more tokens/sec)\n",
        "   - This translates directly to cost savings\n",
        "\n",
        "4. **Request Throughput Chart (Bottom Right)**: Higher is better - shows request capacity\n",
        "   - EAGLE should handle more requests/sec\n",
        "   - Better for production workloads with many concurrent users"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gu5kqxW16ug2"
      },
      "outputs": [],
      "source": [
        "# Set visualization style\n",
        "sns.set_style(\"whitegrid\")\n",
        "plt.rcParams[\"figure.figsize\"] = (16, 12)\n",
        "\n",
        "# Create a 2x2 grid of line charts\n",
        "fig, axes = plt.subplots(2, 2, figsize=(16, 12))\n",
        "fig.suptitle(\n",
        "    \"EAGLE vs Baseline: Performance Across Concurrency Levels\",\n",
        "    fontsize=16,\n",
        "    fontweight=\"bold\",\n",
        "    y=0.995,\n",
        ")\n",
        "\n",
        "# Extract data for plotting\n",
        "concurrency_values = baseline_df[\"Concurrency\"].values\n",
        "baseline_ttft = baseline_df[\"TTFT (ms)\"].values\n",
        "eagle_ttft = eagle_df[\"TTFT (ms)\"].values\n",
        "baseline_tpot = baseline_df[\"TPOT (ms)\"].values\n",
        "eagle_tpot = eagle_df[\"TPOT (ms)\"].values\n",
        "baseline_throughput = baseline_df[\"Throughput (tok/s)\"].values\n",
        "eagle_throughput = eagle_df[\"Throughput (tok/s)\"].values\n",
        "baseline_req_throughput = baseline_df[\"Request Throughput (req/s)\"].values\n",
        "eagle_req_throughput = eagle_df[\"Request Throughput (req/s)\"].values\n",
        "\n",
        "# Plot 1: Time to First Token vs Concurrency\n",
        "ax1 = axes[0, 0]\n",
        "ax1.plot(\n",
        "    concurrency_values,\n",
        "    baseline_ttft,\n",
        "    \"o-\",\n",
        "    color=\"#4285F4\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"Baseline\",\n",
        ")\n",
        "ax1.plot(\n",
        "    concurrency_values,\n",
        "    eagle_ttft,\n",
        "    \"s-\",\n",
        "    color=\"#34A853\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"EAGLE\",\n",
        ")\n",
        "ax1.set_xlabel(\"Max Concurrency\", fontsize=11, fontweight=\"bold\")\n",
        "ax1.set_ylabel(\"Median TTFT (ms)\", fontsize=11, fontweight=\"bold\")\n",
        "ax1.set_title(\n",
        "    \"Time to First Token vs Concurrency\\n(Lower is Better)\",\n",
        "    fontsize=12,\n",
        "    fontweight=\"bold\",\n",
        ")\n",
        "ax1.legend(loc=\"best\", fontsize=10)\n",
        "ax1.grid(True, alpha=0.3)\n",
        "\n",
        "# Plot 2: Time Per Output Token vs Concurrency\n",
        "ax2 = axes[0, 1]\n",
        "ax2.plot(\n",
        "    concurrency_values,\n",
        "    baseline_tpot,\n",
        "    \"o-\",\n",
        "    color=\"#4285F4\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"Baseline\",\n",
        ")\n",
        "ax2.plot(\n",
        "    concurrency_values,\n",
        "    eagle_tpot,\n",
        "    \"s-\",\n",
        "    color=\"#34A853\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"EAGLE\",\n",
        ")\n",
        "ax2.set_xlabel(\"Max Concurrency\", fontsize=11, fontweight=\"bold\")\n",
        "ax2.set_ylabel(\"Median TPOT (ms)\", fontsize=11, fontweight=\"bold\")\n",
        "ax2.set_title(\n",
        "    \"Time Per Output Token vs Concurrency\\n(Lower is Better)\",\n",
        "    fontsize=12,\n",
        "    fontweight=\"bold\",\n",
        ")\n",
        "ax2.legend(loc=\"best\", fontsize=10)\n",
        "ax2.grid(True, alpha=0.3)\n",
        "\n",
        "# Plot 3: Token Throughput vs Concurrency\n",
        "ax3 = axes[1, 0]\n",
        "ax3.plot(\n",
        "    concurrency_values,\n",
        "    baseline_throughput,\n",
        "    \"o-\",\n",
        "    color=\"#4285F4\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"Baseline\",\n",
        ")\n",
        "ax3.plot(\n",
        "    concurrency_values,\n",
        "    eagle_throughput,\n",
        "    \"s-\",\n",
        "    color=\"#34A853\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"EAGLE\",\n",
        ")\n",
        "ax3.set_xlabel(\"Max Concurrency\", fontsize=11, fontweight=\"bold\")\n",
        "ax3.set_ylabel(\"Throughput (tokens/s)\", fontsize=11, fontweight=\"bold\")\n",
        "ax3.set_title(\n",
        "    \"Token Throughput vs Concurrency\\n(Higher is Better)\",\n",
        "    fontsize=12,\n",
        "    fontweight=\"bold\",\n",
        ")\n",
        "ax3.legend(loc=\"best\", fontsize=10)\n",
        "ax3.grid(True, alpha=0.3)\n",
        "\n",
        "# Plot 4: Request Throughput vs Concurrency\n",
        "ax4 = axes[1, 1]\n",
        "ax4.plot(\n",
        "    concurrency_values,\n",
        "    baseline_req_throughput,\n",
        "    \"o-\",\n",
        "    color=\"#4285F4\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"Baseline\",\n",
        ")\n",
        "ax4.plot(\n",
        "    concurrency_values,\n",
        "    eagle_req_throughput,\n",
        "    \"s-\",\n",
        "    color=\"#34A853\",\n",
        "    linewidth=2,\n",
        "    markersize=8,\n",
        "    label=\"EAGLE\",\n",
        ")\n",
        "ax4.set_xlabel(\"Max Concurrency\", fontsize=11, fontweight=\"bold\")\n",
        "ax4.set_ylabel(\"Request Throughput (req/s)\", fontsize=11, fontweight=\"bold\")\n",
        "ax4.set_title(\n",
        "    \"Request Throughput vs Concurrency\\n(Higher is Better)\",\n",
        "    fontsize=12,\n",
        "    fontweight=\"bold\",\n",
        ")\n",
        "ax4.legend(loc=\"best\", fontsize=10)\n",
        "ax4.grid(True, alpha=0.3)\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.savefig(\"eagle_concurrency_analysis.png\", dpi=300, bbox_inches=\"tight\")\n",
        "plt.show()\n",
        "\n",
        "print(\"\\n✅ Visualization saved as 'eagle_concurrency_analysis.png'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mfY7ucr66ug2"
      },
      "source": [
        "## Cleanup\n",
        "\n",
        "**Important:** These endpoints run on expensive hardware (8x H100 GPUs). Make sure to delete them when done to avoid unnecessary costs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hHSKnlF46ug2"
      },
      "outputs": [],
      "source": [
        "# Set to True to delete endpoints and models\n",
        "delete_endpoints = False  # @param {type: \"boolean\"}\n",
        "delete_models = False  # @param {type: \"boolean\"}\n",
        "\n",
        "if delete_endpoints:\n",
        "    print(\"🗑️  Deleting endpoints...\\n\")\n",
        "\n",
        "    try:\n",
        "        # Undeploy and delete baseline endpoint\n",
        "        print(\"   Undeploying baseline endpoint...\")\n",
        "        baseline_endpoint.undeploy_all()\n",
        "        print(\"   Deleting baseline endpoint...\")\n",
        "        baseline_endpoint.delete()\n",
        "        print(\"   ✅ Baseline endpoint deleted\\n\")\n",
        "\n",
        "        # Undeploy and delete EAGLE endpoint\n",
        "        print(\"   Undeploying EAGLE endpoint...\")\n",
        "        eagle_endpoint.undeploy_all()\n",
        "        print(\"   Deleting EAGLE endpoint...\")\n",
        "        eagle_endpoint.delete()\n",
        "        print(\"   ✅ EAGLE endpoint deleted\\n\")\n",
        "\n",
        "        print(\"✅ All endpoints deleted successfully!\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\n❌ Failed to delete endpoints: {e}\")\n",
        "        print(\"   You may need to delete them manually from the console\")\n",
        "else:\n",
        "    print(\"⚠️  Endpoints not deleted (delete_endpoints=False)\")\n",
        "    print(\"   Remember to delete them manually to avoid charges!\")\n",
        "    print(f\"\\n   Baseline endpoint: {baseline_endpoint.resource_name}\")\n",
        "    print(f\"   EAGLE endpoint: {eagle_endpoint.resource_name}\")\n",
        "\n",
        "if delete_models:\n",
        "    print(\"\\n🗑️  Deleting models...\\n\")\n",
        "\n",
        "    try:\n",
        "        baseline_model.delete()\n",
        "        print(\"   ✅ Baseline model deleted\")\n",
        "\n",
        "        eagle_model.delete()\n",
        "        print(\"   ✅ EAGLE model deleted\")\n",
        "\n",
        "        print(\"\\n✅ All models deleted successfully!\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\n❌ Failed to delete models: {e}\")\n",
        "        print(\"   You may need to delete them manually from the console\")\n",
        "else:\n",
        "    print(\"\\n⚠️  Models not deleted (delete_models=False)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1vTE6suu6ug2"
      },
      "source": [
        "## Next Steps\n",
        "\n",
        "Now that you've successfully benchmarked EAGLE on Vertex AI, here are some next steps:\n",
        "\n",
        "### 1. Test with Your Own Data\n",
        "Replace ShareGPT with your production prompts:\n",
        "```python\n",
        "# Save your prompts as JSON in ShareGPT format\n",
        "your_prompts = [\n",
        "    {\"id\": \"1\", \"conversations\": [{\"from\": \"human\", \"value\": \"Your prompt here\"}]},\n",
        "    # ... more prompts\n",
        "]\n",
        "with open(\"/tmp/your_prompts.json\", \"w\") as f:\n",
        "    json.dump(your_prompts, f)\n",
        "```\n",
        "\n",
        "### 2. Optimize EAGLE Parameters\n",
        "Experiment with different EAGLE configurations:\n",
        "- `--speculative-num-steps`: Try 2, 3, 4, 5 (higher = more speculation)\n",
        "- `--speculative-num-draft-tokens`: Try 4, 8, 12 (higher = more tokens per step)\n",
        "- `--speculative-eagle-topk`: Try 3, 4, 5 (higher = more diverse predictions)\n",
        "\n",
        "### 3. Production Deployment\n",
        "For production use:\n",
        "- Enable autoscaling: `min_replica_count=1, max_replica_count=5`\n",
        "- Set up monitoring and alerting\n",
        "- Implement A/B testing between baseline and EAGLE\n",
        "- Configure request/response logging\n",
        "\n",
        "### 4. Cost Optimization\n",
        "- Compare cost per token: `GPU cost / tokens generated`\n",
        "- Calculate break-even point for your workload\n",
        "- Consider using cheaper GPUs (L4, A100) if throughput requirements are lower\n",
        "\n",
        "## Additional Resources\n",
        "\n",
        "- [Vertex AI Model Garden Documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/model-garden/explore-models)\n",
        "- [vLLM Speculative Decoding Guide](https://docs.vllm.ai/en/latest/models/spec_decode.html)\n",
        "- [EAGLE Paper (arXiv)](https://arxiv.org/abs/2401.15077)\n",
        "- [vLLM Benchmark Documentation](https://docs.vllm.ai/en/latest/serving/benchmarking.html)\n",
        "- [Llama 4 Model Card](https://ai.meta.com/llama/)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "benchmarking_eagle_on_vertex_ai.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
