{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d4649c-a8ca-494d-8c11-e2aca8faea64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import plotly.graph_objects as go\n",
    "\n",
    "proj_dir = Path.cwd()\n",
    "proj_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a59f2e07-2505-4ad3-978d-2f2a8d4c7f16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "# Define the directory path where your files are located\n",
    "dir_path = proj_dir/'tgi_benchmark_results/'\n",
    "\n",
    "\n",
    "def build_df():\n",
    "    # Initialize an empty list to store the dataframes\n",
    "    dfs = []\n",
    "\n",
    "    # Iterate through the files in the directory\n",
    "    for tgibs_folder in dir_path.glob(\"*/*_tgibs_*\"):\n",
    "        # Check if the file matches the pattern *_summary.json\n",
    "        summary_file = list(tgibs_folder.glob(\"*_summary.json\"))[0]\n",
    "        # Extract the tgibs value from the filename\n",
    "        hw = tgibs_folder.parts[-2]\n",
    "        tgibs_value = tgibs_folder.name.split('_tgibs_')[1].split('__')[0]\n",
    "\n",
    "        # Load the JSON file\n",
    "        with open(summary_file, 'r') as f:\n",
    "            data = json.load(f)\n",
    "\n",
    "        # Convert the JSON data to a pandas dataframe\n",
    "        df = pd.DataFrame([data])\n",
    "\n",
    "        # Add a column with the tgibs value\n",
    "        df['tgibs'] = int(tgibs_value)\n",
    "        df['hw'] = hw\n",
    "        df['id'] = f\"{hw}_{tgibs_value}\"\n",
    "\n",
    "        # Append the dataframe to the list\n",
    "        dfs.append(df)\n",
    "    df = pd.concat(dfs, ignore_index=True)\n",
    "    df = df.sort_values(by=['tgibs', 'num_concurrent_requests'], ascending=[True, True])\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8508fb9-fa31-4e23-80c1-e77a56d3775e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = build_df()\n",
    "\n",
    "# Create a figure\n",
    "fig = go.Figure()\n",
    "\n",
    "# Group the dataframe by batch_size\n",
    "grouped_df = df.groupby('id')\n",
    "\n",
    "# List of specific batch_sizes to label\n",
    "label_batch_sizes = ['nvidia-a100_8', 'nvidia-h100_8', 'nvidia-a100_8', 'nvidia-h100-fp8_8', 'nvidia-a100_medusa_8']\n",
    "\n",
    "# Iterate over each group\n",
    "for batch_size, group in grouped_df:\n",
    "    # Add a line to the figure\n",
    "    fig.add_trace(go.Scatter(\n",
    "        x=group['results_end_to_end_latency_s_mean'],\n",
    "        y=group['results_num_completed_requests_per_min'],\n",
    "        mode='lines+markers',\n",
    "        name=f\"Batch Size: {batch_size}\",  # Formatting batch size in the legend\n",
    "        hovertemplate=(\n",
    "            f\"<b>Batch Size: {batch_size}</b><br>\"\n",
    "            \"VU: %{text}<br>\"\n",
    "            \"Latency: %{x:.2f}s<br>\"\n",
    "            \"Throughput: %{y:.2f} reqs/min\"\n",
    "        ) + \"<extra></extra>\",\n",
    "        text=[f\"{v} VU\" for v in group['num_concurrent_requests']]  # This will only be visible on hover\n",
    "    ))\n",
    "\n",
    "    # Optionally add annotations only for the first point in the specified batch sizes\n",
    "    if batch_size in label_batch_sizes:\n",
    "        fig.add_annotation(\n",
    "            x=group['results_end_to_end_latency_s_mean'].iloc[0],\n",
    "            y=group['results_num_completed_requests_per_min'].iloc[0],\n",
    "            text=f'{batch_size[:-2].replace(\"nvidia-\", \"\")}',\n",
    "            showarrow=False,\n",
    "            ax=0,\n",
    "            # ay=90,  # Offset to move the text down\n",
    "            xanchor='center',\n",
    "            yanchor='top'\n",
    "        )\n",
    "\n",
    "# Update layout for the figure\n",
    "fig.update_layout(\n",
    "    title_text=\"Requests Throughput vs Latency by Batch Size\",\n",
    "    xaxis_title=\"End to End Latency (seconds)\",\n",
    "    yaxis_title=\"Requests/min\",\n",
    "    showlegend=True,\n",
    ")\n",
    "\n",
    "# Show the figure\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d2719fe-b0b5-400f-83a0-7eaffd8f2254",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2472ada-8215-45cb-9efb-b094f02bb416",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}