src/fmbench/2_deploy_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Deploy Jumpstart and Non Jumpstart Models Asynchronously \n",
    "---------------------\n",
    "*This notebook works best with the conda_python3 kernel on a ml.t3.medium machine*.\n",
    "\n",
    "**This step of our solution design covers setting up the environment, downloading the requirements needed to run the environment, as well as deploying the model endpoints from the config.yml file asychronously.**\n",
    "\n",
    "1. Prerequisite: Navigate to the file: 0_setup.ipynb and Run the cell to import and download the requirements.txt.\n",
    "\n",
    "2. Now you can run this notebook to deploy the models asychronously in different threads. The key components of this notebook for the purposes of understanding are:\n",
    "\n",
    "- Loading the globals.py and config.yml file.\n",
    "\n",
    "- Setting a blocker function deploy_model to deploy the given model endpoint followed by:\n",
    "\n",
    "- A series of async functions to set tasks to deploy the models from the config yml file asynchronously in different threads. View the notebook from the link above.\n",
    "\n",
    "- Once the endpoints are deployed, their model configurations are stored within the endpoints.json file.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import all of the necessary libraries below to run this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if interactive mode is set to no -> pickup fmbench from Python installation path\n",
    "# if interactive mode is set to yes -> pickup fmbench from the current path (one level above this notebook)\n",
    "# if interactive mode is not defined -> pickup fmbench from the current path (one level above this notebook)\n",
    "# the premise is that if run non-interactively then it can only be run through main.py which will set interactive mode to no\n",
    "import os\n",
    "import sys\n",
    "if os.environ.get(\"INTERACTIVE_MODE_SET\", \"yes\") == \"yes\":\n",
    "    sys.path.append(os.path.dirname(os.getcwd()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "import json\n",
    "import boto3\n",
    "import asyncio\n",
    "import logging\n",
    "import importlib.util\n",
    "import fmbench.scripts\n",
    "from pathlib import Path\n",
    "from fmbench.utils import *\n",
    "from fmbench.globals import *\n",
    "from typing import Dict, List, Optional\n",
    "from sagemaker import get_execution_role\n",
    "import importlib.resources as pkg_resources\n",
    "from botocore.exceptions import ClientError\n",
    "from botocore.exceptions import NoCredentialsError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pygmentize globals.py to view and use any of the globally initialized variables "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set up a logger to log all messages while the code runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.basicConfig(format='[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s', level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the config.yml file\n",
    "------\n",
    "\n",
    "The config.yml file contains information that is used across this benchmarking environment, such as information about the aws account, prompts, payloads to be used for invocations, and model configurations like the version of the model, the endpoint name, model_id that needs to be deployed. Configurations also support the gives instance type to be used, for example: \"ml.g5.24xlarge\", the image uri, whether or not to deploy this given model, followed by an inference script \"jumpstart.py\" which supports the inference script for jumpstart models to deploy the model in this deploy notebook. \n",
    "\n",
    "View the contents of the config yml file below and how it is loaded and used throughout this notebook with deploying the model endpoints asynchronously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load the config.yml file referring to the globals.py file\n",
    "config = load_main_config(CONFIG_FILE)\n",
    "\n",
    "## configure the aws region and execution role\n",
    "aws_region = config['aws']['region']\n",
    "\n",
    "\n",
    "# try:\n",
    "#     sagemaker_execution_role = get_execution_role()\n",
    "#     config['aws']['sagemaker_execution_role'] = sagemaker_execution_role\n",
    "#     logger.info(f\"determined SageMaker exeuction role from get_execution_role\")\n",
    "# except Exception as e:\n",
    "#     logger.error(f\"could not determine SageMaker execution role, error={e}\")\n",
    "#     logger.info(f\"going to look for execution role in config file..\")\n",
    "#     sagemaker_execution_role = config['aws'].get('sagemaker_execution_role')\n",
    "#     if sagemaker_execution_role is not None:\n",
    "#         logger.info(f\"found SageMaker execution role in config file..\")\n",
    "\n",
    "logger.info(f\"aws_region={aws_region}, execution_role={config['aws']['sagemaker_execution_role']}\")\n",
    "logger.info(f\"config={json.dumps(config, indent=2)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Deploy a single model: blocking function used for asynchronous deployment\n",
    "\n",
    "This function is designed to deploy a single large language model endpoint. It takes three parameters: experiment_config (a dictionary containing configuration details for the model deployment from the config.yml file), aws_region (the AWS region where the model will be deployed), and role_arn (the AWS role's Amazon Resource Name used for the deployment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to deploy a model\n",
    "def deploy_model(experiment_config: Dict, aws_region: str, role_arn: str) -> Optional[Dict]:\n",
    "\n",
    "    # Log the deployment details\n",
    "    logger.info(f\"going to deploy {experiment_config}, in {aws_region} with {role_arn}\")\n",
    "    model_deployment_result: Optional[Dict] = None\n",
    "\n",
    "    # Check if deployment is enabled in the config; skip if not\n",
    "    deploy = experiment_config.get('deploy', False)\n",
    "    if deploy is False:\n",
    "        logger.error(f\"skipping deployment of {experiment_config['model_id']} because deploy={deploy}\")\n",
    "        model_deployment_result = dict(endpoint_name=experiment_config['ep_name'], experiment_name=experiment_config['name'])\n",
    "        return model_deployment_result\n",
    "\n",
    "    # Initialize the S3 client\n",
    "    s3_client = boto3.client('s3', region_name=aws_region)\n",
    "\n",
    "    # Assuming fmbench is a valid Python package and scripts is a subdirectory within it\n",
    "    scripts_dir = Path(pkg_resources.files('fmbench'), 'scripts')\n",
    "    logger.info(f\"Using fmbench.scripts directory: {scripts_dir}\")\n",
    "\n",
    "    # Proceed with deployment as before\n",
    "    try:\n",
    "        module_name = Path(experiment_config['deployment_script']).stem\n",
    "        logger.info(f\"script provided for deploying this model is --> {module_name}\")\n",
    "        deployment_script_path = scripts_dir / f\"{module_name}.py\"\n",
    "        logger.info(f\"script path is --> {deployment_script_path}\")\n",
    "\n",
    "        # Check and proceed with local script\n",
    "        if not deployment_script_path.exists():\n",
    "            logger.error(f\"Deployment script {deployment_script_path} not found.\")\n",
    "            return None\n",
    "\n",
    "        logger.info(f\"Deploying using local code: {deployment_script_path}\")\n",
    "\n",
    "        spec = importlib.util.spec_from_file_location(module_name, str(deployment_script_path))\n",
    "        module = importlib.util.module_from_spec(spec)\n",
    "        sys.modules[module_name] = module\n",
    "        spec.loader.exec_module(module)\n",
    "        st = time.perf_counter()\n",
    "        model_deployment_result = module.deploy(experiment_config, role_arn)\n",
    "        elapsed_time = time.perf_counter() - st\n",
    "        logger.info(f\"time taken to deploy model_id={experiment_config['model_id']} via \"\n",
    "                    f\"{deployment_script_path} is {elapsed_time:0.2f}\")\n",
    "        return model_deployment_result\n",
    "    except Exception as error:  # Broader exception handling for non-ClientError issues\n",
    "        logger.error(f\"An error occurred during deployment: {error}\")\n",
    "        return model_deployment_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Asynchronous Model Deployment\n",
    "----\n",
    "\n",
    "#### async_deploy_model: \n",
    "\n",
    "- This is an asynchronous wrapper around the deploy_model function. It uses asyncio.to_thread to run the synchronous deploy_model function in a separate thread. This allows the function to be awaited in an asynchronous context, enabling concurrent model deployments without any blocking from the main thread\n",
    "\n",
    "#### async_deploy_all_models Function: \n",
    "\n",
    "- This 'async_deploy_all_models' function is designed to deploy multiple models concurrently. It splits the models into batches and deploys each batch concurrently using asyncio.gather."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Asynchronous wrapper function to allow our deploy_model function to allow concurrent requests for deployment\n",
    "async def async_deploy_model(experiment_config: Dict, role_arn: str, aws_region: str) -> str:\n",
    "    # Run the deploy_model function in a separate thread to deploy the models asychronously\n",
    "    return await asyncio.to_thread(deploy_model, experiment_config, role_arn, aws_region)\n",
    "\n",
    "## Final asychronous function to deploy all of the models concurrently\n",
    "async def async_deploy_all_models(config: Dict) -> List[Dict]:\n",
    "    \n",
    "    ## Extract experiments from the config.yml file (contains information on model configurations)\n",
    "    experiments: List[Dict] = config['experiments']\n",
    "    n: int = 4 # max concurrency so as to not get a throttling exception\n",
    "    \n",
    "    ## Split experiments into smaller batches for concurrent deployment\n",
    "    experiments_splitted = [experiments[i * n:(i + 1) * n] for i in range((len(experiments) + n - 1) // n )]\n",
    "    results = []\n",
    "    for exp_list in experiments_splitted:\n",
    "        \n",
    "        ## send the deployment in batches\n",
    "        result = await asyncio.gather(*[async_deploy_model(m,\n",
    "                                                           config['aws']['region'],\n",
    "                                                           config['aws']['sagemaker_execution_role']) for m in exp_list])\n",
    "        ## Collect and furthermore extend the results from each batch\n",
    "        results.extend(result)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# async version\n",
    "s = time.perf_counter()\n",
    "\n",
    "## Call all of the models for deployment using the config.yml file model configurations\n",
    "endpoint_names = await async_deploy_all_models(config)\n",
    "\n",
    "## Set a timer for model deployment counter\n",
    "elapsed_async = time.perf_counter() - s\n",
    "print(f\"endpoint_names -> {endpoint_names}, deployed in {elapsed_async:0.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Function to get all of the information on the deployed endpoints and store it in a json\n",
    "def get_all_info_for_endpoint(ep: Dict) -> Dict:\n",
    "    \n",
    "    ## extract the endpoint name\n",
    "    ep_name = ep['endpoint_name']\n",
    "    \n",
    "    ## extract the experiment name from the config.yml file\n",
    "    experiment_name = ep['experiment_name']\n",
    "    if ep_name is None:\n",
    "        return None\n",
    "    sm_client = boto3.client('sagemaker')\n",
    "    \n",
    "    # check if the given endpoint is a sagemaker endpoint\n",
    "    next_token = None\n",
    "    ep_names = []\n",
    "    while True:\n",
    "        if next_token:\n",
    "            resp = sm_client.list_endpoints(MaxResults=100, StatusEquals='InService', NextToken=next_token)\n",
    "        else:\n",
    "            resp = sm_client.list_endpoints(MaxResults=100, StatusEquals='InService')\n",
    "        ep_names.extend([e['EndpointName'] for e in resp['Endpoints']])\n",
    "    \n",
    "        next_token = resp.get('NextToken')\n",
    "        if next_token is None:\n",
    "            break\n",
    "    if ep_name in ep_names:\n",
    "        logger.info(f\"ep_name={ep_name} is an InService SageMaker endpoint\")\n",
    "        ## get the description on the configuration of the deployed model\n",
    "        endpoint = sm_client.describe_endpoint(EndpointName=ep_name)\n",
    "        endpoint_config = sm_client.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])\n",
    "        model_config = sm_client.describe_model(ModelName=endpoint_config['ProductionVariants'][0]['ModelName'])\n",
    "        ## Store the experiment name and all of the other model configuration information in the 'info' dict\n",
    "        info = dict(experiment_name=experiment_name,\n",
    "                    endpoint=endpoint,\n",
    "                    endpoint_config=endpoint_config,\n",
    "                    model_config=model_config)\n",
    "    else:\n",
    "        logger.info(f\"ep_name={ep_name} is not an \\\"InService\\\" SageMaker endpoint, \"\n",
    "                    f\"setting all info about it to None\")\n",
    "        info = None\n",
    "    return info\n",
    "\n",
    "all_info = list(filter(None,\n",
    "                  list(map(get_all_info_for_endpoint,\n",
    "                             list(filter(None,\n",
    "                                          endpoint_names))))))\n",
    "\n",
    "## stores information in a dictionary for collectively all of the deployed model endpoints\n",
    "all_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert data to JSON\n",
    "json_data = json.dumps(all_info, indent=2, default=str)\n",
    "\n",
    "# Specify the file name\n",
    "file_name = \"endpoints.json\"\n",
    "\n",
    "# Write to S3\n",
    "endpoint_s3_path = write_to_s3(json_data, config['aws']['bucket'], MODELS_DIR, \"\", file_name)\n",
    "\n",
    "logger.info(f\"The s3 endpoints that are deployed are sent to this file --> {endpoint_s3_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fmbench_python311",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}