Error out when reproducing CUDA graph capture kernels (#359)

Summary:
Pull Request resolved: #359

Fix #276

When a kernel is launched during CUDA graph capture, argument extraction
is skipped (D86722827) and `extracted_args` contains only a `_note` string
instead of per-argument dicts. The reproducer's `_create_arg_from_info()`
then crashes with `AttributeError: 'str' object has no attribute 'get'`.

This change detects the `_note` sentinel in `build_context_bundle()` and
raises a clear `RuntimeError` early, before any downstream code tries to
process the incomplete argument data.

Reviewed By: xuzhao9

Differential Revision: D97514433

fbshipit-source-id: c4e61d50884e0e6e451c581f1d602970247e4849

Author: FindHao

tests/cpu/test_placeholder_replacer.py

@@ -203,5 +203,91 @@ def test_scratch_size_none_when_absent(self):
         self.assertIsNone(bundle.compile["global_scratch_size"])
 
 
+class TestBuildContextBundleCudaGraphCapture(unittest.TestCase):
+    """Tests that build_context_bundle raises RuntimeError for CUDA graph capture launches.
+
+    When a kernel is launched during CUDA graph capture, argument extraction
+    is skipped (see D86722827) and extracted_args contains only a _note string
+    instead of per-argument dicts. build_context_bundle should detect this and
+    raise a clear RuntimeError.
+    """
+
+    def _make_events_with_note(self, note_message):
+        """Create launch + compilation events where extracted_args has a _note sentinel."""
+        launch_event = {
+            "event_type": "launch",
+            "grid": [1, 1, 1],
+            "extracted_args": {
+                "_note": note_message,
+            },
+            "compilation_metadata": {
+                "hash": "abc123",
+                "num_warps": 4,
+                "num_stages": 2,
+            },
+        }
+        comp_event = {
+            "event_type": "compilation",
+            "payload": {
+                "metadata": {"hash": "abc123", "name": "my_kernel"},
+                "python_source": {
+                    "file_path": "/tmp/kernel.py",
+                    "code": "@triton.jit\ndef my_kernel(): pass",
+                },
+            },
+            "stack": [],
+        }
+        return [launch_event, comp_event]
+
+    def test_note_sentinel_raises_runtime_error(self):
+        """Test that _note in extracted_args raises RuntimeError."""
+        from tritonparse.reproducer.ingestion.ndjson import build_context_bundle
+
+        note = "Argument extraction skipped during CUDA graph capture"
+        events = self._make_events_with_note(note)
+
+        with self.assertRaises(RuntimeError) as cm:
+            build_context_bundle(events, line_index=0)
+
+        error_msg = str(cm.exception)
+        self.assertIn("Cannot generate reproducer", error_msg)
+        self.assertIn("my_kernel", error_msg)
+        self.assertIn(note, error_msg)
+
+    def test_no_note_sentinel_succeeds(self):
+        """Test that normal extracted_args without _note works fine."""
+        from tritonparse.reproducer.ingestion.ndjson import build_context_bundle
+
+        launch_event = {
+            "event_type": "launch",
+            "grid": [1, 1, 1],
+            "extracted_args": {
+                "x": {"type": "tensor", "shape": [4, 4], "dtype": "float32"},
+                "BLOCK_SIZE": {"type": "int", "value": 128},
+            },
+            "compilation_metadata": {
+                "hash": "abc123",
+                "num_warps": 4,
+                "num_stages": 2,
+            },
+        }
+        comp_event = {
+            "event_type": "compilation",
+            "payload": {
+                "metadata": {"hash": "abc123", "name": "my_kernel"},
+                "python_source": {
+                    "file_path": "/tmp/kernel.py",
+                    "code": "@triton.jit\ndef my_kernel(): pass",
+                },
+            },
+            "stack": [],
+        }
+        events = [launch_event, comp_event]
+
+        # Should not raise
+        bundle = build_context_bundle(events, line_index=0)
+        self.assertEqual(bundle.kernel_info.function_name, "my_kernel")
+
+
 if __name__ == "__main__":
     unittest.main()

tritonparse/reproducer/ingestion/ndjson.py

@@ -212,6 +212,19 @@ def build_context_bundle(
     kernel_info = get_kernel_info(comp_event)
     grid = launch_event.get("grid")
     extracted_args = launch_event.get("extracted_args", {})
+
+    # Check if this launch event was captured during CUDA graph capture,
+    # which means argument extraction was skipped and we cannot generate
+    # a reproducer. See D86722827 for context on why extraction is skipped.
+    if "_note" in extracted_args:
+        raise RuntimeError(
+            f"Cannot generate reproducer for kernel "
+            f"'{kernel_info.function_name}' at line {line_index}: "
+            f"{extracted_args['_note']}. "
+            f"Kernel launches during CUDA graph capture do not have "
+            f"extracted argument data needed for reproducer generation."
+        )
+
     comp_meta = launch_event.get("compilation_metadata", {})
 
     # Compile metadata subset we care about.