From 3c6803587afd31f53cc52d1c53bdc5a83a1e560a Mon Sep 17 00:00:00 2001 From: Xizhou Feng Date: Tue, 12 May 2026 18:38:34 -0700 Subject: [PATCH] Fix test flakiness and add timeline coverage (#347) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Fix two root causes of flakiness in `test_timeline.py` and add coverage for uncovered functions in `hta/analyzers/timeline.py`. **Flakiness fixes:** - `get_trace(rank)` returns a direct reference to the internal DataFrame, not a copy. All `setUp` methods and inline calls now use `.copy()` to ensure test isolation — preventing shared mutable state between tests that caused non-deterministic failures depending on execution order. - `test_raises_on_partial_columns` had a silent-pass path: when trace data contained all required columns, the test body never executed. Rewritten to always truncate `cols_present` so the ValueError assertion always fires. **Coverage additions (test_coverage_gaps.py):** - `test_simplify_name_truncates_long_names` — covers truncation path in `_simplify_name` - `test_plot_gpu_kernels_delegates` — covers `plot_timeline_gpu_kernels` (previously 0% coverage) - `test_plot_gpu_kernels_from_trace` — covers `plot_timeline_gpu_kernels_from_trace` (previously 0% coverage) Reviewed By: Chenguang-Zhu Differential Revision: D104891668 --- tests/test_analyzers_timeline.py | 62 ++++++++++++++++++++++ tests/test_coverage_gaps.py | 89 ++++++++++++++++++++++++++++++++ tests/test_timeline.py | 28 +++++----- tests/test_trace_parse.py | 47 ++++++++++++++++- 4 files changed, 209 insertions(+), 17 deletions(-) create mode 100644 tests/test_analyzers_timeline.py diff --git a/tests/test_analyzers_timeline.py b/tests/test_analyzers_timeline.py new file mode 100644 index 00000000..41d6135d --- /dev/null +++ b/tests/test_analyzers_timeline.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for hta.analyzers.timeline module.""" + +import os +import unittest +from unittest.mock import Mock, patch + +import pandas as pd +from hta.analyzers.timeline import ( + _simplify_name, + plot_timeline_gpu_kernels, + plot_timeline_gpu_kernels_from_trace, +) +from hta.common.trace import Trace +from hta.common.trace_symbol_table import TraceSymbolTable +from hta.utils.test_utils import get_test_data_dir + + +class TestAnalyzersTimeline(unittest.TestCase): + """Coverage tests for hta/analyzers/timeline.py.""" + + def test_simplify_name_truncates_long_names(self) -> None: + long_name = "a" * 100 + result = _simplify_name(long_name) + self.assertEqual(len(result), 45) + self.assertTrue(result.endswith("...")) + + @patch("hta.analyzers.timeline.px.timeline") + def test_plot_gpu_kernels_delegates(self, mock_px_timeline: Mock) -> None: + mock_px_timeline.return_value = Mock() + sym_table = TraceSymbolTable() + sym_table.add_symbols(["aten::mm", "kernel"]) + sym_id_map = sym_table.get_sym_id_map() + df = pd.DataFrame( + { + "iteration": [1], + "name": [sym_id_map["aten::mm"]], + "cat": [sym_id_map["kernel"]], + "rank": [0], + "stream": [7], + "ts": [1000], + "dur": [500], + } + ) + plot_timeline_gpu_kernels("test", df, sym_table) + mock_px_timeline.assert_called_once() + + @patch("hta.analyzers.timeline.plot_timeline") + def test_plot_gpu_kernels_from_trace(self, mock_plot: Mock) -> None: + trace_path = os.path.join(get_test_data_dir(), "timeline_analysis") + t = Trace(trace_dir=trace_path) + t.parse_traces() + t.decode_symbol_ids(use_shorten_name=False) + plot_timeline_gpu_kernels_from_trace("test", t) + mock_plot.assert_called_once() + + +if __name__ == "__main__": # pragma: no cover + unittest.main() diff --git a/tests/test_coverage_gaps.py b/tests/test_coverage_gaps.py index 872a3315..ad1f7ea0 100644 --- a/tests/test_coverage_gaps.py +++ b/tests/test_coverage_gaps.py @@ -756,6 +756,95 @@ def test_plot_timeline_no_rank_column(self, mock_px_timeline: Mock) -> None: mock_px_timeline.assert_called_once() +# ---- parser_config.py coverage tests ---- + + +class TestParserConfigCoverage(unittest.TestCase): + def test_repr(self) -> None: + from hta.configs.parser_config import ParserConfig + + cfg = ParserConfig.get_default_cfg() + r = repr(cfg) + self.assertIn("ParserConfig(", r) + self.assertIn("parse_all_args=", r) + + def test_get_versioned_cfg(self) -> None: + from hta.configs.parser_config import DEFAULT_PARSE_VERSION, ParserConfig + + cfg = ParserConfig.get_versioned_cfg(DEFAULT_PARSE_VERSION) + self.assertIsInstance(cfg, ParserConfig) + + def test_get_info_args(self) -> None: + from hta.configs.parser_config import ParserConfig + + args = ParserConfig.get_info_args() + self.assertIsInstance(args, list) + self.assertGreater(len(args), 0) + + def test_set_and_get_args_with_selector(self) -> None: + from hta.configs.parser_config import ParserConfig + + cfg = ParserConfig.get_default_cfg() + all_args = cfg.get_args() + if all_args: + cfg.selected_arg_keys = [all_args[0].name] + selected = cfg.get_args() + self.assertEqual(len(selected), 1) + self.assertEqual(selected[0].name, all_args[0].name) + cfg.selected_arg_keys = None + + def test_set_min_required_cols(self) -> None: + from hta.configs.parser_config import ParserConfig + + cfg = ParserConfig.get_default_cfg() + cfg.set_min_required_cols(["a", "b"]) + self.assertEqual(cfg.get_min_required_cols(), ["a", "b"]) + + def test_enable_communication_args(self) -> None: + from hta.configs.parser_config import ParserConfig + + cfg = ParserConfig.get_default_cfg() + before = len(cfg.get_args()) + ParserConfig.enable_communication_args(cfg) + self.assertGreaterEqual(len(cfg.get_args()), before) + + def test_set_global_parser_config_version(self) -> None: + from hta.configs.parser_config import DEFAULT_PARSE_VERSION, ParserConfig + + ParserConfig.set_global_parser_config_version(DEFAULT_PARSE_VERSION) + self.assertIsNotNone(ParserConfig.ARGS_MINIMUM) + self.assertIsNotNone(ParserConfig.ARGS_DEFAULT) + + def test_show_available_args(self) -> None: + from hta.configs.parser_config import ParserConfig + + with patch("builtins.print"): + ParserConfig.show_available_args() + + def test_transform_arg_name_special(self) -> None: + from hta.configs.parser_config import ParserConfig + + self.assertEqual(ParserConfig.transform_arg_name("name"), "arg_name") + self.assertEqual( + ParserConfig.transform_arg_name("Input Dims (MB)"), "input_dims" + ) + + def test_make_attribute_spec_types(self) -> None: + from hta.configs.parser_config import ParserConfig + + spec_int = ParserConfig.make_attribute_spec("count", 42) + self.assertEqual(spec_int.raw_name, "count") + + spec_float = ParserConfig.make_attribute_spec("rate", 3.14) + self.assertIsNotNone(spec_float) + + spec_str = ParserConfig.make_attribute_spec("label", "hello") + self.assertIsNotNone(spec_str) + + spec_obj = ParserConfig.make_attribute_spec("data", [1, 2, 3]) + self.assertIsNotNone(spec_obj) + + # ---- trace_file.py additional tests ---- diff --git a/tests/test_timeline.py b/tests/test_timeline.py index 8f56f038..f5ca34bf 100644 --- a/tests/test_timeline.py +++ b/tests/test_timeline.py @@ -112,7 +112,7 @@ def setUpClass(cls) -> None: cls.t.decode_symbol_ids(use_shorten_name=False) def setUp(self) -> None: - self.df = self.t.get_trace(0) + self.df = self.t.get_trace(0).copy() def test_prepare_returns_required_columns(self) -> None: setting = TimelinePlotSetting(task_height=50, plot_format=PlotFormat.File) @@ -196,7 +196,7 @@ def setUpClass(cls) -> None: cls.t.decode_symbol_ids(use_shorten_name=False) def setUp(self) -> None: - self.df = self.t.get_trace(0) + self.df = self.t.get_trace(0).copy() def test_raises_on_missing_columns(self) -> None: """Should raise ValueError when required columns are missing.""" @@ -205,15 +205,14 @@ def test_raises_on_missing_columns(self) -> None: align_module_with_kernels(minimal_df, ["## forward ##"]) def test_raises_on_partial_columns(self) -> None: - """Should raise ValueError with only the REQUIRED_COLUMNS set (missing kernel columns).""" - cols_present = list( - REQUIRED_COLUMNS_FOR_CPU_GPU_ALIGNMENT & set(self.df.columns) - ) - if len(cols_present) < len(REQUIRED_COLUMNS_FOR_CPU_GPU_ALIGNMENT): - # Not all required columns exist; the function should raise - partial_df = self.df[cols_present].copy() - with self.assertRaises(ValueError): - align_module_with_kernels(partial_df, ["## forward ##"]) + """Should raise ValueError with a subset of required columns.""" + required = REQUIRED_COLUMNS_FOR_CPU_GPU_ALIGNMENT + cols_present = list(required & set(self.df.columns)) + if len(cols_present) >= len(required): + cols_present = cols_present[: len(required) - 1] + partial_df = self.df[cols_present].copy() + with self.assertRaises(ValueError): + align_module_with_kernels(partial_df, ["## forward ##"]) def test_align_with_symbol_table(self) -> None: """Should produce aligned events when all required columns are present.""" @@ -228,8 +227,7 @@ def test_align_with_symbol_table(self) -> None: ) annotations = aligned[aligned["stream"].eq(AnnotationStreamID)] self.assertGreater(aligned.shape[0], 0) - # All annotations should have non-negative num_kernels - if "num_kernels" in annotations.columns and not annotations.empty: + if "num_kernels" in annotations.columns: self.assertEqual(annotations[annotations["num_kernels"].lt(0)].shape[0], 0) def test_align_without_symbol_table(self) -> None: @@ -259,7 +257,7 @@ def setUpClass(cls) -> None: @patch(f"{_MODULE}.px.timeline") def test_plot_calls_plotly(self, mock_timeline: Mock) -> None: - df = self.t.get_trace(0) + df = self.t.get_trace(0).copy() setting = TimelinePlotSetting(task_height=50, plot_format=PlotFormat.File) timeline_events = prepare_timeline_events( df, symbol_table=self.t.symbol_table, setting=setting @@ -302,7 +300,7 @@ def setUpClass(cls) -> None: @patch(f"{_MODULE}.plot_events_timeline") def test_timeline_prepare_and_plot(self, mock_plot: Mock) -> None: - df = self.t.get_trace(0) + df = self.t.get_trace(0).copy() tl = Timeline(df, self.t.symbol_table, filter_func=GPUKernelFilter()) tl.setting.plot_format = PlotFormat.File diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index b2dee438..5fb4e9cd 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -2,12 +2,14 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +import gzip +import json import math import os +import tempfile import unittest from typing import Any, Dict, Optional, Set - -# import unittest.mock as mock +from unittest.mock import patch import pandas as pd from hta.common.trace import parse_trace_dict, Trace @@ -20,6 +22,7 @@ parse_metadata_ijson, parse_trace_dataframe, ParserBackend, + round_down_time_stamps, set_default_trace_parsing_backend, ) from hta.common.trace_symbol_table import TraceSymbolTable @@ -730,5 +733,45 @@ def test_compress_df_mixed_events_filters_nulls(self): self.assertEqual(len(result_df), 1) +class TestTraceParserCoverage(unittest.TestCase): + """Coverage tests for uncovered branches in trace_parser.py.""" + + def test_infer_gpu_type_name_set_none(self) -> None: + self.assertEqual(infer_gpu_type({}, None), "UNKNOWN GPU") + + def test_infer_gpu_type_cuda(self) -> None: + self.assertEqual(infer_gpu_type({}, {"cudaLaunchKernel": 1}), "NVIDIA GPU") + + def test_infer_gpu_type_hip(self) -> None: + self.assertEqual(infer_gpu_type({}, {"hipLaunchKernel": 1}), "AMD GPU") + + def test_auto_detect_parser_backend_no_ijson(self) -> None: + with patch.dict("sys.modules", {"ijson": None}): + result = _auto_detect_parser_backend() + self.assertEqual(result, ParserBackend.JSON) + + def test_parse_trace_dict_gz(self) -> None: + data = {"traceEvents": [{"name": "test", "ph": "X"}]} + with tempfile.NamedTemporaryFile(suffix=".gz", delete=False) as f: + f.write(gzip.compress(json.dumps(data).encode())) + gz_path = f.name + try: + result = parse_trace_dict(gz_path) + self.assertEqual(result, data) + finally: + os.unlink(gz_path) + + def test_parse_trace_dict_invalid_extension(self) -> None: + with self.assertRaises(ValueError): + parse_trace_dict("/tmp/trace.txt") + + def test_round_down_disabled(self) -> None: + df = pd.DataFrame({"ts": [1.5, 2.5], "dur": [0.5, 0.5]}) + with patch("hta.common.trace_parser.hta_options") as mock_opts: + mock_opts.disable_ns_rounding.return_value = True + round_down_time_stamps(df) + self.assertEqual(df["ts"].tolist(), [1.5, 2.5]) + + if __name__ == "__main__": # pragma: no cover unittest.main()