Sample.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure OpenAI Service Load Balancing with Azure API Management\n",
    "\n",
    "This notebook demonstrates how to use Azure API Management to load balance requests to multiple deployed Azure OpenAI services.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "The notebook uses [PowerShell](https://learn.microsoft.com/powershell/scripting/install/installing-powershell) and [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) to deploy all necessary Azure resources. Both tools are available on Windows, macOS and Linux environments.\n",
    "\n",
    "Running this notebook will deploy the following resources in your Azure subscription:\n",
    "- Azure Resource Group\n",
    "- Azure Managed Identity\n",
    "- Azure OpenAI Service (UK South + West US)\n",
    "- Azure API Management"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Login to Azure & set subscription\n",
    "\n",
    "The following will prompt you to login to Azure. Once logged in, the current default subscription in your available subscriptions will be set for deployment.\n",
    "\n",
    "> **Note:** If you have multiple subscriptions, you can change the default subscription by running `az account set --subscription <subscription_id>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "pwsh"
    },
    "polyglot_notebook": {
     "kernelName": "pwsh"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "# Check if you are already logged in\n",
    "$loggedIn = az account show --query \"name\" -o tsv\n",
    "\n",
    "if ($loggedIn -ne $null) {\n",
    "    Write-Host \"Already logged in as $loggedIn\"\n",
    "} else {\n",
    "    Write-Host \"Logging in...\"\n",
    "    az login\n",
    "}\n",
    "\n",
    "# Retrieve ID for current subscription\n",
    "$subscriptionId = (\n",
    "    (\n",
    "        az account list -o json `\n",
    "            --query \"[?isDefault]\"\n",
    "    ) | ConvertFrom-Json\n",
    ").id\n",
    "\n",
    "# Set subscription ID\n",
    "az account set --subscription $subscriptionId\n",
    "Write-Host \"Subscription set to $subscriptionId\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Deploy Azure resources with Bicep\n",
    "\n",
    "The following will deploy all the necessary Azure resources, previously listed, using [Azure Bicep](https://learn.microsoft.com/en-us/azure/azure-resource-manager/bicep/).\n",
    "\n",
    "The deployment occurs at the subscription level, creating a new resource group. The location of the deployment is set to **UK South** and this can be changed, as well as other parameters, in the next code cell.\n",
    "\n",
    "This may take up to 10 minutes due to the Azure API Management resource needing to be activated successfully.\n",
    "\n",
    "### Understanding the deployment\n",
    "\n",
    "#### Managed Identity\n",
    "\n",
    "A [user-assigned Managed Identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/overview) is created for the Azure API Management instance. This is used to authenticate with the Azure OpenAI instances instead of using API keys by using role-based access control (RBAC) permissions.\n",
    "\n",
    "Read more about [how to configure Azure OpenAI Service with managed identities](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity).\n",
    "\n",
    "#### OpenAI Services\n",
    "\n",
    "Two [Azure OpenAI Service](https://learn.microsoft.com/en-us/azure/cognitive-services/openai-service/overview) instances are deployed, one in the UK South region and one in the West US region. These are deployed with the `gpt-35-turbo` models, version 1106 to be used for inference. You can change the the configuration of the instances to deploy by updating the `openAIInstances` parameter in the [`main.bicep`](./infra/main.bicep) file.\n",
    "\n",
    "#### API Management - Backends\n",
    "\n",
    "[API Management backends](https://learn.microsoft.com/en-us/azure/api-management/backends?tabs=bicep) are created for each deployed Azure OpenAI Service instance. Each backend points to the endpoint of the deployed Azure OpenAI Service instances so that we can use them in conjunction with the deployed API Management API.\n",
    "\n",
    "#### API Management - API\n",
    "\n",
    "Azure OpenAI Service has a standard REST API specification that we can use to import the API into API Management. You can [find the latest OpenAPI specifications in the Azure REST API specifications GitHub repository](https://github.com/Azure/azure-rest-api-specs/tree/main/specification/cognitiveservices/data-plane/AzureOpenAI/inference/preview).\n",
    "\n",
    "The API is configured on the `/openai` path for the deployed Azure API Management instance.\n",
    "\n",
    "#### API Management - Subscription\n",
    "\n",
    "In order to access the API, we need to [create a subscription for it](https://learn.microsoft.com/en-us/azure/api-management/api-management-subscriptions). This will generate a subscription key that we can use to make requests to the API.\n",
    "\n",
    "#### API Management - Named Values\n",
    "\n",
    "Based on the deployed Managed Identity, a [API Management Named Value](https://learn.microsoft.com/en-us/azure/api-management/api-management-howto-properties?tabs=azure-portal) is created for the client ID so that it can be used in the API Management policy for retrieving a token to authenticate with the Azure OpenAI Service instances.\n",
    "\n",
    "#### API Management - Policy\n",
    "\n",
    "With everything configured in Azure API Management to take advantage of the multiple Azure OpenAI Service instances, we [deploy a policy to the API Management API](https://learn.microsoft.com/en-us/azure/api-management/api-management-howto-policies) to load balance requests across the backends.\n",
    "\n",
    "The [API Management OpenAI policy](./infra/policies/round-robin-policy.xml) defines a round-robin load balancing strategy based on the following flow:\n",
    "\n",
    "- On an inbound request, we check a cached `backend-counter` variable to determine which backend to use. If this doesn't exist, we create it and immediately set it to `0`.\n",
    "- We then use the `backend-counter` variable to set the backend service to use for the request.\n",
    "- We then retrieve an access token for the Azure OpenAI Service using the Managed Identity associated with the API Management service.\n",
    "- We then change the `backend-counter` to `1` and set it back to the cache.\n",
    "\n",
    "- When processing the request at the backend, we provide a retry policy to handle any transient errors that may occur. This is to ensure that we don't lose any requests to the Azure OpenAI Service, and can swap to another backend if the current one is unavailable.\n",
    "- The process follows the same logic as the inbound request, whereby we check the cached `backend-counter` variable to determine which backend to use and set the Authorization header accordingly.\n",
    "\n",
    "The [retry policy](https://learn.microsoft.com/en-us/azure/api-management/retry-policy) for the backend will only trigger if the status code of the response is `400` or greater, and will retry up to 3 times with a 5 second delay between each retry. If the first request immediately fails, the retry policy will ensure that the request is immediately refired to another backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "pwsh"
    },
    "polyglot_notebook": {
     "kernelName": "pwsh"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "$deploymentName = 'aoai-apim-loadbalancing'\n",
    "$location = 'uksouth'\n",
    "$apiManagementPublisherEmail = 'test@test.com' # Update this to your email address\n",
    "$apiManagementPublisherName = 'Test User' # Update this to your name\n",
    "\n",
    "$deploymentOutputs = (.\\Deploy-Infrastructure.ps1 -DeploymentName $deploymentName -Location $location -PublisherEmail $apiManagementPublisherEmail -PublisherName $apiManagementPublisherName)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get common outputs from Bicep deployment\n",
    "\n",
    "The following will get the outputs from the Bicep deployment and set them as PowerShell variables for use in the Azure CLI commands.\n",
    "\n",
    "> **Note:** All the outputs can be found at the bottom of the [main.bicep](./infra/main.bicep) file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "pwsh"
    },
    "polyglot_notebook": {
     "kernelName": "pwsh"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "$resourceGroup = $deploymentOutputs.resourceGroupInfo.value.name\n",
    "$apiManagement = $deploymentOutputs.apiManagementInfo.value.name\n",
    "$apiManagementGatewayUrl = $deploymentOutputs.apiManagementInfo.value.gatewayUrl\n",
    "$apiManagementSubscriptionName = $deploymentOutputs.apiManagementInfo.value.subscriptionName\n",
    "\n",
    "Write-Host \"Resource group: $resourceGroup\"\n",
    "Write-Host \"API Management: $apiManagement\"\n",
    "Write-Host \"API Management Gateway URL: $apiManagementGatewayUrl\"\n",
    "Write-Host \"API Management Subscription Name: $apiManagementSubscriptionName\"\n",
    "\n",
    "# Retrieve API Management subscription primary key\n",
    "$apimSubscriptionEndpoint = \"https://management.azure.com/subscriptions/$($subscriptionId)/resourceGroups/$($resourceGroup)/providers/Microsoft.ApiManagement/service/$($apiManagement)/subscriptions/$($apiManagementSubscriptionName)/listSecrets?api-version=2023-05-01-preview\"\n",
    "$apiManagementSubscription = (az rest --uri $apimSubscriptionEndpoint --method POST) | ConvertFrom-Json\n",
    "$apiManagementSubscriptionKey = $apiManagementSubscription.primaryKey"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Test API Management\n",
    "\n",
    "Now that everything is configured, you can test the API Management API by making a request to the expected `/openai` endpoints.\n",
    "\n",
    "This mechanism will work for directly communicating with Azure OpenAI Services using standard HTTP requests, as well as using libraries such as [Microsoft Semantic Kernel](https://learn.microsoft.com/en-us/semantic-kernel/overview/) and [LangChain](https://www.langchain.com).\n",
    "\n",
    "The following examples showcase using standard HTTP requests to the API Management API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creates a completion for the chat message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "pwsh"
    },
    "polyglot_notebook": {
     "kernelName": "pwsh"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "$model = \"gpt-35-turbo\"\n",
    "$apiVersion = \"2024-03-01-preview\"\n",
    "$completionsEndpoint = $apiManagementGatewayUrl + \"/openai/deployments/$($model)/chat/completions?api-version=$($apiVersion)\"\n",
    "\n",
    "# Define the request headers\n",
    "$requestHeaders = @{\n",
    "    \"Ocp-Apim-Subscription-Key\" = $apiManagementSubscriptionKey\n",
    "}\n",
    "\n",
    "# Define the request body\n",
    "$requestBody = @{\n",
    "    messages = @(\n",
    "        @{\n",
    "            role = \"system\"\n",
    "            content = \"You are a helpful AI assistant. You always try to provide accurate answers or follow up with another question if not.\"\n",
    "        },\n",
    "        @{\n",
    "            role = \"user\"\n",
    "            content = \"What is the best way to get to London from Berlin?\"\n",
    "        }\n",
    "    )\n",
    "    max_tokens = 200\n",
    "    temperature = 0.7\n",
    "    top_p = 0.95\n",
    "    frequency_penalty = 0\n",
    "    presence_penalty = 0\n",
    "}\n",
    "\n",
    "$requestBodyString = ($requestBody | ConvertTo-Json -Depth 10 -Compress)\n",
    "\n",
    "Write-Host \"Posting request with URI $completionsEndpoint and body $requestBodyString\"\n",
    "\n",
    "$apiManagementResponse = Invoke-WebRequest -Uri $completionsEndpoint -Headers $requestHeaders -Method POST -Body $requestBodyString -ContentType \"application/json\"\n",
    "\n",
    "Write-Host \"Response: $apiManagementResponse\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Cleanup\n",
    "\n",
    "The following will delete all the resources that were deployed as part of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "pwsh"
    },
    "polyglot_notebook": {
     "kernelName": "pwsh"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "az group delete --name $resourceGroup --yes --no-wait"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}