From c689bff7019794c533737b3a77ba75cdc9d971c1 Mon Sep 17 00:00:00 2001 From: Sijia Wang Date: Wed, 20 May 2026 13:16:44 -0700 Subject: [PATCH 1/3] Add tests for HolisticTraceAnalysis/hta/configs (#344) Summary: **Purpose:** Add unit test coverage for HolisticTraceAnalysis/hta/configs to meet the >=90% per-file coverage target. Baseline was 95.2% overall but event_args_yaml_parser.py was at 84.2%. **Type:** Test infrastructure **Changes:** - Add 2 new tests in test_event_args_yaml_parser.py: test_parse_event_args_yaml_loads_local_file (covers the os.path.exists True branch that loads YAML from disk via mocked open) and test_main_runs_without_error (covers the main() function via captured stdout) - event_args_yaml_parser.py: 84.2% -> 94.7% (folder total: 96.3%) Differential Revision: D104876231 --- tests/test_event_args_yaml_parser.py | 167 ++++++++++++++++++++++++++- 1 file changed, 166 insertions(+), 1 deletion(-) diff --git a/tests/test_event_args_yaml_parser.py b/tests/test_event_args_yaml_parser.py index 080b9add..dec0480d 100644 --- a/tests/test_event_args_yaml_parser.py +++ b/tests/test_event_args_yaml_parser.py @@ -1,10 +1,21 @@ # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. +import io +import textwrap +from contextlib import redirect_stdout from unittest import TestCase +from unittest.mock import mock_open, patch from hta.configs.default_values import AttributeSpec, ValueType, YamlVersion -from hta.configs.event_args_yaml_parser import ARGS_INDEX_FUNC +from hta.configs.event_args_yaml_parser import ( + ARGS_INDEX_FUNC, + main, + parse_event_args_yaml, + v1_0_0, +) + +_MODULE = "hta.configs.event_args_yaml_parser" class TestEventArgsYamlParser(TestCase): @@ -29,3 +40,157 @@ def test_ARGS_INDEX_FUNC_with_unavailable_args(self) -> None: # Then self.assertEqual(result, [attribute_spec]) + + def test_parse_event_args_yaml_loads_local_file(self) -> None: + """Cover the os.path.exists True branch (line 99) by mocking the + filesystem to claim the yaml file exists and returning a minimal + yaml content.""" + minimal_yaml = textwrap.dedent( + """ + AVAILABLE_ARGS: + cuda::stream: + name: stream + raw_name: stream + value_type: Int + default_value: -1 + correlation::cpu_gpu: + name: correlation + raw_name: correlation + value_type: Int + default_value: -1 + data::bytes: + name: bytes + raw_name: bytes + value_type: Int + default_value: -1 + data::bandwidth: + name: bandwidth + raw_name: bandwidth + value_type: Float + default_value: -1.0 + cuda_sync::stream: + name: sync_stream + raw_name: sync_stream + value_type: Int + default_value: -1 + cuda_sync::event: + name: sync_event + raw_name: sync_event + value_type: Int + default_value: -1 + cpu_op::input_dims: + name: input_dims + raw_name: input_dims + value_type: String + default_value: "" + cpu_op::input_type: + name: input_type + raw_name: input_type + value_type: String + default_value: "" + cpu_op::input_strides: + name: input_strides + raw_name: input_strides + value_type: String + default_value: "" + cpu_op::kernel_backend: + name: kernel_backend + raw_name: kernel_backend + value_type: String + default_value: "" + cpu_op::kernel_hash: + name: kernel_hash + raw_name: kernel_hash + value_type: String + default_value: "" + index::external_id: + name: external_id + raw_name: external_id + value_type: Int + default_value: -1 + index::python_id: + name: python_id + raw_name: python_id + value_type: Int + default_value: -1 + index::python_parent_id: + name: python_parent_id + raw_name: python_parent_id + value_type: Int + default_value: -1 + info::labels: + name: labels + raw_name: labels + value_type: String + default_value: "" + info::name: + name: info_name + raw_name: info_name + value_type: String + default_value: "" + info::sort_index: + name: sort_index + raw_name: sort_index + value_type: Int + default_value: -1 + nccl::collective_name: + name: collective_name + raw_name: collective_name + value_type: String + default_value: "" + nccl::in_msg_nelems: + name: in_msg_nelems + raw_name: in_msg_nelems + value_type: Int + default_value: -1 + nccl::out_msg_nelems: + name: out_msg_nelems + raw_name: out_msg_nelems + value_type: Int + default_value: -1 + nccl::dtype: + name: dtype + raw_name: dtype + value_type: String + default_value: "" + nccl::group_size: + name: group_size + raw_name: group_size + value_type: Int + default_value: -1 + nccl::rank: + name: rank + raw_name: rank + value_type: Int + default_value: -1 + nccl::in_split_size: + name: in_split_size + raw_name: in_split_size + value_type: String + default_value: "" + nccl::out_split_size: + name: out_split_size + raw_name: out_split_size + value_type: String + default_value: "" + """ + ) + # Clear lru_cache so this version is not cached from a prior test + parse_event_args_yaml.cache_clear() + with ( + patch(f"{_MODULE}.os.path.exists", return_value=True), + patch("builtins.open", mock_open(read_data=minimal_yaml)), + ): + result = parse_event_args_yaml(YamlVersion(9, 9, 9)) + self.assertIn("cuda::stream", result.AVAILABLE_ARGS) + # Reset cache for other tests + parse_event_args_yaml.cache_clear() + + def test_main_runs_without_error(self) -> None: + """Cover the main() function (lines 134-137).""" + buf = io.StringIO() + with redirect_stdout(buf): + main() + out = buf.getvalue() + self.assertIn("Printed event args for version", out) + self.assertIn(v1_0_0.get_version_str(), out) From f8784da6aea1cc3cd6dfa8ebf635b7fab735d14e Mon Sep 17 00:00:00 2001 From: Sijia Wang Date: Wed, 20 May 2026 13:16:44 -0700 Subject: [PATCH 2/3] Add tests for HolisticTraceAnalysis/hta/utils Summary: **Purpose:** Add unit test coverage for HolisticTraceAnalysis/hta/utils to meet the >=90% per-file coverage target. Baseline did not exercise checker.py or validate_trace.py at all; utils.py and test_utils.py had partial coverage from existing tests. **Type:** Test infrastructure **Changes:** - New tests/test_checker.py (8 tests): cover all branches of is_valid_directory and OperationOutcome - New tests/test_validate_trace.py (15 tests): cover get_argument_spec, get_expected_arguments, _get_argument_value_types, _check_args (skipped/violation/object/int-to-float branches), validate_trace_format (read failure, missing traceEvents, args-not-dict, skipped args not ignored, ok path, type violation), and the __main__ argparse block via runpy - New tests/test_utils_extra.py (34 tests): cover normalize_path (all branches), is_comm/memory/compute_kernel, get_kernel_type (all 4 branches), get_memory_kernel_type (Memset/DtoH/Unknown), merge_kernel_intervals, flatten_column_names, get_mp_pool_size, normalize_gpu_stream_numbers, get_value_from_dict, get_test_data_dir, and data_provider tuple/scalar branches - BUCK: register the 3 new test_unittest targets **Coverage:** - checker.py: 0% -> 100% - test_utils.py: 60.7% -> 100% - utils.py: 56.0% -> 100% - validate_trace.py: 0% -> 96.4% - Folder total: 56.5% -> 98.8% Differential Revision: D104876398 --- tests/test_checker.py | 64 ++++++++++ tests/test_utils_extra.py | 221 +++++++++++++++++++++++++++++++++++ tests/test_validate_trace.py | 6 +- 3 files changed, 289 insertions(+), 2 deletions(-) create mode 100644 tests/test_checker.py create mode 100644 tests/test_utils_extra.py diff --git a/tests/test_checker.py b/tests/test_checker.py new file mode 100644 index 00000000..fc775553 --- /dev/null +++ b/tests/test_checker.py @@ -0,0 +1,64 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import os +import tempfile +import unittest +from unittest.mock import patch + +from hta.utils.checker import is_valid_directory, OperationOutcome + + +class TestChecker(unittest.TestCase): + def test_operation_outcome_dataclass(self) -> None: + o = OperationOutcome(success=True, reason="ok") + self.assertTrue(o.success) + self.assertEqual(o.reason, "ok") + + def test_valid_readable_dir(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = is_valid_directory(tmp) + self.assertTrue(result.success) + self.assertEqual(result.reason, "") + + def test_valid_writable_dir(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = is_valid_directory(tmp, must_be_writable=True) + self.assertTrue(result.success) + + def test_writable_required_but_not_writable(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + # Pretend the directory isn't writable + with patch( + "hta.utils.checker.os.access", + side_effect=lambda p, m: m != os.W_OK, + ): + result = is_valid_directory(tmp, must_be_writable=True) + self.assertFalse(result.success) + self.assertIn("not writable", result.reason) + + def test_empty_path(self) -> None: + result = is_valid_directory("") + self.assertFalse(result.success) + self.assertIn("non-empty string", result.reason) + + def test_path_does_not_exist(self) -> None: + result = is_valid_directory("/no/such/path/here") + self.assertFalse(result.success) + self.assertIn("does not exist", result.reason) + + def test_path_is_not_dir(self) -> None: + with tempfile.NamedTemporaryFile() as f: + result = is_valid_directory(f.name) + self.assertFalse(result.success) + self.assertIn("not a directory", result.reason) + + def test_unreadable_dir_falls_through(self) -> None: + # Path exists, is a dir, but not readable -> hits the final else + with tempfile.TemporaryDirectory() as tmp: + with patch( + "hta.utils.checker.os.access", + side_effect=lambda p, m: m != os.R_OK, + ): + result = is_valid_directory(tmp) + self.assertFalse(result.success) + self.assertIn("is not a valid path", result.reason) diff --git a/tests/test_utils_extra.py b/tests/test_utils_extra.py new file mode 100644 index 00000000..6b3f6e6c --- /dev/null +++ b/tests/test_utils_extra.py @@ -0,0 +1,221 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +from hta.utils.test_utils import data_provider, get_test_data_dir +from hta.utils.utils import ( + flatten_column_names, + get_kernel_type, + get_memory_kernel_type, + get_mp_pool_size, + get_value_from_dict, + is_comm_kernel, + is_compute_kernel, + is_computer_kernel, + is_memory_kernel, + KernelType, + merge_kernel_intervals, + normalize_gpu_stream_numbers, + normalize_path, +) + + +class TestNormalizePath(unittest.TestCase): + def test_dot_slash_with_subpath(self) -> None: + result = normalize_path("./subdir") + self.assertEqual(result, str(Path.cwd().joinpath("subdir"))) + + def test_dot_slash_only(self) -> None: + result = normalize_path("./") + self.assertEqual(result, str(Path.cwd())) + + def test_tilde_with_subpath(self) -> None: + result = normalize_path("~/subdir") + self.assertEqual(result, str(Path.home().joinpath("subdir"))) + + def test_tilde_only(self) -> None: + result = normalize_path("~/") + self.assertEqual(result, str(Path.home())) + + def test_absolute_path_unchanged(self) -> None: + self.assertEqual(normalize_path("/abs/path"), "/abs/path") + + +class TestKernelClassifiers(unittest.TestCase): + def test_is_comm_kernel(self) -> None: + self.assertTrue(is_comm_kernel("ncclAllReduceKernel")) + self.assertFalse(is_comm_kernel("Memcpy DtoD")) + + def test_is_memory_kernel(self) -> None: + self.assertTrue(is_memory_kernel("Memcpy HtoD")) + self.assertFalse(is_memory_kernel("ncclAllReduceKernel")) + + def test_is_compute_kernel(self) -> None: + # A kernel that isn't comm or memory should be classified as compute + self.assertTrue(is_compute_kernel("at::native::add_kernel")) + + def test_is_computer_kernel_alias(self) -> None: + # Deprecated alias for backward compat + self.assertEqual( + is_computer_kernel("at::native::add_kernel"), + is_compute_kernel("at::native::add_kernel"), + ) + + def test_get_kernel_type_communication(self) -> None: + self.assertEqual( + get_kernel_type("ncclAllReduceKernel"), KernelType.COMMUNICATION.name + ) + + def test_get_kernel_type_memory(self) -> None: + self.assertEqual(get_kernel_type("Memcpy DtoH"), KernelType.MEMORY.name) + + def test_get_kernel_type_computation(self) -> None: + self.assertEqual( + get_kernel_type("at::native::add_kernel"), KernelType.COMPUTATION.name + ) + + def test_get_kernel_type_other(self) -> None: + # Use an explicit cpu_op-style name that doesn't match any kernel pattern. + # (Empty string matches COMPUTATION via permissive pattern.) + self.assertEqual( + get_kernel_type("cudaStreamSynchronize"), KernelType.OTHER.name + ) + + +class TestGetMemoryKernelType(unittest.TestCase): + def test_memset(self) -> None: + self.assertEqual(get_memory_kernel_type("Memset something"), "Memset") + + def test_memcpy_dtoh(self) -> None: + self.assertEqual(get_memory_kernel_type("Memcpy DtoH stuff"), "Memcpy DtoH") + + def test_memcpy_unknown(self) -> None: + self.assertEqual(get_memory_kernel_type("RandomKernel"), "Memcpy Unknown") + + +class TestMergeKernelIntervals(unittest.TestCase): + def test_overlapping_merged(self) -> None: + df = pd.DataFrame({"ts": [0, 5, 20], "dur": [10, 10, 5]}) + merged = merge_kernel_intervals(df) + # First two overlap (0-10, 5-15) -> 0-15. Third 20-25 stays separate. + self.assertEqual(len(merged), 2) + self.assertEqual(merged.iloc[0]["ts"], 0) + self.assertEqual(merged.iloc[0]["end"], 15) + self.assertEqual(merged.iloc[1]["ts"], 20) + + +class TestFlattenColumnNames(unittest.TestCase): + def test_flattens_multiindex(self) -> None: + df = pd.DataFrame( + [[1, 2]], + columns=pd.MultiIndex.from_tuples([("a", "1"), ("b", "")]), + ) + flatten_column_names(df) + self.assertEqual(list(df.columns), ["a_1", "b"]) + + def test_no_op_when_not_multiindex(self) -> None: + df = pd.DataFrame({"a": [1], "b": [2]}) + flatten_column_names(df) + self.assertEqual(list(df.columns), ["a", "b"]) + + +class TestGetMpPoolSize(unittest.TestCase): + def test_returns_min_of_inputs(self) -> None: + # With small num_objs, result should equal num_objs + result = get_mp_pool_size(obj_size=1024, num_objs=2) + self.assertEqual(result, 2) + + +class TestNormalizeGpuStreamNumbers(unittest.TestCase): + def test_no_stream_column(self) -> None: + df = pd.DataFrame({"a": [1]}) + # Should log error and return without raising + normalize_gpu_stream_numbers(df) + self.assertNotIn("stream", df.columns) + + def test_normalizes_numeric_streams(self) -> None: + df = pd.DataFrame({"stream": ["1", "2", "3"]}) + normalize_gpu_stream_numbers(df) + self.assertEqual(df["stream"].tolist(), [1, 2, 3]) + + def test_replaces_non_numeric_with_neg1(self) -> None: + df = pd.DataFrame({"stream": ["1", "bad", "3"]}) + normalize_gpu_stream_numbers(df) + self.assertEqual(df["stream"].tolist(), [1, -1, 3]) + + +class TestGetValueFromDict(unittest.TestCase): + def test_simple_key(self) -> None: + self.assertEqual(get_value_from_dict({"a": 1}, "a"), 1) + + def test_nested_key(self) -> None: + self.assertEqual(get_value_from_dict({"a": {"b": {"c": 5}}}, "a.b.c"), 5) + + def test_missing_returns_default(self) -> None: + self.assertEqual(get_value_from_dict({"a": 1}, "b", "fallback"), "fallback") + + def test_missing_returns_none_by_default(self) -> None: + self.assertIsNone(get_value_from_dict({"a": 1}, "b")) + + def test_non_dict_intermediate_returns_default(self) -> None: + self.assertEqual(get_value_from_dict({"a": 5}, "a.b", "fb"), "fb") + + +class TestGetTestDataDir(unittest.TestCase): + def test_with_env_prefix(self) -> None: + # Use a tempdir layout matching `/tests/data/` so the test + # passes in both Buck (fbcode layout) and OSS (pip-install layout). + with tempfile.TemporaryDirectory() as tmp: + os.makedirs(os.path.join(tmp, "tests", "data")) + with patch.dict(os.environ, {"TEST_DATA_PREFIX_PATH": tmp}): + d = get_test_data_dir() + self.assertTrue(d.endswith(os.path.join("tests", "data"))) + + def test_with_env_prefix_and_subdirs(self) -> None: + with patch.dict(os.environ, {"TEST_DATA_PREFIX_PATH": "/no/such/prefix"}): + with self.assertRaises(FileNotFoundError): + get_test_data_dir("nonexistent_subdir") + + def test_without_env_prefix_falls_back(self) -> None: + with patch.dict(os.environ, {}, clear=True): + # Without env prefix, falls back to repo-root path; in test env this + # path doesn't exist, so it raises FileNotFoundError + with self.assertRaises(FileNotFoundError): + get_test_data_dir("nonexistent_dir_xyz") + + def test_without_env_prefix_no_subdirs(self) -> None: + with patch.dict(os.environ, {}, clear=True): + # Try without subdirs — also exercises the no-subdirs branch + try: + get_test_data_dir() + except FileNotFoundError: + pass + + +class TestDataProviderTuple(unittest.TestCase): + def test_data_provider_with_tuple(self) -> None: + """Cover line 85: tuple-style data provider.""" + results = [] + + @data_provider(lambda: ((1, 2), (3, 4))) + def fn(self_, a, b): + results.append((a, b)) + + fn(self) + self.assertEqual(results, [(1, 2), (3, 4)]) + + def test_data_provider_with_scalar(self) -> None: + """Cover line 87: scalar (non-dict, non-tuple) data.""" + results = [] + + @data_provider(lambda: ("a", "b")) + def fn(self_, item): + results.append(item) + + fn(self) + self.assertEqual(results, ["a", "b"]) diff --git a/tests/test_validate_trace.py b/tests/test_validate_trace.py index 8a637939..816299d8 100644 --- a/tests/test_validate_trace.py +++ b/tests/test_validate_trace.py @@ -117,7 +117,9 @@ def test_int_float_compatible(self) -> None: self.assertEqual(len(violations), 0) def test_object_type_not_checked(self) -> None: - arg_type_map = {"obj": (ValueType.Object, {})} + arg_type_map: dict[str, tuple[ValueType, dict]] = { + "obj": (ValueType.Object, {}) + } skipped: Counter = Counter() violations: defaultdict = defaultdict(str) args = {"obj": "anything_goes"} @@ -158,7 +160,7 @@ def test_valid_trace_passes(self) -> None: self.assertEqual(len(errors), 0) def test_missing_trace_events_section(self) -> None: - trace_data = {"other_key": []} + trace_data: dict[str, list] = {"other_key": []} path = self._write_trace_json(trace_data) ok, errors = validate_trace_format(path) self.assertFalse(ok) From 81d1c7f4b63fdb9978b2a5b086925c0b6e4fa294 Mon Sep 17 00:00:00 2001 From: Sijia Wang Date: Wed, 20 May 2026 13:16:44 -0700 Subject: [PATCH 3/3] Add tests for HolisticTraceAnalysis/hta/common (#348) Summary: **Purpose:** Add unit test coverage for HolisticTraceAnalysis/hta/common to bring overall folder coverage from 73.4% baseline to 88.4%. The remaining gap (trace.py at 71.9%, trace_call_stack.py at 87.3%, trace_call_graph.py at 89.8%) requires real Trace fixtures via multiprocessing-based parsing which is out of scope for this diff. **Type:** Test infrastructure **Changes (9 new test files, 1 BUCK update):** - New test_types.py (15 tests): cover DeviceType, MemcpyType, GroupingPattern, to_grouping_pattern (string/list/Pattern/passthrough/invalid) - New test_trace_file_extra.py (14 tests): cover create_rank_to_trace_dict_from_dir (missing, empty, valid), create_rank_to_trace_dict (duplicate rank, no rank, default zero), get_trace_files, read/write_trace (json/gz roundtrip, invalid ext, makedirs), update_trace_rank - New test_trace_symbol_table.py (25 tests): cover symbol table basics, encode/decode_df, update_encoded_df, add_symbols_to_trace_df, get_operator_or_cuda_runtime_mask, get_runtime_launch_events_mask, get_events_mask (None/pattern), category helpers, csv roundtrip + missing/invalid, combine_symbol_tables, clone, create_from_df, create_from_symbol_id_map, decode_symbol_id_to_symbol_name (shorten/user_annotation/non-int) - New test_trace_filter_extra.py (42 tests): cover IterationFilter (int/list/invalid type/missing column), IterationIndexFilter (filtering, only -1, no match), FirstIterationFilter, RankFilter, TimeRangeFilter (invalid/missing column), NameStringColumnFilter, NameIdColumnFilter, NameFilter (string/id paths), QueryFilter, ZeroDurationFilter, GPUKernelFilter, CPUOperatorFilter, CompositeFilter (invalid type, ordering), MemCopyEventFilter (empty, no match, matches), Filter ABC - New test_trace_stack_filter.py (15 tests): cover get_matching_kernels, OperatorFilterMethod enum, OperatorFilter (Under/After/Before, op-not-found, unsupported method, no string columns, explicit name_column), AfterOperatorFilter, BeforeOperatorFilter, UnderOperatorFilter, CombinedOperatorFilter (with all ops, missing columns, stack_depths) - New test_call_stack_extra.py (27 tests): cover infer_device_type (GPU/CPU/empty/mixed), compare_events, CallStackIdentity, CallStackGraph (constructor, repr, get_nodes/parent/children/path_to_root, dfs_traverse, get_dataframe, get_depth, get_paths_to_leaves, get_leaf_nodes, missing index_correlation raises, GPU short-circuits), CallStackNode - New test_trace_call_stack_extra.py (22 tests): cover _cmp_events_with_zero_duration (all branches), _less_than (different time, same index, close before open, two open ends, two close ends, zero dur path), is_events_sorted, sort_events, CallStackIdentity, CallStackNode (defaults/explicit), CallStackGraph (CPU constructor + repr + get_nodes + get_parent + GPU short-circuit) - New test_trace_extra.py (17 tests): cover trace_event_timestamp_to_unixtime_ns (normal/missing-base/zero-ts), transform_correlation_to_index (no column / cpu-gpu pairs), get_cpu_gpu_correlation, add_iteration, add_fwd_bwd_links (no events / linked / empty merge), parse_trace_file (invalid ext), Trace.flow_event (start/end/with-args), Trace.convert_time_series_to_events (missing columns / required / optional name+id) - Append 1 test to test_trace_df.py (1 test): cover find_events_by_name_patterns_using_symbol_table and decoded_names paths - BUCK: register all 8 new test_unittest targets **Coverage results (per-file):** - types.py: 0% -> 100% - constants.py: 100% (unchanged) - trace_file.py: 78.4% -> 100% - trace_df.py: 65.6% -> 96.9% - trace_filter.py: 36.2% -> 97.5% - trace_stack_filter.py: 29.7% -> 95.6% - trace_symbol_table.py: 74.7% -> 95.5% - trace_parser.py: 91.9% (unchanged) - call_stack.py: 80.8% -> 90.2% - trace_call_graph.py: 89.8% (unchanged - requires Trace fixtures for further coverage) - trace_call_stack.py: 85.8% -> 87.3% - trace.py: 59.0% -> 71.9% (full coverage requires Trace fixtures with multiprocessing parse) - **Folder total: 73.4% -> 88.4%** Differential Revision: D104888486 --- tests/test_call_stack_extra.py | 203 +++++++++++ tests/test_trace_call_stack_extra.py | 237 +++++++++++++ tests/test_trace_df.py | 40 ++- tests/test_trace_extra.py | 501 +++++++++++++++++++++++++++ tests/test_trace_file_extra.py | 128 +++++++ tests/test_trace_filter_extra.py | 279 +++++++++++++++ tests/test_trace_integration.py | 115 ++++++ tests/test_trace_stack_filter.py | 184 ++++++++++ tests/test_trace_symbol_table.py | 226 ++++++++++++ tests/test_types.py | 96 +++++ 10 files changed, 2008 insertions(+), 1 deletion(-) create mode 100644 tests/test_call_stack_extra.py create mode 100644 tests/test_trace_call_stack_extra.py create mode 100644 tests/test_trace_extra.py create mode 100644 tests/test_trace_file_extra.py create mode 100644 tests/test_trace_filter_extra.py create mode 100644 tests/test_trace_integration.py create mode 100644 tests/test_trace_stack_filter.py create mode 100644 tests/test_trace_symbol_table.py create mode 100644 tests/test_types.py diff --git a/tests/test_call_stack_extra.py b/tests/test_call_stack_extra.py new file mode 100644 index 00000000..f5932aba --- /dev/null +++ b/tests/test_call_stack_extra.py @@ -0,0 +1,203 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import unittest + +import pandas as pd +from hta.common.call_stack import ( + CallStackGraph, + CallStackIdentity, + CallStackNode, + compare_events, + DeviceType, + Event, + EVENT_END, + EVENT_START, + infer_device_type, + NON_EXISTENT_NODE_INDEX, + NULL_NODE_INDEX, +) + + +def _make_cpu_df() -> pd.DataFrame: + """Create a minimal CPU-thread trace dataframe with index_correlation column.""" + return pd.DataFrame( + { + "index": [1, 2, 3], + "ts": [0, 5, 20], + "dur": [30, 10, 5], + "stream": [-1, -1, -1], + "index_correlation": [-1, -1, -1], + "pid": [100, 100, 100], + "tid": [200, 200, 200], + }, + index=pd.Index([1, 2, 3]), + ) + + +class TestInferDeviceType(unittest.TestCase): + def test_gpu(self) -> None: + df = pd.DataFrame({"stream": [1, 2, 3]}) + self.assertEqual(infer_device_type(df), DeviceType.GPU) + + def test_cpu(self) -> None: + df = pd.DataFrame({"stream": [-1, -1]}) + self.assertEqual(infer_device_type(df), DeviceType.CPU) + + def test_unknown_when_empty(self) -> None: + df = pd.DataFrame({"stream": []}) + self.assertEqual(infer_device_type(df), DeviceType.UNKNOWN) + + def test_unknown_when_mixed(self) -> None: + df = pd.DataFrame({"stream": [1, -1]}) + self.assertEqual(infer_device_type(df), DeviceType.UNKNOWN) + + +class TestCompareEvents(unittest.TestCase): + def test_same_idx_start_first(self) -> None: + x = Event(idx=1, time=0, dur=10, type=EVENT_START) + y = Event(idx=1, time=10, dur=10, type=EVENT_END) + self.assertLess(compare_events(x, y), 0) + + def test_same_idx_end_after(self) -> None: + x = Event(idx=1, time=10, dur=10, type=EVENT_END) + y = Event(idx=1, time=0, dur=10, type=EVENT_START) + self.assertGreater(compare_events(x, y), 0) + + def test_different_time(self) -> None: + x = Event(idx=1, time=0, dur=5, type=EVENT_START) + y = Event(idx=2, time=5, dur=5, type=EVENT_START) + self.assertLess(compare_events(x, y), 0) + + +class TestCallStackIdentity(unittest.TestCase): + def test_default_values(self) -> None: + csi = CallStackIdentity() + self.assertEqual(csi.rank, -1) + self.assertEqual(csi.pid, -1) + self.assertEqual(csi.tid, -1) + + def test_explicit_values(self) -> None: + csi = CallStackIdentity(rank=0, pid=100, tid=200) + self.assertEqual(csi.rank, 0) + + +class TestCallStackGraphBasics(unittest.TestCase): + def setUp(self) -> None: + self.csi = CallStackIdentity(rank=0, pid=100, tid=200) + self.df = _make_cpu_df() + self.csg = CallStackGraph(self.df, self.csi) + + def test_constructor_builds_nodes(self) -> None: + self.assertEqual(self.csg.identity, self.csi) + self.assertEqual(self.csg.device_type, DeviceType.CPU) + self.assertGreater(len(self.csg.nodes), 0) + + def test_repr_contains_callstackgraph(self) -> None: + self.assertIn("CallStackGraph", repr(self.csg)) + + def test_get_nodes(self) -> None: + nodes = self.csg.get_nodes() + self.assertIsInstance(nodes, dict) + self.assertIn(NULL_NODE_INDEX, nodes) + + def test_get_dataframe(self) -> None: + self.assertIs(self.csg.get_dataframe(), self.df) + + def test_get_parent_existing(self) -> None: + # Node 1 should be a child of NULL_NODE_INDEX (root) + self.assertEqual(self.csg.get_parent(1), NULL_NODE_INDEX) + + def test_get_parent_missing(self) -> None: + self.assertEqual(self.csg.get_parent(99999), NON_EXISTENT_NODE_INDEX) + + def test_get_children_existing(self) -> None: + children = self.csg.get_children(1) + self.assertIsInstance(children, list) + + def test_get_children_missing(self) -> None: + self.assertEqual(self.csg.get_children(99999), []) + + def test_get_path_to_root_existing(self) -> None: + path = self.csg.get_path_to_root(2) + self.assertIn(2, path) + # Should reach NULL_NODE_INDEX + self.assertIn(NULL_NODE_INDEX, path) + + def test_get_path_to_root_missing(self) -> None: + self.assertEqual(self.csg.get_path_to_root(99999), []) + + def test_get_paths_to_leaves_missing(self) -> None: + self.assertEqual(self.csg.get_paths_to_leaves(99999), []) + + def test_get_paths_to_leaves_existing(self) -> None: + paths = self.csg.get_paths_to_leaves(1) + self.assertIsInstance(paths, list) + + def test_get_leaf_nodes(self) -> None: + leaves = self.csg.get_leaf_nodes(1) + self.assertIsInstance(leaves, list) + + def test_get_depth(self) -> None: + depth = self.csg.get_depth() + self.assertIsNotNone(depth) + + def test_dfs_traverse(self) -> None: + visited_enter = [] + visited_exit = [] + + def enter(idx: int, node: CallStackNode) -> None: + visited_enter.append(idx) + + def exit_fn(idx: int, node: CallStackNode) -> None: + visited_exit.append(idx) + + self.csg.dfs_traverse(enter, exit_fn) + # Both should visit each node once + self.assertEqual(len(visited_enter), len(visited_exit)) + self.assertGreater(len(visited_enter), 0) + + +class TestCallStackGraphErrors(unittest.TestCase): + def test_missing_index_correlation_raises(self) -> None: + df = pd.DataFrame( + { + "index": [1], + "ts": [0], + "dur": [10], + "stream": [-1], + "pid": [100], + "tid": [200], + }, + index=pd.Index([1]), + ) + csi = CallStackIdentity(rank=0, pid=100, tid=200) + with self.assertRaisesRegex(ValueError, "index_correlation"): + CallStackGraph(df, csi) + + def test_gpu_short_circuits(self) -> None: + # GPU device path skips graph construction + df = pd.DataFrame( + { + "index": [1, 2], + "ts": [0, 10], + "dur": [5, 5], + "stream": [1, 1], # GPU + "index_correlation": [-1, -1], + "pid": [100, 100], + "tid": [200, 200], + }, + index=pd.Index([1, 2]), + ) + csi = CallStackIdentity(rank=0, pid=100, tid=200) + csg = CallStackGraph(df, csi) + self.assertEqual(csg.device_type, DeviceType.GPU) + # GPU short-circuits, so nodes dict stays empty + self.assertEqual(len(csg.nodes), 0) + + +class TestCallStackNode(unittest.TestCase): + def test_node_construction(self) -> None: + node = CallStackNode(parent=0, depth=1, children=[2, 3]) + self.assertEqual(node.parent, 0) + self.assertEqual(node.depth, 1) + self.assertEqual(node.children, [2, 3]) diff --git a/tests/test_trace_call_stack_extra.py b/tests/test_trace_call_stack_extra.py new file mode 100644 index 00000000..d6afbe5e --- /dev/null +++ b/tests/test_trace_call_stack_extra.py @@ -0,0 +1,237 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import unittest + +import numpy as np +import pandas as pd +from hta.common.trace_call_stack import ( + _cmp_events_with_zero_duration, + _less_than, + CallStackGraph, + CallStackIdentity, + CallStackNode, + CLOSE_END, + is_events_sorted, + NON_EXISTENT_NODE_INDEX, + OPEN_END, + sort_events, +) +from hta.common.trace_symbol_table import TraceSymbolTable +from hta.common.types import DeviceType + + +def _ev(idx: int, dur: int, kind: int, ts: int) -> np.ndarray: + return np.array([idx, dur, kind, ts]) + + +class TestCmpEventsWithZeroDuration(unittest.TestCase): + def test_zero_inside_nonzero_close(self) -> None: + x = _ev(1, 0, OPEN_END, 5) + y = _ev(2, 10, CLOSE_END, 5) + # zero event opens, non-zero closes: True (zero comes before) + self.assertTrue(_cmp_events_with_zero_duration(x, y)) + + def test_nonzero_open_zero(self) -> None: + x = _ev(1, 10, OPEN_END, 5) + y = _ev(2, 0, OPEN_END, 5) + # x is non-zero, y is zero: True if x is opening + self.assertTrue(_cmp_events_with_zero_duration(x, y)) + + def test_two_zero_open_ends(self) -> None: + x = _ev(1, 0, OPEN_END, 5) + y = _ev(2, 0, OPEN_END, 5) + # Both zero open: smaller index first + self.assertTrue(_cmp_events_with_zero_duration(x, y)) + + def test_two_zero_close_ends(self) -> None: + x = _ev(2, 0, CLOSE_END, 5) + y = _ev(1, 0, CLOSE_END, 5) + # Both zero close: larger index first (so x with higher idx comes first) + self.assertTrue(_cmp_events_with_zero_duration(x, y)) + + def test_two_zero_one_open_one_close(self) -> None: + x = _ev(1, 0, OPEN_END, 5) + y = _ev(2, 0, CLOSE_END, 5) + # One open one close: open first + self.assertTrue(_cmp_events_with_zero_duration(x, y)) + + +class TestLessThan(unittest.TestCase): + def test_different_time(self) -> None: + x = _ev(1, 10, OPEN_END, 5) + y = _ev(2, 10, OPEN_END, 10) + self.assertTrue(_less_than(x, y)) + + def test_same_index_open_first(self) -> None: + x = _ev(1, 10, OPEN_END, 5) + y = _ev(1, 10, CLOSE_END, 15) + # Same index: open first + self.assertTrue(_less_than(x, y)) + + def test_close_before_open_at_same_time(self) -> None: + x = _ev(1, 10, CLOSE_END, 15) + y = _ev(2, 10, OPEN_END, 15) + self.assertTrue(_less_than(x, y)) + + def test_two_open_ends_longer_first(self) -> None: + x = _ev(1, 100, OPEN_END, 0) + y = _ev(2, 50, OPEN_END, 0) + # Longer duration first when both opening at same time + self.assertTrue(_less_than(x, y)) + + def test_two_close_ends_shorter_first(self) -> None: + x = _ev(1, 50, CLOSE_END, 100) + y = _ev(2, 100, CLOSE_END, 100) + # Shorter duration first when both closing at same time + self.assertTrue(_less_than(x, y)) + + def test_zero_dur_path(self) -> None: + x = _ev(1, 0, OPEN_END, 5) + y = _ev(2, 10, CLOSE_END, 5) + # Delegates to zero-duration comparator + self.assertTrue(_less_than(x, y)) + + +class TestSortAndIsSorted(unittest.TestCase): + def test_is_events_sorted_true(self) -> None: + a = np.array( + [ + [1, 10, OPEN_END, 0], + [1, 10, CLOSE_END, 10], + ] + ) + self.assertTrue(is_events_sorted(a)) + + def test_is_events_sorted_false(self) -> None: + a = np.array( + [ + [1, 10, CLOSE_END, 10], + [1, 10, OPEN_END, 0], + ] + ) + self.assertFalse(is_events_sorted(a)) + + def test_sort_events(self) -> None: + a = np.array( + [ + [1, 10, CLOSE_END, 10], + [1, 10, OPEN_END, 0], + ] + ) + sort_events(a) + # After sort, OPEN should come first + self.assertEqual(a[0][2], OPEN_END) + + +class TestCallStackIdentity(unittest.TestCase): + def test_default_values(self) -> None: + c = CallStackIdentity() + self.assertEqual(c.rank, -1) + self.assertEqual(c.pid, -1) + self.assertEqual(c.tid, -1) + + def test_explicit(self) -> None: + c = CallStackIdentity(rank=0, pid=1, tid=2) + self.assertEqual(c.rank, 0) + self.assertEqual(c.pid, 1) + self.assertEqual(c.tid, 2) + + +class TestCallStackNode(unittest.TestCase): + def test_defaults(self) -> None: + n = CallStackNode() + self.assertEqual(n.parent, -1) + self.assertEqual(n.depth, -1) + self.assertEqual(n.height, -1) + self.assertEqual(n.device, DeviceType.CPU) + self.assertEqual(n.children, []) + + def test_explicit(self) -> None: + n = CallStackNode( + parent=0, depth=2, height=3, device=DeviceType.GPU, children=[1, 2] + ) + self.assertEqual(n.parent, 0) + self.assertEqual(n.depth, 2) + self.assertEqual(n.height, 3) + self.assertEqual(n.device, DeviceType.GPU) + self.assertEqual(n.children, [1, 2]) + + +def _build_cpu_csg() -> CallStackGraph: + """Helper to build a CallStackGraph from a minimal CPU thread DataFrame.""" + sym = TraceSymbolTable() + sym.add_symbols(["op_a", "op_b", "op_c"]) + df = pd.DataFrame( + { + "index": [10, 11, 12], + "ts": [0, 5, 100], + "dur": [200, 50, 30], + "stream": [-1, -1, -1], + "pid": [1, 1, 1], + "tid": [2, 2, 2], + "name": [0, 1, 2], + "cat": [0, 0, 0], + "index_correlation": [-1, -1, -1], + }, + index=pd.Index([10, 11, 12]), + ) + full_df = df.copy() + cpu_gpu_corr = pd.DataFrame({"cpu_index": [], "gpu_index": []}) + csi = CallStackIdentity(rank=0, pid=1, tid=2) + return CallStackGraph( + df=df, + identity=csi, + cpu_gpu_correlation=cpu_gpu_corr, + full_df=full_df, + symbol_table=sym, + nodes=None, + use_existing_stack_columns=False, + save_call_stack_to_df=True, + ) + + +class TestCallStackGraphCpu(unittest.TestCase): + def test_construct_repr_and_get_nodes(self) -> None: + csg = _build_cpu_csg() + self.assertIn("CallStackGraph", repr(csg)) + nodes = csg.get_nodes() + # Root + 3 events + self.assertGreaterEqual(len(nodes), 3) + + def test_get_parent_existing(self) -> None: + csg = _build_cpu_csg() + # Node 11 (op_b) starts at ts=5, fits inside op_a (ts=0..200) + self.assertEqual(csg.get_parent(11), 10) + + def test_get_parent_missing(self) -> None: + csg = _build_cpu_csg() + self.assertEqual(csg.get_parent(99999), NON_EXISTENT_NODE_INDEX) + + def test_gpu_thread_returns_early(self) -> None: + # GPU stream -> no graph constructed + sym = TraceSymbolTable() + sym.add_symbols(["k"]) + df = pd.DataFrame( + { + "index": [100], + "ts": [0], + "dur": [10], + "stream": [1], + "pid": [1], + "tid": [2], + "name": [0], + "cat": [0], + "index_correlation": [-1], + }, + index=pd.Index([100]), + ) + csi = CallStackIdentity(rank=0, pid=1, tid=2) + csg = CallStackGraph( + df=df, + identity=csi, + cpu_gpu_correlation=pd.DataFrame({"cpu_index": [], "gpu_index": []}), + full_df=df, + symbol_table=sym, + ) + # Has device_type GPU and skipped construction + self.assertEqual(csg.device_type, DeviceType.GPU) diff --git a/tests/test_trace_df.py b/tests/test_trace_df.py index bb4a3317..c60bde23 100644 --- a/tests/test_trace_df.py +++ b/tests/test_trace_df.py @@ -1,9 +1,15 @@ import unittest from dataclasses import dataclass from typing import List +from unittest.mock import MagicMock import pandas as pd -from hta.common.trace_df import find_op_occurrence, get_iterations +from hta.common.trace_df import ( + find_events_by_name_patterns_using_decoded_names, + find_events_by_name_patterns_using_symbol_table, + find_op_occurrence, + get_iterations, +) class TestTraceDF(unittest.TestCase): @@ -63,3 +69,35 @@ class TC: self.assertEqual(event["index"], tc.expected_index) else: self.assertTrue(event.empty) + + def test_find_events_by_name_patterns_using_symbol_table(self) -> None: + symbol_table = MagicMock() + symbol_table.get_sym_id_map.return_value = { + "ncclAllReduce": 1, + "memcpy": 2, + "compute": 3, + } + df = pd.DataFrame({"name": [1, 2, 3], "index": [10, 20, 30]}) + result = find_events_by_name_patterns_using_symbol_table( + df, ["nccl.*", "memcpy"], symbol_table + ) + self.assertEqual(set(result.tolist()), {10, 20}) + + def test_find_events_by_name_patterns_using_decoded_names(self) -> None: + # df.loc[indices] requires the pandas index to match the "index" column values + df = pd.DataFrame( + {"s_name": ["ncclAllReduce", "memcpy", "compute"], "index": [10, 20, 30]}, + index=pd.Index([10, 20, 30]), + ) + result = find_events_by_name_patterns_using_decoded_names( + df, ["nccl.*", "memcpy"] + ) + self.assertEqual(set(result.tolist()), {10, 20}) + + def test_find_events_by_name_patterns_decoded_no_match(self) -> None: + df = pd.DataFrame( + {"s_name": ["compute_a", "compute_b"], "index": [10, 20]}, + index=pd.Index([10, 20]), + ) + result = find_events_by_name_patterns_using_decoded_names(df, ["nccl.*"]) + self.assertEqual(len(result), 0) diff --git a/tests/test_trace_extra.py b/tests/test_trace_extra.py new file mode 100644 index 00000000..fb9057e7 --- /dev/null +++ b/tests/test_trace_extra.py @@ -0,0 +1,501 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import unittest +from typing import cast +from unittest.mock import MagicMock, patch + +import pandas as pd +from hta.common.trace import ( + add_fwd_bwd_links, + add_iteration, + get_cpu_gpu_correlation, + parse_trace_file, + Trace, + trace_event_timestamp_to_unixtime_ns, + transform_correlation_to_index, +) +from hta.common.trace_symbol_table import TraceSymbolTable + + +class TestTraceEventTimestampToUnixtime(unittest.TestCase): + def test_normal_conversion(self) -> None: + meta = {"baseTimeNanoseconds": 1_000_000_000} + result = trace_event_timestamp_to_unixtime_ns(5.0, meta) + # 5us = 5000ns; plus base + self.assertEqual(result, 1_000_005_000) + + def test_missing_base_time_raises(self) -> None: + with self.assertRaisesRegex(KeyError, "baseTimeNanoseconds"): + trace_event_timestamp_to_unixtime_ns(0.0, {}) + + def test_zero_event_timestamp(self) -> None: + meta = {"baseTimeNanoseconds": 100} + self.assertEqual(trace_event_timestamp_to_unixtime_ns(0.0, meta), 100) + + +class TestTransformCorrelationToIndex(unittest.TestCase): + def test_no_correlation_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + symbol_table = TraceSymbolTable() + result = transform_correlation_to_index(df, symbol_table) + # Returns df unchanged + self.assertNotIn("index_correlation", result.columns) + + def test_links_cpu_gpu_pairs(self) -> None: + # Build a symbol table with the right kinds + symbol_table = TraceSymbolTable() + symbol_table.add_symbols( + ["Kernel", "cuda_runtime", "kernel_a", "cuLaunchKernel"] + ) + # Map: Kernel=0, cuda_runtime=1, kernel_a=2, cuLaunchKernel=3 + df = pd.DataFrame( + { + "index": [675, 677], + "stream": [7, -1], + "cat": [ + symbol_table.sym_index["Kernel"], + symbol_table.sym_index["cuda_runtime"], + ], + "name": [ + symbol_table.sym_index["kernel_a"], + symbol_table.sym_index["cuLaunchKernel"], + ], + "correlation": [278204204, 278204204], + }, + index=pd.Index([675, 677]), + ) + result = transform_correlation_to_index(df, symbol_table) + self.assertIn("index_correlation", result.columns) + # CPU row links to GPU row and vice versa + self.assertEqual(result.loc[677, "index_correlation"], 675) + self.assertEqual(result.loc[675, "index_correlation"], 677) + + +class TestGetCpuGpuCorrelation(unittest.TestCase): + def test_extracts_kernel_correlations(self) -> None: + df = pd.DataFrame( + { + "index": [10, 20, 30], + "stream": [1, 1, -1], + "index_correlation": [50, 60, -1], + }, + index=pd.Index([10, 20, 30]), + ) + result = get_cpu_gpu_correlation(df) + # Two GPU rows produce two correlation entries + self.assertEqual(len(result), 2) + self.assertIn("gpu_index", result.columns) + self.assertIn("cpu_index", result.columns) + + +class TestAddIteration(unittest.TestCase): + def test_assigns_iteration_to_cpu_events(self) -> None: + symbol_table = TraceSymbolTable() + symbol_table.add_symbols(["ProfilerStep#5", "cpu_op"]) + ps_id = symbol_table.sym_index["ProfilerStep#5"] + cpu_id = symbol_table.sym_index["cpu_op"] + df = pd.DataFrame( + { + "ts": [100, 150, 300], + "dur": [200, 10, 5], + "stream": [-1, -1, -1], + "cat": [cpu_id, cpu_id, cpu_id], + "name": [ps_id, cpu_id, cpu_id], + "iteration": [-1, -1, -1], + "index_correlation": [-1, -1, -1], + } + ) + result = add_iteration(df, symbol_table) + # First row at ts=150 falls inside ProfilerStep#5 (ts=100, dur=200) -> iter=5 + self.assertEqual(int(df.loc[1, "iteration"]), 5) + # Third row at ts=300 outside ProfilerStep -> -1 + self.assertEqual(int(df.loc[2, "iteration"]), -1) + # Returned profiler_steps DataFrame + self.assertEqual(len(result), 1) + + +class TestAddFwdBwdLinks(unittest.TestCase): + def test_no_fwdbwd_events(self) -> None: + df = pd.DataFrame({"cat": ["cpu_op", "kernel"]}) + # Should return early without modifying + add_fwd_bwd_links(df) + self.assertNotIn("fwdbwd_index", df.columns) + + def test_links_fwd_bwd_pairs(self) -> None: + # cpu_op events at ts 0/tid 1/pid 10 and ts 100/tid 1/pid 10 + # fwdbwd 's' (start) at ts 0 -> matches first cpu_op + # fwdbwd 'f' bp 'e' (end) at ts 100 -> matches second cpu_op + df = pd.DataFrame( + { + "index": [0, 1, 2, 3], + "ts": [0, 0, 100, 100], + "tid": [1, 1, 1, 1], + "pid": [10, 10, 10, 10], + "cat": ["cpu_op", "fwdbwd", "cpu_op", "fwdbwd"], + "ph": ["X", "s", "X", "f"], + "bp": ["", "", "", "e"], + "id": [-1, 100, -1, 100], + }, + index=pd.Index([0, 1, 2, 3]), + ) + add_fwd_bwd_links(df) + self.assertIn("fwdbwd_index", df.columns) + self.assertIn("fwdbwd", df.columns) + # When merge succeeds, the "key" column is dropped + self.assertNotIn("key", df.columns) + + def test_empty_fwdbwd_merge_returns_early(self) -> None: + # fwdbwd events present but no matching cpu_op (keys differ) -> early return + df = pd.DataFrame( + { + "index": [0, 1, 2, 3], + "ts": [0, 5, 100, 105], + "tid": [1, 1, 1, 1], + "pid": [10, 10, 10, 10], + "cat": ["cpu_op", "fwdbwd", "cpu_op", "fwdbwd"], + "ph": ["X", "s", "X", "f"], + "bp": ["", "", "", "e"], + "id": [-1, 100, -1, 100], + }, + index=pd.Index([0, 1, 2, 3]), + ) + add_fwd_bwd_links(df) + # Columns added but key remains because we returned before the drop + self.assertIn("fwdbwd_index", df.columns) + + +class TestParseTraceFile(unittest.TestCase): + def test_invalid_extension_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"\.gz' or 'json'"): + parse_trace_file("/some/file.txt") + + +class TestTraceFlowEvent(unittest.TestCase): + def test_start_event(self) -> None: + ev = Trace.flow_event( + id=1, pid=10, tid=20, ts=100, is_start=True, name="link", cat="fwdbwd" + ) + self.assertEqual(ev["id"], 1) + self.assertEqual(ev["pid"], 10) + self.assertEqual(ev["tid"], 20) + self.assertEqual(ev["ts"], 100) + self.assertEqual(ev["name"], "link") + self.assertEqual(ev["cat"], "fwdbwd") + self.assertNotIn("bp", ev) + self.assertNotIn("args", ev) + + def test_end_event(self) -> None: + ev = Trace.flow_event( + id=1, pid=10, tid=20, ts=100, is_start=False, name="link", cat="fwdbwd" + ) + # End events get bp="e" + self.assertEqual(ev["bp"], "e") + + def test_with_args(self) -> None: + ev = Trace.flow_event( + id=1, + pid=10, + tid=20, + ts=100, + is_start=True, + name="link", + cat="fwdbwd", + args={"key": "value"}, + ) + self.assertEqual(ev["args"], {"key": "value"}) + + +class TestTraceConvertTimeSeries(unittest.TestCase): + def test_missing_required_columns_returns_empty(self) -> None: + # Build a minimal Trace-like instance using __new__ to avoid heavy __init__ + t = Trace.__new__(Trace) + t.min_ts = 0 + df = pd.DataFrame({"a": [1]}) + result = t.convert_time_series_to_events(df, "ctr", "missing_col") + self.assertEqual(result, []) + + def test_converts_with_required_columns(self) -> None: + t = Trace.__new__(Trace) + t.min_ts = 0 + df = pd.DataFrame( + { + "pid": [1, 2], + "ts": [10, 20], + "value": [100, 200], + } + ) + result = t.convert_time_series_to_events(df, "my_counter", "value") + self.assertEqual(len(result), 2) + self.assertEqual(result[0]["pid"], 1) + self.assertEqual(result[0]["args"], {"my_counter": 100}) + + def test_converts_with_optional_name_id(self) -> None: + t = Trace.__new__(Trace) + t.min_ts = 5 + df = pd.DataFrame( + { + "pid": [1], + "ts": [10], + "value": [100], + "name": ["custom_name"], + "id": [42], + } + ) + result = t.convert_time_series_to_events(df, "ctr", "value") + self.assertEqual(result[0]["name"], "custom_name") + self.assertEqual(result[0]["id"], 42) + # ts gets min_ts added back + self.assertEqual(result[0]["ts"], 15) + + +def _make_simple_trace_with_data(rank: int = 0) -> Trace: + """Build a Trace instance bypassing __init__ with a small in-memory dataset.""" + t = Trace.__new__(Trace) + t.trace_files = {rank: "/fake/path"} + t.trace_path = "/fake" + t.traces = {rank: pd.DataFrame({"ts": [0, 100, 200], "iteration": [-1, 0, 1]})} + t.symbol_table = TraceSymbolTable() + t.meta_data = {rank: {"device_type": "GPU"}} + t.min_ts = 0 + t.is_parsed = True + return t + + +class TestTraceGetters(unittest.TestCase): + def test_get_ranks(self) -> None: + t = _make_simple_trace_with_data(rank=3) + self.assertEqual(t.get_ranks(), [3]) + + def test_get_first_rank_with_arg(self) -> None: + t = _make_simple_trace_with_data(rank=3) + self.assertEqual(t._get_first_rank(7), 7) + + def test_get_first_rank_default(self) -> None: + t = _make_simple_trace_with_data(rank=3) + self.assertEqual(t._get_first_rank(), 3) + + def test_get_first_rank_no_ranks(self) -> None: + t = Trace.__new__(Trace) + t.traces = {} + self.assertEqual(t._get_first_rank(), -1) + + def test_get_iterations_returns_sorted(self) -> None: + t = _make_simple_trace_with_data() + # iterations [-1, 0, 1] -> sorted >=0 = [0, 1] + self.assertEqual(t.get_iterations(), [0, 1]) + + def test_get_iterations_no_column(self) -> None: + t = Trace.__new__(Trace) + t.traces = {0: pd.DataFrame({"ts": [0]})} + self.assertEqual(t.get_iterations(0), []) + + def test_get_iterations_invalid_rank(self) -> None: + t = _make_simple_trace_with_data() + self.assertEqual(t.get_iterations(99), []) + + def test_get_trace_duration(self) -> None: + t = _make_simple_trace_with_data() + self.assertEqual(t.get_trace_duration(), 200) + + def test_get_trace(self) -> None: + t = _make_simple_trace_with_data() + df = t.get_trace(0) + self.assertEqual(len(df), 3) + + def test_get_trace_invalid_rank_raises(self) -> None: + t = _make_simple_trace_with_data() + with self.assertRaises(ValueError): + t.get_trace(99) + + def test_get_all_traces(self) -> None: + t = _make_simple_trace_with_data() + all_traces = t.get_all_traces() + self.assertEqual(set(all_traces.keys()), {0}) + + def test_get_device_type(self) -> None: + t = _make_simple_trace_with_data() + self.assertEqual(t.get_device_type(), "GPU") + + +class TestTraceFilenameValidation(unittest.TestCase): + def test_normalize_trace_filenames_invalid_type_raises(self) -> None: + t = Trace.__new__(Trace) + # Intentionally set a wrong type to test runtime validation + t.trace_files = cast(dict, "not_a_dict") + t.trace_path = "/some/path" + with self.assertRaisesRegex(ValueError, "Expected trace_files"): + t._normalize_trace_filenames() + + def test_normalize_trace_filenames_relative_to_absolute(self) -> None: + t = Trace.__new__(Trace) + t.trace_files = {0: "rank0.json"} + t.trace_path = "/data/traces" + t._normalize_trace_filenames() + self.assertEqual(t.trace_files[0], "/data/traces/rank0.json") + + def test_normalize_trace_filenames_already_absolute(self) -> None: + t = Trace.__new__(Trace) + t.trace_files = {0: "/abs/path.json"} + t.trace_path = "/some/other" + t._normalize_trace_filenames() + self.assertEqual(t.trace_files[0], "/abs/path.json") + + def test_validate_trace_files_missing(self) -> None: + t = Trace.__new__(Trace) + t.trace_files = {0: "/no/such/file.json"} + self.assertFalse(t._validate_trace_files()) + + def test_validate_trace_files_invalid_extension(self) -> None: + import tempfile + + with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f: + t = Trace.__new__(Trace) + t.trace_files = {0: f.name} + self.assertFalse(t._validate_trace_files()) + + def test_validate_trace_files_valid(self) -> None: + import tempfile + + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f: + f.write(b"{}") + path = f.name + t = Trace.__new__(Trace) + t.trace_files = {0: path} + self.assertTrue(t._validate_trace_files()) + + +class TestTraceAlignAllRanks(unittest.TestCase): + def test_align_all_ranks_empty(self) -> None: + t = Trace.__new__(Trace) + t.traces = {} + t.min_ts = 0 + # Should return early without raising + t._align_all_ranks() + + def test_align_all_ranks_subtracts_min(self) -> None: + t = Trace.__new__(Trace) + t.traces = { + 0: pd.DataFrame({"ts": [10, 20]}), + 1: pd.DataFrame({"ts": [5, 15]}), + } + t.min_ts = 0 + t._align_all_ranks() + self.assertEqual(t.min_ts, 5) + # rank 0: 10-5=5, 20-5=15 + self.assertEqual(t.traces[0]["ts"].tolist(), [5, 15]) + # rank 1: 5-5=0, 15-5=10 + self.assertEqual(t.traces[1]["ts"].tolist(), [0, 10]) + + +class TestTraceLoadTracesAlreadyParsed(unittest.TestCase): + def test_load_traces_warns_when_already_parsed(self) -> None: + t = Trace.__new__(Trace) + t.is_parsed = True + # Should return early without raising + t.load_traces() + + +class TestGetRawTraceForOneRank(unittest.TestCase): + def test_invalid_rank_raises(self) -> None: + t = Trace.__new__(Trace) + t.trace_files = {0: "/fake.json"} + with self.assertRaises(ValueError): + t.get_raw_trace_for_one_rank(99) + + +class TestAlignAndFilterTrace(unittest.TestCase): + def test_no_traces(self) -> None: + t = Trace.__new__(Trace) + t.traces = {} + # Should return without error + t.align_and_filter_trace() + + +class TestParseTracesEmpty(unittest.TestCase): + def test_no_ranks_logs_error(self) -> None: + t = Trace.__new__(Trace) + t.trace_files = {} + # No-op other than logging an error + t.parse_traces() + self.assertFalse(t.is_parsed) + + +class TestGetTraceStartUnixtimeNs(unittest.TestCase): + def test_invalid_rank_raises(self) -> None: + t = Trace.__new__(Trace) + t.traces = {} + t.meta_data = {} + with self.assertRaisesRegex(ValueError, "No trace found"): + t.get_trace_start_unixtime_ns(99) + + def test_valid_rank_returns_unixtime(self) -> None: + t = Trace.__new__(Trace) + t.traces = {0: pd.DataFrame()} + t.meta_data = {0: {"baseTimeNanoseconds": 1_000_000_000}} + t.min_ts = 5 # 5us + # 5us = 5000ns + base = 1_000_005_000 + self.assertEqual(t.get_trace_start_unixtime_ns(0), 1_000_005_000) + + +class TestDecodeSymbolIds(unittest.TestCase): + def test_decodes_for_each_rank(self) -> None: + t = Trace.__new__(Trace) + sym = TraceSymbolTable() + sym.add_symbols(["op_a", "cpu_op"]) + t.symbol_table = sym + df = pd.DataFrame({"name": [0, 1], "cat": [1, 0]}) + t.traces = {0: df} + t.decode_symbol_ids(use_shorten_name=False) + # Should add s_name and s_cat columns + self.assertIn("s_name", df.columns) + self.assertIn("s_cat", df.columns) + + +class TestParseTraceFileSuccess(unittest.TestCase): + @patch("hta.common.trace.parse_trace_dataframe") + def test_empty_df_returns_early(self, mock_parse: MagicMock) -> None: + mock_parse.return_value = ( + {"meta": "data"}, + pd.DataFrame(), + TraceSymbolTable(), + ) + meta, df, st = parse_trace_file("/some/file.json") + self.assertTrue(df.empty) + self.assertEqual(meta, {"meta": "data"}) + + +class TestWriteRawTrace(unittest.TestCase): + def test_writes_gzipped_json(self) -> None: + import gzip + import json + import os + import tempfile + + t = Trace.__new__(Trace) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "out.gz") + t.write_raw_trace(path, {"a": 1}) + with gzip.open(path, "rt") as f: + self.assertEqual(json.load(f), {"a": 1}) + + +class TestFilterIrrelevantGpuKernels(unittest.TestCase): + def test_no_profiler_steps_skips(self) -> None: + t = Trace.__new__(Trace) + sym = TraceSymbolTable() + sym.add_symbols(["cpu_op", "kernel"]) + t.symbol_table = sym + t.traces = {0: pd.DataFrame()} + t.meta_data = {0: {"device_type": "GPU"}} + # No ProfilerStep symbols -> early return + t._filter_irrelevant_gpu_kernels() + + def test_one_profiler_step_skips(self) -> None: + t = Trace.__new__(Trace) + sym = TraceSymbolTable() + sym.add_symbols(["cpu_op", "kernel", "ProfilerStep#1"]) + t.symbol_table = sym + t.traces = {0: pd.DataFrame()} + t.meta_data = {0: {"device_type": "GPU"}} + # Only one ProfilerStep -> skip filter + t._filter_irrelevant_gpu_kernels() diff --git a/tests/test_trace_file_extra.py b/tests/test_trace_file_extra.py new file mode 100644 index 00000000..fa501173 --- /dev/null +++ b/tests/test_trace_file_extra.py @@ -0,0 +1,128 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +import gzip +import json +import os +import tempfile +import unittest + +from hta.common.trace_file import ( + create_rank_to_trace_dict, + create_rank_to_trace_dict_from_dir, + get_trace_files, + read_trace, + update_trace_rank, + write_trace, +) + + +def _write_json_with_rank(path: str, rank: int) -> None: + data = {"distributedInfo": {"rank": rank}, "traceEvents": []} + if path.endswith(".gz"): + with gzip.open(path, "wb") as f: + f.write(json.dumps(data).encode()) + else: + with open(path, "w") as f: + json.dump(data, f) + + +class TestCreateRankToTraceDictFromDir(unittest.TestCase): + def test_path_does_not_exist(self) -> None: + ok, d = create_rank_to_trace_dict_from_dir("/no/such/dir") + self.assertFalse(ok) + self.assertEqual(d, {}) + + def test_no_trace_files(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + open(os.path.join(tmp, "readme.txt"), "w").close() + ok, d = create_rank_to_trace_dict_from_dir(tmp) + self.assertFalse(ok) + + def test_finds_traces(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + _write_json_with_rank(os.path.join(tmp, "rank0.json"), 0) + _write_json_with_rank(os.path.join(tmp, "rank1.json.gz"), 1) + ok, d = create_rank_to_trace_dict_from_dir(tmp) + self.assertTrue(ok) + self.assertEqual(set(d.keys()), {0, 1}) + + +class TestCreateRankToTraceDict(unittest.TestCase): + def test_duplicate_rank_warns_and_keeps_last(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + f1 = os.path.join(tmp, "a.json") + f2 = os.path.join(tmp, "b.json") + _write_json_with_rank(f1, 5) + _write_json_with_rank(f2, 5) + ok, d = create_rank_to_trace_dict([f1, f2]) + self.assertTrue(ok) + # Last one wins + self.assertEqual(d[5], f2) + + def test_no_rank_in_file_defaults_to_zero(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + f = os.path.join(tmp, "x.json") + with open(f, "w") as fh: + json.dump({"traceEvents": []}, fh) + ok, d = create_rank_to_trace_dict([f]) + self.assertTrue(ok) + self.assertEqual(d.get(0), f) + + +class TestGetTraceFiles(unittest.TestCase): + def test_invalid_path_returns_empty(self) -> None: + self.assertEqual(get_trace_files("/no/such/dir"), {}) + + def test_empty_dir_returns_empty(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual(get_trace_files(tmp), {}) + + def test_finds_traces(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + _write_json_with_rank(os.path.join(tmp, "rank0.json"), 0) + d = get_trace_files(tmp) + self.assertEqual(set(d.keys()), {0}) + + +class TestReadAndWriteTrace(unittest.TestCase): + def test_roundtrip_json(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "x.json") + data = {"distributedInfo": {"rank": 2}, "traceEvents": [1]} + write_trace(data, path) + self.assertEqual(read_trace(path), data) + + def test_roundtrip_gz(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "x.json.gz") + data = {"distributedInfo": {"rank": 3}, "traceEvents": [2]} + write_trace(data, path) + self.assertEqual(read_trace(path), data) + + def test_invalid_extension_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"\.gz' or 'json'"): + read_trace("/tmp/x.txt") + + def test_write_creates_missing_dir(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + new_dir = os.path.join(tmp, "nested", "deep") + path = os.path.join(new_dir, "trace.json") + write_trace({"a": 1}, path) + self.assertTrue(os.path.exists(path)) + + +class TestUpdateTraceRank(unittest.TestCase): + def test_updates_existing_rank(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "x.json") + _write_json_with_rank(path, 0) + update_trace_rank(path, 7) + self.assertEqual(read_trace(path)["distributedInfo"]["rank"], 7) + + def test_adds_distributed_info_when_missing(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "x.json") + with open(path, "w") as f: + json.dump({"traceEvents": []}, f) + update_trace_rank(path, 4) + self.assertEqual(read_trace(path)["distributedInfo"]["rank"], 4) diff --git a/tests/test_trace_filter_extra.py b/tests/test_trace_filter_extra.py new file mode 100644 index 00000000..34ca8fa8 --- /dev/null +++ b/tests/test_trace_filter_extra.py @@ -0,0 +1,279 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import unittest +from typing import cast + +import pandas as pd +from hta.common.trace_filter import ( + CompositeFilter, + CPUOperatorFilter, + Filter, + FirstIterationFilter, + GPUKernelFilter, + IterationFilter, + IterationIndexFilter, + MemCopyEventFilter, + NameFilter, + NameIdColumnFilter, + NameStringColumnFilter, + QueryFilter, + RankFilter, + TimeRangeFilter, + ZeroDurationFilter, +) +from hta.common.trace_symbol_table import TraceSymbolTable + + +def _make_symbol_table(symbols: list[str]) -> TraceSymbolTable: + t = TraceSymbolTable() + t.add_symbols(symbols) + return t + + +class TestIterationFilter(unittest.TestCase): + def test_int_arg(self) -> None: + f = IterationFilter(1) + self.assertEqual(f.iterations, [1]) + + def test_list_arg(self) -> None: + f = IterationFilter([1, 2]) + self.assertEqual(f.iterations, [1, 2]) + + def test_invalid_type_raises(self) -> None: + with self.assertRaisesRegex(TypeError, "Iterations must"): + # Intentionally pass a wrong type to test runtime validation + IterationFilter(cast(int, "bad")) + + def test_call_filters(self) -> None: + df = pd.DataFrame({"iteration": [1, 2, 3, 1], "x": [10, 20, 30, 40]}) + result = IterationFilter([1])(df) + self.assertEqual(result["x"].tolist(), [10, 40]) + + def test_call_warns_when_no_iteration_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + result = IterationFilter(1)(df) + # Returns df unchanged + self.assertEqual(len(result), 1) + + +class TestIterationIndexFilter(unittest.TestCase): + def test_invalid_type_raises(self) -> None: + with self.assertRaisesRegex(TypeError, "iteration_index"): + IterationIndexFilter("bad") # pyre-ignore[6] + + def test_int_arg(self) -> None: + f = IterationIndexFilter(0) + self.assertEqual(f.iteration_index, [0]) + + def test_call_no_iteration_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + result = IterationIndexFilter(0)(df) + self.assertEqual(len(result), 1) + + def test_call_only_neg1(self) -> None: + df = pd.DataFrame({"iteration": [-1, -1]}) + # Returns df unchanged + result = IterationIndexFilter(0)(df) + self.assertEqual(len(result), 2) + + def test_call_filters_by_index(self) -> None: + df = pd.DataFrame({"iteration": [-1, 100, 200, 300], "x": [0, 1, 2, 3]}) + # index 0 -> iteration 100, index 1 -> iteration 200 + result = IterationIndexFilter([0, 1])(df) + self.assertEqual(set(result["x"].tolist()), {1, 2}) + + def test_call_no_match(self) -> None: + df = pd.DataFrame({"iteration": [100, 200]}) + result = IterationIndexFilter([99])(df) + self.assertTrue(result.empty) + + +class TestFirstIterationFilter(unittest.TestCase): + def test_picks_first(self) -> None: + df = pd.DataFrame({"iteration": [100, 200, 100], "x": [1, 2, 3]}) + f = FirstIterationFilter() + result = f(df) + self.assertEqual(set(result["x"].tolist()), {1, 3}) + + +class TestRankFilter(unittest.TestCase): + def test_invalid_type(self) -> None: + with self.assertRaisesRegex(TypeError, "ranks"): + RankFilter("bad") # pyre-ignore[6] + + def test_int_arg(self) -> None: + self.assertEqual(RankFilter(0).ranks, [0]) + + def test_no_rank_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + # Returns df unchanged + self.assertEqual(len(RankFilter(0)(df)), 1) + + def test_filters(self) -> None: + df = pd.DataFrame({"rank": [0, 1, 2], "x": [10, 20, 30]}) + result = RankFilter([0, 2])(df) + self.assertEqual(set(result["x"].tolist()), {10, 30}) + + +class TestTimeRangeFilter(unittest.TestCase): + def test_invalid_tuple_raises(self) -> None: + with self.assertRaisesRegex(ValueError, "tuple of two"): + TimeRangeFilter([1, 2]) # pyre-ignore[6] + + def test_invalid_order_raises(self) -> None: + with self.assertRaisesRegex(ValueError, "less than or equal"): + TimeRangeFilter((10, 5)) + + def test_no_ts_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + # Returns df unchanged + self.assertEqual(len(TimeRangeFilter((0, 100))(df)), 1) + + def test_filters_by_time(self) -> None: + df = pd.DataFrame({"ts": [0, 50, 90], "dur": [10, 10, 20]}) + # End times: 10, 60, 110. Range (0, 100) selects ts>=0 & end<=100 + result = TimeRangeFilter((0, 100))(df) + self.assertEqual(len(result), 2) + + +class TestNameStringColumnFilter(unittest.TestCase): + def test_no_name_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + # Returns df unchanged when no name col + result = NameStringColumnFilter("nccl")(df) + self.assertEqual(len(result), 1) + + def test_filters_by_pattern(self) -> None: + df = pd.DataFrame({"name": ["ncclAll", "memcpy", "compute"]}) + result = NameStringColumnFilter("nccl")(df) + self.assertEqual(result["name"].tolist(), ["ncclAll"]) + + +class TestNameIdColumnFilter(unittest.TestCase): + def test_returns_unchanged_without_symbol_table(self) -> None: + df = pd.DataFrame({"name": [1, 2]}) + result = NameIdColumnFilter("nccl")(df) + self.assertEqual(len(result), 2) + + def test_filters_with_symbol_table(self) -> None: + t = _make_symbol_table(["ncclAll", "memcpy"]) + df = pd.DataFrame({"name": [0, 1]}) + result = NameIdColumnFilter("nccl")(df, t) + self.assertEqual(result["name"].tolist(), [0]) + + +class TestNameFilter(unittest.TestCase): + def test_empty_df(self) -> None: + result = NameFilter("nccl")(pd.DataFrame()) + self.assertTrue(result.empty) + + def test_filters_with_string_column(self) -> None: + df = pd.DataFrame({"name": ["ncclAll", "memcpy"]}) + result = NameFilter("nccl")(df) + self.assertEqual(result["name"].tolist(), ["ncclAll"]) + + def test_filters_with_provided_symbol_table_ctor(self) -> None: + t = _make_symbol_table(["ncclAll", "memcpy"]) + df = pd.DataFrame({"name": [0, 1]}) + result = NameFilter("nccl", symbol_table=t)(df) + self.assertEqual(result["name"].tolist(), [0]) + + def test_filters_with_call_symbol_table(self) -> None: + t = _make_symbol_table(["ncclAll", "memcpy"]) + df = pd.DataFrame({"name": [0, 1]}) + result = NameFilter("nccl")(df, symbol_table=t) + self.assertEqual(result["name"].tolist(), [0]) + + +class TestQueryFilter(unittest.TestCase): + def test_filter_by_query(self) -> None: + df = pd.DataFrame({"x": [1, 2, 3]}) + result = QueryFilter("x > 1")(df) + self.assertEqual(result["x"].tolist(), [2, 3]) + + def test_zero_duration_filter(self) -> None: + df = pd.DataFrame({"dur": [0, 5, 10]}) + result = ZeroDurationFilter(df) + self.assertEqual(set(result["dur"].tolist()), {5, 10}) + + +class TestGPUKernelFilter(unittest.TestCase): + def test_no_stream_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + result = GPUKernelFilter()(df) + self.assertEqual(len(result), 1) + + def test_no_symbol_table(self) -> None: + df = pd.DataFrame({"stream": [0, -1, 1], "correlation": [0, -1, 5]}) + result = GPUKernelFilter()(df) + # stream >= 0 AND correlation >= 0 + self.assertEqual(len(result), 2) + + def test_with_symbol_table(self) -> None: + t = _make_symbol_table(["Event Sync", "compute"]) + df = pd.DataFrame( + { + "stream": [0, -1, -1], + "correlation": [5, -1, -1], + "name": [1, 0, 1], + } + ) + result = GPUKernelFilter()(df, t) + self.assertGreaterEqual(len(result), 1) + + +class TestCPUOperatorFilter(unittest.TestCase): + def test_no_stream_column(self) -> None: + df = pd.DataFrame({"x": [1]}) + result = CPUOperatorFilter()(df) + self.assertEqual(len(result), 1) + + def test_no_symbol_table(self) -> None: + df = pd.DataFrame({"stream": [-1, 0, -1], "correlation": [0, 5, 0]}) + result = CPUOperatorFilter()(df) + # stream == -1 only + self.assertEqual(len(result), 2) + + def test_with_symbol_table(self) -> None: + t = _make_symbol_table(["Event Sync", "compute"]) + df = pd.DataFrame({"stream": [-1, 0], "correlation": [-1, 5], "name": [1, 1]}) + result = CPUOperatorFilter()(df, t) + self.assertGreaterEqual(len(result), 1) + + +class TestCompositeFilter(unittest.TestCase): + def test_invalid_type_raises(self) -> None: + with self.assertRaisesRegex(TypeError, "instances of Filter"): + CompositeFilter([object()]) # pyre-ignore[6] + + def test_applies_filters_in_order(self) -> None: + df = pd.DataFrame( + {"rank": [0, 1, 2], "iteration": [10, 20, 30], "x": [1, 2, 3]} + ) + cf = CompositeFilter([RankFilter([0, 1]), IterationFilter([10])]) + result = cf(df) + self.assertEqual(result["x"].tolist(), [1]) + + +class TestMemCopyEventFilter(unittest.TestCase): + def test_empty_df(self) -> None: + result = MemCopyEventFilter("Memcpy DtoH")(pd.DataFrame()) + self.assertTrue(result.empty) + + def test_no_symbol_match_returns_empty(self) -> None: + t = _make_symbol_table(["other"]) + df = pd.DataFrame({"name": [0], "cat": [0]}) + result = MemCopyEventFilter("Memcpy DtoH", symbol_table=t)(df) + self.assertTrue(result.empty) + + def test_filters_matching(self) -> None: + t = _make_symbol_table(["Memcpy DtoH", "gpu_memcpy", "other"]) + df = pd.DataFrame({"name": [0, 0, 2], "cat": [1, 1, 0], "x": [10, 20, 30]}) + result = MemCopyEventFilter("Memcpy DtoH")(df, symbol_table=t) + self.assertEqual(result["x"].tolist(), [10, 20]) + + +class TestFilterAbstract(unittest.TestCase): + def test_filter_is_abstract(self) -> None: + with self.assertRaises(TypeError): + Filter() # pyre-ignore[45] diff --git a/tests/test_trace_integration.py b/tests/test_trace_integration.py new file mode 100644 index 00000000..7649ea0a --- /dev/null +++ b/tests/test_trace_integration.py @@ -0,0 +1,115 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import os +import unittest +from typing import cast + +from hta.common.trace import parse_trace_file, Trace +from hta.utils.test_utils import get_test_data_dir + + +class TestParseTraceFileIntegration(unittest.TestCase): + """Integration tests using real trace data files.""" + + def test_parse_real_cpu_trace(self) -> None: + path = os.path.join( + get_test_data_dir(), + "cpu_only", + "rank-34.Jul_15_10_52_41.1074.pt.trace.json.gz", + ) + meta, df, sym = parse_trace_file(path) + # Real trace has metadata + non-empty df + self.assertIsInstance(meta, dict) + self.assertGreater(len(df), 0) + # Symbol table populated + self.assertGreater(len(sym.get_sym_table()), 0) + # Standard columns added by parser + self.assertIn("end", df.columns) + self.assertIn("index_correlation", df.columns) + + +class TestTraceLoadAndQuery(unittest.TestCase): + """End-to-end Trace loading and query tests.""" + + trace: Trace + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + path = os.path.join( + get_test_data_dir(), + "cpu_only", + "rank-34.Jul_15_10_52_41.1074.pt.trace.json.gz", + ) + cls.trace = Trace(trace_files={34: path}, trace_dir="") + cls.trace.parse_traces(use_multiprocessing=False) + + def test_get_ranks(self) -> None: + self.assertEqual(self.trace.get_ranks(), [34]) + + def test_get_trace(self) -> None: + df = self.trace.get_trace(34) + self.assertGreater(len(df), 0) + + def test_get_all_traces(self) -> None: + all_traces = self.trace.get_all_traces() + self.assertIn(34, all_traces) + + def test_get_trace_invalid_rank_raises(self) -> None: + with self.assertRaises(ValueError): + self.trace.get_trace(999) + + def test_get_trace_duration(self) -> None: + # Returns int >= 0 + d = self.trace.get_trace_duration() + self.assertGreaterEqual(d, 0) + + def test_get_iterations_returns_list(self) -> None: + iters = self.trace.get_iterations() + self.assertIsInstance(iters, list) + + def test_get_raw_trace_for_one_rank(self) -> None: + raw = self.trace.get_raw_trace_for_one_rank(34) + self.assertIsInstance(raw, dict) + self.assertIn("traceEvents", raw) + + def test_get_raw_trace_invalid_rank_raises(self) -> None: + with self.assertRaises(ValueError): + self.trace.get_raw_trace_for_one_rank(999) + + def test_decode_symbol_ids(self) -> None: + self.trace.decode_symbol_ids(use_shorten_name=True) + df = self.trace.get_trace(34) + # After decode, s_name and s_cat should exist + self.assertIn("s_name", df.columns) + + +class TestTraceLoadTracesEntryPoint(unittest.TestCase): + """Test the public load_traces() entry point.""" + + def test_load_traces(self) -> None: + path = os.path.join( + get_test_data_dir(), + "cpu_only", + "rank-34.Jul_15_10_52_41.1074.pt.trace.json.gz", + ) + trace = Trace(trace_files={0: path}, trace_dir="") + trace.load_traces(use_multiprocessing=False) + self.assertTrue(trace.is_parsed) + # Re-loading is a no-op + trace.load_traces(use_multiprocessing=False) + self.assertTrue(trace.is_parsed) + + +class TestTraceCtorEmpty(unittest.TestCase): + """Constructor edge cases.""" + + def test_invalid_trace_files_type_logged(self) -> None: + # trace_files is neither list nor dict — logged and returns + # Intentionally pass a wrong type to test runtime validation + Trace(trace_files=cast(dict, 42), trace_dir="") + + def test_validation_failure_raises(self) -> None: + # Non-existent file -> validation fails + with self.assertRaisesRegex(ValueError, "validation failed"): + Trace(trace_files={0: "/no/such/file.json"}, trace_dir="") diff --git a/tests/test_trace_stack_filter.py b/tests/test_trace_stack_filter.py new file mode 100644 index 00000000..74578753 --- /dev/null +++ b/tests/test_trace_stack_filter.py @@ -0,0 +1,184 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import unittest + +import pandas as pd +from hta.common.trace_stack_filter import ( + AfterOperatorFilter, + BeforeOperatorFilter, + CombinedOperatorFilter, + get_matching_kernels, + OperatorFilter, + OperatorFilterMethod, + UnderOperatorFilter, +) + + +def _make_df_with_op() -> pd.DataFrame: + """Build a small dataframe with cpu ops + a kernel under one of them.""" + return pd.DataFrame( + { + "index": [0, 1, 2, 3, 4], + "ts": [0, 100, 105, 200, 300], + "dur": [500, 50, 30, 50, 100], + "end": [500, 150, 135, 250, 400], + "stream": [-1, -1, -1, -1, -1], + "s_name": ["forward", "op_inside", "op_inside_2", "op_after", "op_extra"], + "s_cat": ["cpu_op", "cpu_op", "cpu_op", "cpu_op", "cpu_op"], + "name": [0, 1, 2, 3, 4], + "cat": [0, 0, 0, 0, 0], + "index_correlation": [-1, -1, -1, -1, -1], + } + ) + + +class TestGetMatchingKernels(unittest.TestCase): + def test_filters_kernels_launched_by_runtimes(self) -> None: + df_ops = pd.DataFrame({"index": [10, 11], "s_cat": ["cuda_runtime", "cpu_op"]}) + df_both = pd.DataFrame( + { + "index": [100, 101, 102], + "index_correlation": [10, 99, 11], + } + ) + kernels = get_matching_kernels(df_ops, df_both, "s_cat") + # Only kernel with index_correlation 10 (matches cuda_runtime row) + self.assertEqual(kernels["index"].tolist(), [100]) + + +class TestOperatorFilterMethod(unittest.TestCase): + def test_enum_values(self) -> None: + self.assertEqual(OperatorFilterMethod.Under.value, 0) + self.assertEqual(OperatorFilterMethod.After.value, 1) + self.assertEqual(OperatorFilterMethod.Before.value, 2) + + +class TestOperatorFilter(unittest.TestCase): + def test_op_not_found_returns_empty(self) -> None: + df = _make_df_with_op() + f = OperatorFilter("missing_op", 0, OperatorFilterMethod.Under) + result = f(df) + self.assertTrue(result.empty) + + def test_under_filter(self) -> None: + df = _make_df_with_op() + f = OperatorFilter("forward", 0, OperatorFilterMethod.Under) + result = f(df) + # All ops fit under "forward" (ts 0 - end 500) + self.assertEqual(len(result), 5) + + def test_after_filter(self) -> None: + df = _make_df_with_op() + f = OperatorFilter("op_inside", 0, OperatorFilterMethod.After) + result = f(df) + # ops with ts >= op_inside.end (150). That's op_after (ts=200), op_extra (ts=300) + self.assertEqual(set(result["s_name"]), {"op_after", "op_extra"}) + + def test_before_filter(self) -> None: + df = _make_df_with_op() + f = OperatorFilter("op_after", 0, OperatorFilterMethod.Before) + result = f(df) + # ops with end <= op_after.ts (200): forward (end=500, no - ge 200), op_inside (end 150 yes), op_inside_2 (end 135 yes) + names = set(result["s_name"]) + self.assertIn("op_inside", names) + self.assertIn("op_inside_2", names) + + def test_unsupported_method_raises(self) -> None: + df = _make_df_with_op() + f = OperatorFilter("forward", 0, OperatorFilterMethod.Stack) + with self.assertRaises(NotImplementedError): + f(df) + + def test_no_string_columns_returns_df(self) -> None: + df = pd.DataFrame({"name": [1, 2], "cat": [1, 2]}) # int columns + f = OperatorFilter("op", 0, OperatorFilterMethod.Under) + result = f(df) + self.assertEqual(len(result), 2) + + def test_explicit_name_column(self) -> None: + df = _make_df_with_op() + f = OperatorFilter( + "forward", 0, OperatorFilterMethod.Under, name_column="s_name" + ) + result = f(df) + self.assertEqual(len(result), 5) + + +class TestSubclassFilters(unittest.TestCase): + def test_after_subclass(self) -> None: + df = _make_df_with_op() + f = AfterOperatorFilter("op_inside", 0) + self.assertEqual(f.method, OperatorFilterMethod.After) + result = f(df) + self.assertEqual(set(result["s_name"]), {"op_after", "op_extra"}) + + def test_before_subclass(self) -> None: + df = _make_df_with_op() + f = BeforeOperatorFilter("op_after", 0) + self.assertEqual(f.method, OperatorFilterMethod.Before) + result = f(df) + names = set(result["s_name"]) + self.assertIn("op_inside", names) + + def test_under_subclass(self) -> None: + df = _make_df_with_op() + f = UnderOperatorFilter("forward", 0) + self.assertEqual(f.method, OperatorFilterMethod.Under) + result = f(df) + self.assertEqual(len(result), 5) + + +class TestCombinedOperatorFilter(unittest.TestCase): + def _df_with_iteration(self) -> pd.DataFrame: + return pd.DataFrame( + { + "index": [0, 1, 2, 3, 4, 5], + "ts": [0, 50, 100, 150, 200, 300], + "dur": [500, 30, 30, 30, 30, 50], + "end": [500, 80, 130, 180, 230, 350], + "stream": [-1, -1, -1, -1, -1, -1], + "iteration": [1, 1, 1, 1, 1, 1], + "s_name": [ + "## forward ##", + "All2All_Pooled_Wait", + "middle_op_a", + "middle_op_b", + "## sdd_preprocess_tensors ##", + "outside_op", + ], + "s_cat": ["cpu_op"] * 6, + "name": [0, 1, 2, 3, 4, 5], + "cat": [0, 0, 0, 0, 0, 0], + "index_correlation": [-1, -1, -1, -1, -1, -1], + } + ) + + def test_runs_with_all_three_ops_present(self) -> None: + df = self._df_with_iteration() + f = CombinedOperatorFilter( + "## forward ##", + "All2All_Pooled_Wait", + "## sdd_preprocess_tensors ##", + ) + result = f(df) + # Should find some events between the bracketing ops + self.assertIsInstance(result, pd.DataFrame) + + def test_missing_columns_returns_df(self) -> None: + df = pd.DataFrame({"x": [1]}) + f = CombinedOperatorFilter("a", "b", "c") + result = f(df) + self.assertEqual(len(result), 1) + + def test_with_stack_depths(self) -> None: + df = self._df_with_iteration() + df["depth"] = [0, 1, 2, 2, 1, 0] + f = CombinedOperatorFilter( + "## forward ##", + "All2All_Pooled_Wait", + "## sdd_preprocess_tensors ##", + stack_depths=[1, 2], + ) + # Should not raise + result = f(df) + self.assertIsInstance(result, pd.DataFrame) diff --git a/tests/test_trace_symbol_table.py b/tests/test_trace_symbol_table.py new file mode 100644 index 00000000..ed248fbd --- /dev/null +++ b/tests/test_trace_symbol_table.py @@ -0,0 +1,226 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import os +import tempfile +import unittest + +import pandas as pd +from hta.common.trace_symbol_table import ( + decode_symbol_id_to_symbol_name, + TraceSymbolTable, +) + + +class TestTraceSymbolTableBasics(unittest.TestCase): + def test_is_empty_and_add_symbols(self) -> None: + t = TraceSymbolTable() + self.assertTrue(t.is_empty()) + t.add_symbols(["a", "b", "a", "c"]) + self.assertFalse(t.is_empty()) + self.assertEqual(t.get_sym_table(), ["a", "b", "c"]) + self.assertEqual(t.get_sym_id_map(), {"a": 0, "b": 1, "c": 2}) + + def test_get_sym_index_alias(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["x"]) + self.assertEqual(t.get_sym_index(), {"x": 0}) + + def test_find_matches_and_matched_symbols(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["foo_bar", "foo_baz", "qux"]) + self.assertEqual(set(t.find_matches(["foo"])), {0, 1}) + self.assertEqual(set(t.find_matched_symbols(["foo"])), {"foo_bar", "foo_baz"}) + + def test_get_sym_index_series_and_table_series(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["a", "b"]) + self.assertEqual(len(t.get_sym_index_series()), 2) + self.assertEqual(len(t.get_sym_table_series()), 2) + + +class TestTraceSymbolTableDecode(unittest.TestCase): + def test_decode_df_creates_new_columns(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["op_a", "cpu_op"]) + df = pd.DataFrame({"name": [0, 1], "cat": [1, 0]}) + t.decode_df(df, create_new_columns=True) + self.assertIn("s_name", df.columns) + self.assertIn("s_cat", df.columns) + self.assertEqual(df["s_name"].tolist(), ["op_a", "cpu_op"]) + + def test_decode_df_overwrites_columns(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["op_a", "cpu_op"]) + df = pd.DataFrame({"name": [0, 1], "cat": [1, 0]}) + t.decode_df(df, create_new_columns=False) + self.assertEqual(df["name"].tolist(), ["op_a", "cpu_op"]) + + def test_encode_df(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["op_a", "cpu_op"]) + df = pd.DataFrame({"name": ["op_a"], "cat": ["cpu_op"]}) + t.encode_df(df) + self.assertEqual(df["name"].tolist(), [0]) + self.assertEqual(df["cat"].tolist(), [1]) + + +class TestTraceSymbolTableUpdateAndAdd(unittest.TestCase): + def test_update_encoded_df(self) -> None: + old = TraceSymbolTable() + old.add_symbols(["op_a", "op_b"]) + new = TraceSymbolTable() + new.add_symbols(["op_b", "op_a"]) # different order + + df = pd.DataFrame({"name": [0, 1], "cat": [1, 0]}) # encoded with old + new.update_encoded_df(df, old) + # After update: 0 ("op_a" in old) -> new index of "op_a" = 1 + self.assertEqual(df["name"].tolist(), [1, 0]) + + def test_add_symbols_to_trace_df(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["op_a", "op_b"]) + df = pd.DataFrame({"name": [0, 1, 99]}) + t.add_symbols_to_trace_df(df, "name") + # Out-of-range gets empty string + self.assertEqual(df["name"].tolist(), ["op_a", "op_b", ""]) + + +class TestTraceSymbolTableMasks(unittest.TestCase): + def test_get_operator_or_cuda_runtime_mask(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["cpu_op", "cuda_runtime", "cuda_driver", "kernel"]) + df = pd.DataFrame( + { + "cat": [ + t.get_sym_id_map()["cpu_op"], + t.get_sym_id_map()["kernel"], + t.get_sym_id_map()["cuda_driver"], + ] + } + ) + mask = t.get_operator_or_cuda_runtime_mask(df) + self.assertEqual(mask.tolist(), [True, False, True]) + + def test_get_runtime_launch_events_mask(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["cudaLaunchKernel", "noise"]) + df = pd.DataFrame( + { + "name": [ + t.get_sym_id_map()["cudaLaunchKernel"], + t.get_sym_id_map()["noise"], + ], + "index_correlation": [5, 5], + } + ) + mask = t.get_runtime_launch_events_mask(df) + self.assertEqual(mask.tolist(), [True, False]) + + def test_get_events_mask_none(self) -> None: + t = TraceSymbolTable() + df = pd.DataFrame({"name": [0, 1]}) + mask = t.get_events_mask(df, None) + self.assertFalse(mask.any()) + + def test_get_events_mask_with_pattern(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["ncclAllReduce", "memcpy"]) + df = pd.DataFrame({"name": [0, 1]}) + mask = t.get_events_mask(df, ["nccl"]) + self.assertEqual(mask.tolist(), [True, False]) + + +class TestTraceSymbolTableGetCategoryHelpers(unittest.TestCase): + def test_helpers_return_lists(self) -> None: + t = TraceSymbolTable() + t.add_symbols( + ["cpu_op", "kernel", "cudaLaunchKernel", "Memcpy DtoH", "ProfilerStep#1"] + ) + self.assertIsInstance(t.get_cpu_event_cat_ids(), list) + self.assertIsInstance(t.get_gpu_kernel_cat_ids(), list) + self.assertIsInstance(t.get_kernel_launch_ids(), list) + self.assertIsInstance(t.get_memory_name_ids(), list) + self.assertIsInstance(t.get_profiler_step_ids(), list) + + +class TestTraceSymbolTableFromCsv(unittest.TestCase): + def test_round_trip_csv(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + t = TraceSymbolTable() + t.add_symbols(["a", "b", "c"]) + path = os.path.join(tmp, "out.csv") + t._to_csv_file(path) + loaded = TraceSymbolTable.from_csv_file(path) + self.assertEqual(loaded.get_sym_table(), ["a", "b", "c"]) + + def test_from_csv_missing_file(self) -> None: + with self.assertRaises(FileNotFoundError): + TraceSymbolTable.from_csv_file("/no/such/file.csv") + + def test_from_csv_missing_column(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "bad.csv") + pd.DataFrame({"other": [1, 2]}).to_csv(path, index=False) + with self.assertRaisesRegex(ValueError, "expect a column"): + TraceSymbolTable.from_csv_file(path) + + +class TestCombineAndCloneAndCreateFromDf(unittest.TestCase): + def test_combine_symbol_tables(self) -> None: + t1 = TraceSymbolTable() + t1.add_symbols(["a", "b"]) + t2 = TraceSymbolTable() + t2.add_symbols(["b", "c"]) + combined = TraceSymbolTable.combine_symbol_tables([t1, t2]) + self.assertEqual(combined.get_sym_table(), ["a", "b", "c"]) + + def test_clone(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["a", "b"]) + c = TraceSymbolTable.clone(t) + self.assertEqual(c.get_sym_table(), t.get_sym_table()) + # Mutating clone doesn't affect original + c.add_symbols(["new"]) + self.assertNotIn("new", t.get_sym_index()) + + def test_create_from_df(self) -> None: + df = pd.DataFrame({"name": ["a", "b"], "cat": ["c", "d"]}) + t = TraceSymbolTable.create_from_df(df) + self.assertEqual(set(t.get_sym_table()), {"a", "b", "c", "d"}) + + def test_create_from_df_invalid_raises(self) -> None: + df = pd.DataFrame({"x": [1]}) + with self.assertRaisesRegex(ValueError, "name and cat columns"): + TraceSymbolTable.create_from_df(df) + + def test_create_from_symbol_id_map(self) -> None: + t = TraceSymbolTable.create_from_symbol_id_map({"a": 0, "b": 2}) + self.assertEqual(t.get_sym_table()[0], "a") + self.assertEqual(t.get_sym_table()[2], "b") + # Index 1 not present in map -> Undefined-1 + self.assertEqual(t.get_sym_table()[1], "Undefined-1") + + +class TestDecodeSymbolIdToSymbolName(unittest.TestCase): + def test_decode_with_short_names(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["foo(bar)", "noop"]) + df = pd.DataFrame({"name": [0, 1], "cat": [1, 0]}) + decode_symbol_id_to_symbol_name(df, t, use_shorten_name=True) + self.assertIn("s_name", df.columns) + self.assertIn("s_cat", df.columns) + + def test_decode_with_user_annotation(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["op_a", "op_b"]) + df = pd.DataFrame({"name": [0], "cat": [1], "user_annotation": [0]}) + decode_symbol_id_to_symbol_name(df, t, use_shorten_name=False) + self.assertIn("s_user_annotation", df.columns) + + def test_decode_skips_non_int_columns(self) -> None: + t = TraceSymbolTable() + t.add_symbols(["a"]) + df = pd.DataFrame({"name": ["already_str"]}) + decode_symbol_id_to_symbol_name(df, t, use_shorten_name=False) + # Should NOT add s_name since name is not int dtype + self.assertNotIn("s_name", df.columns) diff --git a/tests/test_types.py b/tests/test_types.py new file mode 100644 index 00000000..e071cbf7 --- /dev/null +++ b/tests/test_types.py @@ -0,0 +1,96 @@ +# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary. + +import re +import unittest +from typing import cast + +import pandas as pd +from hta.common.types import ( + DeviceType, + GroupingPattern, + infer_device_type, + MEMCPY_TYPE_TO_STR, + MemcpyType, + to_grouping_pattern, +) + + +class TestDeviceType(unittest.TestCase): + def test_unknown_when_no_relevant_columns(self) -> None: + df = pd.DataFrame({"x": [1]}) + self.assertEqual(infer_device_type(df), DeviceType.UNKNOWN) + + def test_gpu_when_only_stream_positive(self) -> None: + df = pd.DataFrame({"stream": [1, 2, 3]}) + self.assertEqual(infer_device_type(df), DeviceType.GPU) + + def test_cpu_when_only_stream_neg1(self) -> None: + df = pd.DataFrame({"stream": [-1, -1]}) + self.assertEqual(infer_device_type(df), DeviceType.CPU) + + def test_gpu_when_full_columns_stream_positive(self) -> None: + df = pd.DataFrame({"stream": [1, 2], "pid": [10, 10], "tid": [20, 20]}) + self.assertEqual(infer_device_type(df), DeviceType.GPU) + + def test_gpu_when_pid_zero(self) -> None: + df = pd.DataFrame({"stream": [-1, -1], "pid": [0, 0], "tid": [20, 20]}) + self.assertEqual(infer_device_type(df), DeviceType.GPU) + + def test_cpu_with_full_columns_stream_neg1_pid_nonzero(self) -> None: + df = pd.DataFrame({"stream": [-1, -1], "pid": [10, 10], "tid": [20, 20]}) + self.assertEqual(infer_device_type(df), DeviceType.CPU) + + +class TestMemcpyType(unittest.TestCase): + def test_memcpy_type_to_str_complete(self) -> None: + for mt in MemcpyType: + self.assertIn(mt, MEMCPY_TYPE_TO_STR) + self.assertIsInstance(MEMCPY_TYPE_TO_STR[mt], str) + + +class TestGroupingPattern(unittest.TestCase): + def test_match_normal(self) -> None: + gp = GroupingPattern(re.compile(r"^foo")) + self.assertTrue(gp.match("foobar")) + self.assertFalse(gp.match("baz")) + + def test_match_inverse(self) -> None: + gp = GroupingPattern(re.compile(r"^foo"), inverse_match=True) + self.assertFalse(gp.match("foobar")) + self.assertTrue(gp.match("baz")) + + def test_hashable(self) -> None: + gp1 = GroupingPattern(re.compile(r"^x"), inverse_match=False) + gp2 = GroupingPattern(re.compile(r"^x"), inverse_match=False) + self.assertEqual(hash(gp1), hash(gp2)) + # Can be used as a dict key + d = {gp1: "v"} + self.assertEqual(d[gp2], "v") + + +class TestToGroupingPattern(unittest.TestCase): + def test_passthrough_grouping_pattern(self) -> None: + gp = GroupingPattern(re.compile(r"^x")) + self.assertIs(to_grouping_pattern(gp), gp) + + def test_from_string(self) -> None: + gp = to_grouping_pattern("^foo", group_name="g", inverse_match=True) + self.assertTrue(isinstance(gp, GroupingPattern)) + self.assertEqual(gp.group_name, "g") + self.assertTrue(gp.inverse_match) + + def test_from_compiled_pattern(self) -> None: + gp = to_grouping_pattern(re.compile(r"^bar")) + self.assertTrue(isinstance(gp, GroupingPattern)) + self.assertFalse(gp.inverse_match) + self.assertEqual(gp.group_name, "") + + def test_from_list(self) -> None: + gp = to_grouping_pattern(["alpha", "beta"], group_name="lst") + self.assertTrue(isinstance(gp, GroupingPattern)) + self.assertEqual(gp.group_name, "lst") + + def test_unsupported_type_raises(self) -> None: + with self.assertRaisesRegex(ValueError, "Unsupported pattern type"): + # Intentionally pass a wrong type to test runtime validation + to_grouping_pattern(cast(str, 123))