Emit annotate() label paths in the debug sidecar

OnnxDebugSession had no way to recover the annotate()/@Annotated label
hierarchy a node carried at compile time — annotations lived only on the
in-process graph. Key them by canonical node id (matching placements and
the residual states) so they round-trip onto a fresh rebuild via the new
OnnxDebugSession.annotation(node) accessor.

Test Suite Status: Unknown
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

tests/debug/test_onnx_debug_session.py

@@ -125,6 +125,43 @@ def test_sidecar_written_with_expected_schema(token_artifact):
     assert sidecar["assert_coverage"]["n_asserts"] > 0
 
 
+def _build_annotated_graph():
+    """Small graph whose nodes carry nested ``annotate()`` label paths."""
+    from torchwright.graph.node import annotate
+
+    a = create_input("a", 1)
+    b = create_input("b", 1)
+    with annotate("scene"):
+        prod = signed_multiply(a, b, max_abs1=10, max_abs2=10)
+        with annotate("scene/sum"):
+            out = add(prod, multiply_const(a, 2.0))
+    return out, create_pos_encoding(), prod, out
+
+
+def test_sidecar_carries_annotations(tmp_path):
+    """The ``annotate()`` label path round-trips through the debug sidecar
+    and onto OnnxDebugSession.annotation() against a fresh rebuild."""
+    out, pos, prod, top = _build_annotated_graph()
+    onnx_path = str(tmp_path / "annotated.onnx")
+    compile_headless_to_onnx(
+        out, pos, onnx_path, d=D, d_head=D_HEAD, max_seq_len=16, verbose=False
+    )
+
+    with open(debug_meta_path_for(onnx_path)) as f:
+        sidecar = json.load(f)
+    # The annotated nodes' paths are present; "scene/sum" is deduped, not
+    # "scene/scene/sum".
+    labels = set(sidecar["annotations"].values())
+    assert "scene" in labels
+    assert "scene/sum" in labels
+
+    # Round-trip onto a fresh rebuild (new node ids) via the session.
+    out2, pos2, prod2, top2 = _build_annotated_graph()
+    sess = OnnxDebugSession(onnx_path, out2, pos2)
+    assert sess.annotation(prod2) == "scene"
+    assert sess.annotation(top2) == "scene/sum"
+
+
 # ---------------------------------------------------------------------------
 # Parity with CompiledHeadless
 # ---------------------------------------------------------------------------

torchwright/compiler/export.py

@@ -372,10 +372,12 @@ def _write_debug_sidecar(
     The sidecar carries the residual assignment keyed by CANONICAL node
     id (see :mod:`torchwright.compiler.graph_identity`) per capture
     state, a structural fingerprint of the compiled graph for rebuild
-    validation, and the Assert/DebugWatch coverage present at compile
-    time (so the loader can warn when a rebuilt graph carries fewer
-    checks than the compiled one did — the fingerprint is deliberately
-    wrapper-transparent and cannot see that).
+    validation, the ``annotate()`` label path for every reachable node
+    that carries one (also keyed by canonical id), and the
+    Assert/DebugWatch coverage present at compile time (so the loader can
+    warn when a rebuilt graph carries fewer checks than the compiled one
+    did — the fingerprint is deliberately wrapper-transparent and cannot
+    see that).
 
     State keys correspond one-to-one with the per-layer residual tensor
     names in the emitted ONNX graph: ``"input"`` ↔ ``res_0``,
@@ -392,6 +394,7 @@ def _write_debug_sidecar(
         canonical_ids,
         debug_fingerprint,
         encode_cols,
+        nodes_by_canonical_id,
         unwrap_debug,
     )
 
@@ -437,6 +440,16 @@ def _write_debug_sidecar(
             if unwrap_debug(a.inputs[0]).node_id in canon
         }
     )
+    # Annotation paths (the "/"-separated label hierarchy set by
+    # ``annotate()`` / ``@annotated``) keyed by canonical node id, for
+    # every reachable node that carries one.  Wrappers are stepped
+    # through by ``_canonical_walk``, so this reads the annotation on the
+    # wrapped node — the same node the residual states key by.
+    annotations = {
+        str(cid): node.annotation
+        for cid, node in nodes_by_canonical_id(out).items()
+        if node.annotation is not None
+    }
     matrices, placements, n_heads, d_hidden_per_layer = _build_matrix_occupancy(
         compiled, canon, d, d_head
     )
@@ -451,6 +464,7 @@ def _write_debug_sidecar(
         "n_layers": len(compiled.layers),
         "matrices": matrices,
         "placements": placements,
+        "annotations": annotations,
         "input_specs": [list(spec) for spec in input_specs],
         "cache_stride": int(cache_stride),
         "cache_window": int(cache_window) if cache_window is not None else None,

torchwright/debug/onnx_debug.py

@@ -220,6 +220,16 @@ def __init__(
                 ra.assign(st, node, decode_cols(runs))
         self._ra = ra
         self._key_to_state = key_to_state
+
+        # --- Annotation paths from the sidecar, keyed by canonical id and
+        # mapped onto the rebuilt graph's live node objects so callers can
+        # look one up by node (mirroring node.annotation on the in-process
+        # backend).  Reachable-but-unannotated nodes are simply absent.
+        self._annotation_by_node_id: Dict[int, str] = {}
+        for cid_str, label in (sidecar.get("annotations") or {}).items():
+            node = by_canon.get(int(cid_str))
+            if node is not None:
+                self._annotation_by_node_id[node.node_id] = label
         # Post-MLP triples — the ordered states every check/probe scans,
         # mirroring the in-process backend.
         self._ordered: List[tuple] = [
@@ -525,6 +535,16 @@ def debug_value(self, node: Node) -> Optional[torch.Tensor]:
         res_tensor, _ = tensor_pair
         return extract_compiled_value(node, ds.ra, state, res_tensor)
 
+    def annotation(self, node: Node) -> Optional[str]:
+        """The ``annotate()`` path recorded for ``node`` at compile time.
+
+        Looks the rebuilt ``node`` up against the annotations the sidecar
+        captured (keyed by canonical id).  Returns ``None`` for nodes that
+        carried no annotation when compiled, or that aren't reachable from
+        the output.  Assert/DebugWatch wrappers are unwrapped first.
+        """
+        return self._annotation_by_node_id.get(unwrap_debug(node).node_id)
+
     # ---- DebugRuntime protocol surface (probe entry points) ----------------
 
     def debug_layout(self) -> tuple: