Filter trace events with corrupted CUPTI durations in _compress_df (#325)

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

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

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}")
 

hta/configs/parser_config.py

@@ -1,11 +1,13 @@
 # pyre-strict
+from __future__ import annotations
 
 import copy
 import re
 from enum import Enum
 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)
 

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:

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