From 6df2498715b891c57f09f3769ba29df90c60fc8a Mon Sep 17 00:00:00 2001
From: Shen Xu <shenchenxu@meta.com>
Date: Tue, 1 Oct 2024 19:30:26 -0700
Subject: [PATCH] Allow Inspector to accept ETDump bytes directly (#5657)

Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/5657

Reviewed By: tarun292

Differential Revision: D63400412
---
 devtools/inspector/_inspector.py           | 13 ++++++++++---
 devtools/inspector/_inspector_utils.py     | 19 +++++++++++++------
 devtools/inspector/tests/inspector_test.py |  4 +++-
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/devtools/inspector/_inspector.py b/devtools/inspector/_inspector.py
index 0539d4f5e4b..3691cd0234d 100644
--- a/devtools/inspector/_inspector.py
+++ b/devtools/inspector/_inspector.py
@@ -967,6 +967,7 @@ class Inspector:
     def __init__(
         self,
         etdump_path: Optional[str] = None,
+        etdump_data: Optional[bytes] = None,
         etrecord: Optional[Union[ETRecord, str]] = None,
         source_time_scale: TimeScale = TimeScale.NS,
         target_time_scale: TimeScale = TimeScale.MS,
@@ -980,11 +981,12 @@ def __init__(
         enable_module_hierarchy: bool = False,
     ) -> None:
         r"""
-        Initialize an `Inspector` instance with the underlying `EventBlock`\ s populated with data from the provided ETDump path
+        Initialize an `Inspector` instance with the underlying `EventBlock`\ s populated with data from the provided ETDump path or binary,
         and optional ETRecord path.
 
         Args:
-            etdump_path: Path to the ETDump file.
+            etdump_path: Path to the ETDump file. Either this parameter or etdump_data should be provided.
+            etdump_data: ETDump binary. Either this parameter or etdump_path should be provided.
             etrecord: Optional ETRecord object or path to the ETRecord file.
             source_time_scale: The time scale of the performance data retrieved from the runtime. The default time hook implentation in the runtime returns NS.
             target_time_scale: The target time scale to which the users want their performance data converted to. Defaults to MS.
@@ -1025,8 +1027,13 @@ def __init__(
         else:
             raise TypeError("Unsupported ETRecord type")
 
+        if (etdump_path is None) == (etdump_data is None):
+            raise ValueError(
+                "Expecting exactly one of etdump_path or etdump_data to be specified."
+            )
+
         # Create EventBlocks from ETDump
-        etdump = gen_etdump_object(etdump_path=etdump_path)
+        etdump = gen_etdump_object(etdump_path=etdump_path, etdump_data=etdump_data)
         if debug_buffer_path is not None:
             with open(debug_buffer_path, "rb") as f:
                 output_buffer = f.read()
diff --git a/devtools/inspector/_inspector_utils.py b/devtools/inspector/_inspector_utils.py
index 5f04e2d0413..a2989c224e1 100644
--- a/devtools/inspector/_inspector_utils.py
+++ b/devtools/inspector/_inspector_utils.py
@@ -279,13 +279,20 @@ def _extract_debug_handles(graph: OperatorGraph):
     return debug_handle_to_op_node_map
 
 
-def gen_etdump_object(etdump_path: Optional[str] = None) -> ETDumpFlatCC:
+def gen_etdump_object(
+    etdump_path: Optional[str] = None, etdump_data: Optional[bytes] = None
+) -> ETDumpFlatCC:
     # Gen event blocks from etdump
-    if etdump_path is None:
-        raise ValueError("Etdump_path must be specified.")
-    with open(etdump_path, "rb") as buff:
-        etdump = deserialize_from_etdump_flatcc(buff.read())
-        return etdump
+    if etdump_data is None and etdump_path is not None:
+        with open(etdump_path, "rb") as buff:
+            etdump_data = buff.read()
+
+    if etdump_data is None:
+        raise ValueError(
+            "Unable to get ETDump data. One and only one of etdump_path and etdump_data must be specified."
+        )
+
+    return deserialize_from_etdump_flatcc(etdump_data)
 
 
 def plot_metric(result: List[float], metric_name: str):
diff --git a/devtools/inspector/tests/inspector_test.py b/devtools/inspector/tests/inspector_test.py
index 34c96eef534..4b3f8075d8e 100644
--- a/devtools/inspector/tests/inspector_test.py
+++ b/devtools/inspector/tests/inspector_test.py
@@ -86,7 +86,9 @@ def test_inspector_constructor(self):
 
             # Assert that expected functions are called
             mock_parse_etrecord.assert_called_once_with(etrecord_path=ETRECORD_PATH)
-            mock_gen_etdump.assert_called_once_with(etdump_path=ETDUMP_PATH)
+            mock_gen_etdump.assert_called_once_with(
+                etdump_path=ETDUMP_PATH, etdump_data=None
+            )
             mock_gen_from_etdump.assert_called_once()
             # Because we mocked parse_etrecord() to return None, this method shouldn't be called
             mock_gen_graphs_from_etrecord.assert_not_called()