phi3/2_serving.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "796cf06d-e4d3-422e-8301-a961e9520f52",
   "metadata": {},
   "source": [
    "# Open Source LLM serving using the Azure ML Python SDK\n",
    "\n",
    "[Note] Please use `Python 3.10 - SDK v2 (azureml_py310_sdkv2)` conda environment.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538d1c11-a30d-4342-85d5-ba59e1890b76",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%store -r job_name\n",
    "try:\n",
    "    job_name\n",
    "except NameError:\n",
    "    print(\"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\")\n",
    "    print(\"[ERROR] Please run the previous notebook (model training) again.\")\n",
    "    print(\"++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb139b1d-500c-4ef2-9af3-728f2a5ea05f",
   "metadata": {},
   "source": [
    "## 1. Load config file\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7f6ccf1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os, sys\n",
    "module_path = \"../../0_lab_preparation\"\n",
    "sys.path.append(os.path.abspath(module_path))\n",
    "\n",
    "from common import check_kernel\n",
    "check_kernel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5234c47-b3e5-4218-8a98-3988c8991643",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import yaml\n",
    "from logger import logger\n",
    "from datetime import datetime\n",
    "snapshot_date = datetime.now().strftime(\"%Y-%m-%d\")\n",
    "\n",
    "with open('config.yml') as f:\n",
    "    d = yaml.load(f, Loader=yaml.FullLoader)\n",
    "    \n",
    "AZURE_SUBSCRIPTION_ID = d['config']['AZURE_SUBSCRIPTION_ID']\n",
    "AZURE_RESOURCE_GROUP = d['config']['AZURE_RESOURCE_GROUP']\n",
    "AZURE_WORKSPACE = d['config']['AZURE_WORKSPACE']\n",
    "AZURE_DATA_NAME = d['config']['AZURE_DATA_NAME']    \n",
    "DATA_DIR = d['config']['DATA_DIR']\n",
    "CLOUD_DIR = d['config']['CLOUD_DIR']\n",
    "HF_MODEL_NAME_OR_PATH = d['config']['HF_MODEL_NAME_OR_PATH']\n",
    "IS_DEBUG = d['config']['IS_DEBUG']\n",
    "\n",
    "azure_env_name = d['serve']['azure_env_name']\n",
    "azure_model_name = d['serve']['azure_model_name']\n",
    "azure_endpoint_name = d['serve']['azure_endpoint_name']\n",
    "azure_deployment_name = d['serve']['azure_deployment_name']\n",
    "azure_serving_cluster_size = d['serve']['azure_serving_cluster_size']\n",
    "\n",
    "\n",
    "logger.info(\"===== 0. Azure ML Deployment Info =====\")\n",
    "logger.info(f\"AZURE_SUBSCRIPTION_ID={AZURE_SUBSCRIPTION_ID}\")\n",
    "logger.info(f\"AZURE_RESOURCE_GROUP={AZURE_RESOURCE_GROUP}\")\n",
    "logger.info(f\"AZURE_WORKSPACE={AZURE_WORKSPACE}\")\n",
    "logger.info(f\"AZURE_DATA_NAME={AZURE_DATA_NAME}\")\n",
    "logger.info(f\"DATA_DIR={DATA_DIR}\")\n",
    "logger.info(f\"CLOUD_DIR={CLOUD_DIR}\")\n",
    "logger.info(f\"HF_MODEL_NAME_OR_PATH={HF_MODEL_NAME_OR_PATH}\")\n",
    "logger.info(f\"IS_DEBUG={IS_DEBUG}\")\n",
    "\n",
    "logger.info(f\"azure_env_name={azure_env_name}\")\n",
    "logger.info(f\"azure_model_name={azure_model_name}\")\n",
    "logger.info(f\"azure_endpoint_name={azure_endpoint_name}\")\n",
    "logger.info(f\"azure_deployment_name={azure_deployment_name}\")\n",
    "logger.info(f\"azure_serving_cluster_size={azure_serving_cluster_size}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9843e0f-3cf1-4e86-abb7-a49919fac8d4",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 2. Serving preparation\n",
    "\n",
    "---\n",
    "\n",
    "### 2.1. Configure workspace details\n",
    "\n",
    "To connect to a workspace, we need identifying parameters - a subscription, a resource group, and a workspace name. We will use these details in the MLClient from azure.ai.ml to get a handle on the Azure Machine Learning workspace we need. We will use the default Azure authentication for this hands-on.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccb4a273-ba31-4f47-a2fd-dc8cdea390f6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import required libraries\n",
    "import time\n",
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "from azure.ai.ml import MLClient, Input\n",
    "from azure.ai.ml import command\n",
    "from azure.ai.ml.entities import Model\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "from azure.core.exceptions import ResourceNotFoundError, ResourceExistsError\n",
    "\n",
    "logger.info(f\"===== 2. Serving preparation =====\")\n",
    "logger.info(f\"Calling DefaultAzureCredential.\")\n",
    "credential = DefaultAzureCredential()\n",
    "ml_client = None\n",
    "try:\n",
    "    ml_client = MLClient.from_config(credential)\n",
    "except Exception as ex:\n",
    "    print(ex)\n",
    "    ml_client = MLClient(credential, AZURE_SUBSCRIPTION_ID, AZURE_RESOURCE_GROUP, AZURE_WORKSPACE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba2a377b-d10c-413e-a67b-2c11a3cff7fd",
   "metadata": {},
   "source": [
    "### 2.2. Create model asset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73532c39-3fdd-40a7-b2be-9f5a2f22443a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_or_create_model_asset(ml_client, model_name, job_name, model_dir=\"outputs\", model_type=\"custom_model\", update=False):\n",
    "    \n",
    "    try:\n",
    "        latest_model_version = max([int(m.version) for m in ml_client.models.list(name=model_name)])\n",
    "        if update:\n",
    "            raise ResourceExistsError('Found Model asset, but will update the Model.')\n",
    "        else:\n",
    "            model_asset = ml_client.models.get(name=model_name, version=latest_model_version)\n",
    "            logger.info(f\"Found Model asset: {model_name}. Will not create again\")\n",
    "    except (ResourceNotFoundError, ResourceExistsError) as e:\n",
    "        logger.info(f\"Exception: {e}\")        \n",
    "        model_path = f\"azureml://jobs/{job_name}/outputs/artifacts/paths/{model_dir}/\"    \n",
    "        run_model = Model(\n",
    "            name=model_name,        \n",
    "            path=model_path,\n",
    "            description=\"Model created from run.\",\n",
    "            type=model_type # mlflow_model, custom_model, triton_model\n",
    "        )\n",
    "        model_asset = ml_client.models.create_or_update(run_model)\n",
    "        logger.info(f\"Created Model asset: {model_name}\")\n",
    "\n",
    "    return model_asset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aaa9671-fc98-4e5e-a70c-7771caa1c7d6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_dir = d['train']['model_dir']\n",
    "model = get_or_create_model_asset(ml_client, azure_model_name, job_name, model_dir, model_type=\"custom_model\", update=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0df561a-7846-4450-a2dd-4af6396b1719",
   "metadata": {},
   "source": [
    "### 2.3. Create AzureML environment\n",
    "\n",
    "Azure ML defines containers (called environment asset) in which your code will run. We can use the built-in environment or build a custom environment (Docker container, conda). This hands-on uses Docker container.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08ac357a-3944-4702-a338-b9f4b67dadc9",
   "metadata": {},
   "source": [
    "#### Docker environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ace0e7-e49c-4d8a-ac2f-604b6e00b145",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile {CLOUD_DIR}/serve/Dockerfile\n",
    "FROM mcr.microsoft.com/aifx/acpt/stable-ubuntu2004-cu124-py310-torch241:biweekly.202410.2\n",
    "\n",
    "# Install pip dependencies\n",
    "COPY requirements.txt .\n",
    "RUN pip install -r requirements.txt --no-cache-dir\n",
    "\n",
    "# Inference requirements\n",
    "COPY --from=mcr.microsoft.com/azureml/o16n-base/python-assets:20230419.v1 /artifacts /var/\n",
    "\n",
    "RUN /var/requirements/install_system_requirements.sh && \\\n",
    "    cp /var/configuration/rsyslog.conf /etc/rsyslog.conf && \\\n",
    "    cp /var/configuration/nginx.conf /etc/nginx/sites-available/app && \\\n",
    "    ln -sf /etc/nginx/sites-available/app /etc/nginx/sites-enabled/app && \\\n",
    "    rm -f /etc/nginx/sites-enabled/default\n",
    "ENV SVDIR=/var/runit\n",
    "ENV WORKER_TIMEOUT=400\n",
    "EXPOSE 5001 8883 8888\n",
    "\n",
    "# support Deepspeed launcher requirement of passwordless ssh login\n",
    "RUN apt-get update\n",
    "RUN apt-get install -y openssh-server openssh-client\n",
    "\n",
    "RUN MAX_JOBS=4 pip install flash-attn==2.6.3 --no-build-isolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e4b20f8-d76c-46ca-bea8-b501c4210210",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile {CLOUD_DIR}/serve/requirements.txt\n",
    "azureml-core==1.58.0\n",
    "azureml-dataset-runtime==1.58.0\n",
    "azureml-defaults==1.58.0\n",
    "azure-ml==0.0.1\n",
    "azure-ml-component==0.9.18.post2\n",
    "azureml-mlflow==1.58.0\n",
    "azureml-contrib-services==1.58.0\n",
    "azureml-contrib-services==1.58.0\n",
    "azureml-automl-common-tools==1.58.0\n",
    "torch-tb-profiler==0.4.3\n",
    "azureml-inference-server-http~=1.3\n",
    "inference-schema==1.8.0\n",
    "MarkupSafe==3.0.2\n",
    "regex\n",
    "pybind11\n",
    "bitsandbytes==0.44.1\n",
    "transformers==4.46.1\n",
    "peft==0.13.2\n",
    "accelerate==1.1.0\n",
    "datasets\n",
    "scipy\n",
    "azure-identity\n",
    "packaging==24.1\n",
    "timm==1.0.11\n",
    "einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7b19bc6-dc43-4948-8c71-97a033a5f5a3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Environment, BuildContext\n",
    "\n",
    "def get_or_create_docker_environment_asset(ml_client, env_name, docker_dir, update=False):\n",
    "    \n",
    "    try:\n",
    "        latest_env_version = max([int(e.version) for e in ml_client.environments.list(name=env_name)])\n",
    "        if update:\n",
    "            raise ResourceExistsError('Found Environment asset, but will update the Environment.')\n",
    "        else:\n",
    "            env_asset = ml_client.environments.get(name=env_name, version=latest_env_version)\n",
    "            logger.info(f\"Found Environment asset: {env_name}. Will not create again\")\n",
    "    except (ResourceNotFoundError, ResourceExistsError) as e:\n",
    "        logger.info(f\"Exception: {e}\")\n",
    "        env_docker_image = Environment(\n",
    "            build=BuildContext(path=docker_dir),\n",
    "            name=env_name,\n",
    "            description=\"Environment created from a Docker context.\",\n",
    "        )\n",
    "        env_asset = ml_client.environments.create_or_update(env_docker_image)\n",
    "        logger.info(f\"Created Environment asset: {env_name}\")\n",
    "    \n",
    "    return env_asset\n",
    "\n",
    "env = get_or_create_docker_environment_asset(ml_client, azure_env_name, f\"{CLOUD_DIR}/serve\", update=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e496e364-c8ed-43c1-a248-2530c1eb44a4",
   "metadata": {},
   "source": [
    "### 2.4. Serving script\n",
    "\n",
    "If you are not serving with MLflow but with a custom model, you are free to write your own code.The `score.py` example below shows how to write the code.\n",
    "\n",
    "-   `init()`: This function is the place to write logic for global initialization operations like loading the LLM model.\n",
    "-   `run()`: Inference logic is called for every invocation of the endpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d2a7975-8473-4f48-9d15-5180e3276b11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile src_serve/score.py\n",
    "import os\n",
    "import logging\n",
    "import json\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline\n",
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "def init():\n",
    "    \"\"\"\n",
    "    This function is called when the container is initialized/started, typically after create/update of the deployment.\n",
    "    You can write the logic here to perform init operations like caching the model in memory\n",
    "    \"\"\"\n",
    "    global model\n",
    "    global tokenizer\n",
    "    # AZUREML_MODEL_DIR is an environment variable created during deployment.\n",
    "    # It is the path to the model folder (./azureml-models/$MODEL_NAME/$VERSION)\n",
    "    # Please provide your model's folder name if there is one\n",
    "    model_path = os.path.join(\n",
    "        os.getenv(\"AZUREML_MODEL_DIR\"), \"{{score_model_dir}}\"\n",
    "    )\n",
    "    model_id = \"{{hf_model_name_or_path}}\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "    model = AutoModelForCausalLM.from_pretrained(model_id, device_map={\"\":0}, torch_dtype=\"auto\", trust_remote_code=True)\n",
    "\n",
    "    model.load_adapter(model_path)\n",
    "    logging.info(\"Loaded model.\")\n",
    "    \n",
    "def run(json_data: str):\n",
    "    logging.info(\"Request received\")\n",
    "    data = json.loads(json_data)\n",
    "    input_data= data[\"input_data\"]\n",
    "    params = data['params']\n",
    "    \n",
    "    pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
    "    output = pipe(input_data, **params)\n",
    "    generated_text = output[0]['generated_text']\n",
    "    logging.info(\"Output Response: \" + generated_text)\n",
    "    json_result = {\"result\": str(generated_text)}\n",
    "    \n",
    "    return json_result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20399372-3db8-4ca9-86ba-5d3123e108cf",
   "metadata": {},
   "source": [
    "Plug in the appropriate variables in the model inference script.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68daf5a2-4aaa-4649-bb58-eec6abc942bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import jinja2\n",
    "from pathlib import Path\n",
    "TRAINED_MLFLOW = False\n",
    "\n",
    "jinja_env = jinja2.Environment()  \n",
    "\n",
    "template = jinja_env.from_string(Path(\"src_serve/score.py\").open().read())\n",
    "score_model_dir = os.path.join(model_dir, \"peft\") if TRAINED_MLFLOW else model_dir    \n",
    "\n",
    "Path(\"src_serve/score.py\").open(\"w\").write(\n",
    "    template.render(score_model_dir=score_model_dir, hf_model_name_or_path=HF_MODEL_NAME_OR_PATH)\n",
    ")\n",
    "\n",
    "!pygmentize src_serve/score.py | cat -n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aafb187e-370f-4481-9d82-a38ae982c1e3",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 3. Serving\n",
    "\n",
    "---\n",
    "\n",
    "### 3.1. Create endpoint\n",
    "\n",
    "Create an endpoint. This process does not provision a GPU cluster yet.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c22433-4ab8-4db9-956b-7f437b86dfa6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "logger.info(f\"===== 3. Serving =====\")\n",
    "\n",
    "from azure.ai.ml.entities import (\n",
    "    ManagedOnlineEndpoint,\n",
    "    IdentityConfiguration,\n",
    "    ManagedIdentityConfiguration,\n",
    ")\n",
    "\n",
    "# Check if the endpoint already exists in the workspace\n",
    "try:\n",
    "    endpoint = ml_client.online_endpoints.get(azure_endpoint_name)\n",
    "    logger.info(\"---Endpoint already exists---\")\n",
    "except:\n",
    "    # Create an online endpoint if it doesn't exist\n",
    "\n",
    "    # Define the endpoint\n",
    "    endpoint = ManagedOnlineEndpoint(\n",
    "        name=azure_endpoint_name,\n",
    "        description=f\"Test endpoint for {model.name}\",\n",
    "        # identity=IdentityConfiguration(\n",
    "        #     type=\"user_assigned\",\n",
    "        #     user_assigned_identities=[ManagedIdentityConfiguration(resource_id=uai_id)],\n",
    "        # )\n",
    "        # if uai_id != \"\"\n",
    "        # else None,\n",
    "    )\n",
    "\n",
    "# Trigger the endpoint creation\n",
    "try:\n",
    "    ml_client.begin_create_or_update(endpoint).wait()\n",
    "    logger.info(\"\\n---Endpoint created successfully---\\n\")\n",
    "except Exception as err:\n",
    "    raise RuntimeError(\n",
    "        f\"Endpoint creation failed. Detailed Response:\\n{err}\"\n",
    "    ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f8a05af-0aca-4d36-9c0f-bfa4dcc6203b",
   "metadata": {},
   "source": [
    "### 3.2. Create Deployment\n",
    "\n",
    "Create a Deployment. This takes a lot of time as GPU clusters must be provisioned and the serving environment must be built."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3afaa8b-5af1-49d1-990f-414da0effe8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "from azure.ai.ml.entities import (    \n",
    "    OnlineRequestSettings,\n",
    "    CodeConfiguration,\n",
    "    ManagedOnlineDeployment,\n",
    "    ProbeSettings,\n",
    "    Environment\n",
    ")\n",
    "\n",
    "deployment = ManagedOnlineDeployment(\n",
    "    name=azure_deployment_name,\n",
    "    endpoint_name=azure_endpoint_name,\n",
    "    model=model,\n",
    "    instance_type=azure_serving_cluster_size,\n",
    "    instance_count=1,\n",
    "    #code_configuration=code_configuration,\n",
    "    environment=env,\n",
    "    scoring_script=\"score.py\",\n",
    "    code_path=\"./src_serve\",\n",
    "    #environment_variables=deployment_env_vars,\n",
    "    request_settings=OnlineRequestSettings(\n",
    "        max_concurrent_requests_per_instance=3,\n",
    "        request_timeout_ms=90000, \n",
    "        max_queue_wait_ms=60000\n",
    "    ),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Trigger the deployment creation\n",
    "try:\n",
    "    ml_client.begin_create_or_update(deployment).wait()\n",
    "    logger.info(\"\\n---Deployment created successfully---\\n\")\n",
    "except Exception as err:\n",
    "    raise RuntimeError(\n",
    "        f\"Deployment creation failed. Detailed Response:\\n{err}\"\n",
    "    ) from err\n",
    "    \n",
    "endpoint.traffic = {azure_deployment_name: 100}\n",
    "endpoint_poller = ml_client.online_endpoints.begin_create_or_update(endpoint)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bc7b788-01f1-47d8-8142-5fdfb0014063",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "endpoint_poller.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72f437f6-153a-42d5-ab22-0011d0fe2481",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 4. Test\n",
    "\n",
    "---\n",
    "\n",
    "### 4.1. Invocation\n",
    "\n",
    "Try calling the endpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf37c63-4d1e-44a3-8c94-268ffc716da9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "sample = {\n",
    "    \"input_data\": \n",
    "        [\n",
    "            {\"role\": \"user\", \"content\": \"Tell me Microsoft's brief history.\"},\n",
    "            {\"role\": \"assistant\", \"content\": \"Microsoft was founded by Bill Gates and Paul Allen on April 4, 1975, to develop and sell a BASIC interpreter for the Altair 8800.\"},\n",
    "            {\"role\": \"user\", \"content\": \"What about Amazon's history?\"}\n",
    "        ],\n",
    "    \"params\": {\n",
    "        \"temperature\": 0.1,\n",
    "        \"max_new_tokens\": 128,\n",
    "        \"do_sample\": True,\n",
    "        \"return_full_text\": False\n",
    "    }\n",
    "}\n",
    "\n",
    "test_src_dir = \"./phi3-inference-test\"\n",
    "os.makedirs(test_src_dir, exist_ok=True)\n",
    "logger.info(f\"Test script directory: {test_src_dir}\")\n",
    "sample_data_path = os.path.join(test_src_dir, \"sample-request.json\")\n",
    "\n",
    "with open(sample_data_path, \"w\") as f:\n",
    "    json.dump(sample, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53b3079d-a412-497d-aa37-795a1683ac78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=azure_endpoint_name,\n",
    "    deployment_name=azure_deployment_name,\n",
    "    request_file=sample_data_path,\n",
    ")\n",
    "\n",
    "result_json = json.loads(result)\n",
    "print(result_json['result'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a9d96bd-75da-4c10-923c-edad899fc4d3",
   "metadata": {},
   "source": [
    "### 4.2. LLM latency/throughput benchmarking\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20bf4fff-915b-45a8-9f46-bff0e80373df",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from time import perf_counter\n",
    "\n",
    "def benchmark_latency(endpoint_name, deployment_name, sample_data_path, num_warmups=1, num_infers=5):\n",
    "    print(f\"Measuring latency for Endpoint '{endpoint_name}' and Deployment '{deployment_name}', num_infers={num_infers}\")\n",
    "\n",
    "    latencies = []\n",
    "    # warm up\n",
    "    for _ in range(num_warmups):\n",
    "        result = ml_client.online_endpoints.invoke(\n",
    "            endpoint_name=endpoint_name,\n",
    "            deployment_name=deployment_name,\n",
    "            request_file=sample_data_path,\n",
    "        ) \n",
    "        \n",
    "    begin = time.time()        \n",
    "    # Timed run\n",
    "    for _ in range(num_infers):\n",
    "        start_time = perf_counter()\n",
    "        result = ml_client.online_endpoints.invoke(\n",
    "            endpoint_name=endpoint_name,\n",
    "            deployment_name=deployment_name,\n",
    "            request_file=sample_data_path,\n",
    "        )\n",
    "        latency = perf_counter() - start_time\n",
    "        latencies.append(latency)\n",
    "    end = time.time() \n",
    "        \n",
    "    # Compute run statistics\n",
    "    duration = end - begin    \n",
    "    time_avg_sec = np.mean(latencies)\n",
    "    time_std_sec = np.std(latencies)\n",
    "    time_p95_sec = np.percentile(latencies, 95)\n",
    "    time_p99_sec = np.percentile(latencies, 99)\n",
    "    \n",
    "    # Metrics\n",
    "    metrics = {\n",
    "        'duration': duration,\n",
    "        'avg_sec': time_avg_sec,\n",
    "        'std_sec': time_std_sec,        \n",
    "        'p95_sec': time_p95_sec,\n",
    "        'p99_sec': time_p99_sec    \n",
    "    }\n",
    "    \n",
    "    return metrics\n",
    "\n",
    "def benchmark_latency_multicore(endpoint_name, deployment_name, sample_data_path, num_warmups=1, num_infers=5, num_threads=2):\n",
    "    import time\n",
    "    import concurrent.futures\n",
    "\n",
    "    # Warmup\n",
    "    for _ in range(num_warmups):\n",
    "        result = ml_client.online_endpoints.invoke(\n",
    "            endpoint_name=endpoint_name,\n",
    "            deployment_name=deployment_name,\n",
    "            request_file=sample_data_path,\n",
    "        )        \n",
    "                \n",
    "    latencies = []\n",
    "\n",
    "    # Thread task: Each of these thread tasks executes in a serial loop for a single model.\n",
    "    #              Multiple of these threads are launched to achieve parallelism.\n",
    "    def task(model):\n",
    "        for _ in range(num_infers):\n",
    "            start = time.time()\n",
    "            result = ml_client.online_endpoints.invoke(\n",
    "                endpoint_name=endpoint_name,\n",
    "                deployment_name=deployment_name,\n",
    "                request_file=sample_data_path,\n",
    "            )   \n",
    "            finish = time.time()\n",
    "            latencies.append(finish - start)\n",
    "            \n",
    "    # Submit tasks\n",
    "    begin = time.time()\n",
    "    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as pool:\n",
    "        for i in range(num_threads):\n",
    "            pool.submit(task, model)\n",
    "    end = time.time()\n",
    "\n",
    "    # Compute metrics\n",
    "    duration = end - begin\n",
    "    inferences = len(latencies)\n",
    "    throughput = inferences / duration\n",
    "    avg_latency = sum(latencies) / len(latencies)\n",
    "    \n",
    "    # Compute run statistics\n",
    "    time_avg_sec = np.mean(latencies)\n",
    "    time_std_sec = np.std(latencies)\n",
    "    time_p95_sec = np.percentile(latencies, 95)\n",
    "    time_p99_sec = np.percentile(latencies, 99)\n",
    "    \n",
    "    time_std_sec = np.std(latencies)\n",
    "    time_p95_sec = np.percentile(latencies, 95)\n",
    "    time_p99_sec = np.percentile(latencies, 99)\n",
    "\n",
    "    # Metrics\n",
    "    metrics = {\n",
    "        'threads': num_threads,\n",
    "        'duration': duration,\n",
    "        'throughput': throughput,\n",
    "        'avg_sec': avg_latency,\n",
    "        'std_sec': time_std_sec,        \n",
    "        'p95_sec': time_p95_sec,\n",
    "        'p99_sec': time_p99_sec    \n",
    "    }\n",
    "    \n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8827fb45-c7ab-4aa8-b9f4-1c4e610d695a",
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_result = benchmark_latency(azure_endpoint_name, azure_deployment_name, sample_data_path, num_warmups=1, num_infers=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c62e9289-fe65-4cbf-abc8-f19a4f486bde",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(benchmark_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5834a237-e751-446a-ac21-7272c29b0c2c",
   "metadata": {},
   "source": [
    "## Clean up\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59508aad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf {test_src_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc72663d-d773-435b-871f-cc51b1e51763",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ml_client.online_endpoints.begin_delete(azure_endpoint_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK v2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "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.10.14"
  },
  "microsoft": {
   "ms_spell_check": {
    "ms_spell_check_language": "en"
   }
  },
  "nteract": {
   "version": "nteract-front-end@1.0.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}