labs/backend-pool-load-balancing/backend-pool-load-balancing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# APIM ❤️ OpenAI\n",
    "\n",
    "## Backend pool Load Balancing lab\n",
    "![flow](../../images/backend-pool-load-balancing.gif)\n",
    "\n",
    "Playground to try the built-in load balancing [backend pool functionality of APIM](https://learn.microsoft.com/azure/api-management/backends?tabs=bicep) to a list of Azure OpenAI endpoints.\n",
    "\n",
    "Notes:\n",
    "- **This is a typical prioritized PTU with fallback consumption scenario**. The lab specifically showcases how a priority 1 (highest) backend is exhausted before gracefully falling back to two equally-weighted priority 2 backends.\n",
    "- The backend pool uses round-robin by default.\n",
    "- Priority and weight-based routing are supported and can be adjusted by modifying `priority` (the lower the number, the higher the priority) and `weight` parameters in the `openai_resources` variable below.\n",
    "- The `retry` API Management policy initiates a retry to an available backend if an HTTP 429 status code is encountered. This is transparent to the caller.\n",
    "\n",
    "### Result\n",
    "![result](result.png)\n",
    "\n",
    "### Prerequisites\n",
    "- [Python 3.12 or later version](https://www.python.org/) installed\n",
    "- [Pandas Library](https://pandas.pydata.org/) and matplotlib installed\n",
    "- [VS Code](https://code.visualstudio.com/) installed with the [Jupyter notebook extension](https://marketplace.visualstudio.com/items?itemName=ms-toolsai.jupyter) enabled\n",
    "- [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) installed\n",
    "- [An Azure Subscription](https://azure.microsoft.com/free/) with Contributor permissions\n",
    "- [Access granted to Azure OpenAI](https://aka.ms/oai/access) or just enable the mock service\n",
    "- [Sign in to Azure with Azure CLI](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='0'></a>\n",
    "### 0️⃣ Initialize notebook variables\n",
    "\n",
    "- Resources will be suffixed by a unique string based on your subscription id.\n",
    "- Adjust the location parameters according your preferences and on the [product availability by Azure region.](https://azure.microsoft.com/explore/global-infrastructure/products-by-region/?cdn=disable&products=cognitive-services,api-management)\n",
    "- Adjust the OpenAI model and version according the [availability by region.](https://learn.microsoft.com/azure/ai-services/openai/concepts/models) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, json\n",
    "sys.path.insert(1, '../../shared')  # add the shared directory to the Python path\n",
    "import utils\n",
    "\n",
    "deployment_name = os.path.basename(os.path.dirname(globals()['__vsc_ipynb_file__']))\n",
    "resource_group_name = f\"lab-{deployment_name}\" # change the name to match your naming style\n",
    "resource_group_location = \"westeurope\"\n",
    "\n",
    "apim_sku = 'Basicv2'\n",
    "\n",
    "# Prioritize East US until exhaustion (simulate PTU with TPM), then equally distribute between Sweden and West US (consumption fallback)\n",
    "openai_resources = [\n",
    "    {\"name\": \"openai1\", \"location\": \"eastus\", \"priority\": 1},\n",
    "    {\"name\": \"openai2\", \"location\": \"swedencentral\", \"priority\": 2, \"weight\": 50},\n",
    "    {\"name\": \"openai3\", \"location\": \"westus\", \"priority\": 2, \"weight\": 50}\n",
    "]\n",
    "\n",
    "openai_deployment_name = \"gpt-4o-mini\"\n",
    "openai_model_name = \"gpt-4o-mini\"\n",
    "openai_model_version = \"2024-07-18\"\n",
    "openai_model_capacity = 8\n",
    "openai_model_sku = 'Standard'\n",
    "openai_api_version = \"2024-02-01\"\n",
    "\n",
    "utils.print_ok('Notebook initialized')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='1'></a>\n",
    "### 1️⃣ Verify the Azure CLI and the connected Azure subscription\n",
    "\n",
    "The following commands ensure that you have the latest version of the Azure CLI and that the Azure CLI is connected to your Azure subscription."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = utils.run(\"az account show\", \"Retrieved az account\", \"Failed to get the current az account\")\n",
    "\n",
    "if output.success and output.json_data:\n",
    "    current_user = output.json_data['user']['name']\n",
    "    tenant_id = output.json_data['tenantId']\n",
    "    subscription_id = output.json_data['id']\n",
    "\n",
    "    utils.print_info(f\"Current user: {current_user}\")\n",
    "    utils.print_info(f\"Tenant ID: {tenant_id}\")\n",
    "    utils.print_info(f\"Subscription ID: {subscription_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2'></a>\n",
    "### 2️⃣ Create deployment using 🦾 Bicep\n",
    "\n",
    "This lab uses [Bicep](https://learn.microsoft.com/azure/azure-resource-manager/bicep/overview?tabs=bicep) to declarative define all the resources that will be deployed in the specified resource group. Change the parameters or the [main.bicep](main.bicep) directly to try different configurations.\n",
    "\n",
    "`openAIModelCapacity` is set intentionally low to `8` (8k tokens per minute) to trigger the retry logic in the load balancer (transparent to the user) as well as the priority failover from priority 1 to 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the resource group if doesn't exist\n",
    "utils.create_resource_group(resource_group_name, resource_group_location)\n",
    "\n",
    "# Define the Bicep parameters\n",
    "bicep_parameters = {\n",
    "    \"$schema\": \"https://schema.management.azure.com/schemas/2019-04-01/deploymentParameters.json#\",\n",
    "    \"contentVersion\": \"1.0.0.0\",\n",
    "    \"parameters\": {\n",
    "        \"apimSku\": { \"value\": apim_sku },\n",
    "        \"openAIConfig\": { \"value\": openai_resources },\n",
    "        \"openAIDeploymentName\": { \"value\": openai_deployment_name },\n",
    "        \"openAIModelName\": { \"value\": openai_model_name },\n",
    "        \"openAIModelVersion\": { \"value\": openai_model_version },\n",
    "        \"openAIModelCapacity\": { \"value\": openai_model_capacity },\n",
    "        \"openAIModelSKU\": { \"value\": openai_model_sku },\n",
    "        \"openAIAPIVersion\": { \"value\": openai_api_version }\n",
    "    }\n",
    "}\n",
    "\n",
    "# Write the parameters to the params.json file\n",
    "with open('params.json', 'w') as bicep_parameters_file:\n",
    "    bicep_parameters_file.write(json.dumps(bicep_parameters))\n",
    "\n",
    "# Run the deployment\n",
    "output = utils.run(f\"az deployment group create --name {deployment_name} --resource-group {resource_group_name} --template-file main.bicep --parameters params.json\",\n",
    "    f\"Deployment '{deployment_name}' succeeded\", f\"Deployment '{deployment_name}' failed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3'></a>\n",
    "### 3️⃣ Get the deployment outputs\n",
    "\n",
    "Retrieve the required outputs from the Bicep deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain all of the outputs from the deployment\n",
    "output = utils.run(f\"az deployment group show --name {deployment_name} -g {resource_group_name}\", f\"Retrieved deployment: {deployment_name}\", f\"Failed to retrieve deployment: {deployment_name}\")\n",
    "\n",
    "if output.success and output.json_data:\n",
    "    apim_service_id = utils.get_deployment_output(output, 'apimServiceId', 'APIM Service Id')\n",
    "    apim_resource_gateway_url = utils.get_deployment_output(output, 'apimResourceGatewayURL', 'APIM API Gateway URL')\n",
    "    apim_subscription_key = utils.get_deployment_output(output, 'apimSubscriptionKey', 'APIM Subscription Key (masked)', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='requests'></a>\n",
    "### 🧪 Test the API using a direct HTTP call\n",
    "Requests is an elegant and simple HTTP library for Python that will be used here to make raw API requests and inspect the responses. \n",
    "\n",
    "You will not see HTTP 429s returned as API Management's `retry` policy will select an available backend. If no backends are viable, an HTTP 503 will be returned.\n",
    "\n",
    "Tip: Use the [tracing tool](../../tools/tracing.ipynb) to track the behavior of the backend pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests, time\n",
    "\n",
    "runs = 20\n",
    "sleep_time_ms = 100\n",
    "url = f\"{apim_resource_gateway_url}/openai/deployments/{openai_deployment_name}/chat/completions?api-version={openai_api_version}\"\n",
    "messages = {\"messages\": [\n",
    "    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "]}\n",
    "api_runs = []\n",
    "\n",
    "# Initialize a session for connection pooling and set any default headers\n",
    "session = requests.Session()\n",
    "session.headers.update({'api-key': apim_subscription_key})\n",
    "\n",
    "try:\n",
    "    for i in range(runs):\n",
    "        print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "        start_time = time.time()\n",
    "        response = session.post(url, json = messages)\n",
    "        response_time = time.time() - start_time\n",
    "        print(f\"⌚ {response_time:.2f} seconds\")\n",
    "\n",
    "        utils.print_response_code(response)\n",
    "\n",
    "        if \"x-ms-region\" in response.headers:\n",
    "            print(f\"x-ms-region: \\x1b[1;32m{response.headers.get(\"x-ms-region\")}\\x1b[0m\") # this header is useful to determine the region of the backend that served the request\n",
    "            api_runs.append((response_time, response.headers.get(\"x-ms-region\")))\n",
    "\n",
    "        if (response.status_code == 200):\n",
    "            data = json.loads(response.text)\n",
    "            print(f\"Token usage: {json.dumps(dict(data.get(\"usage\")), indent = 4)}\\n\")\n",
    "            print(f\"💬 {data.get(\"choices\")[0].get(\"message\").get(\"content\")}\\n\")\n",
    "        else:\n",
    "            print(f\"{response.text}\\n\")\n",
    "\n",
    "        time.sleep(sleep_time_ms/1000)\n",
    "finally:\n",
    "    # Close the session to release the connection\n",
    "    session.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='plot'></a>\n",
    "### 🔍 Analyze Load Balancing results\n",
    "\n",
    "The priority 1 backend will be used until TPM exhaustion sets in, then distribution will occur near equally across the two priority 2 backends with 50/50 weights.  \n",
    "\n",
    "Please note that the first request of the lab can take a bit longer and should be discounted in terms of duration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle as pltRectangle\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.figsize'] = [15, 7]\n",
    "df = pd.DataFrame(api_runs, columns = ['Response Time', 'Region'])\n",
    "df['Run'] = range(1, len(df) + 1)\n",
    "\n",
    "# Define a color map for each region\n",
    "color_map = {'East US': 'lightpink', 'Sweden Central': 'lightyellow', 'West US': 'lightblue'}  # Add more regions and colors as needed\n",
    "\n",
    "# Plot the dataframe with colored bars\n",
    "ax = df.plot(kind = 'bar', x = 'Run', y = 'Response Time', color = [color_map.get(region, 'gray') for region in df['Region']], legend = False)\n",
    "\n",
    "# Add legend\n",
    "legend_labels = [pltRectangle((0, 0), 1, 1, color = color_map.get(region, 'gray')) for region in df['Region'].unique()]\n",
    "ax.legend(legend_labels, df['Region'].unique())\n",
    "\n",
    "plt.title('Load Balancing results')\n",
    "plt.xlabel('Run #')\n",
    "plt.ylabel('Response Time')\n",
    "plt.xticks(rotation = 0)\n",
    "\n",
    "average = df['Response Time'].mean()\n",
    "plt.axhline(y = average, color = 'r', linestyle = '--', label = f'Average: {average:.2f}')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='sdk'></a>\n",
    "### 🧪 Test the API using the Azure OpenAI Python SDK\n",
    "\n",
    "Repeat the same test using the Python SDK to ensure compatibility. Note that we do not know what region served the response; we only see that we obtained a response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "runs = 20\n",
    "sleep_time_ms = 100\n",
    "\n",
    "client = AzureOpenAI(\n",
    "    azure_endpoint = apim_resource_gateway_url,\n",
    "    api_key = apim_subscription_key,\n",
    "    api_version = openai_api_version\n",
    ")\n",
    "\n",
    "for i in range(runs):\n",
    "    print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "    start_time = time.time()\n",
    "    raw_response = client.chat.completions.with_raw_response.create(\n",
    "        model = openai_model_name,\n",
    "        messages = [\n",
    "            {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "        ])\n",
    "    response_time = time.time() - start_time\n",
    "\n",
    "    print(f\"⌚ {response_time:.2f} seconds\")\n",
    "    print(f\"x-ms-region: \\x1b[1;32m{raw_response.headers.get(\"x-ms-region\")}\\x1b[0m\") # this header is useful to determine the region of the backend that served the request\n",
    "\n",
    "    response = raw_response.parse()\n",
    "\n",
    "    if response.usage:\n",
    "        print(f\"Token usage:\\n   Total tokens: {response.usage.total_tokens}\\n   Prompt tokens: {response.usage.prompt_tokens}\\n   Completion tokens: {response.usage.completion_tokens}\\n\")\n",
    "\n",
    "    print(f\"💬 {response.choices[0].message.content}\\n\")\n",
    "\n",
    "    time.sleep(sleep_time_ms/1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='clean'></a>\n",
    "### 🗑️ Clean up resources\n",
    "\n",
    "When you're finished with the lab, you should remove all your deployed resources from Azure to avoid extra charges and keep your Azure subscription uncluttered.\n",
    "Use the [clean-up-resources notebook](clean-up-resources.ipynb) for that."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}