Fix test flakiness and add timeline coverage (#347)

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()

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 ----
 
 

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
 

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()