Adding support for a scale factor when parsing ETDump results

sdk/etdb/inspector.py

@@ -153,12 +153,16 @@ class Event:
 
     @staticmethod
     def _gen_from_profile_events(
-        signature: ProfileEventSignature, events: List[ProfileEvent]
+        signature: ProfileEventSignature,
+        events: List[ProfileEvent],
+        scale_factor: int = 1,
     ) -> "Event":
         """
         Given a ProfileEventSignature and a list of ProfileEvents with that signature,
         return an Event object matching the ProfileEventSignature, with perf_data
         populated from the list of ProfileEvents
+
+        An optional inverse scale factor can be provided to adjust the event timestamps
         """
         if signature.delegate_id is not None:  # 0 is a valid value
             delegate_debug_identifier = signature.delegate_id
@@ -170,7 +174,10 @@ def _gen_from_profile_events(
         name = signature.name if not is_delegated_op else str(delegate_debug_identifier)
 
         perf_data = PerfData(
-            [float(event.end_time - event.start_time) / 1000 for event in events]
+            [
+                float(event.end_time - event.start_time) / scale_factor
+                for event in events
+            ]
         )
 
         return Event(
@@ -247,10 +254,15 @@ def to_dataframe(self) -> pd.DataFrame:
         return df
 
     @staticmethod
-    def _gen_from_etdump(etdump: ETDumpFlatCC) -> List["EventBlock"]:
+    def _gen_from_etdump(
+        etdump: ETDumpFlatCC, scale_factor: int = 1
+    ) -> List["EventBlock"]:
         """
         Given an etdump, generate a list of EventBlocks corresponding to the
-        contents
+        contents.
+
+        An optional (inverse) scale factor can be provided to adjust the
+        etdump timestamps associated with each EventBlocks
         """
 
         # Group all the RunData by the set of profile events
@@ -286,7 +298,7 @@ def _gen_from_etdump(etdump: ETDumpFlatCC) -> List["EventBlock"]:
             EventBlock(
                 name=str(index),
                 events=[
-                    Event._gen_from_profile_events(signature, event)
+                    Event._gen_from_profile_events(signature, event, scale_factor)
                     for signature, event in profile_events.items()
                 ],
             )
@@ -360,16 +372,25 @@ class Inspector:
     """
 
     def __init__(
-        self, etdump_path: Optional[str] = None, etrecord_path: Optional[str] = None
+        self,
+        etdump_path: Optional[str] = None,
+        etrecord_path: Optional[str] = None,
+        etdump_scale: int = 1000,
     ) -> None:
         """
         Create an inspector instance from the provided ETDump/ETRecord
+
+        Args:
+            etdump_path: Path to the ETDump file.
+            etrecord_path: Path to the ETRecord file.
+            etdump_scale: Inverse Scale Factor used to cast the timestamps in ETDump
+                defaults to milli (1000ms = 1s).
         """
 
         # TODO: etrecord_path can be optional, so need to support the case when it is not present
         self._etrecord = gen_etrecord_object(etrecord_path=etrecord_path)
         etdump = gen_etdump_object(etdump_path=etdump_path)
-        self.event_blocks = EventBlock._gen_from_etdump(etdump)
+        self.event_blocks = EventBlock._gen_from_etdump(etdump, etdump_scale)
 
         self._op_graph_dict: Mapping[
             str, OperatorGraphWithStats

sdk/etdb/tests/event_blocks_test.py

@@ -128,7 +128,7 @@ def test_gen_from_etdump(self) -> None:
         - Correct number of Events and Raw Data values (iterations)
         """
         etdump: ETDumpFlatCC = TestEventBlock._get_sample_etdump_flatcc()
-        blocks: List[EventBlock] = EventBlock._gen_from_etdump(etdump)
+        blocks: List[EventBlock] = EventBlock._gen_from_etdump(etdump, 1000)
 
         self.assertEqual(len(blocks), 2, f"Expected 2 runs, got {len(blocks)}")
 
@@ -152,6 +152,7 @@ def _test_profile_event_generation(
             instruction_id: int,
             delegate_debug_id_int: Optional[int] = None,
             delegate_debug_id_str: Optional[str] = None,
+            scale_factor: int = 1000,
         ) -> None:
             """
             Helper function for testing that the provided ProfileEvent fields are
@@ -180,17 +181,19 @@ def _test_profile_event_generation(
                 TestEventBlock._gen_sample_profile_event(
                     name,
                     instruction_id,
-                    (0, 10),
+                    (0, time),
                     delegate_debug_id,
                 )
                 for time in durations
             ]
-            event = Event._gen_from_profile_events(signature, profile_events)
+            event = Event._gen_from_profile_events(
+                signature, profile_events, scale_factor=scale_factor
+            )
 
             is_delegated = delegate_debug_id is not None
             expected_event = Event(
                 str(delegate_debug_id) if is_delegated else name,
-                PerfData([float(duration) / 1000 for duration in durations]),
+                PerfData([float(duration) / scale_factor for duration in durations]),
                 instruction_id=signature.instruction_id,
                 delegate_debug_identifier=delegate_debug_id,
                 is_delegated_op=is_delegated,
@@ -206,6 +209,11 @@ def _test_profile_event_generation(
         # Delegate with String Debug ID
         _test_profile_event_generation("delegate", 1, None, "identifier")
 
+        # Manipulating the scale factor
+        _test_profile_event_generation(
+            "delegate", 1, None, "identifier", scale_factor=10000
+        )
+
     def test_gen_resolve_debug_handles(self) -> None:
         """
         Test that gen_resolve_debug_handles() correctly populates the EventBlock