From 1cea8d47317af76ca66c99f7d83b645dd379fb9b Mon Sep 17 00:00:00 2001 From: Benedykt Bela Date: Fri, 7 Nov 2025 14:27:33 +0200 Subject: [PATCH 1/4] Rename Trace class to TraceCollection Rename Trace class to TraceCollection and its file trace.py to trace_collection.py. Adjusted the whole code to work with changed names. Signed-off-by: Benedykt Bela --- benchmarks/trace_load_benchmark.py | 4 +-- examples/symbol_table_demo.py | 19 +++++----- hta/analyzers/breakdown_analysis.py | 22 ++++++------ hta/analyzers/communication_analysis.py | 6 ++-- hta/analyzers/critical_path_analysis.py | 24 ++++++------- hta/analyzers/cuda_kernel_analysis.py | 14 ++++---- hta/analyzers/cupti_counter_analysis.py | 8 ++--- hta/analyzers/straggler.py | 6 ++-- hta/analyzers/straggler_analysis.py | 6 ++-- hta/analyzers/timeline.py | 6 ++-- hta/analyzers/trace_counters.py | 30 ++++++++-------- hta/common/call_stack.py | 12 +++---- hta/common/execution_trace.py | 10 +++--- hta/common/trace_call_graph.py | 16 +++++---- hta/common/{trace.py => trace_collection.py} | 12 +++---- hta/trace_analysis.py | 6 ++-- hta/trace_diff.py | 36 +++++++++---------- tests/test_call_stack.py | 2 +- tests/test_correlation.py | 2 +- tests/test_custom_trace_parser.py | 6 ++-- tests/test_timeline.py | 4 +-- tests/test_trace_analysis.py | 6 ++-- tests/test_trace_call_graph.py | 10 +++--- tests/test_trace_filter.py | 6 ++-- tests/test_trace_parse.py | 38 ++++++++++++-------- 25 files changed, 166 insertions(+), 145 deletions(-) rename hta/common/{trace.py => trace_collection.py} (98%) diff --git a/benchmarks/trace_load_benchmark.py b/benchmarks/trace_load_benchmark.py index 8e8c9307..96728462 100644 --- a/benchmarks/trace_load_benchmark.py +++ b/benchmarks/trace_load_benchmark.py @@ -7,7 +7,7 @@ import pyperf -from hta.common.trace import parse_trace_file, Trace +from hta.common.trace_collection import parse_trace_file, TraceCollection from hta.common.trace_parser import set_default_trace_parsing_backend from hta.configs.config import logger @@ -40,7 +40,7 @@ def load_and_parse_trace( range_it = range(loops) t0 = pyperf.perf_counter() for _ in range_it: - trace = Trace(trace_dir=trace_dir) + trace = TraceCollection(trace_dir=trace_dir) trace.parse_traces(max_ranks=max_ranks, use_multiprocessing=use_multiprocessing) return pyperf.perf_counter() - t0 diff --git a/examples/symbol_table_demo.py b/examples/symbol_table_demo.py index e5856647..f5e63d70 100644 --- a/examples/symbol_table_demo.py +++ b/examples/symbol_table_demo.py @@ -18,7 +18,7 @@ import pandas as pd import plotly.express as px -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection path_to_hta = "~/HolisticTraceAnalysis" trace_dir: str = path_to_hta + "/tests/data/vision_transformer" @@ -32,13 +32,15 @@ def set_pandas_display_options(): pd.set_option("display.float_format", "{:.2f}".format) -def demo_statistics(trace: Trace, rank: int, k: Optional[int] = None) -> pd.DataFrame: +def demo_statistics( + trace: TraceCollection, rank: int, k: Optional[int] = None +) -> pd.DataFrame: """ Show the first k items of the kernels by duration in a specific rank's trace. Args: - trace: a Trace instance. + trace: a TraceCollection instance. rank: the rank to be analyzed. k: how many items to show in the output; If None, then show all items. @@ -98,8 +100,8 @@ def demo_visualization(df: pd.DataFrame, title: str, visualize: bool = False) -> logging.info(f"{title}\n{df}\n") -def load_trace(trace_dir, max_ranks) -> Trace: - trace = Trace(trace_dir=trace_dir) +def load_trace(trace_dir, max_ranks) -> TraceCollection: + trace = TraceCollection(trace_dir=trace_dir) trace.parse_traces(max_ranks=max_ranks, use_multiprocessing=True) return trace @@ -107,7 +109,7 @@ def load_trace(trace_dir, max_ranks) -> Trace: def run_demo( trace_dir: str, max_ranks: int, - preloaded_trace: Optional[Trace] = None, + preloaded_trace: Optional[TraceCollection] = None, ): """_summary_ @@ -115,7 +117,8 @@ def run_demo( trace_name (str): name of the trace base_trace_dir (str): the base path of the traces max_ranks (int): maximum number of ranks to be analyzed - preloaded_trace (Optional[Trace], optional): a preloaded trace. Defaults to None. + preloaded_trace (Optional[TraceCollection], optional): a preloaded collection of traces. + Defaults to None. """ # load the trace if preloaded_trace is None: @@ -171,7 +174,7 @@ def run_demo( ) -def trace_info(trace: Trace): +def trace_info(trace: TraceCollection): rank = next(iter(trace.traces)) df = trace.get_trace(rank) logging.info(f"\n===Dataframe of Rank {rank}") diff --git a/hta/analyzers/breakdown_analysis.py b/hta/analyzers/breakdown_analysis.py index c9708ec8..75d4871a 100644 --- a/hta/analyzers/breakdown_analysis.py +++ b/hta/analyzers/breakdown_analysis.py @@ -23,7 +23,7 @@ # import statement used without the "if TYPE_CHECKING" guard will cause a circular # dependency with trace_analysis.py causing mypy to fail and should not be removed. if TYPE_CHECKING: - from hta.common.trace import Trace + from hta.common.trace_collection import TraceCollection # This configures the threshold under which we consider gaps between # kernels to be due to realistic delays in launching back-back kernels on the GPU @@ -36,7 +36,7 @@ def __init__(self): @classmethod def get_gpu_kernel_breakdown( cls, - t: "Trace", + t: "TraceCollection", visualize: bool = True, duration_ratio: float = 0.8, num_kernels: int = 10, @@ -216,13 +216,13 @@ def get_gpu_kernel_breakdown( def _get_gpu_kernel_interval_dataframe( cls, trace_df: pd.DataFrame, - t: "Trace", + t: "TraceCollection", ) -> pd.DataFrame: """Obtains all GPU kernels in the trace dataframe and assigns them an interval index that can be used for analyzing overlap. @args: trace_df (pd.DataFrame) : trace df for specific rank Please make sure this includes "end" column. - @args: t (Trace) : trace object + @args: t (TraceCollection) : object representing collection of traces. Returns: pd.DataFrame with GPU kernels subset with an interval index of [start, end) intervals. @@ -238,13 +238,13 @@ def _get_gpu_kernel_interval_dataframe( def _get_gpu_user_anno_interval_dataframe( cls, trace_df: pd.DataFrame, - t: "Trace", + t: "TraceCollection", ) -> Optional[pd.DataFrame]: """Obtains all GPU user annotations in the trace dataframe and assigns them an interval index that can be used for analyzing overlap. @args: trace_df (pd.DataFrame) : trace df for specific rank Please make sure this includes "end" column. - @args: t (Trace) : trace object + @args: t (TraceCollection) : object representing collection of traces. Returns: pd.DataFrame with GPU kernels subset with an interval index of [start, end) intervals. @@ -309,7 +309,7 @@ def _associate_gpu_kernels_with_user_annotations( @classmethod def get_gpu_kernels_with_user_annotations( cls, - t: "Trace", + t: "TraceCollection", rank: int, expand_names: bool = True, shortern_names: bool = True, @@ -345,7 +345,7 @@ def get_gpu_kernels_with_user_annotations( @classmethod def get_gpu_user_annotation_breakdown( cls, - t: "Trace", + t: "TraceCollection", use_gpu_annotation: bool = True, visualize: bool = True, duration_ratio: float = 0.8, @@ -639,7 +639,9 @@ def _get_idle_time_for_kernels(cls, kernels_df: pd.DataFrame) -> Tuple[int, int] return kernel_time - kernel_run_time, kernel_time @classmethod - def get_temporal_breakdown(cls, t: "Trace", visualize: bool = True) -> pd.DataFrame: + def get_temporal_breakdown( + cls, t: "TraceCollection", visualize: bool = True + ) -> pd.DataFrame: """ Temporal breakdown implementation. See `get_temporal_breakdown` in `trace_analysis.py` for details. """ @@ -802,7 +804,7 @@ def _analyze_idle_time_for_stream( @classmethod def get_idle_time_breakdown( cls, - t: "Trace", + t: "TraceCollection", consecutive_kernel_delay: int, rank: int = 0, streams: Optional[List[int]] = None, diff --git a/hta/analyzers/communication_analysis.py b/hta/analyzers/communication_analysis.py index 74031119..109ec1c2 100644 --- a/hta/analyzers/communication_analysis.py +++ b/hta/analyzers/communication_analysis.py @@ -13,7 +13,7 @@ # import statement used without the "if TYPE_CHECKING" guard will cause a circular # dependency with trace_analysis.py causing mypy to fail and should not be removed. if TYPE_CHECKING: - from hta.trace import Trace + from hta.common.trace_collection import TraceCollection class CommunicationAnalysis: @@ -21,7 +21,9 @@ def __init__(self): pass @classmethod - def get_comm_comp_overlap(cls, t: "Trace", visualize: bool = True) -> pd.DataFrame: + def get_comm_comp_overlap( + cls, t: "TraceCollection", visualize: bool = True + ) -> pd.DataFrame: """ Communication analysis implementation. See `get_comm_comp_overlap` in `trace_analysis.py` for details. """ diff --git a/hta/analyzers/critical_path_analysis.py b/hta/analyzers/critical_path_analysis.py index 9444271e..d38e7015 100644 --- a/hta/analyzers/critical_path_analysis.py +++ b/hta/analyzers/critical_path_analysis.py @@ -27,7 +27,7 @@ # from hta.common.trace_call_graph import CallGraph, CallStackGraph, DeviceType from hta.common.call_stack import CallGraph, CallStackGraph, DeviceType -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_symbol_table import decode_symbol_id_to_symbol_name from hta.configs.config import logger from hta.utils.utils import is_comm_kernel @@ -185,8 +185,8 @@ def _add_zero_weight_launch_edges(self) -> bool: def __init__( self, - t: Optional["Trace"], - t_full: "Trace", + t: Optional["TraceCollection"], + t_full: "TraceCollection", rank: int, G=None, data_load_events=None, @@ -200,8 +200,8 @@ def __init__( For (2) the networkx.DiGraph object G is utilized, see restore_cpgraph() function in this file. Args: - t (Trace): Clipped trace object focussing on region of interest. - t_full (Trace): Full Trace object. + t (TraceCollection): Clipped collection of traces focussing on region of interest. + t_full (TraceCollection): Full TraceCollection object. rank (int): Rank to perform analysis on. G (networkx.DiGraph): An optional DiGraph object. data_load_events (List[str]): List of events (regex) to be considered as @@ -1699,7 +1699,7 @@ def save(self, out_dir: str) -> str: return zip_filename -def restore_cpgraph(zip_filename: str, t_full: "Trace", rank: int) -> CPGraph: +def restore_cpgraph(zip_filename: str, t_full: "TraceCollection", rank: int) -> CPGraph: """Restores the critical path graph object from a zip file. The graph will already be constructed in this case. You can however run critical_path() again and modify the graph etc. @@ -1757,7 +1757,7 @@ def __init__(self): @classmethod def critical_path_analysis( cls, - t: "Trace", + t: "TraceCollection", rank: int, annotation: str, instance_id: Union[Optional[int], Tuple[int, int]], @@ -1772,7 +1772,7 @@ def critical_path_analysis( by passing annotation='ProfilerStep'. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. rank (int): rank to analyze for the critical path. annotation (str): a trace annotation to limit the analysis to, for example "ProfilerStep" would match all annotations that @@ -1871,7 +1871,7 @@ def critical_path_analysis( logger.info(f"Clipped dataframe has {len(clipped_df)} events") # XXX This is a bit hacky but CallGraph does not really support a way - # to specify a dataframe, just supports passing Trace object + # to specify a dataframe, just supports passing TraceCollection object t_copy = deepcopy(t) t_copy.traces[rank] = clipped_df t1 = time.perf_counter() @@ -1893,7 +1893,7 @@ def _is_zero_weight_launch_edge(e: CPEdge) -> bool: @classmethod def overlay_critical_path_analysis( cls, - t: "Trace", + t: "TraceCollection", rank: int, critical_path_graph: CPGraph, output_dir: str, @@ -1906,7 +1906,7 @@ def overlay_critical_path_analysis( for visualization. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. rank (int): rank to generate the time series for. critical_path_graph: Critical Path Graph object generated previously. output_dir (str): Output directory to store overlaid trace. @@ -1972,7 +1972,7 @@ def get_flow_event( is_critical = e in critical_path_graph.critical_path_edges_set - return Trace.flow_event( + return TraceCollection.flow_event( id=flow_id, pid=event["pid"], tid=event["tid"], diff --git a/hta/analyzers/cuda_kernel_analysis.py b/hta/analyzers/cuda_kernel_analysis.py index 5270510e..eccfd37f 100644 --- a/hta/analyzers/cuda_kernel_analysis.py +++ b/hta/analyzers/cuda_kernel_analysis.py @@ -9,9 +9,9 @@ import pandas as pd import plotly.express as px -from hta.common.trace import Trace from hta.common.trace_call_graph import CallGraph +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger from hta.utils.checker import is_valid_directory from plotly import graph_objects as go @@ -24,7 +24,7 @@ def __init__(self): @classmethod def get_frequent_cuda_kernel_sequences( cls, - t: Trace, + t: TraceCollection, operator_name: str, output_dir: str, min_pattern_len: int = 3, @@ -43,7 +43,7 @@ def get_frequent_cuda_kernel_sequences( (3) Overlay the top_k identified repeated patterns back to the trace file. Args: - t (Trace): A trace object containing the trace data. + t (TraceCollection): A trace collection object containing data of multiple traces. operator_name (str): Name of the cpu operators which launch the cuda kernels. output_dir (str): Path to the folder containing the new trace file with overlaid top k frequent patterns. min_pattern_len (int): Minimum length of the cuda kernel sequences that should be identified. @@ -134,7 +134,7 @@ def get_frequent_cuda_kernel_sequences( @classmethod def _generate_frequent_pattern_results( cls, - t: "Trace", + t: "TraceCollection", pattern_counts: Dict[Tuple[int, ...], int], pattern_durations: Dict[Tuple[int, ...], List[int]], pattern_occurrences: Dict[Tuple[int, ...], Set[int]], @@ -226,7 +226,7 @@ def _generate_frequent_pattern_results( @classmethod def _overlay_frequent_patterns_with_trace( cls, - t: "Trace", + t: "TraceCollection", patterns_df: pd.DataFrame, rank: int = 0, compress_other_kernels: bool = True, @@ -404,7 +404,7 @@ def visualize_cuda_launch_kernel_info( @classmethod def get_aten_op_kernels_and_delay( cls, - trace: Trace, + trace: TraceCollection, ranks: Optional[List[int]] = None, sort_by: Optional[List[str]] = None, ) -> Dict[int, pd.DataFrame]: @@ -536,7 +536,7 @@ def find_deepest_aten_op( @classmethod def cuda_kernel_launch_stats( cls, - t: Trace, + t: TraceCollection, ranks: Optional[List[int]] = None, runtime_cutoff: int = 50, launch_delay_cutoff: int = 100, diff --git a/hta/analyzers/cupti_counter_analysis.py b/hta/analyzers/cupti_counter_analysis.py index bc5b08e7..d5c095cc 100644 --- a/hta/analyzers/cupti_counter_analysis.py +++ b/hta/analyzers/cupti_counter_analysis.py @@ -5,9 +5,9 @@ from typing import Dict, List, Optional import pandas as pd -from hta.common.trace import Trace from hta.common.trace_call_graph import CallGraph +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger CUDA_SASS_INSTRUCTION_COUNTER_FLOPS: Dict[str, float] = { @@ -25,7 +25,7 @@ def __init__(self): @classmethod def _get_counter_data_with_operators_for_rank( cls, - t: Trace, + t: TraceCollection, rank: int, cg: CallGraph, ) -> Optional[pd.DataFrame]: @@ -90,13 +90,13 @@ def stringify_op_stack(ops: List[int]) -> List[str]: @classmethod def get_counter_data_with_operators( cls, - t: Trace, + t: TraceCollection, ranks: Optional[List[int]] = None, ) -> List[pd.DataFrame]: """Correlates the Kernel counter events with pytorch operators using the callgraph. Args: - t (Trace): trace object + t (TraceCollection): Collection of multiple traces. ranks (List[int]): List of ranks on which to run the analysis. Default = [0]. Returns: A list of dataframes, one per rank, containing kernel name, diff --git a/hta/analyzers/straggler.py b/hta/analyzers/straggler.py index 48e83831..9dec3422 100644 --- a/hta/analyzers/straggler.py +++ b/hta/analyzers/straggler.py @@ -9,15 +9,15 @@ import plotly.express as px from hta.analyzers.timeline import _get_unique_values, plot_timeline_gpu_kernels -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_symbol_table import TraceSymbolTable -def extract_iteration_info(trace: Trace) -> pd.DataFrame: +def extract_iteration_info(trace: TraceCollection) -> pd.DataFrame: """Extract the iteration information from a trace. Args: - trace (Trace) : a trace object + trace (TraceCollection) : Collection of multiple traces. Returns: a DataFrame that contains all the iteration information with rank as a column. diff --git a/hta/analyzers/straggler_analysis.py b/hta/analyzers/straggler_analysis.py index cc9a30c2..933e3002 100644 --- a/hta/analyzers/straggler_analysis.py +++ b/hta/analyzers/straggler_analysis.py @@ -15,7 +15,7 @@ # import statement used without the "if TYPE_CHECKING" guard will cause a circular # dependency with trace_analysis.py causing mypy to fail and should not be removed. if TYPE_CHECKING: - from hta.common.trace import Trace + from hta.common.trace_collection import TraceCollection class StragglerAnalysis: @@ -23,7 +23,7 @@ def __init__(self): pass @classmethod - def get_profiler_steps(cls, t: "Trace") -> List[int]: + def get_profiler_steps(cls, t: "TraceCollection") -> List[int]: """ Profiler steps implementation. Returns the list of profiler steps. """ @@ -34,7 +34,7 @@ def get_profiler_steps(cls, t: "Trace") -> List[int]: @classmethod def get_potential_stragglers( cls, - t: "Trace", + t: "TraceCollection", profiler_steps: Optional[List[int]] = None, num_candidates: int = 2, visualize: bool = False, diff --git a/hta/analyzers/timeline.py b/hta/analyzers/timeline.py index 26dc5045..e8452b2b 100644 --- a/hta/analyzers/timeline.py +++ b/hta/analyzers/timeline.py @@ -9,7 +9,7 @@ import pandas as pd import plotly.express as px -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_symbol_table import TraceSymbolTable from hta.configs.config import logger @@ -222,7 +222,7 @@ def plot_timeline_gpu_kernels( def plot_timeline_gpu_kernels_from_trace( title: str, - trace_data: Trace, + trace_data: TraceCollection, ranks: Optional[List[int]] = None, iterations: Optional[List[int]] = None, streams: Optional[List[int]] = None, @@ -233,7 +233,7 @@ def plot_timeline_gpu_kernels_from_trace( Args: title (str): a title for the timeline plot - trace_data (Trace): a Trace object + trace_data (TraceCollection): a TraceCollection object ranks (List[int]): filter the input DataFrame with the given set of ranks; use all ranks if None. iterations (List[int]): filter the input DataFrame with the given set of iterations; use all iterations if None. streams (List[int]): filter the input DataFrame with the given set of streams; use all streams if None. diff --git a/hta/analyzers/trace_counters.py b/hta/analyzers/trace_counters.py index babf1a9d..cd13406e 100644 --- a/hta/analyzers/trace_counters.py +++ b/hta/analyzers/trace_counters.py @@ -6,7 +6,7 @@ import pandas as pd -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger from hta.utils.utils import get_kernel_type, get_memory_kernel_type, KernelType @@ -18,7 +18,7 @@ def __init__(self): @classmethod def _get_queue_length_time_series_for_rank( cls, - t: "Trace", + t: "TraceCollection", rank: int, ) -> Optional[pd.DataFrame]: """ @@ -31,7 +31,7 @@ def _get_queue_length_time_series_for_rank( 2. Decremented when a CUDA kernel/memcopy operation executes on a stream. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. rank (int): rank to generate the time series for. Returns: @@ -96,7 +96,7 @@ def _get_queue_length_time_series_for_rank( @classmethod def get_queue_length_time_series( cls, - t: "Trace", + t: "TraceCollection", ranks: Optional[List[int]] = None, ) -> Dict[int, pd.DataFrame]: """ @@ -108,7 +108,7 @@ def get_queue_length_time_series( 2. Decremented when a CUDA kernel/memcopy operation executes on a stream. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. rank (int): rank to perform this analysis for. Returns: @@ -169,14 +169,14 @@ def get_queue_length_summary_from_time_series( @classmethod def get_queue_length_summary( cls, - t: "Trace", + t: "TraceCollection", ranks: Optional[List[int]] = None, ) -> Optional[pd.DataFrame]: """ Returns an (optional) dataframewith queue length statistics per CUDA stream and rank. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. ranks (list of int): ranks to perform this analysis. Returns: @@ -194,7 +194,7 @@ def get_queue_length_summary( @classmethod def get_time_spent_blocked_on_full_queue( cls, - t: "Trace", + t: "TraceCollection", queue_length_dict: Dict[int, pd.DataFrame], max_queue_length: int, ) -> Optional[pd.DataFrame]: @@ -205,7 +205,7 @@ def get_time_spent_blocked_on_full_queue( up the time spent on all streams where the queue is full (see max_queue_length) Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. queue_length_dict (Dict[int, pd.DataFrame]): A dictionary of rank -> time series with the queue length of each CUDA stream. This is the output of get_queue_length_time_series(). max_queue_length (int): Max kernel launch queue length. @@ -257,14 +257,14 @@ def get_time_spent_blocked_on_full_queue( @classmethod def _get_memory_bw_time_series_for_rank( - cls, t: "Trace", rank: int + cls, t: "TraceCollection", rank: int ) -> Optional[pd.DataFrame]: """ Returns time series for the memory bandwidth of memory copy and memory set operations for specified rank. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. rank (int): rank to generate the time series for. Returns: @@ -325,14 +325,14 @@ def _get_memory_bw_time_series_for_rank( @classmethod def get_memory_bw_time_series( cls, - t: "Trace", + t: "TraceCollection", ranks: Optional[List[int]] = None, ) -> Dict[int, pd.DataFrame]: """ Returns a dictionary of rank -> time series for the memory bandwidth. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. ranks (list of int): ranks to perform this analysis for. Returns: @@ -359,7 +359,7 @@ def get_memory_bw_time_series( @classmethod def get_memory_bw_summary( cls, - t: "Trace", + t: "TraceCollection", ranks: Optional[List[int]] = None, ) -> Optional[pd.DataFrame]: """ @@ -367,7 +367,7 @@ def get_memory_bw_summary( tracked memory ops are MemcpyDtoH, MemcpyHtoD, MemcpyDtoD and MemSet. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input collection of traces data structure. ranks (list of int): ranks to perform this analysis for. Returns: diff --git a/hta/common/call_stack.py b/hta/common/call_stack.py index 9996246c..3e4b305f 100644 --- a/hta/common/call_stack.py +++ b/hta/common/call_stack.py @@ -13,7 +13,7 @@ import numpy as np import pandas as pd -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_filter import Filter NON_EXISTENT_NODE_INDEX = -2 @@ -450,7 +450,7 @@ class CallGraph: includes all CallStackGraph objects in the trace and provides further further query and statistic APIs. Attributes: - trace_data (Trace) : the trace data represented in a Trace object, which + trace_data (TraceCollection) : the traces data represented in a TraceCollection object, which contains multiple DataFrame objects mapping the traces of each trainer. call_stacks (List[CallStackGraph]) : a list of per-thread CallStackGraph objects. mapping (pd.DataFrame) : the mapping from CallStackIdentity to CallStackGraph using a DataFrame @@ -458,16 +458,16 @@ class CallGraph: def __init__( self, - trace: Trace, + trace: TraceCollection, 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, ) -> None: - """Construct a CallGraph from a Trace object + """Construct a CallGraph from a TraceCollection object Args: - trace (Trace): the trace data used to construct this CallGraph object. + trace (TraceCollection): the traces data used to construct this CallGraph object. ranks (List[int]) : filter the traces using the given set of ranks. Using all ranks if None. filter_func (Callable) : used to preprocess the trace events and filter events out. Please see filters in hta/common/trace_filter.py for details. remapped_tids (Dict[Dict[int, int]]) : a dictionary that stores per-rank thread ID remappings. @@ -477,7 +477,7 @@ def __init__( Raises: ValueError: the trace data is invalid. """ - self.trace_data: Trace = trace + self.trace_data: TraceCollection = trace self.mapping: pd.DataFrame = pd.DataFrame() self.call_stacks: List[CallStackGraph] = [] self.pre_process_trace_data = pre_process_trace_data diff --git a/hta/common/execution_trace.py b/hta/common/execution_trace.py index 7d7e11e3..08d20b71 100644 --- a/hta/common/execution_trace.py +++ b/hta/common/execution_trace.py @@ -13,7 +13,7 @@ import numpy as np import pandas as pd -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger from hta.utils.utils import normalize_path @@ -96,12 +96,14 @@ def _et_has_overlap(trace_df: pd.DataFrame, et: ExecutionTrace) -> bool: return has_overlap -def correlate_execution_trace(trace: Trace, rank: int, et: ExecutionTrace) -> None: +def correlate_execution_trace( + trace: TraceCollection, rank: int, et: ExecutionTrace +) -> None: """Correlate the trace from a specific rank with Execution Trace object. Args: - trace (Trace): Trace object loaded using `TraceAnalysis(trace_dir=trace_dir)` - or other method. + trace (TraceCollection): TraceCollection object loaded using + `TraceAnalysis(trace_dir=trace_dir)` or other method. rank (int): Rank to correlate with. et (ExecutionTrace): An Execution Trace object to correlate with. diff --git a/hta/common/trace_call_graph.py b/hta/common/trace_call_graph.py index 0925bd09..c67e4592 100644 --- a/hta/common/trace_call_graph.py +++ b/hta/common/trace_call_graph.py @@ -3,8 +3,8 @@ import pandas as pd -from hta.common.trace import get_cpu_gpu_correlation, Trace from hta.common.trace_call_stack import CallStackGraph, CallStackIdentity, CallStackNode +from hta.common.trace_collection import get_cpu_gpu_correlation, TraceCollection from hta.common.trace_symbol_table import TraceSymbolTable from hta.common.types import DeviceType, infer_device_type from hta.configs.config import logger @@ -27,7 +27,7 @@ class CallGraph: launches, AllToAll communications, etc. Attributes: - trace_data (Trace) : A container consisting of a mapping from each trainer to a Trace DataFrame. + trace_data (TraceCollection) : A container consisting of a mapping from each trainer to a Trace object. ranks (List[int]) : A list of trainer IDs (i.e., ranks). rank_to_stacks (Dict[int, Dict[CallStackIdentity, CallStackGraph]]): a map from ranks to their CallStackGraph objects, which are represented as another map from CallStackIdentity to CallStackGraph objects. @@ -54,11 +54,13 @@ class CallGraph: "kernel_span", ] - def __init__(self, trace: Trace, ranks: Optional[List[int]] = None) -> None: - """Construct a CallGraph object from a Trace object. + def __init__( + self, trace: TraceCollection, ranks: Optional[List[int]] = None + ) -> None: + """Construct a CallGraph object from a TraceCollection object. Args: - trace (Trace): The Trace object used to construct this CallGraph object. + trace (TraceCollection): The TraceCollection object used to construct this CallGraph object. ranks (List[int]) : Only construct the CallGraph objects for the given set of ranks. When not provided, uses all available ranks in . Caution: this might be time-consuming. @@ -66,7 +68,7 @@ def __init__(self, trace: Trace, ranks: Optional[List[int]] = None) -> None: Raises: ValueError: If the trace data is invalid. """ - self.trace_data: Trace = trace + self.trace_data: TraceCollection = trace _ranks: List[int] = self.trace_data.get_ranks() self.ranks: List[int] = [r for r in ranks if r in _ranks] if ranks else _ranks if (len(self.ranks)) == 0: @@ -115,7 +117,7 @@ def from_dataframe( Returns: A CallGraph object. """ - t = Trace(trace_files={}, trace_dir="") + t = TraceCollection(trace_files={}, trace_dir="") t.symbol_table = ( symbol_table if symbol_table else TraceSymbolTable.create_from_df(df) ) diff --git a/hta/common/trace.py b/hta/common/trace_collection.py similarity index 98% rename from hta/common/trace.py rename to hta/common/trace_collection.py index 450d2aac..664ac0df 100644 --- a/hta/common/trace.py +++ b/hta/common/trace_collection.py @@ -338,12 +338,12 @@ def add_fwd_bwd_links(df: pd.DataFrame) -> None: logger.debug(f"Time taken to add fwd_bwd links: {t1 - t0 :.2f} seconds") -class Trace: +class TraceCollection: """ A container for the traces collected for a distributed ML training job. An ML training job can have multiple trace collections. Each of those trace collections maps to - one Trace object. + one TraceCollection object. Attributes: @@ -364,7 +364,7 @@ def __init__( parser_config: Optional[ParserConfig] = None, ) -> None: """ - The constructor of a Trace object. + The constructor of a TraceCollection object. Args: trace_files: Optional[Union[List[str], Dict[int, str]]]: either a list of trace file names or a map from rank to trace file names. When a list is provided, HTA will infer the ranks by reading the trace file metadata. @@ -629,7 +629,7 @@ def get_trace(self, rank: int) -> pd.DataFrame: The DataFrame for the given rank. Raises: - ValueError when this Trace object doesn't have trace for the given rank. + ValueError when this TraceCollection object doesn't have trace for the given rank. """ if rank not in self.traces: logger.error(f"get_rank_trace - no trace for rank {rank}") @@ -656,7 +656,7 @@ def get_raw_trace_for_one_rank(self, rank: int = 0) -> Dict[str, Any]: The raw content of the trace file of the given rank. Raises: - ValueError when this Trace object doesn't have trace for the given rank. + ValueError when this TraceCollection object doesn't have trace for the given rank. """ if rank not in self.trace_files: logger.error(f"get_rank_trace - no trace for rank {rank}") @@ -956,7 +956,7 @@ def get_trace_start_unixtime_ns(self, rank: int) -> int: Returns: Unixtime (nanoseconds) of trace start (int) - Raises: ValueError when this Trace object doesn't have trace for the given rank. + Raises: ValueError when this TraceCollection object doesn't have trace for the given rank. """ if rank not in self.traces: err_msg: str = f"No trace found for rank {rank}" diff --git a/hta/trace_analysis.py b/hta/trace_analysis.py index 56704af1..893f639c 100644 --- a/hta/trace_analysis.py +++ b/hta/trace_analysis.py @@ -17,7 +17,7 @@ from hta.analyzers.straggler_analysis import StragglerAnalysis from hta.analyzers.trace_counters import TraceCounters from hta.common.constants import CUDA_MAX_LAUNCH_QUEUE_PER_STREAM -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger from hta.configs.default_values import DEFAULT_TRACE_DIR @@ -34,7 +34,7 @@ def __init__( trace_dir: str = DEFAULT_TRACE_DIR, include_last_profiler_step: Optional[bool] = False, ): - self.t = Trace(trace_files, trace_dir) + self.t = TraceCollection(trace_files, trace_dir) self.t.load_traces(include_last_profiler_step) assert self.t.is_parsed is True @@ -688,7 +688,7 @@ def critical_path_analysis( by passing annotation='ProfilerStep'. See notes for how to pick the iteration. Args: - t (Trace): Input trace data structure. + t (TraceCollection): Input trace collection data structure. rank (int): rank to analyze for the critical path. annotation (str): a trace annotation to limit the analysis to, for example "ProfilerStep" would match all annotations that diff --git a/hta/trace_diff.py b/hta/trace_diff.py index 830ed398..99f38989 100644 --- a/hta/trace_diff.py +++ b/hta/trace_diff.py @@ -10,7 +10,7 @@ import pandas as pd import plotly.graph_objects as go -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.configs.config import logger from hta.utils.utils import flatten_column_names, shorten_name @@ -26,26 +26,26 @@ class DeviceType(Enum): class LabeledTrace: - """A wrapper class for the Trace class which assigns a label to each trace object. + """A wrapper class for the TraceCollection class which assigns a label to each trace object. Attributes: label (str): a label attached to the trace. - t (Trace): a Trace object that contains the trace data. + t (TraceCollection): a TraceCollection object that contains data of multiple traces. iteration_df (pd.DataFrame): a DataFrame that contains the """ def __init__( self, label: str = None, - t: Optional[Trace] = None, + t: Optional[TraceCollection] = None, trace_dir: Optional[str] = None, ): - """Construct a LabeledTrace from either a Trace object or trace files in trace_dir.""" + """Construct a LabeledTrace from either a TraceCollection object or trace files in trace_dir.""" self.label = label if label else f"t{random.randint(0,10)}" if t is not None: self.t = t elif trace_dir is not None and os.path.isdir(trace_dir): - self.t = Trace(trace_dir=trace_dir) + self.t = TraceCollection(trace_dir=trace_dir) else: raise ValueError( "either a trace object or a valid trace dir must be provided in LabeledTrace.__init__()" @@ -213,13 +213,13 @@ def get_ops_summary(self, ops: pd.DataFrame) -> pd.DataFrame: def _trace_argument_adapter( - t: Union[LabeledTrace, Trace, TraceDir], default_label: str + t: Union[LabeledTrace, TraceCollection, TraceDir], default_label: str ) -> LabeledTrace: """A helper function to construct a LabeledTrace from several argument types.""" lt: LabeledTrace if isinstance(t, TraceDir): lt = LabeledTrace(label=default_label, trace_dir=t) - elif isinstance(t, Trace): + elif isinstance(t, TraceCollection): lt = LabeledTrace(label=default_label, t=t) elif isinstance(t, LabeledTrace): lt = t @@ -245,13 +245,13 @@ def compare_traces( Compare the operators/kernels counts and total duration of two traces. Args: - control (Union[LabeledTrace, Trace, TraceDir]): the control trace. - A string or Trace object that defines the control trace. Possible values can be: + control (Union[LabeledTrace, TraceCollection, TraceDir]): the control trace. + A string or TraceCollection object that defines the control trace. Possible values can be: 1. a str (TraceDir) that points to parent path of the trace files. - 2. a Trace object that contains the trace records and metadata. - 3. a LabeledTrace object, which is a wrapper of the Trace object with a label to identify the trace. + 2. a TraceCollection object that contains the trace records and metadata. + 3. a LabeledTrace object, which is a wrapper of the TraceCollection object with a label to identify the trace. - test (Union[LabeledTrace, Trace, TraceDir]): the test trace. + test (Union[LabeledTrace, TraceCollection, TraceDir]): the test trace. Similar to the control trace except it defines the test trace. control_rank (Optional[Union[int, List[int]]]): Specify which ranks of the control trace to use. @@ -363,13 +363,13 @@ def ops_diff( Get the operator difference between two traces. Args: - control (Union[LabeledTrace, Trace, TraceDir]): The control trace. - A string or Trace object that defines the control trace. A possible value can be: + control (Union[LabeledTrace, TraceCollection, TraceDir]): The control trace. + A string or TraceCollection object that defines the control trace. A possible value can be: 1. a str (TraceDir) that points to parent path of the trace files. - 2. a Trace object that contains the trace records and metadata. - 3. a LabeledTrace object, which is a wrapper of the Trace object with a label to identify the trace. + 2. a TraceCollection object that contains the trace records and metadata. + 3. a LabeledTrace object, which is a wrapper of the TraceCollection object with a label to identify the trace. - test (Union[LabeledTrace, Trace, TraceDir]): The test trace. + test (Union[LabeledTrace, TraceCollection, TraceDir]): The test trace. Similar to the control trace except it defines the test trace. control_rank (Optional[Union[int, List[int]]]): Specify which ranks diff --git a/tests/test_call_stack.py b/tests/test_call_stack.py index fba55254..96eb85b9 100644 --- a/tests/test_call_stack.py +++ b/tests/test_call_stack.py @@ -144,7 +144,7 @@ class CallGraphTestCase(unittest.TestCase): def setUp(self) -> None: super().setUp() - # Mock Trace class for testing + # Mock TraceCollection class for testing class MockTrace: def __init__(self, traces): self.traces = traces diff --git a/tests/test_correlation.py b/tests/test_correlation.py index 084392cd..13b13d25 100644 --- a/tests/test_correlation.py +++ b/tests/test_correlation.py @@ -6,7 +6,7 @@ import pandas as pd -from hta.common.trace import transform_correlation_to_index +from hta.common.trace_collection import transform_correlation_to_index from hta.common.trace_symbol_table import TraceSymbolTable diff --git a/tests/test_custom_trace_parser.py b/tests/test_custom_trace_parser.py index daaa614b..20e7d10b 100644 --- a/tests/test_custom_trace_parser.py +++ b/tests/test_custom_trace_parser.py @@ -2,7 +2,7 @@ import tempfile import unittest -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_file import write_trace from hta.configs.parser_config import AVAILABLE_ARGS, ParserConfig @@ -76,9 +76,9 @@ def setUp(self) -> None: ], } - def _create_and_load_trace(self, trace_dir: str) -> Trace: + def _create_and_load_trace(self, trace_dir: str) -> TraceCollection: write_trace(self.trace_data, os.path.join(trace_dir, "trace_0.json.gz")) - t = Trace(trace_dir=trace_dir) + t = TraceCollection(trace_dir=trace_dir) t.parse_traces(use_multiprocessing=False) return t diff --git a/tests/test_timeline.py b/tests/test_timeline.py index 52fdb834..46f9196f 100644 --- a/tests/test_timeline.py +++ b/tests/test_timeline.py @@ -21,8 +21,8 @@ Timeline, TimelinePlotSetting, ) -from hta.common.trace import Trace from hta.common.trace_call_graph import CallGraph +from hta.common.trace_collection import TraceCollection from hta.common.trace_filter import CPUOperatorFilter, GPUKernelFilter from hta.common.trace_symbol_table import TraceSymbolTable @@ -32,7 +32,7 @@ class TestTimelineAnalysis(unittest.TestCase): base_data_dir = str(Path(hta.__file__).parent.parent.joinpath("tests/data")) trace_path: str = os.path.join(base_data_dir, "timeline_analysis") - t = Trace(trace_dir=trace_path) + t = TraceCollection(trace_dir=trace_path) t.parse_traces() t.decode_symbol_ids(use_shorten_name=False) cg = CallGraph(t, ranks=[0]) diff --git a/tests/test_trace_analysis.py b/tests/test_trace_analysis.py index 5d115032..d94e8bba 100644 --- a/tests/test_trace_analysis.py +++ b/tests/test_trace_analysis.py @@ -13,7 +13,7 @@ import hta import pandas as pd -from hta.common.trace import PHASE_COUNTER +from hta.common.trace_collection import PHASE_COUNTER from hta.trace_analysis import TimeSeriesTypes, TraceAnalysis @@ -77,7 +77,7 @@ def setUp(self): str(Path(self.overlaid_trace_dir)), "overlaid_rank-0.json.gz" ) - @patch.object(hta.common.trace.Trace, "write_raw_trace") + @patch.object(hta.common.trace_collection.TraceCollection, "write_raw_trace") def test_frequent_cuda_kernel_sequences(self, mock_write_trace): frequent_patterns_dfs = ( self.vision_transformer_t.get_frequent_cuda_kernel_sequences( @@ -497,7 +497,7 @@ def test_get_mtia_queue_length_stats(self): msg=f"queue_full_df = {queue_full_df}", ) - @patch.object(hta.common.trace.Trace, "write_raw_trace") + @patch.object(hta.common.trace_collection.TraceCollection, "write_raw_trace") def test_generate_trace_with_counters(self, mock_write_trace): # Use a trace with some kernels missing attribution to operators # to check if our logic is robust and does not lead to negative values. diff --git a/tests/test_trace_call_graph.py b/tests/test_trace_call_graph.py index bc92ae4e..8c4e0197 100644 --- a/tests/test_trace_call_graph.py +++ b/tests/test_trace_call_graph.py @@ -3,10 +3,10 @@ from pathlib import Path import pandas as pd -from hta.common.trace import Trace from hta.common.trace_call_graph import CallGraph from hta.common.trace_call_stack import CallStackGraph, CallStackIdentity +from hta.common.trace_collection import TraceCollection class TraceCallGraphTestCase(unittest.TestCase): @@ -16,7 +16,7 @@ def setUp(self) -> None: "tests/data/call_stack/backward_thread.json" ) self.test_trace_backward_threads: str = str(test_data_path) - self.t_backward_threads: Trace = Trace( + self.t_backward_threads: TraceCollection = TraceCollection( trace_files={0: self.test_trace_backward_threads}, trace_dir="", ) @@ -90,7 +90,7 @@ def test_link_main_and_bwd_stacks_has_bwd_annotation(self) -> None: self.assertListEqual(main_stack_root, bwd_stack_root) def test_link_main_and_bwd_stacks_no_bwd_annotation(self) -> None: - t: Trace = self.t_backward_threads + t: TraceCollection = self.t_backward_threads # remove backward annotation for _, df in t.get_all_traces().items(): df.drop(df.loc[df["s_name"].eq("## backward ##")].index, inplace=True) @@ -108,7 +108,9 @@ def test_link_main_and_bwd_stacks_no_bwd_annotation(self) -> None: def test_skip_gpu_threads(self) -> None: trace_file = self.test_trace_backward_threads - t: Trace = Trace(trace_files={i: trace_file for i in range(4)}) + t: TraceCollection = TraceCollection( + trace_files={i: trace_file for i in range(4)} + ) t.parse_traces() # set a new pid for the traces for rank, df in t.get_all_traces().items(): diff --git a/tests/test_trace_filter.py b/tests/test_trace_filter.py index cb519cd1..11ab9f0b 100644 --- a/tests/test_trace_filter.py +++ b/tests/test_trace_filter.py @@ -9,7 +9,7 @@ import numpy as np import pandas as pd -from hta.common.trace import Trace +from hta.common.trace_collection import TraceCollection from hta.common.trace_filter import ( CompositeFilter, CPUOperatorFilter, @@ -27,7 +27,7 @@ class TestTraceFilters(unittest.TestCase): base_data_dir = str(Path(hta.__file__).parent.parent.joinpath("tests/data")) trace_dir: str = os.path.join(base_data_dir, "trace_filter") - htaTrace: Trace = Trace(trace_dir=trace_dir) + htaTrace: TraceCollection = TraceCollection(trace_dir=trace_dir) htaTrace.parse_traces() def setUp(self): @@ -307,7 +307,7 @@ class TestTraceFiltersSyncEvents(unittest.TestCase): base_data_dir = str(Path(hta.__file__).parent.parent.joinpath("tests/data")) trace_dir: str = os.path.join(base_data_dir, "critical_path/cuda_event_sync") - htaTrace: Trace = Trace(trace_dir=trace_dir) + htaTrace: TraceCollection = TraceCollection(trace_dir=trace_dir) htaTrace.parse_traces() def setUp(self): diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index e64a2680..6aefa1eb 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -11,7 +11,7 @@ # import unittest.mock as mock import pandas as pd -from hta.common.trace import parse_trace_dict, Trace +from hta.common.trace_collection import parse_trace_dict, TraceCollection from hta.common.trace_parser import ( _auto_detect_parser_backend, _open_trace_file, @@ -91,11 +91,11 @@ def prepare_ground_truth_df(trace_dir, rank_0_file) -> pd.DataFrame: class TraceParseTestCase(unittest.TestCase): - vision_transformer_t: Trace + vision_transformer_t: TraceCollection vision_transformer_raw_df: pd.DataFrame - inference_t: Trace + inference_t: TraceCollection inference_raw_df: pd.DataFrame - triton_t: Trace + triton_t: TraceCollection triton_raw_df: pd.DataFrame @classmethod @@ -119,21 +119,23 @@ def setUpClass(cls): ] max_ranks = 8 - cls.vision_transformer_t: Trace = Trace(trace_dir=vision_transformer_trace_dir) + cls.vision_transformer_t: TraceCollection = TraceCollection( + trace_dir=vision_transformer_trace_dir + ) cls.vision_transformer_t.parse_traces( max_ranks=max_ranks, use_multiprocessing=True ) cls.vision_transformer_raw_df = prepare_ground_truth_df( vision_transformer_trace_dir, vision_transformer_rank_0_file ) - cls.inference_t: Trace = Trace( + cls.inference_t: TraceCollection = TraceCollection( trace_files=inference_trace_files, trace_dir=os.getcwd() ) cls.inference_t.parse_traces(max_ranks=max_ranks, use_multiprocessing=True) cls.inference_raw_df = prepare_ground_truth_df( inference_trace_dir, inference_rank_0_file ) - cls.triton_t: Trace = Trace(trace_dir=triton_trace_dir) + cls.triton_t: TraceCollection = TraceCollection(trace_dir=triton_trace_dir) cls.triton_t.parse_traces(max_ranks=max_ranks, use_multiprocessing=True) cls.triton_t.align_and_filter_trace(include_last_profiler_step=True) cls.triton_raw_df = prepare_ground_truth_df( @@ -302,7 +304,9 @@ def setUpClass(cls): def test_ijson_parser(self): set_default_trace_parsing_backend(ParserBackend.IJSON) - inference_t: Trace = Trace(trace_dir=self.inference_trace_dir) + inference_t: TraceCollection = TraceCollection( + trace_dir=self.inference_trace_dir + ) inference_t.parse_traces(max_ranks=1) self.assertEqual(len(inference_t.traces), 1) @@ -311,7 +315,9 @@ def test_ijson_parser(self): def test_ijson_batched_parser(self): set_default_trace_parsing_backend(ParserBackend.IJSON_BATCHED) - inference_t: Trace = Trace(trace_dir=self.inference_trace_dir) + inference_t: TraceCollection = TraceCollection( + trace_dir=self.inference_trace_dir + ) inference_t.parse_traces(max_ranks=1) self.assertEqual(len(inference_t.traces), 1) @@ -320,7 +326,9 @@ def test_ijson_batched_parser(self): def test_ijson_batch_and_compress_parser(self): set_default_trace_parsing_backend(ParserBackend.IJSON_BATCH_AND_COMPRESS) - inference_t: Trace = Trace(trace_dir=self.inference_trace_dir) + inference_t: TraceCollection = TraceCollection( + trace_dir=self.inference_trace_dir + ) inference_t.parse_traces(max_ranks=1) self.assertEqual(len(inference_t.traces), 1) @@ -376,7 +384,7 @@ def test_align_and_filter_mtia(self) -> None: mtia_trace_dir: str = add_test_data_path_prefix_if_exists( "tests/data/mtia_trace_single_rank" ) - t: Trace = Trace(trace_dir=mtia_trace_dir) + t: TraceCollection = TraceCollection(trace_dir=mtia_trace_dir) t.parse_traces() t.align_and_filter_trace() t.decode_symbol_ids(use_shorten_name=False) @@ -424,7 +432,7 @@ def setUp(self) -> None: self.resnet_nccl_trace: str = add_test_data_path_prefix_if_exists( "tests/data/nccl_parser_config" ) - self.resnet_nccl_t: Trace = Trace( + self.resnet_nccl_t: TraceCollection = TraceCollection( trace_dir=self.resnet_nccl_trace, parser_config=self.custom_cfg ) @@ -432,7 +440,7 @@ def setUp(self) -> None: triton_trace: str = add_test_data_path_prefix_if_exists( "tests/data/triton_example" ) - self.triton_t: Trace = Trace( + self.triton_t: TraceCollection = TraceCollection( trace_dir=triton_trace, parser_config=self.custom_cfg ) @@ -607,8 +615,8 @@ def test_fix_mtia_memory_kernels(self) -> None: "aten::add": 2004, } ) - # Create a Trace object - t = Trace(trace_dir="", trace_files={}) + # Create a TraceCollection object + t = TraceCollection(trace_dir="", trace_files={}) t.traces[0] = df.copy() t.symbol_table = symbol_table From f5f3bd091d2df80b5d42c3f6a057494f4067ac27 Mon Sep 17 00:00:00 2001 From: Benedykt Bela Date: Fri, 7 Nov 2025 14:32:39 +0200 Subject: [PATCH 2/4] Add singletrace module. Signed-off-by: Benedykt Bela --- hta/common/singletrace.py | 43 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 hta/common/singletrace.py diff --git a/hta/common/singletrace.py b/hta/common/singletrace.py new file mode 100644 index 00000000..956d1f8f --- /dev/null +++ b/hta/common/singletrace.py @@ -0,0 +1,43 @@ +from typing import Any, Dict + +import pandas as pd + +from hta.common.trace_symbol_table import TraceSymbolTable + +MetaData = Dict[str, Any] + + +class _SingleTrace: + """Class representing a single rank in a trace.""" + + device: str = "GENERIC" + + def __init__( + self, meta: MetaData, df: pd.DataFrame, symbol_table: TraceSymbolTable + ): + self.meta = meta + self.df = df + self.symbol_table = symbol_table + + +class _XPUSingleTrace(_SingleTrace): + """Class representing a single XPU rank in a trace.""" + + device: str = "INTEL GPU" + + +Trace = _SingleTrace + + +def create_default(meta=None, df=None, symbol_table=None) -> Trace: + """Factory method to create default Trace object.""" + return _SingleTrace(meta, df, symbol_table) + + +def create(device_type: str, meta, df, symbol_table) -> Trace: + """Factory method to create Trace object based on device type.""" + + if device_type == "INTEL GPU": + return _XPUSingleTrace(meta, df, symbol_table) + else: + return _SingleTrace(meta, df, symbol_table) From 4807b1154df77ba5c459be87a70882f1cb3a18bf Mon Sep 17 00:00:00 2001 From: Benedykt Bela Date: Fri, 7 Nov 2025 16:25:03 +0200 Subject: [PATCH 3/4] Adjust code to use the new singletrace module. The approach is to get the trace DataFrame directly from TraceCollection object instead of getting Trace object first and then accessing its DataFrame attribute. My current approach is to use the Trace objects directly when: - A method has more than one argument that is aggregated by Trace object. This way we reduce the number of parameters in that method. - A method passes the whole Trace object to another method. - A method returns more than one attribute that is aggregated by Trace object. This way we reduce the number of return values. With this approach we don't need to modify all the methods that use pd.DataFrame as input/output parameters. When there is a need to pass the Trace object to use its methods or attributes, we can simply adjust the methods that need it. Signed-off-by: Benedykt Bela --- examples/symbol_table_demo.py | 8 +- hta/analyzers/breakdown_analysis.py | 21 +-- hta/analyzers/communication_analysis.py | 9 +- hta/analyzers/critical_path_analysis.py | 8 +- hta/analyzers/cuda_kernel_analysis.py | 8 +- hta/analyzers/cupti_counter_analysis.py | 2 +- hta/analyzers/straggler.py | 2 +- hta/analyzers/straggler_analysis.py | 5 +- hta/analyzers/timeline.py | 2 +- hta/analyzers/trace_counters.py | 8 +- hta/common/call_stack.py | 8 +- hta/common/execution_trace.py | 2 +- hta/common/singletrace.py | 10 +- hta/common/trace_call_graph.py | 9 +- hta/common/trace_collection.py | 179 ++++++++++++------------ hta/common/trace_parser.py | 22 ++- hta/trace_diff.py | 4 +- tests/test_call_stack.py | 3 + tests/test_correlation.py | 4 +- tests/test_critical_path_analysis.py | 2 +- tests/test_custom_trace_parser.py | 4 +- tests/test_timeline.py | 4 +- tests/test_trace_analysis.py | 2 +- tests/test_trace_call_graph.py | 15 +- tests/test_trace_filter.py | 26 ++-- tests/test_trace_parse.py | 34 +++-- 26 files changed, 222 insertions(+), 179 deletions(-) diff --git a/examples/symbol_table_demo.py b/examples/symbol_table_demo.py index f5e63d70..e7455333 100644 --- a/examples/symbol_table_demo.py +++ b/examples/symbol_table_demo.py @@ -47,7 +47,7 @@ def demo_statistics( Returns: The resulted dataframe from this analysis. """ - df = trace.get_trace(rank) + df = trace.get_trace_df(rank) sym_id_map = trace.symbol_table.get_sym_id_map() sym_table = trace.symbol_table.get_sym_table() @@ -151,13 +151,13 @@ def run_demo( break logging.info("\n===End of Symbol to ID Map\n") - df = demo_trace.get_trace(0) + df = demo_trace.get_trace_df(0) logging.info(f"\n===Data Frame of Rank-0===\ntype={type(df)}\n") logging.info(f"\n{df}\n") logging.info("\n===End of Data Frame\n") logging.info(f"===Data Frame Info===\ntype={type(df)}\n") - demo_trace.get_trace(0).info() + demo_trace.get_trace_df(0).info() logging.info("\n===Kernel Statistics===\n") top_k: int = 10 @@ -176,7 +176,7 @@ def run_demo( def trace_info(trace: TraceCollection): rank = next(iter(trace.traces)) - df = trace.get_trace(rank) + df = trace.get_trace_df(rank) logging.info(f"\n===Dataframe of Rank {rank}") df.info() diff --git a/hta/analyzers/breakdown_analysis.py b/hta/analyzers/breakdown_analysis.py index 75d4871a..12c18c0c 100644 --- a/hta/analyzers/breakdown_analysis.py +++ b/hta/analyzers/breakdown_analysis.py @@ -8,6 +8,7 @@ import pandas as pd import plotly.express as px import plotly.graph_objects as go +from hta.common.singletrace import Trace from hta.common.trace_filter import GPUKernelFilter from hta.common.trace_symbol_table import decode_symbol_id_to_symbol_name @@ -75,8 +76,8 @@ def get_gpu_kernel_breakdown( kernel_type_to_analysis.append(KernelType.MEMORY.name) kernel_per_rank: Dict[str, Dict] = defaultdict(dict) - for rank, trace_df in t.traces.items(): - gpu_kernels = trace_df[trace_df["stream"].ne(-1)].copy() + for rank, trace in t.traces.items(): + gpu_kernels = trace.df[trace.df["stream"].ne(-1)].copy() gpu_kernels["kernel_type"] = gpu_kernels[["name"]].apply( lambda x: get_kernel_type(sym_table[x["name"]]), axis=1 ) @@ -318,7 +319,7 @@ def get_gpu_kernels_with_user_annotations( GPU user annotation. If the kernel overlaps with multiple user annotations, we will pick the lowest/leaf annotation in the stack to attribute to. Read more in get_gpu_kernels_with_user_annotations in hta/trace_analysis.py.""" - trace_df = t.get_trace(rank) + trace_df = t.get_trace_df(rank) trace_df["end"] = trace_df["ts"] + trace_df["dur"] trace_df["user_annotation"] = -1 @@ -401,8 +402,8 @@ def get_gpu_user_annotation_breakdown( kernel_per_rank: Dict[int, pd.DataFrame] = {} - for rank, trace_df in t.traces.items(): - gpu_user_annotation_kernels = trace_df[trace_df["cat"].eq(idx)].copy() + for rank, trace in t.traces.items(): + gpu_user_annotation_kernels = trace.df[trace.df["cat"].eq(idx)].copy() t.symbol_table.add_symbols_to_trace_df(gpu_user_annotation_kernels, "name") logger.info( f"rank = {rank}, num {annotation}s = {len(gpu_user_annotation_kernels)}" @@ -647,9 +648,9 @@ def get_temporal_breakdown( """ sym_table = t.symbol_table.get_sym_table() - def idle_time_per_rank(trace_df: pd.DataFrame) -> Tuple[int, int, int, int]: + def idle_time_per_rank(trace: Trace) -> Tuple[int, int, int, int]: """returns idle_time (us) , compute_time (us), non_compute_time (us), total_time (us)""" - gpu_kernels = trace_df[trace_df["stream"].ne(-1)].copy() + gpu_kernels = trace.df[trace.df["stream"].ne(-1)].copy() idle_time, kernel_time = cls._get_idle_time_for_kernels(gpu_kernels) gpu_kernels["kernel_type"] = gpu_kernels[["name"]].apply( @@ -672,10 +673,10 @@ def idle_time_per_rank(trace_df: pd.DataFrame) -> Tuple[int, int, int, int]: return idle_time, compute_time, non_compute_time, kernel_time result: Dict[str, List[float]] = defaultdict(list) - for rank, trace_df in t.traces.items(): + for rank, trace in t.traces.items(): result["rank"].append(rank) idle_time, compute_time, non_compute_time, kernel_time = idle_time_per_rank( - trace_df + trace ) result["idle_time(us)"].append(idle_time) result["compute_time(us)"].append(compute_time) @@ -826,7 +827,7 @@ def get_idle_time_breakdown( and median of idle intervals between kernels on a CUDA stream, also broken down by the idleness category (default = False). """ - trace_df: pd.DataFrame = t.get_trace(rank) + trace_df: pd.DataFrame = t.get_trace_df(rank) # Need to filter out events with `cuda_sync` category kernel_cats = [ diff --git a/hta/analyzers/communication_analysis.py b/hta/analyzers/communication_analysis.py index 109ec1c2..e4fd91c6 100644 --- a/hta/analyzers/communication_analysis.py +++ b/hta/analyzers/communication_analysis.py @@ -8,6 +8,7 @@ import pandas as pd import plotly.express as px +from hta.common.singletrace import Trace from hta.utils.utils import get_kernel_type, KernelType, merge_kernel_intervals # import statement used without the "if TYPE_CHECKING" guard will cause a circular @@ -29,11 +30,11 @@ def get_comm_comp_overlap( """ sym_table = t.symbol_table.get_sym_table() - def get_comm_comp_overlap_value(trace_df: pd.DataFrame) -> float: + def get_comm_comp_overlap_value(trace: Trace) -> float: """ Compute the overlap percentage between communication and computation kernels for one rank. """ - gpu_kernels = trace_df[trace_df["stream"].ne(-1)].copy() + gpu_kernels = trace.df[trace.df["stream"].ne(-1)].copy() gpu_kernels["kernel_type"] = gpu_kernels[["name"]].apply( lambda x: get_kernel_type(sym_table[x["name"]]), axis=1 ) @@ -77,10 +78,10 @@ def get_comm_comp_overlap_value(trace_df: pd.DataFrame) -> float: ).sum() result: Dict[str, List[float]] = defaultdict(list) - for rank, trace_df in t.traces.items(): + for rank, trace in t.traces.items(): result["rank"].append(rank) result["comp_comm_overlap_ratio"].append( - get_comm_comp_overlap_value(trace_df) + get_comm_comp_overlap_value(trace) ) result_df = pd.DataFrame(result) result_df["comp_comm_overlap_pctg"] = round( diff --git a/hta/analyzers/critical_path_analysis.py b/hta/analyzers/critical_path_analysis.py index d38e7015..1b8fb065 100644 --- a/hta/analyzers/critical_path_analysis.py +++ b/hta/analyzers/critical_path_analysis.py @@ -211,7 +211,7 @@ def __init__( self.rank: int = rank self.t = t self.t_full = t_full - self.full_trace_df: pd.DataFrame = self.t_full.get_trace(rank) + self.full_trace_df: pd.DataFrame = self.t_full.get_trace_df(rank) self.symbol_table = t_full.symbol_table self.data_load_events = data_load_events @@ -221,7 +221,7 @@ def __init__( if t is None: return - self.trace_df: pd.DataFrame = t.get_trace(rank) + self.trace_df: pd.DataFrame = t.get_trace_df(rank) self.critical_path_nodes: List[int] = [] self.critical_path_events_set: Set[int] = set() @@ -1798,7 +1798,7 @@ def critical_path_analysis( """ global PROFILE_TIMES t0 = time.perf_counter() - trace_df: pd.DataFrame = t.get_trace(rank) + trace_df: pd.DataFrame = t.get_trace_df(rank) sym_index = t.symbol_table.get_sym_id_map() if "cuda_sync" not in sym_index: @@ -1873,7 +1873,7 @@ def critical_path_analysis( # XXX This is a bit hacky but CallGraph does not really support a way # to specify a dataframe, just supports passing TraceCollection object t_copy = deepcopy(t) - t_copy.traces[rank] = clipped_df + t_copy.traces[rank].df = clipped_df t1 = time.perf_counter() logger.info(f"Preprocessing took {t1 - t0:.2f} seconds") diff --git a/hta/analyzers/cuda_kernel_analysis.py b/hta/analyzers/cuda_kernel_analysis.py index eccfd37f..79e81e7a 100644 --- a/hta/analyzers/cuda_kernel_analysis.py +++ b/hta/analyzers/cuda_kernel_analysis.py @@ -71,7 +71,7 @@ def get_frequent_cuda_kernel_sequences( sym_index = t.symbol_table.get_sym_id_map() cg = CallGraph(t, ranks=[rank]) - trace_df = t.get_trace(rank) + trace_df = t.get_trace_df(rank) # cpu_kernels = trace_df[trace_df["stream"].eq(-1)].copy() # gpu_kernels = trace_df[trace_df["stream"].ne(-1)].copy() @@ -444,8 +444,8 @@ def find_deepest_aten_op( result_dict: Dict = {} for rank in ranks: - trace_data = trace.get_trace(rank) - symbol_table.decode_df(trace.traces[rank], create_new_columns=True) + trace_data = trace.get_trace_df(rank) + symbol_table.decode_df(trace_data, create_new_columns=True) aten_operations = trace_data[ trace_data["s_name"].str.startswith("aten::") ].copy() @@ -556,7 +556,7 @@ def cuda_kernel_launch_stats( for rank in ranks: # get trace for a rank - trace_df: pd.DataFrame = t.get_trace(rank) + trace_df: pd.DataFrame = t.get_trace_df(rank) # filter out events which have correlation value matching to # cudaLaunchKernel, cudaLaunchKernelExC, cudaMemcpyAsync, cudaMemsetAsync diff --git a/hta/analyzers/cupti_counter_analysis.py b/hta/analyzers/cupti_counter_analysis.py index d5c095cc..c538ccdb 100644 --- a/hta/analyzers/cupti_counter_analysis.py +++ b/hta/analyzers/cupti_counter_analysis.py @@ -31,7 +31,7 @@ def _get_counter_data_with_operators_for_rank( ) -> Optional[pd.DataFrame]: sym_table = t.symbol_table.get_sym_table() t.decode_symbol_ids(use_shorten_name=False) - df = t.get_trace(rank) + df = t.get_trace_df(rank) # Get valid cuda kernels gpu_kernels = ( diff --git a/hta/analyzers/straggler.py b/hta/analyzers/straggler.py index 9dec3422..feddd3a7 100644 --- a/hta/analyzers/straggler.py +++ b/hta/analyzers/straggler.py @@ -51,7 +51,7 @@ def _extract_one_rank(df) -> pd.DataFrame: return p_steps df_iterations = pd.concat( - [_extract_one_rank(trace.get_trace(r)) for r in ranks], + [_extract_one_rank(trace.get_trace_df(r)) for r in ranks], keys=ranks, names=["rank", "iter"], ).reset_index() diff --git a/hta/analyzers/straggler_analysis.py b/hta/analyzers/straggler_analysis.py index 933e3002..dcf70f60 100644 --- a/hta/analyzers/straggler_analysis.py +++ b/hta/analyzers/straggler_analysis.py @@ -62,7 +62,10 @@ def get_potential_stragglers( ranks = list(t.get_all_traces().keys()) df_all = pd.concat( - [t.get_trace(r) for r in ranks], axis=0, keys=ranks, names=["rank", "idx"] + [t.get_trace_df(r) for r in ranks], + axis=0, + keys=ranks, + names=["rank", "idx"], ).reset_index() df_selected_profiler_steps = df_all.loc[ diff --git a/hta/analyzers/timeline.py b/hta/analyzers/timeline.py index e8452b2b..5be73034 100644 --- a/hta/analyzers/timeline.py +++ b/hta/analyzers/timeline.py @@ -244,7 +244,7 @@ def plot_timeline_gpu_kernels_from_trace( else: _ranks = list(trace_data.get_all_traces().keys()) df = pd.concat( - [trace_data.get_trace(r) for r in _ranks], + [trace_data.get_trace_df(r) for r in _ranks], axis=0, keys=_ranks, names=["rank", "idx"], diff --git a/hta/analyzers/trace_counters.py b/hta/analyzers/trace_counters.py index cd13406e..ff80bb78 100644 --- a/hta/analyzers/trace_counters.py +++ b/hta/analyzers/trace_counters.py @@ -44,8 +44,8 @@ def _get_queue_length_time_series_for_rank( time series. The value remains constant until the next timestamp. In essence, it can be thought of as a step function. """ - # get trace for a rank - trace_df: pd.DataFrame = t.get_trace(rank) + # get dataframe of trace for a rank + trace_df: pd.DataFrame = t.get_trace_df(rank) # CUDA Runtime events that may launch kernels # - filter events that have a correlated kernel event only. @@ -274,8 +274,8 @@ def _get_memory_bw_time_series_for_rank( ts (timestamp), pid (of corresponding GPU), name of memory copy type and memory_bw_gbps (memory bandwidth in GB/sec). """ - # get trace for a rank - trace_df: pd.DataFrame = t.get_trace(rank) + # get dataframe of trace for a rank + trace_df: pd.DataFrame = t.get_trace_df(rank) sym_table = t.symbol_table.get_sym_table() gpu_kernels = trace_df[trace_df["stream"].ne(-1)].copy() diff --git a/hta/common/call_stack.py b/hta/common/call_stack.py index 3e4b305f..a094f744 100644 --- a/hta/common/call_stack.py +++ b/hta/common/call_stack.py @@ -522,9 +522,9 @@ def _construct_call_graph( for rank in ranks: if self.pre_process_trace_data: self.constrain_child_time_withinin_parent( - self.trace_data.get_trace(rank) + self.trace_data.get_trace_df(rank) ) - df = self.trace_data.get_trace(rank).copy() + df = self.trace_data.get_trace_df(rank).copy() if remapped_tids and rank in remapped_tids: df = self._remap_tids(df, remapped_tids[rank]) for (pid, tid), df_thread in df.groupby(["pid", "tid"]): @@ -568,7 +568,7 @@ def _construct_call_graph( if node_id >= 0 } ) - df = self.trace_data.get_trace(rank) + df = self.trace_data.get_trace_df(rank) if ( not hta_options.disable_call_graph_depth() and call_stack_indices.size > 0 @@ -600,7 +600,7 @@ def get_stack_of_node( Raises: ValueError when the index is not in the DataFrame. """ - df = self.trace_data.get_trace(rank) + df = self.trace_data.get_trace_df(rank) # If it is a GPU kernel, get the stack from the launch event. if df.loc[node_id]["stream"] > -1: diff --git a/hta/common/execution_trace.py b/hta/common/execution_trace.py index 08d20b71..d4fe07e3 100644 --- a/hta/common/execution_trace.py +++ b/hta/common/execution_trace.py @@ -121,7 +121,7 @@ def correlate_execution_trace( Please note (2) is not supported yet and will come in future PRs. """ - trace_df = trace.get_trace(rank) + trace_df = trace.get_trace_df(rank) if not _et_has_overlap(trace_df, et): logging.error( diff --git a/hta/common/singletrace.py b/hta/common/singletrace.py index 956d1f8f..33a6e6b9 100644 --- a/hta/common/singletrace.py +++ b/hta/common/singletrace.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import Any, Dict, List import pandas as pd @@ -19,6 +19,14 @@ def __init__( self.df = df self.symbol_table = symbol_table + def get_sym_table(self) -> List[str]: + """Get the list of symbols from the symbol table. + + Returns: + List of symbol strings in order of their IDs. + """ + return self.symbol_table.get_sym_table() + class _XPUSingleTrace(_SingleTrace): """Class representing a single XPU rank in a trace.""" diff --git a/hta/common/trace_call_graph.py b/hta/common/trace_call_graph.py index c67e4592..bba26aef 100644 --- a/hta/common/trace_call_graph.py +++ b/hta/common/trace_call_graph.py @@ -3,6 +3,7 @@ import pandas as pd +from hta.common import singletrace from hta.common.trace_call_stack import CallStackGraph, CallStackIdentity, CallStackNode from hta.common.trace_collection import get_cpu_gpu_correlation, TraceCollection from hta.common.trace_symbol_table import TraceSymbolTable @@ -96,7 +97,7 @@ def __init__( self._cached_nodes: Dict[int, CallStackNode] = self.rank_to_nodes[ self._cached_rank ] - self._cached_df: pd.DataFrame = self.trace_data.get_trace(self._cached_rank) + self._cached_df: pd.DataFrame = self.trace_data.get_trace_df(self._cached_rank) self._cached_gpu_kernels: pd.DataFrame = self._cached_df.loc[ self._cached_df["stream"].ne(-1) ] @@ -121,7 +122,7 @@ def from_dataframe( t.symbol_table = ( symbol_table if symbol_table else TraceSymbolTable.create_from_df(df) ) - t.traces[rank] = df.copy() + t.traces[rank] = singletrace.create_default(df=df.copy()) t.is_parsed = True cg = CallGraph(t) @@ -138,7 +139,7 @@ def _construct_call_graph(self) -> None: t0 = perf_counter() self.rank_to_nodes[rank] = {} self.rank_to_stacks[rank] = {} - df = self.trace_data.get_trace(rank) + df = self.trace_data.get_trace_df(rank) # add an "end" column for time interval based filtering if "end" not in df.columns: df["end"] = df["ts"] + df["dur"] @@ -364,7 +365,7 @@ def _update_cached_data(self, rank: int) -> None: if rank in self.ranks and rank != self._cached_rank: self._cached_rank = rank self._cached_nodes = self.rank_to_nodes[self._cached_rank] - self._cached_df = self.trace_data.get_trace(self._cached_rank) + self._cached_df = self.trace_data.get_trace_df(self._cached_rank) self._cached_gpu_kernels = self._cached_df.loc[ self._cached_df["stream"].ne(-1) ] diff --git a/hta/common/trace_collection.py b/hta/common/trace_collection.py index 664ac0df..cfdde1d0 100644 --- a/hta/common/trace_collection.py +++ b/hta/common/trace_collection.py @@ -13,12 +13,13 @@ import time import tracemalloc import warnings -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Union import numpy as np import pandas as pd +from hta.common.singletrace import Trace from hta.common.trace_file import create_rank_to_trace_dict, get_trace_files from hta.common.trace_filter import CPUOperatorFilter, GPUKernelFilter from hta.common.trace_parser import parse_trace_dataframe, parse_trace_dict @@ -60,10 +61,8 @@ def trace_event_timestamp_to_unixtime_ns( return trace_event_ts_ns + base_time_nanoseconds -def transform_correlation_to_index( - df: pd.DataFrame, symbol_table: TraceSymbolTable -) -> pd.DataFrame: - """Transform correlation to index_correlation and add a index_correlation column to df. +def transform_correlation_to_index(trace: Trace) -> pd.DataFrame: + """Transform correlation to index_correlation and add a index_correlation column to trace's df. The correlation in the trace is a reference ID which links a Cuda kernel launch and the cuda kernel. Because the correlation is not an index, using `correction` to find @@ -78,21 +77,21 @@ def transform_correlation_to_index( The transform can be illustrated as follows: - Given the following input DataFrame : + Given the following input trace's DataFrame: | index | stream | cat | s_cat | correlation | =============================================== | 675 | 7 | 248 | Kernel | 278204204 | | 677 | -1 | 9 | Runtime| 278204204 | - After calling `transform_correlation_to_index(df)`, will become: + After calling `transform_correlation_to_index(df)`, trace's df will become: | index | stream | cat | s_cat | correlation | index_correlation | ==================================================================== | 675 | 7 | 248 | Kernel | 278204204 | 677 | | 677 | -1 | 9 | Runtime| 278204204 | 675 | - This function changes the input DataFrame by adding a new column `index_correlation`. + This function changes the input trace's DataFrame by adding a new column `index_correlation`. Example use: find the kernels of all Cuda Launch cuda_launch_events = df.loc[df['cat'].eq(9) & df['index_correlation'].gt(0)] @@ -100,15 +99,15 @@ def transform_correlation_to_index( df.loc[cuda_kernel_indices] Args: - df (pd.DataFrame): the input DataFrame - symbol_table: the TraceSymbolTable for the trace + trace (Trace): the input Trace object, containing both the DataFrame and the symbol table. Returns: pd.DataFrame: the transformed DataFrame with a index_correlation column. Affects: - This function adds a index_correlation column to df. + This function adds a index_correlation column to trace's df. """ + df = trace.df if "correlation" not in df.columns: return df @@ -119,8 +118,8 @@ def transform_correlation_to_index( df["correlation"].ne(-1), ["index", "correlation", "stream", "name"] ] - on_cpu = CPUOperatorFilter()(corr_df, symbol_table) - on_gpu = GPUKernelFilter()(corr_df, symbol_table) + on_cpu = CPUOperatorFilter()(corr_df, trace.symbol_table) + on_gpu = GPUKernelFilter()(corr_df, trace.symbol_table) # We only need to merge once. # index_x --> index_y will be cpu to gpu mapping @@ -154,9 +153,9 @@ def get_cpu_gpu_correlation(df: pd.DataFrame) -> pd.DataFrame: return cpu_gpu_correlation -def add_iteration(df: pd.DataFrame, symbol_table: TraceSymbolTable) -> pd.DataFrame: +def add_iteration(trace: Trace) -> pd.DataFrame: """ - Add an iteration column to the DataFrame . + Add an iteration column to the trace's DataFrame. This function extracts the trace iteration number from the `ProfilerStep` annotation in the name column and then apply the following logic to determine which iteration @@ -169,17 +168,18 @@ def add_iteration(df: pd.DataFrame, symbol_table: TraceSymbolTable) -> pd.DataFr steps, set the iteration number to -1. Args: - df: a DataFrame representation of a trace. - symbol_table: the TraceSymbolTable for the trace + trace (Trace): a Trace object containing both the DataFrame and the symbol table. Returns: A DataFrame with the profiler steps information. Note: - This function will change the input DataFrame by adding the iteration number column. + This function will change the input trace's DataFrame by adding the iteration number column. """ - s_map = pd.Series(symbol_table.sym_index) - s_tab = pd.Series(symbol_table.sym_table) + df = trace.df + + s_map = pd.Series(trace.symbol_table.sym_index) + s_tab = pd.Series(trace.symbol_table.sym_table) profiler_step_ids = s_map[s_map.index.str.startswith("ProfilerStep")] profiler_step_ids.sort_index() @@ -228,20 +228,17 @@ def _get_profiler_step(ts: int) -> int: return profiler_steps -def parse_trace_file( - trace_file_path: str, - cfg: Optional[ParserConfig] = None, -) -> Tuple[MetaData, pd.DataFrame, TraceSymbolTable]: - """parse a single trace file into a meat test_data dictionary and a dataframe of events. +def parse_trace_file(trace_file_path: str, cfg: Optional[ParserConfig] = None) -> Trace: + """parse a single trace file into a trace (Trace) object. Args: trace_file_path (str): The path to a trace file. When the trace_file is a relative path. This method combines the object's trace_path with trace_file to get the full path of the trace file. cfg (ParserConfig, Optional): A ParserConfig object controls how to parse the trace file. Returns: - Tuple[MetaData, pd.DataFrame, TraceSymbolTable] - The first item is the trace's metadata; - The second item is the dataframe representation of the trace's events. - The third item is the symbol table to encode the symbols of the trace. + Trace object that contains: + Trace's metadata. + DataFrame representation of the trace's events. + Symbol table to encode the symbols of the trace. Raises: OSError when the trace file doesn't exist or current process has no permission to access it. @@ -257,21 +254,21 @@ def parse_trace_file( t_start = time.perf_counter() cfg = cfg or ParserConfig.get_default_cfg() - meta, df, local_symbol_table = parse_trace_dataframe(trace_file_path, cfg) + trace: Trace = parse_trace_dataframe(trace_file_path, cfg) # add fwd bwd links between CPU ops - add_fwd_bwd_links(df) - - df = transform_correlation_to_index(df, local_symbol_table) + add_fwd_bwd_links(trace.df) + trace.df = transform_correlation_to_index(trace) + add_iteration(trace) - add_iteration(df, local_symbol_table) - df["end"] = df["ts"] + df["dur"] + trace.df["end"] = trace.df["ts"] + trace.df["dur"] t_end = time.perf_counter() logger.warning( f"Overall parsing of {trace_file_path} in {(t_end - t_start):.2f} seconds; current PID:{os. getpid()}" ) - return meta, df, local_symbol_table + + return trace class _TraceFileParserWrapper: @@ -280,9 +277,7 @@ class _TraceFileParserWrapper: def __init__(self, cfg: ParserConfig) -> None: self.cfg = cfg - def __call__( - self, trace_file: str - ) -> Tuple[MetaData, pd.DataFrame, TraceSymbolTable]: + def __call__(self, trace_file: str) -> Trace: return parse_trace_file(trace_file, self.cfg) @@ -399,7 +394,7 @@ def __init__( return logger.debug(self.trace_files) - self.traces: Dict[int, pd.DataFrame] = {} + self.traces: Dict[int, Trace] = {} self.symbol_table = TraceSymbolTable() self.meta_data: Dict[int, MetaData] = {} self.min_ts: int = 0 @@ -428,10 +423,11 @@ def load_traces( use_memory_profiling=use_memory_profiling, ) self.align_and_filter_trace(include_last_profiler_step) - for rank, df in self.traces.items(): - df = self.traces[rank].set_index("index", drop=False) + for trace in self.traces.values(): + df = trace.df + df = df.set_index("index", drop=False) df.index.names = [None] - self.traces[rank] = df + trace.df = df self.is_parsed = True def parse_single_rank(self, rank: int) -> None: @@ -443,20 +439,17 @@ def parse_single_rank(self, rank: int) -> None: """ if rank in self.trace_files: trace_filepath = self.trace_files[rank] - ( - self.meta_data[rank], - self.traces[rank], - local_symbol_table, - ) = parse_trace_file(trace_filepath, self.parser_config) + trace = parse_trace_file(trace_filepath, self.parser_config) # update the global symbol table - self.symbol_table.add_symbols(local_symbol_table.get_sym_table()) + local_symbol_table = trace.get_sym_table() + self.symbol_table.add_symbols(local_symbol_table) # fix the encoding of the data frame - local_table = local_symbol_table.get_sym_table() global_map = self.symbol_table.get_sym_id_map() for col in ["cat", "name"]: - self.traces[rank][col] = self.traces[rank][col].apply( - lambda idx: global_map[local_table[idx]] + trace.df[col] = trace.df[col].apply( + lambda idx: global_map[local_symbol_table[idx]] ) + self.traces[rank] = trace def parse_multiple_ranks( self, @@ -482,13 +475,9 @@ def parse_multiple_ranks( if not use_multiprocessing: for rank in ranks: logger.debug(f"parsing trace for rank-{rank}") - result = parse_trace_file(self.trace_files[rank], self.parser_config) - self.meta_data[rank], self.traces[rank], local_symbol_tables[rank] = ( - result[0], - result[1], - result[2], - ) - self.symbol_table.add_symbols(local_symbol_tables[rank].get_sym_table()) + trace = parse_trace_file(self.trace_files[rank], self.parser_config) + self.traces[rank] = trace + self.symbol_table.add_symbols(trace.get_sym_table()) logger.debug(f"finished parsing for all {len(ranks)} ranks") else: num_procs = min(mp.cpu_count(), len(ranks)) @@ -504,24 +493,20 @@ def parse_multiple_ranks( _parser = _TraceFileParserWrapper(self.parser_config) with mp.get_context("fork").Pool(num_procs) as pool: - results = pool.map(_parser, trace_paths, chunksize=1) + trace_list = pool.map(_parser, trace_paths, chunksize=1) logger.debug(f"finished parallel parsing using {num_procs} processes.") # collect the results - for rank, result in zip(ranks, results): - self.meta_data[rank], self.traces[rank], local_symbol_tables[rank] = ( - result[0], - result[1], - result[2], - ) - self.symbol_table.add_symbols(local_symbol_tables[rank].get_sym_table()) + for rank, trace in zip(ranks, trace_list): + self.traces[rank] = trace + self.symbol_table.add_symbols(trace.get_sym_table()) # Now we update the IDs in the Dataframe using the global symbols table. global_map = self.symbol_table.get_sym_id_map() - for rank in ranks: - local_table = local_symbol_tables[rank].get_sym_table() + for trace in self.traces.values(): + local_table = trace.get_sym_table() for col in ["cat", "name"]: - self.traces[rank][col] = self.traces[rank][col].apply( + trace.df[col] = trace.df[col].apply( lambda idx: global_map[local_table[idx]] ) @@ -600,11 +585,30 @@ def get_iterations(self, rank: Optional[int] = None) -> List[int]: rank = self._get_first_rank(rank) if rank in self.get_ranks(): - df = self.traces[rank] + df = self.get_trace_df(rank) if "iteration" in df.columns: return sorted([i for i in df["iteration"].unique() if i >= 0]) return [] + def get_trace_df(self, rank: int) -> pd.DataFrame: + """ + Get the trace's DataFrame for a given rank. + + Args: + rank (int) : the rank of the trainer. + + Returns: + The trace's DataFrame for the given rank. + + Raises: + ValueError when this TraceCollection object doesn't have trace for the given rank. + """ + if rank not in self.traces: + logger.error(f"get_rank_trace - no trace for rank {rank}") + raise ValueError + + return self.traces[rank].df + def get_trace_duration(self, rank: Optional[int] = None) -> int: """Get the duration of specified rank. @@ -615,10 +619,11 @@ def get_trace_duration(self, rank: Optional[int] = None) -> int: Returns: duration of trace (int) """ rank = self._get_first_rank(rank) - trace_df = self.get_trace(rank) + trace_df = self.get_trace_df(rank) + return trace_df.ts.max() - trace_df.ts.min() - def get_trace(self, rank: int) -> pd.DataFrame: + def get_trace(self, rank: int) -> Trace: """ Get the trace for a given rank. @@ -626,7 +631,7 @@ def get_trace(self, rank: int) -> pd.DataFrame: rank (int) : the rank of the trainer whose trace is to be returned. Returns: - The DataFrame for the given rank. + The trace for the given rank. Raises: ValueError when this TraceCollection object doesn't have trace for the given rank. @@ -636,11 +641,11 @@ def get_trace(self, rank: int) -> pd.DataFrame: raise ValueError return self.traces[rank] - def get_all_traces(self) -> Dict[int, pd.DataFrame]: + def get_all_traces(self) -> Dict[int, Trace]: """ Get the traces of all ranks. Returns: - A dictionary with rank as key and its trace test_data as value. + A dictionary with rank as key and its trace as value. """ return self.traces @@ -709,10 +714,9 @@ def _align_all_ranks(self) -> None: """ Align dataframes for all ranks such that the earliest event starts at time 0. """ - self.min_ts = min(trace_df["ts"].min() for trace_df in self.traces.values()) - for rank, trace_df in self.traces.items(): - trace_df["ts"] = trace_df["ts"] - self.min_ts - self.traces[rank] = trace_df + self.min_ts = min(trace.df["ts"].min() for trace in self.traces.values()) + for _, trace in self.traces.items(): + trace.df["ts"] = trace.df["ts"] - self.min_ts def _fix_mtia_memory_kernels(self, trace_df: pd.DataFrame) -> pd.DataFrame: """ @@ -850,11 +854,11 @@ def filter_mtia_kernels_for_one_rank(trace_df: pd.DataFrame) -> pd.DataFrame: "There is only one iteration in the trace. The analysis result may not be accurate." ) include_last_profiler_step = True - for rank, trace_df in self.traces.items(): + for trace in self.traces.values(): if device_type != "MTIA": - self.traces[rank] = filter_gpu_kernels_with_cpu_correlation(trace_df) + trace.df = filter_gpu_kernels_with_cpu_correlation(trace.df) else: - self.traces[rank] = filter_mtia_kernels_for_one_rank(trace_df) + trace.df = filter_mtia_kernels_for_one_rank(trace.df) def decode_symbol_ids(self, use_shorten_name: bool = True) -> None: """Decode the name and cat column to show the original string names. @@ -866,7 +870,7 @@ def decode_symbol_ids(self, use_shorten_name: bool = True) -> None: for rank in self.traces: decode_symbol_id_to_symbol_name( - self.get_trace(rank), self.symbol_table, use_shorten_name + self.get_trace_df(rank), self.symbol_table, use_shorten_name ) def get_device_type(self) -> str: @@ -876,7 +880,8 @@ def get_device_type(self) -> str: The device type parsed from the trace metadata. If the device type is not found, return "UNKNOWN". """ rank = next(iter(self.traces)) - device_type = self.meta_data[rank].get("device_type", "UNKNOWN") + trace: Trace = self.traces[rank] + device_type = trace.meta.get("device_type", "UNKNOWN") return device_type def convert_time_series_to_events( @@ -963,4 +968,6 @@ def get_trace_start_unixtime_ns(self, rank: int) -> int: logger.warning(err_msg) raise ValueError(err_msg) - return trace_event_timestamp_to_unixtime_ns(self.min_ts, self.meta_data[rank]) + trace: Trace = self.get_trace(rank) + + return trace_event_timestamp_to_unixtime_ns(self.min_ts, trace.meta) diff --git a/hta/common/trace_parser.py b/hta/common/trace_parser.py index 780531b8..d0ca8da7 100644 --- a/hta/common/trace_parser.py +++ b/hta/common/trace_parser.py @@ -21,6 +21,8 @@ import numpy as np import pandas as pd +from hta.common import singletrace +from hta.common.singletrace import Trace from hta.common.trace_symbol_table import TraceSymbolTable from hta.configs.config import logger @@ -87,6 +89,8 @@ def infer_gpu_type( return "AMD GPU" if "runFunction - job_prep_and_submit_for_execution" in name_set: return "MTIA" + if "urEnqueueKernelLaunch" in name_set: + return "INTEL GPU" return "UNKNOWN GPU" @@ -459,17 +463,17 @@ def _parse_trace_dataframe_ijson( def parse_trace_dataframe( trace_file_path: str, cfg: ParserConfig, -) -> Tuple[MetaData, pd.DataFrame, TraceSymbolTable]: - """Parse a single trace file into a meat test_data dictionary and a dataframe of events. +) -> Trace: + """Parse a single trace file into a trace (Trace) object. Args: trace_file_path (str): The path to a trace file. When the trace_file is a relative path. This method combines the object's trace_path with trace_file to get the full path of the trace file. cfg (ParserConfig, Optional): A ParserConfig object controls how to parse the trace file. Returns: - Tuple[MetaData, pd.DataFrame, TraceSymbolTable] - The first item is the trace's metadata; - The second item is the dataframe representation of the trace's events. - The third item is the symbol table to encode the symbols of the trace. + Trace object that contains: + Trace's metadata. + DataFrame representation of the trace's events. + Symbol table to encode the symbols of the trace. Raises: JSONDecodeError when the trace file is not a valid JSON document. @@ -509,6 +513,9 @@ def parse_trace_dataframe( device_type = infer_gpu_type(meta, local_symbol_table.get_sym_id_map()) meta[str(MetaDataKey.DEVICE_TYPE)] = device_type + # Create Trace object representing single rank from the trace collection. + trace = singletrace.create(device_type, meta, df, local_symbol_table) + t_end = time.perf_counter() logger.warning( f"Parsed {trace_file_path} backend={parser_backend} in {(t_end - t_start):.2f} seconds; current PID:{os. getpid()}" @@ -519,7 +526,8 @@ def parse_trace_dataframe( logger.warning( f"Parser Memory usage peak = {(peak/1024/1024):.2f} MB, current = {(current/1024/1024):.2f} MB" ) - return meta, df, local_symbol_table + + return trace # --- Trace metadata reader --- diff --git a/hta/trace_diff.py b/hta/trace_diff.py index 99f38989..c479030a 100644 --- a/hta/trace_diff.py +++ b/hta/trace_diff.py @@ -139,10 +139,10 @@ def extract_ops( # Filter the trace by ranks if len(ranks) == 1: - df_rank = self.t.get_trace(ranks[0]) + df_rank = self.t.get_trace_df(ranks[0]) else: df_rank = pd.concat( - [self.t.get_trace(r) for r in ranks], + [self.t.get_trace_df(r) for r in ranks], axis=0, keys=_ranks, names=["rank", "idx"], diff --git a/tests/test_call_stack.py b/tests/test_call_stack.py index 96eb85b9..e8199338 100644 --- a/tests/test_call_stack.py +++ b/tests/test_call_stack.py @@ -155,6 +155,9 @@ def get_all_traces(self): def get_trace(self, rank): return self.traces[rank] + def get_trace_df(self, rank): + return self.traces[rank] + # Create test data for trim_trace_events self.df_trim = pd.DataFrame( { diff --git a/tests/test_correlation.py b/tests/test_correlation.py index 13b13d25..a7dd568f 100644 --- a/tests/test_correlation.py +++ b/tests/test_correlation.py @@ -6,6 +6,7 @@ import pandas as pd +from hta.common import singletrace from hta.common.trace_collection import transform_correlation_to_index from hta.common.trace_symbol_table import TraceSymbolTable @@ -35,7 +36,8 @@ def test_something(self): mock_symbol_table = TraceSymbolTable() mock_symbol_table.sym_index = sym_index expected_index_correlation = [5, 6, 0, -1, 1, 2, 0, -1] - df2 = transform_correlation_to_index(df, mock_symbol_table) + trace = singletrace.create_default(df=df, symbol_table=mock_symbol_table) + df2 = transform_correlation_to_index(trace) self.assertListEqual( expected_index_correlation, df2["index_correlation"].tolist() ) diff --git a/tests/test_critical_path_analysis.py b/tests/test_critical_path_analysis.py index 2bac7462..9dc118bf 100644 --- a/tests/test_critical_path_analysis.py +++ b/tests/test_critical_path_analysis.py @@ -68,7 +68,7 @@ def _critical_path_on_simple_add_trace(self) -> CPGraph: def test_critical_path_basic_add(self): critical_path_t = self.simple_add_trace cp_graph = self._critical_path_on_simple_add_trace() - trace_df = critical_path_t.t.get_trace(0) + trace_df = critical_path_t.t.get_trace_df(0) # Check the graph construction for the aten::relu_ operator # There are 3 stacked operators/runtime events here; diff --git a/tests/test_custom_trace_parser.py b/tests/test_custom_trace_parser.py index 20e7d10b..d73afac2 100644 --- a/tests/test_custom_trace_parser.py +++ b/tests/test_custom_trace_parser.py @@ -86,7 +86,7 @@ def test_default_config(self) -> None: with tempfile.TemporaryDirectory() as t_dir: t = self._create_and_load_trace(t_dir) cfg = ParserConfig.get_default_cfg() - df = t.get_trace(self.rank) + df = t.get_trace_df(self.rank) self.assertTrue(all(arg.name in df.columns for arg in cfg.get_args())) def test_custom_config(self) -> None: @@ -119,6 +119,6 @@ def test_custom_config(self) -> None: with tempfile.TemporaryDirectory() as t_dir: t = self._create_and_load_trace(t_dir) cfg = ParserConfig.get_default_cfg() - df = t.get_trace(self.rank) + df = t.get_trace_df(self.rank) self.assertTrue(all(arg.name in df.columns for arg in cfg.get_args())) self.assertFalse(all(arg.name in df.columns for arg in removed_args)) diff --git a/tests/test_timeline.py b/tests/test_timeline.py index 46f9196f..0c7a549e 100644 --- a/tests/test_timeline.py +++ b/tests/test_timeline.py @@ -40,7 +40,7 @@ class TestTimelineAnalysis(unittest.TestCase): def setUp(self) -> None: self.trace_path: str = TestTimelineAnalysis.trace_path self.t = TestTimelineAnalysis.t - self.df = self.t.get_trace(0) + self.df = self.t.get_trace_df(0) @patch(f"{_MODULE}.px.timeline") def test_plot_timeline(self, mock_timeline: Mock) -> None: @@ -253,7 +253,7 @@ class _TCase(NamedTuple): @patch(f"{_MODULE}.plot_events_timeline") def test_timeline_class(self, mock_plot: Mock) -> None: - df = self.t.get_trace(0) + df = self.t.get_trace_df(0) save_path = os.path.join(self.trace_path, "timeline_class.html") tl = Timeline(df, self.t.symbol_table, filter_func=GPUKernelFilter()) tl.setting.plot_format = PlotFormat.File diff --git a/tests/test_trace_analysis.py b/tests/test_trace_analysis.py index d94e8bba..c139c946 100644 --- a/tests/test_trace_analysis.py +++ b/tests/test_trace_analysis.py @@ -327,7 +327,7 @@ def __test_gpu_user_annotation_common( annotation = "gpu_user_annotation" if use_gpu_annotation else "user_annotation" idx = analyzer.t.symbol_table.sym_index[annotation] - trace_df = analyzer.t.get_trace(0) + trace_df = analyzer.t.get_trace_df(0) analyzer.t.symbol_table.add_symbols_to_trace_df(trace_df, "name") ref_sum_df = ( trace_df[trace_df.cat == idx][["name", "dur"]] diff --git a/tests/test_trace_call_graph.py b/tests/test_trace_call_graph.py index 8c4e0197..471e5730 100644 --- a/tests/test_trace_call_graph.py +++ b/tests/test_trace_call_graph.py @@ -27,7 +27,7 @@ def setUp(self) -> None: self.t_backward_threads, ranks=[0] ) self.df_backward_threads: pd.DataFrame = ( - self.cg_backward_threads.trace_data.get_trace(0) + self.cg_backward_threads.trace_data.get_trace_df(0) ) @staticmethod @@ -92,14 +92,17 @@ def test_link_main_and_bwd_stacks_has_bwd_annotation(self) -> None: def test_link_main_and_bwd_stacks_no_bwd_annotation(self) -> None: t: TraceCollection = self.t_backward_threads # remove backward annotation - for _, df in t.get_all_traces().items(): + for _, trace in t.get_all_traces().items(): + df = trace.df df.drop(df.loc[df["s_name"].eq("## backward ##")].index, inplace=True) cg: CallGraph = CallGraph(t) autograd_index = self._get_first_index( - t.get_trace(0), "autograd::engine::evaluate_function: AddmmBackward0" + t.get_trace_df(0), "autograd::engine::evaluate_function: AddmmBackward0" + ) + profiler_step_index = self._get_first_index( + t.get_trace_df(0), "ProfilerStep#552" ) - profiler_step_index = self._get_first_index(t.get_trace(0), "ProfilerStep#552") self.assertTrue(autograd_index > -1) self.assertTrue(profiler_step_index > -1) csg: CallStackGraph = cg.get_csg_of_node(autograd_index, 0) @@ -113,8 +116,8 @@ def test_skip_gpu_threads(self) -> None: ) t.parse_traces() # set a new pid for the traces - for rank, df in t.get_all_traces().items(): - df["pid"] = rank + 1 + for rank, trace in t.get_all_traces().items(): + trace.df["pid"] = rank + 1 # For the test traces, there should be exactly two stacks for each rank. cg: CallGraph = CallGraph(t) self.assertListEqual(cg.mapping.groupby("rank").size().unique().tolist(), [2]) diff --git a/tests/test_trace_filter.py b/tests/test_trace_filter.py index 11ab9f0b..af592a98 100644 --- a/tests/test_trace_filter.py +++ b/tests/test_trace_filter.py @@ -32,7 +32,7 @@ class TestTraceFilters(unittest.TestCase): def setUp(self): self.htaTrace = TestTraceFilters.htaTrace - self.df = self.htaTrace.get_trace(0) + self.df = self.htaTrace.get_trace_df(0) def testIterationFilter(self) -> None: f = IterationFilter([551]) @@ -47,10 +47,10 @@ def testRankFilter(self) -> None: f = RankFilter([0]) ranks_df = [] - for rank, df in self.htaTrace.traces.items(): + for rank, trace in self.htaTrace.traces.items(): # add rank column - df["rank"] = rank - ranks_df.append(df) + trace.df["rank"] = rank + ranks_df.append(trace.df) # combine both ranks combined_df = pd.concat(ranks_df) @@ -63,7 +63,7 @@ def testTimeRangeFilter(self) -> None: start_time = 1682725898237042 end_time = 1682725898240570 f = TimeRangeFilter((start_time, end_time)) - filtered_df = f(self.htaTrace.traces[0]) + filtered_df = f(self.htaTrace.get_trace_df(0)) # rows are present in time range self.assertEqual(filtered_df.shape[0], 93) @@ -109,7 +109,7 @@ def testNameFilterWithoutSymbolTable(self) -> None: def testGPUKernelFilter(self) -> None: f = GPUKernelFilter() - filtered_df = f(self.htaTrace.traces[0]) + filtered_df = f(self.htaTrace.get_trace_df(0)) # GPU kernel is present, note we are not reading CUDA sync events. self.assertTrue(filtered_df[(filtered_df["stream"] > 0)].size > 0) @@ -119,7 +119,7 @@ def testGPUKernelFilter(self) -> None: def testCPUOperatorFilter(self) -> None: f = CPUOperatorFilter() - filtered_df = f(self.htaTrace.traces[0]) + filtered_df = f(self.htaTrace.get_trace_df(0)) # CPU operator is present self.assertTrue(filtered_df[(filtered_df["stream"] < 0)].size > 0) @@ -135,7 +135,7 @@ def testCompositeFilter(self) -> None: f2 = TimeRangeFilter((start_time, end_time)) cf = CompositeFilter([f1, f2]) - filtered_df = cf(self.htaTrace.traces[0]) + filtered_df = cf(self.htaTrace.get_trace_df(0)) self.assertTrue(filtered_df[(filtered_df["iteration"] == 551)].size > 0) self.assertTrue( filtered_df[ @@ -270,7 +270,7 @@ class TC: ] self.htaTrace.decode_symbol_ids(use_shorten_name=False) - df = self.htaTrace.traces[0] + df = self.htaTrace.get_trace_df(0) for i, tc in enumerate(test_cases): f = CombinedOperatorFilter( tc.root_op_name, @@ -294,7 +294,7 @@ class TC: def testUnderOperatorFilter(self) -> None: op_name = "forward" self.htaTrace.decode_symbol_ids(use_shorten_name=False) - df = FirstIterationFilter()(self.htaTrace.traces[0]) + df = FirstIterationFilter()(self.htaTrace.get_trace_df(0)) f = UnderOperatorFilter(op_name=op_name, position=0, include_gpu_kernels=True) self.assertEqual(f(df).shape[0], 146) @@ -312,12 +312,12 @@ class TestTraceFiltersSyncEvents(unittest.TestCase): def setUp(self): self.htaTrace = TestTraceFiltersSyncEvents.htaTrace - self.df = self.htaTrace.get_trace(0) + self.df = self.htaTrace.get_trace_df(0) def testGPUKernelFilter(self) -> None: f = GPUKernelFilter() filtered_df = f( - self.htaTrace.traces[0], symbol_table=self.htaTrace.symbol_table + self.htaTrace.get_trace_df(0), symbol_table=self.htaTrace.symbol_table ) # GPU kernel is present @@ -330,7 +330,7 @@ def testGPUKernelFilter(self) -> None: def testCPUOperatorFilter(self) -> None: f = CPUOperatorFilter() filtered_df = f( - self.htaTrace.traces[0], symbol_table=self.htaTrace.symbol_table + self.htaTrace.get_trace_df(0), symbol_table=self.htaTrace.symbol_table ) # CPU operator is present diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index 6aefa1eb..0367cce8 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -11,6 +11,8 @@ # import unittest.mock as mock import pandas as pd +from hta.common import singletrace +from hta.common.singletrace import Trace from hta.common.trace_collection import parse_trace_dict, TraceCollection from hta.common.trace_parser import ( _auto_detect_parser_backend, @@ -167,8 +169,8 @@ def test_trace_load(self) -> None: sym_id_map = t.symbol_table.get_sym_id_map() sym_table = t.symbol_table.get_sym_table() - rank_0_df_name_id = t.traces[0]["name"] - rank_0_df_name = t.traces[0]["name"].apply(lambda x: sym_table[x]) + rank_0_df_name_id = t.get_trace_df(0)["name"] + rank_0_df_name = t.get_trace_df(0)["name"].apply(lambda x: sym_table[x]) ground_truth_name = raw_df["name"] ground_truth_name_id = raw_df["name"].apply(lambda x: sym_id_map[x]) @@ -188,19 +190,21 @@ def test_trace_load(self) -> None: sym_id_map = t.symbol_table.get_sym_id_map() profiler_steps = [v for k, v in sym_id_map.items() if "ProfilerStep" in k] - filtered_profiler_steps = t.traces[0]["name"].isin(profiler_steps).sum() + filtered_profiler_steps = ( + t.get_trace_df(0)["name"].isin(profiler_steps).sum() + ) self.assertEqual( filtered_profiler_steps + int(raw_profiler_steps > 1), raw_profiler_steps, ) - self.assertLessEqual(len(t.traces[0]), len(raw_df)) - self.assertGreaterEqual(t.traces[0]["ts"].min(), 0) + self.assertLessEqual(len(t.get_trace_df(0)), len(raw_df)) + self.assertGreaterEqual(t.get_trace_df(0)["ts"].min(), 0) def test_trace_iteration(self) -> None: # run tests for each collection of traces for t in self.traces: - df = t.traces[0] + df = t.get_trace_df(0) sym_id_map = t.symbol_table.get_sym_id_map() iterations = { f"ProfilerStep#{i}" @@ -228,7 +232,8 @@ def test_trace_iteration(self) -> None: ) def test_trace_metadata(self) -> None: - trace_meta = self.vision_transformer_t.meta_data[0] + trace: Trace = self.vision_transformer_t.get_trace(0) + trace_meta = trace.meta exp_meta = EXPECTED_META_VISION_TRANFORMER self.assertEqual(trace_meta["schemaVersion"], exp_meta["schemaVersion"]) self.assertEqual(trace_meta["distributedInfo"], exp_meta["distributedInfo"]) @@ -394,7 +399,7 @@ def test_align_and_filter_mtia(self) -> None: self.assertGreaterEqual(len(t.get_ranks()), 1) # Ensure that the trace has the correct iterations - result_df = t.get_trace(t.get_ranks()[0]) + result_df = t.get_trace_df(t.get_ranks()[0]) self.assertTrue(result_df["ts"].ge(0).all()) self.assertTrue(result_df["iteration"].ge(0).all()) @@ -452,7 +457,7 @@ def test_nccl_parser_config(self) -> None: self.resnet_nccl_t.parse_traces(max_ranks=1, use_multiprocessing=False) self.resnet_nccl_t.decode_symbol_ids(use_shorten_name=False) - trace_df = self.resnet_nccl_t.get_trace(0) + trace_df = self.resnet_nccl_t.get_trace_df(0) self.assertGreater(len(trace_df), 0) nccl_kernels = trace_df.query( @@ -480,7 +485,7 @@ def test_triton_trace(self) -> None: self.triton_t.parse_traces(max_ranks=1, use_multiprocessing=False) self.triton_t.decode_symbol_ids(use_shorten_name=False) - trace_df = self.triton_t.get_trace(0) + trace_df = self.triton_t.get_trace_df(0) self.assertGreater(len(trace_df), 0) self.assertTrue("kernel_backend" in trace_df.columns) self.assertTrue("kernel_hash" in trace_df.columns) @@ -546,7 +551,8 @@ def test_parse_all_args( trace_file = os.path.join(self.resnet_nccl_trace, "nccl_data.json.gz") cfg = ParserConfig(ParserConfig.get_minimum_args()) cfg.set_parse_all_args(parse_all_args) - _, df, _ = parse_trace_dataframe(trace_file, cfg) + trace: Trace = parse_trace_dataframe(trace_file, cfg) + df = trace.df self.assertTrue(expected_columns.issubset(set(df.columns))) self.assertTrue(expected_missing_columns.isdisjoint(set(df.columns))) @@ -617,7 +623,7 @@ def test_fix_mtia_memory_kernels(self) -> None: ) # Create a TraceCollection object t = TraceCollection(trace_dir="", trace_files={}) - t.traces[0] = df.copy() + t.traces[0] = singletrace.create_default(df=df.copy()) t.symbol_table = symbol_table # Expected result after applying fix @@ -625,8 +631,8 @@ def test_fix_mtia_memory_kernels(self) -> None: expected_df.loc[[1, 3], "iteration"] = 1 expected_df.loc[[2], "stream"] = expected_df.loc[[2], "tid"] - t._fix_mtia_memory_kernels(t.get_trace(0)) - fixed_df = t.get_trace(0) + t._fix_mtia_memory_kernels(t.get_trace_df(0)) + fixed_df = t.get_trace_df(0) # Validate results pd.testing.assert_frame_equal(fixed_df, expected_df) From 4870b9926096156dd5c404360857a90aa505b6d0 Mon Sep 17 00:00:00 2001 From: Benedykt Bela Date: Mon, 17 Nov 2025 15:21:37 +0200 Subject: [PATCH 4/4] Address review comments. Signed-off-by: Benedykt Bela --- hta/common/singletrace.py | 5 ----- hta/common/trace_call_graph.py | 2 +- hta/common/trace_collection.py | 31 ++++++++++++++++++++++++------- tests/test_call_stack.py | 3 --- tests/test_correlation.py | 2 +- tests/test_trace_parse.py | 5 ++--- 6 files changed, 28 insertions(+), 20 deletions(-) diff --git a/hta/common/singletrace.py b/hta/common/singletrace.py index 33a6e6b9..52c9f83a 100644 --- a/hta/common/singletrace.py +++ b/hta/common/singletrace.py @@ -37,11 +37,6 @@ class _XPUSingleTrace(_SingleTrace): Trace = _SingleTrace -def create_default(meta=None, df=None, symbol_table=None) -> Trace: - """Factory method to create default Trace object.""" - return _SingleTrace(meta, df, symbol_table) - - def create(device_type: str, meta, df, symbol_table) -> Trace: """Factory method to create Trace object based on device type.""" diff --git a/hta/common/trace_call_graph.py b/hta/common/trace_call_graph.py index bba26aef..b54ac668 100644 --- a/hta/common/trace_call_graph.py +++ b/hta/common/trace_call_graph.py @@ -122,7 +122,7 @@ def from_dataframe( t.symbol_table = ( symbol_table if symbol_table else TraceSymbolTable.create_from_df(df) ) - t.traces[rank] = singletrace.create_default(df=df.copy()) + t.traces[rank] = singletrace.create(None, None, df.copy(), None) t.is_parsed = True cg = CallGraph(t) diff --git a/hta/common/trace_collection.py b/hta/common/trace_collection.py index cfdde1d0..8ca76fe8 100644 --- a/hta/common/trace_collection.py +++ b/hta/common/trace_collection.py @@ -345,8 +345,7 @@ class TraceCollection: trace_path (str) : the path to the folder where the collected raw traces are stored. In other words, `trace_path = normalize_path(base_trace_dir)`. trace_files (Dict[int, str]) : a dictionary that maps the rank of a job's trainer to its trace file. - traces (Dict[int, pd.DataFrame]) : a dictionary that maps the rank of a job's trainer to its trace data. - meta_data (Dict[int, MetaData]) : a dictionary that maps the rank of a job's trainer to its meta_data. + traces (Dict[int, Trace]) : a dictionary that maps the rank of a job's trainer to its trace data. symbol_table (TraceSymbolTable) : a symbol table used to encode the symbols in the trace. is_parsed (bool) : a flag indicting whether the trace is parsed or not. parser_config (ParserConfig) : a configuration object for customizing the parser. @@ -396,7 +395,6 @@ def __init__( logger.debug(self.trace_files) self.traces: Dict[int, Trace] = {} self.symbol_table = TraceSymbolTable() - self.meta_data: Dict[int, MetaData] = {} self.min_ts: int = 0 self._normalize_trace_filenames() @@ -590,6 +588,25 @@ def get_iterations(self, rank: Optional[int] = None) -> List[int]: return sorted([i for i in df["iteration"].unique() if i >= 0]) return [] + def get_trace_meta(self, rank: int) -> MetaData: + """ + Get the trace's MetaData for a given rank. + + Args: + rank (int) : the rank of the trainer. + + Returns: + The trace's MetaData for the given rank. + + Raises: + ValueError when this TraceCollection object doesn't have trace for the given rank. + """ + if rank not in self.traces: + logger.error(f"get_rank_trace - no trace for rank {rank}") + raise ValueError + + return self.traces[rank].meta + def get_trace_df(self, rank: int) -> pd.DataFrame: """ Get the trace's DataFrame for a given rank. @@ -880,8 +897,8 @@ def get_device_type(self) -> str: The device type parsed from the trace metadata. If the device type is not found, return "UNKNOWN". """ rank = next(iter(self.traces)) - trace: Trace = self.traces[rank] - device_type = trace.meta.get("device_type", "UNKNOWN") + meta = self.get_trace_meta(rank) + device_type = meta.get("device_type", "UNKNOWN") return device_type def convert_time_series_to_events( @@ -968,6 +985,6 @@ def get_trace_start_unixtime_ns(self, rank: int) -> int: logger.warning(err_msg) raise ValueError(err_msg) - trace: Trace = self.get_trace(rank) + meta = self.get_trace_meta(rank) - return trace_event_timestamp_to_unixtime_ns(self.min_ts, trace.meta) + return trace_event_timestamp_to_unixtime_ns(self.min_ts, meta) diff --git a/tests/test_call_stack.py b/tests/test_call_stack.py index e8199338..cf03614a 100644 --- a/tests/test_call_stack.py +++ b/tests/test_call_stack.py @@ -152,9 +152,6 @@ def __init__(self, traces): def get_all_traces(self): return self.traces.keys() - def get_trace(self, rank): - return self.traces[rank] - def get_trace_df(self, rank): return self.traces[rank] diff --git a/tests/test_correlation.py b/tests/test_correlation.py index a7dd568f..53cd52ca 100644 --- a/tests/test_correlation.py +++ b/tests/test_correlation.py @@ -36,7 +36,7 @@ def test_something(self): mock_symbol_table = TraceSymbolTable() mock_symbol_table.sym_index = sym_index expected_index_correlation = [5, 6, 0, -1, 1, 2, 0, -1] - trace = singletrace.create_default(df=df, symbol_table=mock_symbol_table) + trace = singletrace.create(None, None, df, mock_symbol_table) df2 = transform_correlation_to_index(trace) self.assertListEqual( expected_index_correlation, df2["index_correlation"].tolist() diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index 0367cce8..2407c55a 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -232,8 +232,7 @@ def test_trace_iteration(self) -> None: ) def test_trace_metadata(self) -> None: - trace: Trace = self.vision_transformer_t.get_trace(0) - trace_meta = trace.meta + trace_meta = self.vision_transformer_t.get_trace_meta(0) exp_meta = EXPECTED_META_VISION_TRANFORMER self.assertEqual(trace_meta["schemaVersion"], exp_meta["schemaVersion"]) self.assertEqual(trace_meta["distributedInfo"], exp_meta["distributedInfo"]) @@ -623,7 +622,7 @@ def test_fix_mtia_memory_kernels(self) -> None: ) # Create a TraceCollection object t = TraceCollection(trace_dir="", trace_files={}) - t.traces[0] = singletrace.create_default(df=df.copy()) + t.traces[0] = singletrace.create(None, None, df.copy(), None) t.symbol_table = symbol_table # Expected result after applying fix