notebook_examples/aqua-batch-inferencing.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "fac705fa",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "@notebook{aqua-batch-inferencing.ipynb,\n",
    "    title: AI Quick Action - Batch inferencing,\n",
    "    summary: Perform batch inferencing on LLMs using AI Quick Actions.,\n",
    "    developed_on: generalml_p311_cpu_x86_64_v1,\n",
    "    keywords: quick action, batch, inferencing, llm\n",
    "    license: Universal Permissive License v 1.0\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9990553-e6f8-4293-b89e-946f3e4c17e3",
   "metadata": {},
   "source": [
    "# AI Quick Action - Batch Inferencing "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46694130-657e-46ec-ba4c-dc75353b7cd6",
   "metadata": {},
   "source": [
    "#### This notebook offers a detailed, step-by-step guide for performing batch inference on LLMs using AI Quick Action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb3da79-4abf-4fde-9465-c0f85efaa978",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from ads.jobs import Job, DataScienceJob, ContainerRuntime\n",
    "import ads\n",
    "import json\n",
    "\n",
    "ads.set_auth(\"resource_principal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b09b398b-fbf8-44ef-a3c5-6727125bc931",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "##  Update following variables as required\n",
    "\n",
    "compartment_id = os.environ[\"PROJECT_COMPARTMENT_OCID\"]\n",
    "project_id = os.environ[\"PROJECT_OCID\"]\n",
    "\n",
    "log_group_id = \"ocid1.loggroup.oc1.xxx.xxxxx\"\n",
    "log_id = \"cid1.log.oc1.xxx.xxxxx\"\n",
    "\n",
    "instance_shape = \"VM.GPU.A10.2\"\n",
    "region = \"us-ashburn-1\"\n",
    "\n",
    "container_image = \"dsmc://odsc-vllm-serving:0.6.0\"\n",
    "\n",
    "bucket = \"<bucket_name>\"  # this should be a versioned bucket\n",
    "namespace = \"<bucket_namespace>\"\n",
    "model_name = \"meta-llama/Meta-Llama-3.1-8B-Instruct\"\n",
    "hf_token = \"<your-huggingface-token>\"\n",
    "oci_iam_type = \"resource_principal\"\n",
    "prefix = \"batch-inference\"\n",
    "\n",
    "job_artifacts = \"job-artifacts\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cda36d73-4a00-4f2e-ab64-b779d8da27ce",
   "metadata": {},
   "source": [
    "## Prepare The JOB Artifacts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b700f73d-c52e-44ff-8151-577de0f62132",
   "metadata": {},
   "source": [
    "### 1. Create a **job_artifacts** folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afaca820-c3ae-4ce0-8743-2384bdf4332f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.makedirs(job_artifacts, exist_ok=True)\n",
    "artifacts_path = os.path.expanduser(job_artifacts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "227e0cb0-cd12-4c78-ba94-04fce49c096b",
   "metadata": {},
   "source": [
    "### 2. Add [**input.json**](batch-inferencing-data/input.json) to **job_artifacts** folder "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf618122",
   "metadata": {},
   "source": [
    "**Sample input.json**\n",
    "```json\n",
    "\n",
    "    { \n",
    "      \"vllm_engine_config\": {\n",
    "        \"tensor_parallel_size\": 2,\n",
    "        \"disable_custom_all_reduce\": true\n",
    "      },\n",
    "      \"sampling_config\": {\n",
    "        \"max_tokens\": 250,\n",
    "        \"temperature\": 0.7,\n",
    "        \"top_p\": 0.85\n",
    "      },\n",
    "      \"data\": [\n",
    "        [\n",
    "          {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are a friendly chatbot who is a great story teller.\"\n",
    "          },\n",
    "          {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"Tell me a 1000 words story\"\n",
    "          }\n",
    "        ]\n",
    "      ]\n",
    "    }\n",
    "\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adeade9c-475e-4204-9985-61afad889617",
   "metadata": {},
   "source": [
    "This file defines the configuration and data for batch inferenceing. It contains the following key sections:\n",
    "\n",
    "#### 1.1. vLLM Engine Configuration (`vllm_engine_config`)\n",
    "The vLLM engine config from official [vllm documentation](https://docs.vllm.ai/en/latest/dev/offline_inference/llm.html) can be added here.\n",
    "\n",
    "#### 1.2. Sampling Configuration (`sampling_config`)\n",
    "- **`max_tokens`**: Specifies the maximum number of tokens the model can generate for each prompt, set to `250`.\n",
    "- **`temperature`**: Controls the randomness of the sampling process, with a value of `0.7`, balancing between randomness and determinism.\n",
    "- **`top_p`**: Controls the diversity of the generated text using nucleus sampling, set to `0.85`, meaning only the top 85% probable next tokens will be considered.\n",
    "\n",
    "More details can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html#vllm.SamplingParams).\n",
    "\n",
    "#### 1.3. Data Section (`data`)\n",
    "The `data` section contains a list of messages. Each message object has two required fields:\n",
    "- **`role`** : the role of the messenger (either system, user, assistant or tool)\n",
    "- **`content`** : the content of the message (e.g., Write me a beautiful poem)\n",
    "\n",
    "For more information on prompt engineering, please refer to the OpenAI prompt engineering [documentation](https://platform.openai.com/docs/guides/prompt-engineering/strategy-write-clear-instructions). \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30fc8340-4e20-40d4-a033-3f0c520328dd",
   "metadata": {},
   "source": [
    "### 2. Creating the inferencing script within **job_artifacts** folder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a11327a6-d85a-45f1-a2e3-ba5014f55eb6",
   "metadata": {},
   "source": [
    "This script is designed for batch inferencing using a VLLM engine, performing the following key tasks:\n",
    "\n",
    "- Loads model configurations and Hugging Face credentials from environment variables.\n",
    "- Initializes the model and tokenizer for inference.\n",
    "- Processes input prompts in batches, generates responses, and uploads the generated results to an Object Storage bucket.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1236d3a4-b408-403f-be02-73a5c8be69ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile job-artifacts/vllm_batch_inferencing.py\n",
    "\n",
    "from typing import Any, Dict, List, Union\n",
    "import os\n",
    "import json\n",
    "import logging\n",
    "import time\n",
    "from ads.model.datascience_model import DataScienceModel\n",
    "from vllm import LLM, SamplingParams\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from huggingface_hub import login\n",
    "from ads.common.utils import ObjectStorageDetails\n",
    "from ads.aqua.common.utils import upload_folder\n",
    "\n",
    "\n",
    "## Constants class for centralized variable management\n",
    "class Constants:\n",
    "    import os\n",
    "\n",
    "    # Environment variables used in the script\n",
    "    OCI_IAM_TYPE = (\n",
    "        \"OCI_IAM_TYPE\"  # Authentication type for OCI (Oracle Cloud Infrastructure)\n",
    "    )\n",
    "    INPUT_PROMPT = \"input_prompt\"  # Key for input prompt in data structure\n",
    "    OUTPUT = \"generated_output\"  # Key for generated output in data structure\n",
    "    VLLM_ENGINE_CONFIG = (\n",
    "        \"vllm_engine_config\"  # Key for VLLM engine configuration in data\n",
    "    )\n",
    "    SAMPLING_CONFIG = \"sampling_config\"  # Key for sampling parameters in data\n",
    "    MODEL = \"MODEL\"  # Environment variable for model selection\n",
    "    HF_TOKEN = \"HF_TOKEN\"  # Environment variable for Hugging Face token\n",
    "    OUTPUT_FOLDER_PATH = os.path.expanduser(\n",
    "        \"~/outputs\"\n",
    "    )  # Output folder path for generated data\n",
    "    OUTPUT_FILE_PATH = os.path.join(\n",
    "        OUTPUT_FOLDER_PATH, \"output_prompts.json\"\n",
    "    )  # Path to save output prompts\n",
    "    OS_OBJECT = \"batch-inference\"  # Object name used for storing in OCI\n",
    "\n",
    "\n",
    "class Deployment:\n",
    "    def __init__(self, vllm_config: List[Dict[str, Any]], **kwargs: Any) -> None:\n",
    "        \"\"\"\n",
    "        Initializes the Deployment object.\n",
    "        Loads the model and tokenizer using the provided VLLM configuration.\n",
    "        Also handles authentication for Hugging Face using the token from environment variables.\n",
    "        \"\"\"\n",
    "        print(\"Initializing VLLM with the provided model\")\n",
    "\n",
    "        # Login to Hugging Face if token is provided\n",
    "        hf_token = os.environ.get(Constants.HF_TOKEN, \"\")\n",
    "        if hf_token:\n",
    "            login(token=hf_token)\n",
    "\n",
    "        # Get model name from environment variables,\n",
    "        try:\n",
    "            model_name = os.environ.get(\n",
    "                Constants.MODEL, \"meta-llama/Meta-Llama-3.1-8B-Instruct\"\n",
    "            )\n",
    "            if not model_name:\n",
    "                raise ValueError\n",
    "        except ValueError:\n",
    "            print(f\"Invalid model name in {Constants.MODEL}\")\n",
    "\n",
    "        self.model = os.environ[Constants.MODEL]\n",
    "\n",
    "        # Store the VLLM configuration\n",
    "        self.vllm_config = vllm_config\n",
    "\n",
    "        # Initialize the tokenizer and model using the Hugging Face `transformers` library\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(self.model)\n",
    "        self.llm = LLM(model=self.model, **self.vllm_config)\n",
    "\n",
    "        print(\"Initialization complete.\")\n",
    "\n",
    "    def requests(\n",
    "        self, data: List[Dict[str, Any]], sampling_config: Dict[str, Any] = None\n",
    "    ) -> List[Dict[str, Any]]:\n",
    "        \"\"\"\n",
    "        Processes a batch of requests (prompts) and returns the generated responses.\n",
    "\n",
    "        Args:\n",
    "            data: A list of dictionaries, each containing input prompts.\n",
    "            sampling_config: Optional configuration for controlling sampling behavior (e.g., max tokens, temperature).\n",
    "\n",
    "        Returns:\n",
    "            A list of dictionaries containing the input prompts and their corresponding generated outputs.\n",
    "        \"\"\"\n",
    "        # Convert input messages to token IDs for the model using a chat template\n",
    "        prompt_token_ids = [\n",
    "            self.tokenizer.apply_chat_template(messages, add_generation_prompt=True)\n",
    "            for messages in data\n",
    "        ]\n",
    "\n",
    "        # Set sampling parameters (e.g., max tokens, temperature, top-p)\n",
    "        sampling_params = (\n",
    "            SamplingParams(**sampling_config)\n",
    "            if sampling_config\n",
    "            else SamplingParams(max_tokens=250, temperature=0.6, top_p=0.9)\n",
    "        )\n",
    "\n",
    "        # Generate outputs using the VLLM model\n",
    "        outputs = self.llm.generate(\n",
    "            prompt_token_ids=prompt_token_ids, sampling_params=sampling_params\n",
    "        )\n",
    "\n",
    "        processed_outputs = []\n",
    "        # Loop through each output and decode the token IDs into human-readable text\n",
    "        for output in outputs:\n",
    "            input_prompt = self.tokenizer.decode(\n",
    "                output.prompt_token_ids\n",
    "            )  # Decode the input prompt\n",
    "            generated_text = output.outputs[\n",
    "                0\n",
    "            ].text  # Extract the first generated output\n",
    "            # Store the input prompt, output, and configurations in the result\n",
    "            processed_outputs.append(\n",
    "                {\n",
    "                    Constants.INPUT_PROMPT: input_prompt,\n",
    "                    Constants.OUTPUT: generated_text,\n",
    "                    Constants.VLLM_ENGINE_CONFIG: self.vllm_config,\n",
    "                    Constants.SAMPLING_CONFIG: sampling_config,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        # Save the generated output to a JSON file\n",
    "        os.makedirs(\n",
    "            Constants.OUTPUT_FOLDER_PATH, exist_ok=True\n",
    "        )  # Ensure the output folder exists\n",
    "        with open(Constants.OUTPUT_FILE_PATH, \"w\") as f:\n",
    "            json.dump(processed_outputs, f)  # Write the output data to a JSON file\n",
    "\n",
    "        return processed_outputs  # Return the list of processed outputs\n",
    "\n",
    "\n",
    "def main():\n",
    "    \"\"\"\n",
    "    Main function to run the batch inference pipeline.\n",
    "    Reads input data from a file, initializes the deployment, runs the inference, and uploads results.\n",
    "    \"\"\"\n",
    "    # Load input data from the specified JSON file\n",
    "    with open(\"job-artifacts/input.json\") as f:\n",
    "        input_data = json.load(f)\n",
    "\n",
    "    # Initialize the deployment with the VLLM engine configuration from input data\n",
    "    deployment = Deployment(input_data[Constants.VLLM_ENGINE_CONFIG])\n",
    "\n",
    "    # Run the batch inference and generate responses based on the input prompts\n",
    "    response = deployment.requests(\n",
    "        input_data[\"data\"], input_data.get(Constants.SAMPLING_CONFIG)\n",
    "    )\n",
    "\n",
    "    # Upload the result to Object Storage in Oracle Cloud\n",
    "    prefix = os.environ.get(\"PREFIX\", Constants.OS_OBJECT)\n",
    "    os_path = ObjectStorageDetails(\n",
    "        os.environ[\"BUCKET\"], os.environ[\"NAMESPACE\"], prefix\n",
    "    ).path\n",
    "\n",
    "    # Upload the folder containing the output data to Object Storage\n",
    "    model_artifact_path = upload_folder(\n",
    "        os_path=os_path, local_dir=Constants.OUTPUT_FOLDER_PATH, model_name=\"outputs\"\n",
    "    )\n",
    "\n",
    "    # Print the path to the uploaded model artifact\n",
    "    print(model_artifact_path)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940cc7dd-7cea-4e5a-a220-12fd14ae135a",
   "metadata": {},
   "source": [
    "## Define and Run the inferencing Job"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56a7c157-5411-4c26-b4e2-61af18bc89e4",
   "metadata": {},
   "source": [
    "### Setup Job Infrastructure\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6782a837-4793-4597-98c9-ee124d572a72",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "infrastructure = (\n",
    "    DataScienceJob()\n",
    "    # Configure logging for getting the job run outputs.\n",
    "    .with_log_group_id(log_group_id)\n",
    "    # Log resource will be auto-generated if log ID is not specified.\n",
    "    .with_log_id(log_id)\n",
    "    .with_job_infrastructure_type(\"ME_STANDALONE\")\n",
    "    # If you are in an OCI data science notebook session,\n",
    "    # the following configurations are not required.\n",
    "    # Configurations from the notebook session will be used as defaults.\n",
    "    .with_compartment_id(compartment_id)\n",
    "    .with_project_id(project_id)\n",
    "    .with_shape_name(instance_shape)\n",
    "    # Minimum/Default block storage size is 50 (GB).\n",
    "    .with_block_storage_size(80)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d61b80b-ea95-4501-a8ed-441a7916cfa7",
   "metadata": {},
   "source": [
    "### Configure Job Conatiner Runtime\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6c65fcf-44f0-4dc8-8dd7-ae70d5e7be0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "conatiner_runtime = (\n",
    "    ContainerRuntime()\n",
    "    # Specify the service conda environment by slug name.\n",
    "    .with_image(container_image)\n",
    "    # Environment variable\n",
    "    .with_environment_variable(\n",
    "        HF_TOKEN=hf_token,\n",
    "        MODEL=model_name,\n",
    "        BUCKET=bucket,\n",
    "        NAMESPACE=namespace,\n",
    "        OCI_IAM_TYPE=oci_iam_type,\n",
    "        PREFIX=prefix,\n",
    "    )\n",
    "    # Command line argument\n",
    "    .with_entrypoint([\"bash\", \"-c\"])\n",
    "    .with_cmd(\n",
    "        \"microdnf install -y unzip &&\"\n",
    "        + \"pip install oracle-ads[opctl] &&\"\n",
    "        + \"cd /home/datascience/ &&\"\n",
    "        + \"ls -lt &&\"\n",
    "        + \"unzip job-artifacts.zip -d . && \"\n",
    "        + \"chmod +x job-artifacts/vllm_batch_inferencing.py &&\"\n",
    "        + \"python job-artifacts/vllm_batch_inferencing.py\"\n",
    "    )\n",
    "    .with_artifact(artifacts_path)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d455458-4b77-460f-87dd-0b269fbe71ab",
   "metadata": {},
   "source": [
    "### Run a Job and Monitor outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92834d05-a498-423b-9c4b-ba2a46abddfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = (\n",
    "    Job(name=f\"Batch-inferencing of {model_name} using AI Quick Action.\")\n",
    "    .with_infrastructure(infrastructure)\n",
    "    .with_runtime(conatiner_runtime)\n",
    ")\n",
    "# Create the job on OCI Data Science\n",
    "job.create()\n",
    "# Start a job run\n",
    "run = job.run()\n",
    "\n",
    "# Stream the job run outputs\n",
    "run.watch()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "499313e1-5afc-497a-ac5c-2fec658233d2",
   "metadata": {},
   "source": [
    "## Download the output file from OS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "724d3b9e-5573-4145-add1-acd1430efdc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!oci os object bulk-download  --bucket-name $bucket --namespace $namespace --prefix prefix --dest-dir . --auth resource_principal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26fa1d8-98a1-4398-8454-92d55619ab0e",
   "metadata": {},
   "source": [
    "### Display downloaded output file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e5c0181-16f3-425a-8e90-053ec2765f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file = os.path.join(os.path.abspath(prefix), \"outputs\", \"output_prompts.json\")\n",
    "\n",
    "!python -m json.tool $output_file"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch24_p310_gpu_x86_64_v1]",
   "language": "python",
   "name": "conda-env-pytorch24_p310_gpu_x86_64_v1-py"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}