diff --git a/hta/common/constants.py b/hta/common/constants.py index 52e4bd97..95031668 100644 --- a/hta/common/constants.py +++ b/hta/common/constants.py @@ -6,3 +6,6 @@ # https://forums.developer.nvidia.com/t/maximum-number-of-operations-in-a-stream/255260 # We anecdotally see 1022 - 1024 to cause blocking of runtime operations. CUDA_MAX_LAUNCH_QUEUE_PER_STREAM: int = 1024 + +# Max valid trace event duration. Events beyond this are corrupted (CUPTI timestamp overflow). +MAX_EVENT_DURATION_US: int = 604_800_000_000 # 7 days in microseconds diff --git a/hta/common/trace_call_stack.py b/hta/common/trace_call_stack.py index 90dda343..d1593ec1 100644 --- a/hta/common/trace_call_stack.py +++ b/hta/common/trace_call_stack.py @@ -356,6 +356,7 @@ def _construct_call_stack_graph(self, df: pd.DataFrame) -> None: value_name="time", ) .replace({"ts": -1, "end": 1}) + .assign(kind=lambda x: x["kind"].astype(int)) .sort_values("time") ).to_numpy() diff --git a/hta/common/trace_parser.py b/hta/common/trace_parser.py index 6e0f4705..0e22fe8b 100644 --- a/hta/common/trace_parser.py +++ b/hta/common/trace_parser.py @@ -138,10 +138,26 @@ def parse_trace_dict(trace_file_path: str) -> Dict[str, Any]: trace_record: Dict[str, Any] = {} if trace_file_path.endswith(".gz"): with gzip.open(trace_file_path, "rb") as fh: - trace_record = json.loads(fh.read()) + data = fh.read() + try: + trace_record = json.loads(data) + except UnicodeDecodeError as e: + logger.error( + f"Trace file {trace_file_path} contains invalid UTF-8 " + f"bytes and cannot be parsed: {e}" + ) + raise elif trace_file_path.endswith(".json"): - with open(trace_file_path, "r") as fh2: - trace_record = json.loads(fh2.read()) + with open(trace_file_path, "rb") as fh2: + data = fh2.read() + try: + trace_record = json.loads(data) + except UnicodeDecodeError as e: + logger.error( + f"Trace file {trace_file_path} contains invalid UTF-8 " + f"bytes and cannot be parsed: {e}" + ) + raise else: raise ValueError( f"expect the value of trace_file ({trace_file_path}) ends with '.gz' or 'json'" @@ -292,6 +308,19 @@ def _compress_df( df.dropna(axis=0, subset=["dur", "cat"], inplace=True) df.drop(df[df["cat"] == "Trace"].index, inplace=True) + # drop events with invalid durations (negative or exceeding threshold). + # Corrupted CUPTI timestamps can produce durations of 100+ years. + max_dur = cfg.max_event_duration_us + if max_dur is not None: + invalid_dur = (df["dur"] < 0) | (df["dur"] > max_dur) + n_invalid = invalid_dur.sum() + if n_invalid > 0: + logger.warning( + f"Dropping {n_invalid} trace events with invalid duration " + f"(dur < 0 or dur > {max_dur} us)." + ) + df.drop(df[invalid_dur].index, inplace=True) + # drop columns columns_to_drop = {"ph", "id", "bp", "s"}.intersection(set(df.columns)) df.drop(list(columns_to_drop), axis=1, inplace=True) @@ -484,20 +513,28 @@ def parse_trace_dataframe( if parser_backend == ParserBackend.JSON: meta, df, local_symbol_table = _parse_trace_dataframe_json(trace_file_path, cfg) - elif parser_backend == ParserBackend.IJSON: - meta, df, local_symbol_table = _parse_trace_dataframe_ijson( - trace_file_path, cfg - ) - elif parser_backend == ParserBackend.IJSON_BATCHED: - meta, df, local_symbol_table = _parse_trace_dataframe_ijson( - trace_file_path, - cfg, - batched=True, - ) - elif parser_backend == ParserBackend.IJSON_BATCH_AND_COMPRESS: - meta, df, local_symbol_table = _parse_trace_dataframe_ijson( - trace_file_path, cfg, batched=True, compress_on_fly=True + elif parser_backend in ( + ParserBackend.IJSON, + ParserBackend.IJSON_BATCHED, + ParserBackend.IJSON_BATCH_AND_COMPRESS, + ): + batched = parser_backend in ( + ParserBackend.IJSON_BATCHED, + ParserBackend.IJSON_BATCH_AND_COMPRESS, ) + compress = parser_backend == ParserBackend.IJSON_BATCH_AND_COMPRESS + try: + meta, df, local_symbol_table = _parse_trace_dataframe_ijson( + trace_file_path, cfg, batched=batched, compress_on_fly=compress + ) + except Exception as e: + logger.warning( + f"ijson (yajl) parser failed for {trace_file_path}: {e}. " + f"Falling back to json parser." + ) + meta, df, local_symbol_table = _parse_trace_dataframe_json( + trace_file_path, cfg + ) else: raise ValueError(f"unexpected or unsupported parser = {parser_backend}") diff --git a/hta/configs/parser_config.py b/hta/configs/parser_config.py index 4016b401..e046311b 100644 --- a/hta/configs/parser_config.py +++ b/hta/configs/parser_config.py @@ -1,4 +1,5 @@ # pyre-strict +from __future__ import annotations import copy import re @@ -6,6 +7,7 @@ from typing import Dict, List, Optional, Set, Union import pandas as pd +from hta.common.constants import MAX_EVENT_DURATION_US from hta.configs.default_values import AttributeSpec, EventArgs, ValueType from hta.configs.event_args_yaml_parser import ( parse_event_args_yaml, @@ -113,6 +115,10 @@ def __init__( # config to convert ts to nearest integer self.convert_ts_to_integer = convert_ts_to_integer + # Max valid event duration in microseconds. Events exceeding this are + # dropped as corrupted (e.g. CUPTI timestamp overflow). Set to None to disable. + self.max_event_duration_us: Optional[int] = MAX_EVENT_DURATION_US + def clone(self) -> "ParserConfig": return copy.deepcopy(self) diff --git a/tests/data/large_integer_trace/rank-70.Mar_25_20_48_12.2861.pt.trace.json.gz b/tests/data/large_integer_trace/rank-70.Mar_25_20_48_12.2861.pt.trace.json.gz new file mode 100644 index 00000000..409ef991 Binary files /dev/null and b/tests/data/large_integer_trace/rank-70.Mar_25_20_48_12.2861.pt.trace.json.gz differ diff --git a/tests/test_call_stack.py b/tests/test_call_stack.py index 34bc13f5..13242399 100644 --- a/tests/test_call_stack.py +++ b/tests/test_call_stack.py @@ -5,6 +5,7 @@ import unittest +import numpy as np import pandas as pd from hta.common.call_stack import ( CallGraph, @@ -138,6 +139,28 @@ def test_node_depth(self): # Verify df is used self.assertIsNotNone(df) + def test_events_array_numeric_dtype(self): + """astype(int) ensures kind column is numeric for np.unique(axis=0).""" + df = pd.DataFrame( + {"index": [0, 1], "ts": [100, 200], "dur": [10, 20], "stream": [-1, -1]} + ) + df["end"] = df["ts"] + df["dur"] + events = ( + df.melt( + id_vars=["index", "dur"], + value_vars=["ts", "end"], + var_name="kind", + value_name="time", + ) + .replace({"ts": -1, "end": 1}) + .assign(kind=lambda x: x["kind"].astype(int)) + .to_numpy() + ) + # kind column must be numeric, not object, for np.unique(axis=0) + # to work on newer numpy versions. + self.assertTrue(events.dtype.kind in ("i", "f")) # int or float + self.assertIsNotNone(np.unique(events, axis=0)) + class CallGraphTestCase(unittest.TestCase): def setUp(self) -> None: diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index 8ce43286..054fd66b 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -10,9 +10,11 @@ # import unittest.mock as mock import pandas as pd +from hta.common.constants import MAX_EVENT_DURATION_US from hta.common.trace import parse_trace_dict, Trace from hta.common.trace_parser import ( _auto_detect_parser_backend, + _compress_df, _open_trace_file, get_default_trace_parsing_backend, infer_gpu_type, @@ -670,5 +672,96 @@ def test_fix_mtia_memory_kernels(self) -> None: pd.testing.assert_frame_equal(fixed_df, expected_df) +class InvalidDurationFilterTest(unittest.TestCase): + """Tests that _compress_df drops events with invalid durations.""" + + @data_provider( + lambda: ( + {"durations": [100, -1, 200], "expected_count": 2}, + {"durations": [100, MAX_EVENT_DURATION_US + 1, 200], "expected_count": 2}, + {"durations": [0, 1, MAX_EVENT_DURATION_US], "expected_count": 3}, + {"durations": [-5, MAX_EVENT_DURATION_US + 1], "expected_count": 0}, + ) + ) + def test_invalid_duration_filtering( + self, durations: list, expected_count: int + ) -> None: + df = pd.DataFrame( + { + "cat": ["kernel"] * len(durations), + "name": ["k"] * len(durations), + "ts": list(range(len(durations))), + "dur": durations, + "pid": [0] * len(durations), + "tid": [0] * len(durations), + } + ) + result, _ = _compress_df(df) + self.assertEqual(len(result), expected_count) + + +class TraceParserErrorRecoveryTest(unittest.TestCase): + """Tests for graceful error recovery in trace parsing.""" + + def _create_temp_trace(self, content: bytes, suffix: str = ".json") -> str: + """Write content to a temp file and return the path.""" + import tempfile + + fd, path = tempfile.mkstemp(suffix=suffix) + os.write(fd, content) + os.close(fd) + self.addCleanup(os.remove, path) + return path + + def test_parse_trace_dict_raises_on_invalid_utf8(self) -> None: + """parse_trace_dict raises UnicodeDecodeError on corrupted trace files.""" + import json as json_mod + + valid_json = json_mod.dumps( + { + "schemaVersion": 1, + "traceEvents": [ + { + "ph": "X", + "cat": "kernel", + "name": "k", + "pid": 0, + "tid": 0, + "ts": 100, + "dur": 50, + }, + ], + } + ) + # Inject invalid UTF-8 bytes before the closing brace + bad_content = valid_json.encode("utf-8") + bad_content = bad_content[:-1] + b"\xff\xfe" + bad_content[-1:] + path = self._create_temp_trace(bad_content) + + with self.assertRaises(UnicodeDecodeError): + parse_trace_dict(path) + + def test_parse_trace_dataframe_ijson_fallback_on_large_integer(self) -> None: + """parse_trace_dataframe falls back to JSON when ijson hits large integers. + + Uses a real trace containing Comms Id = 10159970886766084423 which + exceeds yajl's signed int64 max (9223372036854775807). This is the + exact pattern from Kineto traces that causes PPS parsing failures. + Source: mvai_gpu_traces fire-longwg-gwatch-ifr-mtml-opt-v37 rank-70. + """ + trace_dir = add_test_data_path_prefix_if_exists( + "tests/data/large_integer_trace" + ) + trace_file = os.path.join( + trace_dir, "rank-70.Mar_25_20_48_12.2861.pt.trace.json.gz" + ) + cfg = ParserConfig(ParserConfig.get_minimum_args()) + cfg.parser_backend = ParserBackend.IJSON_BATCH_AND_COMPRESS + + meta, df, sym_table = parse_trace_dataframe(trace_file, cfg) + self.assertFalse(df.empty) + self.assertIn("device_type", meta) + + if __name__ == "__main__": # pragma: no cover unittest.main()