examples/experimental/critical_path_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "10a4e2e8",
   "metadata": {},
   "source": [
    "__<font color='red'>\n",
    "Note: This feature is currently under development so it will not be fully functional yet. \n",
    "Also, the APIs here are experimental and subject to change.\n",
    "</font>__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f114ee",
   "metadata": {},
   "source": [
    "## Setup and loading traces\n",
    "To run this demo notebook on your laptop\n",
    "1. Clone the repo `git clone https://github.com/facebookresearch/HolisticTraceAnalysis.git`\n",
    "1. [Optional and recommended] Setup a conda environment. See README for details.\n",
    "1. Set the `trace_dir` parameter in the next cell to the location of `tests/data/critical_path/simple_add` folder in your local `HolisticTraceAnalysis` installation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c8e32f5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hta:Parsed /Users/bcoutinho/Work/hta/HolisticTraceAnalysis2/tests/data/critical_path/simple_add/benchmark_result_493459_1694039959_trace.json.gz time = 0.01 seconds \n",
      "WARNING:hta:Parsed /Users/bcoutinho/Work/hta/HolisticTraceAnalysis2/tests/data/critical_path/simple_add/benchmark_result_493459_1694039959_trace.json.gz backend=json in 0.02 seconds; current PID:54853\n",
      "WARNING:hta:Overall parsing of /Users/bcoutinho/Work/hta/HolisticTraceAnalysis2/tests/data/critical_path/simple_add/benchmark_result_493459_1694039959_trace.json.gz in 0.03 seconds; current PID:54853\n",
      "WARNING:hta:leaving parse_multiple_ranks duration=0.04 seconds\n",
      "WARNING:hta:leaving parse_traces duration=0.04 seconds\n",
      "WARNING:hta:ProfilerStep not found in the trace. The analysis result may not be accurate.\n"
     ]
    }
   ],
   "source": [
    "from hta.trace_analysis import TraceAnalysis\n",
    "base_dir = <update base path here> #\"~/Work/hta/HolisticTraceAnalysis2\"\n",
    "trace_dir = base_dir + \"/tests/data/critical_path/simple_add/\"\n",
    "\n",
    "analyzer = TraceAnalysis(trace_dir=trace_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0da2d07a-b4d3-4e7b-87d6-229b32773617",
   "metadata": {},
   "source": [
    "# Critical Path Analysis\n",
    "\n",
    "Critical path analysis is a commonly applied technique in HPC and AI/ML optimization. It can be leveraged in two ways:\n",
    "1. **Efficiency opportunities** Operations/kernels on critical path should be the target of performance analysis and optimizations. They can provide the “best bang for buck” for performance improvements.This is not limited to just CPU/GPU kernels. Delays in launching or executing CUDA kernels can constitute a significant portion of the critical path as well. This could be optimized by operator fusion and CUDA graphs for example.\n",
    "2. **Simulating improvements / gains:** After identifying critical path we can estimate improvements to the execution time without actually running anything. This helps us estimate the possible gains from various performance optimization.\n",
    "\n",
    "## Approach\n",
    "In a nutshell, computing the critical path involves 1) constructing a weighted DAG connecting all the operations, 2) finding the longest path in this DAG. The challenging part is constructing the DAG here. \n",
    "\n",
    "**Nodes**: The Nodes in the critical path graph represent points in time. Each operator/kernel thus has two nodes viz. a begin and an end node. In case of nested operators we also link the nodes in the order they appear in the callstack. An example of stacked operators is shown below-\n",
    "```\n",
    "          |----------------------- Op A ----------------------|\n",
    "                    |--- Op B ---|        |-- Op C--|\n",
    "Critical Path Graph\n",
    "         (OpA.b)--->(ObB.b)----->(OpB.e)->(OpC.b)-->(OpC.e)->(OpA.e)\n",
    "```\n",
    "\n",
    "**Edges** in this DAG can be one of two types\n",
    "1. Timing edges (weight = time): include durations for the operators/kernels as well as delays to launch operators between CPU and GPU.\n",
    "1. Dependency edges (weight = 0): do not have a time component but show a dependency between operations themselves. This includes data dependencies and synchronization between CPU and GPU.\n",
    "\n",
    "## Lightweight Critical Path\n",
    "Our initial implementation of Critical Path analysis does not consider dependency between PyTorch operators, this requires combining the trace with Chakra Execution Traces and processing Tensor information.\n",
    "\n",
    "The key idea here is to simplify the dependency analysis between PyTorch operators\n",
    "- We assume that all CPU operators are dependent serially on the last operator that ran on a CPU thread. This frees us from attempting to correlate with Execution Traces.\n",
    "- The operator dependency part can be added back later.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "868efbc5-f90b-479d-8c2f-d4a9e26f5ced",
   "metadata": {},
   "source": [
    "## Using this notebook\n",
    "\n",
    "We can demonstrate critical path analysis on most PyTorch traces. However, currently traces are missing\n",
    "information regarding CUDA synchronization. \n",
    "\n",
    "The above is fixed in this [PR](https://github.com/pytorch/pytorch/pull/105187) but not enabled by default. Please follow the documentation in this [PR](https://github.com/pytorch/pytorch/pull/105187) to enable CUDA synchronization events to get best results from this analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ca7b1630-9905-437a-92f2-34d3b84d048d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m\n",
       "\u001b[0manalyzer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcritical_path_analysis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mrank\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mannotation\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0minstance_id\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNoneType\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mTuple\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTuple\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mhta\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0manalyzers\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcritical_path_analysis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCPGraph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbool\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Perform critical path analysis for trace events within a rank.\n",
       "We further reduce the region of interest by selecting\n",
       "a trace annotation and instance id. This will\n",
       "limit the analysis to events within the time range of that annoation.\n",
       "This will include GPU kernels launched by the cpu operators in that\n",
       "time duration.\n",
       "For example, you can use this to limit the analysis to one iteration\n",
       "by passing annotation='ProfilerStep'. See notes for how to pick the iteration.\n",
       "\n",
       "Args:\n",
       "    t (Trace): Input trace data structure.\n",
       "    rank (int): rank to analyze for the critical path.\n",
       "    annotation (str): a trace annotation to limit the analysis to,\n",
       "        for example \"ProfilerStep\" would match all annotations that\n",
       "        match this string (ProfilerStep#100, ProfilerStep#101 etc)\n",
       "    instance_id: can be either of the following\n",
       "        (int) - specify which instance of the annotation to consider.\n",
       "                Defaults to the first instance.\n",
       "        (Tuple(int, int)) - considers a range of annotation instances start to end,\n",
       "                inclusive of both start and end instance.\n",
       "Returns:\n",
       "    Tuple[CPGraph, bool]\n",
       "        A tuple of CPGraph object and a success or fail boolean value.\n",
       "        True indicates that the critical path analysis algorithm succeeded.\n",
       "\n",
       "CPGraph object that can be used to obtain statistics and further\n",
       "visualize the critical path. CPGraph is also a subinstance of a networkx.DiGraph.\n",
       "Run 'CPGraph?' for more info and APIs.\n",
       "\n",
       "Notes:\n",
       "\n",
       "1. Avoid using the first step / iteration in a trace as it usually\n",
       "   has some missing events.\n",
       "2. The analysis requires CUDA synchronization events in the GPU trace,\n",
       "   that were added in https://github.com/pytorch/pytorch/pull/105187\n",
       "   Please see the documentation of this PR on how to enable CUDA sync events in the trace.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/Work/hta/HolisticTraceAnalysis2/hta/trace_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m      method"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "analyzer.critical_path_analysis?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1debaba3-bbac-498b-8f6f-a0b69658e6b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "annotation = \"[param|pytorch.model.alex_net|0|0|0|measure|forward]\"\n",
    "instance_id = 1  # note this is zero based\n",
    "\n",
    "# You can also pass a tuple specifying range of annotation instances\n",
    "# annotation = \"ProfilerStep\"  # will match multiple ProfilerStepXXX annotations\n",
    "# instance_id = (2,3)          # include 3rd to 4th iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "4e8affb1-74f2-4b88-a6eb-85d3a1b67c79",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5170\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5183\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5194\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5196\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5198\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5200\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5202\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5204\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5206\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5208\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5586\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5599\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5610\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5612\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5614\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5616\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5618\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5620\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5622\n",
      "WARNING:hta:CUDA Stream Wait event was not matched to a cudaRecordEvent, name = Stream Wait Event, correlation = 5624\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cp_graph, success = analyzer.critical_path_analysis(\n",
    "    rank = 0, annotation=annotation, instance_id=instance_id)\n",
    "success                                             "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f05f9c-437b-441e-9353-a87307b25e88",
   "metadata": {},
   "source": [
    "## Getting a high level overview of the Critical Path\n",
    "\n",
    "We can use the `summary()` method of the critical path graph object to get an overview of the bottlenecks\n",
    "that are on the critical path. The summary contains a \"bounded_by\" column that shows the % of total\n",
    "time on critical path a particular bottleneck uses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "ba5e8a12-0319-4453-9d7c-d4a18dfa4dec",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mcp_graph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msummary\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mpandas\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mseries\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSeries\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Displays a summary or breakdown of the critical path into one of the following\n",
       "- cpu_bound\n",
       "- gpu_compute_bound\n",
       "- gpu_communication_bound (NCCL Collectives)\n",
       "- gpu_kernel_kernel_overhead (Gaps between GPU kernels)\n",
       "- gpu_kernel_launch_overhead (Delay launching kernels from CPU to GPU)\n",
       "\n",
       "Returns:\n",
       "    A summary pandas series with bottleneck type -> % of duration on critical path.\n",
       "Also see get_critical_path_breakdown().\n",
       "\u001b[0;31mFile:\u001b[0m      ~/Work/hta/HolisticTraceAnalysis2/hta/analyzers/critical_path_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m      method"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cp_graph.summary?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "893f8617-1c23-43bc-abea-3c258ecb8a46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Critical Path broken down by boundedness = (in % of duration)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "bound_by\n",
       "                               0.000000\n",
       "cpu_bound                     99.779087\n",
       "gpu_compute_bound              0.210157\n",
       "gpu_kernel_kernel_overhead     0.003723\n",
       "gpu_kernel_launch_overhead     0.007033\n",
       "Name: duration, dtype: float64"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cp_graph.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212636bb-8481-4b4b-8de2-5c36cb947b20",
   "metadata": {},
   "source": [
    "In this scenario, the CPU is the primary constraint. This may be because the CPU is not efficiently scheduling GPU kernels, resulting in it becoming a bottleneck. For more information on the specific events and edges involved in the critical path, you can refer to the dataframe generated by `cp_graph.get_critical_path_breakdown()`.\n",
    "\n",
    "To debug further we can visualize the critical path as shown in the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c121d9b8-72c8-45a4-972c-7c29c163212b",
   "metadata": {},
   "source": [
    "## Overlay and visualize the Critical Path\n",
    "The `overlay_critical_path_analysis()` function exposes the critical path on the original trace file.\n",
    "There are two modes for the output:\n",
    "1. When `only_show_critical_events=True` (default value) the output trace only contains CPU operators and GPU events on the critical path. One can compare it with the original trace to contrast the critical path identified by the algorithm.\n",
    "1. When `only_show_critical_events=False` in the output trace file search for \"critical\" to highlight events on the critical path.\n",
    "\n",
    "Edges in the critical path graph will be shown using arrows or flow events. Critical edges are marked using the \"critical\" flag in the args property.\n",
    "\n",
    "The category names of flow events begin with \"critical_path_\". They have the following meaning:\n",
    "* `critical_path_dependency`: an inter operator dependency.\n",
    "* `critical_path_operator`: an edge between start and end of a CPU operator or GPU kernel.\n",
    "* `critical_path_kernel_launch_delay`: delay to launch a GPU kernel that is likely to be on the critical path.\n",
    "* `critical_path_kernel_kernel_delay`: delay between running successive GPU kernels.\n",
    "* `critical_path_sync_dependency`: these edges denote synchronization or control dependencies between events such as GPU->CPU, GPU->GPU synchronization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "19434b4f-9799-40f3-b6df-ec99fce37246",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m\n",
       "\u001b[0manalyzer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moverlay_critical_path_analysis\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mrank\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mcritical_path_graph\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mhta\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0manalyzers\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcritical_path_analysis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCPGraph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0moutput_dir\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0monly_show_critical_events\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mshow_all_edges\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Overlay the identified critical path on top of the trace file\n",
       "for visualization.\n",
       "\n",
       "Args:\n",
       "    rank (int): rank to generate the time series for.\n",
       "    critical_path_graph: Critical Path Graph object generated previously\n",
       "    output_dir (str): Output directory to store overlaid trace.\n",
       "    only_show_critical_events (bool): When set the output trace will only\n",
       "        have operators and GPU kernels on the critical path. It will\n",
       "        still retain the user annotations.\n",
       "    show_all_edges (bool): When set this will add edge events for\n",
       "        all types of edges in the critical path graph. This is useful\n",
       "        for debugging the algorithm.\n",
       "\n",
       "Returns: the overlaid trace file path. The generated trace file will\n",
       "have a prefix of \"overlaid_critical_path\\_\" in its name compared\n",
       "to the original trace file.\n",
       "\n",
       "Note: In case of kernel launches that are not on the critical path the graph\n",
       "still has a 0 weight edge between CUDA runtime and kernel. These 0 weight\n",
       "edges are not shown in the overlaid trace by default. Set the environment\n",
       "variable CRITICAL_PATH_SHOW_ZERO_WEIGHT_LAUNCH_EDGE=1 to enable adding this\n",
       "to the overlaid trace. Add this to your notebook\n",
       "`os.environ[\"CRITICAL_PATH_SHOW_ZERO_WEIGHT_LAUNCH_EDGE\"] = 1`\n",
       "\u001b[0;31mFile:\u001b[0m      ~/Work/hta/HolisticTraceAnalysis2/hta/trace_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m      method"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "analyzer.overlay_critical_path_analysis?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "65e20368-2e85-48f1-af48-60d334ec90a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR:hta:The path /Users/bcoutinho/HolisticTraceAnalysis/tests/data/critical_path/simple_add/overlaid does not exist.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "''"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Make sure '~/HolisticTraceAnalysis/tests/data/critical_path/simple_add/overlaid' exists or add a different directory\n",
    "# to write your trace too.\n",
    "analyzer.overlay_critical_path_analysis(\n",
    "    0, cp_graph, output_dir='~/HolisticTraceAnalysis/tests/data/critical_path/simple_add/overlaid')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b9a9ff-2141-4568-8ae3-625d770ac215",
   "metadata": {},
   "source": [
    "## Saving and Restoring CPGraph object\n",
    "\n",
    "Constructing the critical path graph can often take time. However, if a user simply wants to simulate new changes\n",
    "on an existing CPGRaph this can make iteration a lot slower. Note that this save() operation can only be done after\n",
    "the graph has been constructed!\n",
    "To speed things up we add a `save()` and `restore()` feature to CPGraph object.\n",
    "\n",
    "You can save and `CPGraph` object as follows. The output is a zip file with several different saved datastructure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "1eb1a4ff-f198-4b25-a15e-713478442bbb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mcp_graph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout_dir\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Saves the critical path graph object to a zip file.\n",
       "Note that this save() operation can only be done after\n",
       "the graph has been constructed!\n",
       "\n",
       "Args:\n",
       "     out_dir(str) is the directory used to first dump a\n",
       "        collection of files used to save the state of CPGraph object.\n",
       "        The directory is then compressed into a zipfile with the same name.\n",
       "Returns:\n",
       "    trace_zip_file (str): path to the zip file containing the object.\n",
       "\n",
       "Note:\n",
       "    Internally it saves the data into 3 files that are zipped up.\n",
       "    1. Is the trace_data saved as a csv.\n",
       "    2. The networkx graph is serialized using python pickle.\n",
       "    2. Other object data is serialized using python pickle.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/Work/hta/HolisticTraceAnalysis2/hta/analyzers/critical_path_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m      method"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cp_graph.save?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "687b8d9e-f064-48e3-ab28-5580099d536b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hta:Saved files to /tmp/my_saved_cp_graph.zip\n"
     ]
    }
   ],
   "source": [
    "zip_file = cp_graph.save(out_dir=\"/tmp/my_saved_cp_graph\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "860de8ef-51f0-4c92-965a-7e8744f16e19",
   "metadata": {},
   "source": [
    "### Restoring saved graph object\n",
    "The saved graph can easily be restored as follows. This does required the original trace dataframe and rank identifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6498661d-91bf-4069-87da-d7a334e8278c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mrestore_cpgraph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip_filename\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt_full\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Trace'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrank\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mhta\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0manalyzers\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcritical_path_analysis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCPGraph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Restores the critical path graph object from a zip file.\n",
       "The graph will already be constructed in this case. You can\n",
       "however run critical_path() again and modify the graph etc.\n",
       "\n",
       "Returns:\n",
       "    trace_zip_file (str): path to the zip file containing the object.\n",
       "Returns:\n",
       "    CPGraph object restored from the file.\n",
       "\u001b[0;31mFile:\u001b[0m      ~/Work/hta/HolisticTraceAnalysis2/hta/analyzers/critical_path_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m      function"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from hta.analyzers.critical_path_analysis import restore_cpgraph\n",
    "restore_cpgraph?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "05ee0618-e32b-4112-a38c-c6b1741281c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:hta:Extracted zip archive to /tmp/tmp/my_saved_cp_graph\n",
      "WARNING:hta:Restoring graph from /tmp/tmp/my_saved_cp_graph/cp_graph.pkl\n",
      "WARNING:hta:Loading cp_graph from /tmp/tmp/my_saved_cp_graph\n"
     ]
    }
   ],
   "source": [
    "rest_graph = restore_cpgraph(zip_filename=zip_file, t_full=analyzer.t, rank=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "721551b4-de2f-4b76-bfc4-554526faff36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(rest_graph.critical_path_edges_set) == len(cp_graph.critical_path_edges_set)\n",
    "len(rest_graph.nodes) == len(cp_graph.nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "07b9a4ff-74e0-42d5-b31e-1a31efc67226",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Critical Path broken down by boundedness = (in % of duration)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "bound_by\n",
       "                               0.000000\n",
       "cpu_bound                     99.779087\n",
       "gpu_compute_bound              0.210157\n",
       "gpu_kernel_kernel_overhead     0.003723\n",
       "gpu_kernel_launch_overhead     0.007033\n",
       "Name: duration, dtype: float64"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rest_graph.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "d81d2125-f25a-4e70-883f-45cfdbaaa326",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Critical Path broken down by boundedness = (in % of duration)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "bound_by\n",
       "                               0.000000\n",
       "cpu_bound                     99.779087\n",
       "gpu_compute_bound              0.210157\n",
       "gpu_kernel_kernel_overhead     0.003723\n",
       "gpu_kernel_launch_overhead     0.007033\n",
       "Name: duration, dtype: float64"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## You can also re-compute the critical path object\n",
    "rest_graph.critical_path()\n",
    "rest_graph.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2982bca5-b1d9-4673-96e2-9a50e14de9e4",
   "metadata": {},
   "source": [
    "## More details on the CPGraph object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "8b70c7e2-103c-475a-b030-356fc6f4d4ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mInit signature:\u001b[0m\n",
       "\u001b[0mCPGraph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mt\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mForwardRef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Trace'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mt_full\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Trace'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mrank\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mG\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m     \n",
       "Critical path analysis graph representation for trace from one rank.\n",
       "This object constructs a graph that can be analyzed using networkx library.\n",
       "\n",
       "We maintain a mapping between node ids -> CPNode objects\n",
       "and use the integer as a node in the networkx graph datastructure.\n",
       "Edges are directly used as the type is hashable.\n",
       "\n",
       "Attributes:\n",
       "    trace_df (pd.DataFrame): dataframe of trace events used to construct this graph.\n",
       "    symbol_table (TraceSymbolTable): a symbol table used to encode the symbols in the trace.\n",
       "    node_list (List[int]): list of critical path node objects, index in this list is always the node id..\n",
       "    critical_path_nodes (List[int]): list of node ids on the critical path.\n",
       "    critical_path_events_set (Set[int]): set of event ids corresponding to the critical path nodes.\n",
       "    critical_path_edges_set (Set[CPEdge]): set of edge objects that are on the critical path.\n",
       "\u001b[0;31mInit docstring:\u001b[0m\n",
       "Initialize a critical path graph object. This can be done in two\n",
       "ways\n",
       "    1) Generate critical path analysis graph from scatch.\n",
       "    2) Restore a serialized CPGraph object.\n",
       "\n",
       "For (2) the networkx.DiGraph object G is utilized, see restore_cpgraph() function in this file.\n",
       "\n",
       "Args:\n",
       "    t (Trace): Clipped trace object focussing on region of interest.\n",
       "    t_full (Trace): Full Trace object.\n",
       "    rank (int): Rank to perform analysis on.\n",
       "    G (networkx.DiGraph): An optional DiGraph object.\n",
       "\u001b[0;31mFile:\u001b[0m           ~/Work/hta/HolisticTraceAnalysis2/hta/analyzers/critical_path_analysis.py\n",
       "\u001b[0;31mType:\u001b[0m           type\n",
       "\u001b[0;31mSubclasses:\u001b[0m     "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from hta.analyzers.critical_path_analysis import CPGraph\n",
    "CPGraph?"
   ]
  }
 ],
 "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}