From e0f5b72007307bc0f2d78a232813e655b3ad08d6 Mon Sep 17 00:00:00 2001 From: Firoz Syed Date: Fri, 3 Apr 2026 08:02:44 -0700 Subject: [PATCH 1/3] Fix numpy compat issue with np.unique on object dtype arrays Summary: `np.unique(events, axis=0)` in `trace_call_stack.py:367` fails on newer numpy versions when the `events` array has `dtype object`. **Root cause**: DataFrame.replace({"ts": -1, "end": 1}) converts the kind column values from strings ("ts", "end") to ints (-1, 1), but on newer pandas versions the column dtype remains object instead of being inferred as int64. When to_numpy() creates the array, the object dtype causes np.unique(axis=0) to throw TypeError. Fix: add `.assign(kind=lambda x: x["kind"].astype(int))` to the chain, forcing `kind` to int dtype before `to_numpy()`. This fixes the OSS CI failures: - [facebookresearch/HolisticTraceAnalysis: CI / build (3.10)](https://github.com/facebookresearch/HolisticTraceAnalysis/actions/runs/23827081316/job/69452407257?pr=324) Differential Revision: D99093340 --- hta/common/trace_call_stack.py | 1 + hta/configs/parser_config.py | 1 + tests/test_call_stack.py | 23 +++++++++++++++++++++++ 3 files changed, 25 insertions(+) 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/configs/parser_config.py b/hta/configs/parser_config.py index 4016b401..c39bd010 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 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: From d1fa67bf13a19475891662877803aadcc5b4e887 Mon Sep 17 00:00:00 2001 From: Firoz Syed Date: Fri, 3 Apr 2026 08:02:45 -0700 Subject: [PATCH 2/3] Add graceful error recovery for trace parsing (#324) Summary: Pull Request resolved: https://github.com/facebookresearch/HolisticTraceAnalysis/pull/324 OSS HTA uses ijson with the yajl C backend for memory-efficient streaming of large traces. However, yajl has a signed int64 limit and strict character validation that causes parsing failures on traces with: - Large integers (e.g. Comms Id values exceeding int64 max) - Invalid UTF-8 bytes - Unescaped quotes from old Kineto versions Internal HTA does not hit these issues because it uses Python json.loads (arbitrary precision integers) with a strip_invalid_bytes fallback. This diff adds two error handling mechanisms to OSS HTA: 1. parse_trace_dict: logs a clear error and raises on UnicodeDecodeError, telling the user the trace file is corrupted. 2. parse_trace_dataframe: wraps ijson parsing in try/except and falls back to the JSON backend when yajl fails. This preserves ijson memory efficiency for normal traces while gracefully handling traces with valid JSON that yajl cannot parse (e.g. large integers). Differential Revision: D98693551 --- hta/common/trace_parser.py | 56 +++++++++++----- ...k-70.Mar_25_20_48_12.2861.pt.trace.json.gz | Bin 0 -> 238 bytes tests/test_trace_parse.py | 63 ++++++++++++++++++ 3 files changed, 103 insertions(+), 16 deletions(-) create mode 100644 tests/data/large_integer_trace/rank-70.Mar_25_20_48_12.2861.pt.trace.json.gz diff --git a/hta/common/trace_parser.py b/hta/common/trace_parser.py index 6e0f4705..326744af 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'" @@ -484,20 +500,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/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 0000000000000000000000000000000000000000..409ef9917ee5fcdbf1c88a4c3588934608e1939a GIT binary patch literal 238 zcmVIA1X`E;2YaF)na)E_8BX zV`VOCb8l_{ZH>DM!axv3e`T0D#YW9*I}r=9vk*ZfER(^A+0C*^Y(oBBCL*S=JKQ^S z&xjq>Y|WWN_h=nzh$Fy07~=hlUEiUG>xjNa4S2x_kW4T9ZGxBvhE literal 0 HcmV?d00001 diff --git a/tests/test_trace_parse.py b/tests/test_trace_parse.py index 8ce43286..f1653d42 100644 --- a/tests/test_trace_parse.py +++ b/tests/test_trace_parse.py @@ -670,5 +670,68 @@ def test_fix_mtia_memory_kernels(self) -> None: pd.testing.assert_frame_equal(fixed_df, expected_df) +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() From a39086e925d31f59e0ea7a10c90eba87dd0ecc8a Mon Sep 17 00:00:00 2001 From: Firoz Syed Date: Fri, 3 Apr 2026 09:45:35 -0700 Subject: [PATCH 3/3] Filter trace events with corrupted CUPTI durations in _compress_df (#325) Summary: Pull Request resolved: https://github.com/facebookresearch/HolisticTraceAnalysis/pull/325 Some GPU traces have corrupted timestamps from CUPTI, producing nonsensical event durations (e.g., 236 years). These corrupted values propagate through HTA's DataFrame and skew all downstream analysis (GPU hours, iteration times, comm/IO breakdowns). This diff adds an acceptance check in `_compress_df()`, the single chokepoint that all trace parser backends (json, ijson, ijson_batched) flow through. Events with negative or excessively large durations are dropped at parse time, before any downstream consumer sees them. **Threshold: 7 days (604,800,000,000 us)** Query over last 30 days on `ai_infra.gpu_trace_stats`: | Duration Bucket | Actual Range (ms) | Trace Count | |-----------------|------------------------------------------|-------------| | 0-1s | 0 - 1,000 | 13,404,751 | | 1s-10s | 1,000 - 10,000 | 67,420,645 | | 10s-1min | 10,000 - 60,000 | 7,364,988 | | 1min-1hr | 60,000 - 3,503,332 | 792,591 | | 1hr-1day | 4,007,737 - 53,174,381 | 5 | | 1day-7days | 201,830,019 - 286,093,774 | 2 | | 7days-1year | 610,791,801 - 29,085,041,461 | 14 | | 1year+ | 55,971,542,156 - 9,223,372,033,022 | 86 | 99.999% of traces are under 1 hour (max 3.5M ms). Everything above 7 days is clearly corrupted (max is 9.2e15 ms = INT64_MAX/1000, classic CUPTI overflow). Related: D96421400 (downstream fix in Durin post-processing) Differential Revision: D98705186 --- hta/common/constants.py | 3 +++ hta/common/trace_parser.py | 13 +++++++++++++ hta/configs/parser_config.py | 5 +++++ tests/test_trace_parse.py | 30 ++++++++++++++++++++++++++++++ 4 files changed, 51 insertions(+) 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_parser.py b/hta/common/trace_parser.py index 326744af..0e22fe8b 100644 --- a/hta/common/trace_parser.py +++ b/hta/common/trace_parser.py @@ -308,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) diff --git a/hta/configs/parser_config.py b/hta/configs/parser_config.py index c39bd010..e046311b 100644 --- a/hta/configs/parser_config.py +++ b/hta/configs/parser_config.py @@ -7,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, @@ -114,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/test_trace_parse.py b/tests/test_trace_parse.py index f1653d42..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,6 +672,34 @@ 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."""