Revert D78594843

Summary:
This diff reverts D78594843
S545236

Depends on D78594843

Reviewed By: hjli-creator

Differential Revision: D79185506

Author: Dark Knight

hta/analyzers/critical_path_analysis.py

@@ -622,7 +622,7 @@ def _construct_graph_from_cuda_runtime_events(self) -> None:
                         self._add_edge_helper(
                             previous_node, node, CPEdgeType.DEPENDENCY
                         )
-                if last_node and not last_node.is_start:
+                if not last_node.is_start:
                     self._add_edge_helper(last_node, node, CPEdgeType.DEPENDENCY)
 
             last_node = node
@@ -652,7 +652,7 @@ def _construct_graph_from_call_stacks(self) -> None:
         for csg in cpu_call_stacks:
             self._construct_graph_from_call_stack(csg)
 
-    def _get_bwd_tid(self, trace_df: pd.DataFrame) -> Optional[int]:
+    def _get_bwd_tid(self, trace_df: pd.DataFrame) -> int | None:
         """Get the thread id for the backward pass, or None is one cannot be identified.
 
         We identify the backward pass as the thread which contains "autograd" events. If
@@ -668,7 +668,7 @@ def _get_bwd_tid(self, trace_df: pd.DataFrame) -> Optional[int]:
 
         return self._get_tid_for_event(trace_df, "autograd")
 
-    def _get_fwd_tid(self, trace_df: pd.DataFrame) -> Optional[int]:
+    def _get_fwd_tid(self, trace_df: pd.DataFrame) -> int | None:
         """Get the thread id for the forward pass, or None is one cannot be identified.
 
         We identify the forward pass as the thread which contains
@@ -683,7 +683,7 @@ def _get_fwd_tid(self, trace_df: pd.DataFrame) -> Optional[int]:
         """
         return self._get_tid_for_event(trace_df, "forward")
 
-    def _get_tid_for_event(self, trace_df: pd.DataFrame, ev_name: str) -> Optional[int]:
+    def _get_tid_for_event(self, trace_df: pd.DataFrame, ev_name: str) -> int | None:
         events = [sym for sym in self.symbol_table.sym_table if ev_name in sym]
         event_ids = [self.symbol_table.sym_index[candidate] for candidate in events]
         cpu_op_id = self.symbol_table.sym_index.get("cpu_op", -1)

hta/common/call_stack.py

@@ -4,7 +4,6 @@
 
 import functools
 import logging
-import queue
 from collections import namedtuple
 from enum import Enum
 from time import perf_counter
@@ -462,8 +461,7 @@ def __init__(
         trace: Trace,
         ranks: Optional[List[int]] = None,
         filter_func: Optional[Filter] = None,
-        remapped_tids: Optional[Dict[int, Dict[int, int]]] = None,
-        pre_process_trace_data: bool = False,
+        remapped_tids: Dict[int, Dict[int, int]] | None = None,
     ) -> None:
         """Construct a CallGraph from a Trace object <trace_data>
 
@@ -474,14 +472,12 @@ def __init__(
             remapped_tids (Dict[Dict[int, int]]) : a dictionary that stores per-rank thread ID remappings.
                 For example: { 0: { 300: 400}} means that on rank 0, thread ID 300 is remapped to 400.
                 This is useful for training jobs, where the backward thread can typically be merged into the main trainer thread.
-            pre_process_trace_data (bool) : whether to pre-process the trace data. If True, the event duration from trace data will be trimmed to not exceed the duration of the parent event if exist.
         Raises:
             ValueError: the trace data is invalid.
         """
         self.trace_data: Trace = trace
         self.mapping: pd.DataFrame = pd.DataFrame()
         self.call_stacks: List[CallStackGraph] = []
-        self.pre_process_trace_data = pre_process_trace_data
 
         _ranks = [k for k in trace.get_all_traces()] if ranks is None else ranks
         self._construct_call_graph(_ranks, filter_func, remapped_tids)
@@ -505,7 +501,7 @@ def _construct_call_graph(
         self,
         ranks: List[int],
         filter_func: Optional[Filter],
-        remapped_tids: Optional[Dict[int, Dict[int, int]]] = None,
+        remapped_tids: Dict[int, Dict[int, int]] | None = None,
     ) -> None:
         """
         Construct the call graph from the traces of a distributed training job.
@@ -521,8 +517,6 @@ def _construct_call_graph(
         t0 = perf_counter()
         # construct a call stack graph for each thread/stream
         for rank in ranks:
-            if self.pre_process_trace_data:
-                self.trim_trace_events(self.trace_data.get_trace(rank))
             df = self.trace_data.get_trace(rank).copy()
             if remapped_tids and rank in remapped_tids:
                 df = self._remap_tids(df, remapped_tids[rank])
@@ -531,11 +525,7 @@ def _construct_call_graph(
                     # Filter out gpu annotations and sync events
                     df_thread = df_thread[df_thread["stream"].gt(0)]
                 csi = CallStackIdentity(rank, pid, tid)
-                csg = CallStackGraph(
-                    df_thread,
-                    csi,
-                    filter_func,
-                )
+                csg = CallStackGraph(df_thread, csi, filter_func)
                 self.call_stacks.append(csg)
                 call_stack_ids.append(csi)
         t1 = perf_counter()
@@ -617,44 +607,3 @@ def get_stack_of_node(
         stack_nodes = np.array(leaf_nodes + parent_nodes + kernel_nodes)
         df_stack = df.reindex(stack_nodes)
         return df_stack
-
-    def trim_trace_events(self, df: pd.DataFrame) -> None:
-        """Trim the trace events to not exceed the duration of the parent event if exist. The original DataFrame is modified in place.
-
-        Args:
-            df (pd.DataFrame): the trace data frame.
-        """
-        adj: Dict[int, List[int]] = {}
-        python_id_event_idx_map: Dict[int, int] = {}
-        roots = []
-        for row in df.itertuples():
-            if hasattr(row, "python_id") and row.python_id > -1:
-                python_id_event_idx_map[row.python_id] = row.index
-                if (
-                    hasattr(row, "python_parent_id")
-                    and row.python_parent_id > -1
-                    and df.loc[python_id_event_idx_map[row.python_parent_id]]["tid"]
-                    == row.tid
-                ):
-                    children = adj.get(row.python_parent_id, [])
-                    children.append(row.python_id)
-                    adj[row.python_parent_id] = children
-                else:
-                    roots.append(row.python_id)
-        for root in roots:
-            # BFS trim event duration
-            q: queue.Queue[int] = queue.Queue()
-            q.put(root)
-            while not q.empty():
-                cur_py_id = q.get()
-                if cur_py_id in adj:
-                    cur_id = python_id_event_idx_map[cur_py_id]
-                    for child_py_id in adj[cur_py_id]:
-                        child_id = python_id_event_idx_map[child_py_id]
-                        df.at[child_id, "dur"] = min(
-                            df.at[child_id, "dur"],
-                            df.at[cur_id, "ts"]
-                            + df.at[cur_id, "dur"]
-                            - df.at[child_id, "ts"],
-                        )
-                        q.put(child_py_id)

hta/configs/event_args_formats/event_args_1.0.0.yaml

@@ -11,16 +11,6 @@ AVAILABLE_ARGS:
     raw_name: External id
     value_type: Int
     default_value: -1
-  index::python_id:
-    name: python_id
-    raw_name: Python id
-    value_type: Int
-    default_value: -1
-  index::python_parent_id:
-    name: python_parent_id
-    raw_name: Python parent id
-    value_type: Int
-    default_value: -1
   cpu_op::concrete_inputs:
     name: concrete_inputs
     raw_name: Concrete Inputs

hta/configs/event_args_yaml_parser.py

@@ -17,11 +17,6 @@
 v1_0_0: YamlVersion = YamlVersion(1, 0, 0)
 
 
-ARGS_INDEX_FUNC: Callable[[Dict[str, AttributeSpec]], List[AttributeSpec]] = (
-    lambda available_args: [
-        available_args[k] for k in ["index::external_id", "index::python_id"]
-    ]
-)
 ARGS_INPUT_SHAPE_FUNC: Callable[[Dict[str, AttributeSpec]], List[AttributeSpec]] = (
     lambda available_args: [
         available_args[k]
@@ -79,7 +74,7 @@
         + ARGS_BANDWIDTH_FUNC(available_args)
         + ARGS_SYNC_FUNC(available_args)
         + ARGS_INPUT_SHAPE_FUNC(available_args)
-        + ARGS_INDEX_FUNC(available_args)
+        + [available_args["index::external_id"]]
     )
 )
 

tests/test_call_stack.py

@@ -1,18 +1,12 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-# Unit tests for call_stack.py
 
 import unittest
 
 import pandas as pd
 
-from hta.common.call_stack import (
-    CallGraph,
-    CallStackGraph,
-    CallStackIdentity,
-    CallStackNode,
-)
+from hta.common.call_stack import CallStackGraph, CallStackIdentity, CallStackNode
 from hta.common.trace_filter import ZeroDurationFilter
 
 
@@ -87,7 +81,7 @@ def test_construct_call_graph_0_dur(self):
         self.assertDictEqual(nodes, self.nodes2)
 
     def test_sort_events(self):
-        index = [0, 1, 2, 3]
+        index = [1, 2, 3, 4]
         start = [0, 0, 5, 5]
         dur = [10, 5, 1, 5]
         stream = [-1, -1, -1, -1]
@@ -102,11 +96,11 @@ def test_sort_events(self):
             }
         )
         nodes = {
-            -1: CallStackNode(parent=-1, depth=-1, children=[0]),
-            0: CallStackNode(parent=-1, depth=0, children=[1, 3]),
-            1: CallStackNode(parent=0, depth=1, children=[]),
-            2: CallStackNode(parent=3, depth=2, children=[]),
-            3: CallStackNode(parent=0, depth=1, children=[2]),
+            -1: CallStackNode(parent=-1, depth=-1, children=[1]),
+            1: CallStackNode(parent=-1, depth=0, children=[2, 4]),
+            2: CallStackNode(parent=1, depth=1, children=[]),
+            4: CallStackNode(parent=1, depth=1, children=[3]),
+            3: CallStackNode(parent=4, depth=2, children=[]),
         }
         csg = CallStackGraph(df, self.csi)
         self.assertDictEqual(nodes, csg.get_nodes())
@@ -136,164 +130,6 @@ def test_node_depth(self):
         depth_from_csg = csg.get_depth().to_dict()
         depth_from_nodes = {idx: node.depth for idx, node in nodes.items() if idx >= 0}
         self.assertDictEqual(depth_from_csg, depth_from_nodes)
-        # Verify df is used
-        self.assertIsNotNone(df)
-
-
-class CallGraphTestCase(unittest.TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-
-        # Mock Trace class for testing
-        class MockTrace:
-            def __init__(self, traces):
-                self.traces = traces
-
-            def get_all_traces(self):
-                return self.traces.keys()
-
-            def get_trace(self, rank):
-                return self.traces[rank]
-
-        # Create test data for trim_trace_events
-        self.df_trim = pd.DataFrame(
-            {
-                "index": [0, 1, 2, 3],
-                "ts": [0, 2, 4, 6],
-                "dur": [10, 6, 8, 2],  # Child event 2 exceeds parent event 1's duration
-                "pid": [1, 1, 1, 1],
-                "tid": [1, 1, 1, 1],
-                "stream": [-1, -1, -1, -1],
-                "index_correlation": [-1, -1, -1, -1],
-                "python_id": [100, 101, 102, 103],
-                "python_parent_id": [-1, 100, 101, 102],
-            }
-        )
-
-        self.trace_mock = MockTrace({0: self.df_trim})
-
-    def test_trim_trace_events_basic(self):
-        # Given & when
-        CallGraph(self.trace_mock, pre_process_trace_data=True)
-
-        # Then
-        # Event 1: ts=2, dur=6, end=8
-        # Event 2: ts=4, dur=8 (original), should be trimmed to dur=4 to end at 8
-        self.assertEqual(self.df_trim.at[2, "dur"], 4)
-
-    def test_trim_trace_events_complex_hierarchy(self):
-        # Given
-        df_complex = pd.DataFrame(
-            {
-                "index": [0, 1, 2, 3, 4],
-                "ts": [0, 2, 4, 6, 7],
-                "dur": [15, 10, 10, 20, 8],
-                "pid": [1, 1, 1, 1, 1],
-                "tid": [1, 1, 1, 1, 1],
-                "stream": [-1, -1, -1, -1, -1],
-                "index_correlation": [-1, -1, -1, -1, -1],
-                "python_id": [100, 101, 102, 103, 104],
-                "python_parent_id": [-1, 100, 101, 102, 102],
-            }
-        )
-
-        trace_complex = type(self.trace_mock)({0: df_complex})
-
-        # When
-        CallGraph(trace_complex, pre_process_trace_data=True)
-
-        # Then
-        # Event 1: ts=2, dur=10, end=12
-        # Event 2: ts=4, dur=10, should be trimmed to dur=8 to end at 12
-        # Event 3: ts=6, dur=20, should be trimmed to dur=6 to end at 12
-        # Event 4: ts=7, dur=8, should be trimmed to dur=5 to end at 12
-        self.assertEqual(df_complex.at[2, "dur"], 8)
-        self.assertEqual(df_complex.at[3, "dur"], 6)
-        self.assertEqual(df_complex.at[4, "dur"], 5)
-
-    def test_trim_trace_events_different_threads(self):
-        # Given
-        df_threads = pd.DataFrame(
-            {
-                "index": [0, 1, 2, 3],
-                "ts": [0, 2, 4, 6],
-                "dur": [10, 6, 8, 2],
-                "pid": [1, 1, 1, 1],
-                "tid": [1, 1, 2, 2],  # Events 2 and 3 are in a different thread
-                "stream": [-1, -1, -1, -1],
-                "index_correlation": [-1, -1, -1, -1],
-                "python_id": [100, 101, 102, 103],
-                "python_parent_id": [-1, 100, 101, 102],
-            }
-        )
-
-        trace_threads = type(self.trace_mock)({0: df_threads})
-
-        # When
-        CallGraph(trace_threads, pre_process_trace_data=True)
-
-        # Then
-        # Event 2 should not be trimmed because it's in a different thread than its parent
-        self.assertEqual(df_threads.at[2, "dur"], 8)
-
-    def test_trim_trace_events_no_trimming_needed(self):
-        # Given
-        df_no_trim = pd.DataFrame(
-            {
-                "index": [0, 1, 2, 3],
-                "ts": [0, 2, 4, 6],
-                "dur": [10, 6, 3, 1],  # All child events end before their parents
-                "pid": [1, 1, 1, 1],
-                "tid": [1, 1, 1, 1],
-                "stream": [-1, -1, -1, -1],
-                "index_correlation": [-1, -1, -1, -1],
-                "python_id": [100, 101, 102, 103],
-                "python_parent_id": [-1, 100, 101, 102],
-            }
-        )
-
-        trace_no_trim = type(self.trace_mock)({0: df_no_trim})
-
-        # When
-        CallGraph(trace_no_trim, pre_process_trace_data=True)
-
-        # Then
-        # Durations should remain unchanged
-        self.assertEqual(df_no_trim.at[1, "dur"], 6)
-        self.assertEqual(df_no_trim.at[2, "dur"], 3)
-        self.assertEqual(df_no_trim.at[3, "dur"], 1)
-
-    def test_trim_trace_events_multiple_children(self):
-        # Given
-        df_multi_children = pd.DataFrame(
-            {
-                "index": [0, 1, 2, 3, 4],
-                "ts": [0, 2, 3, 6, 8],
-                "dur": [
-                    10,
-                    8,
-                    3,
-                    2,
-                    5,
-                ],  # Child event 4 exceeds parent event 1's duration
-                "pid": [1, 1, 1, 1, 1],
-                "tid": [1, 1, 1, 1, 1],
-                "stream": [-1, -1, -1, -1, -1],
-                "index_correlation": [-1, -1, -1, -1, -1],
-                "python_id": [100, 101, 102, 103, 104],
-                "python_parent_id": [-1, 100, 101, 101, 101],
-            }
-        )
-
-        trace_multi = type(self.trace_mock)({0: df_multi_children})
-
-        # When
-        CallGraph(trace_multi, pre_process_trace_data=True)
-
-        # Then
-        # Event 1: ts=2, dur=8, end=10
-        # Event 4: ts=8, dur=5, should be trimmed to dur=2 to end at 10
-        self.assertEqual(df_multi_children.at[4, "dur"], 2)
 
 
 if __name__ == "__main__":  # pragma: no cover