Exporters return OnnxArtifact; token extra_metadata; max_layers=400, verbose=False defaults

- OnnxArtifact (export.py, next to the sidecar plumbing): frozen
  dataclass of paths + build scalars (kind, n_layers,
  per_layer_n_heads, d, d_head, cache_stride, cache_window; token kind
  adds d_embed/vocab_size). Built strictly AFTER export from
  paths/scalars — holds no graph/weights/exporter state (the streaming
  memory bound is a stated hard invariant). .load() and
  .debug_session(output_node, pos_encoding) lazy-import their targets
  to avoid the onnx_load/onnx_debug import cycle.
- compile_to_onnx and compile_headless_to_onnx now return it (were
  -> None; nothing consumed the old return).
- compile_to_onnx gains extra_metadata, written nested under the
  sidecar 'extra' key (mirrors the headless exporter; top-level keys
  unchanged — additive only). Surfaced via OnnxTokenModule.metadata.
- compile_headless_to_onnx gains optimize=0, threaded to
  forward_compile (compile_to_onnx already had it).
- Defaults flipped on both exporters: max_layers 200->400 (pure cap at
  optimize=0; no optimize>=1 caller relies on the old default — doom
  passes it explicitly), verbose True->False (user decision).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: physicsrob

torchwright/compiler/export.py

@@ -1,17 +1,20 @@
 """Compile a torchwright graph to a KV-cached ONNX model.
 
-Two exporters, symmetric:
+Two exporters, symmetric, both returning an :class:`OnnxArtifact`
+(paths + small build metadata; ``artifact.load()`` /
+``artifact.debug_session(...)``):
 
     compile_to_onnx(output_node, pos_encoding, embedding, path, ...)
         Token I/O: token_ids -> logits.  Sidecar format
         ``torchwright.token.v1`` carries the vocab.  Consumer:
-        :mod:`torchwright.compiler.repl`.
+        ``OnnxTokenModule`` via
+        :func:`torchwright.compiler.onnx_load.load_onnx`.
 
     compile_headless_to_onnx(output_node, pos_encoding, path, ...)
         Float I/O: inputs -> outputs.  Sidecar format
         ``torchwright.headless.v1`` carries the alphabetically-ordered
-        input column names.  Consumer:
-        :mod:`torchwright.compiler.onnx_load`.
+        input column names.  Consumer: ``OnnxHeadlessModule`` via
+        :func:`torchwright.compiler.onnx_load.load_onnx`.
 
 Both speak the STATIC-cache prefill/decode protocol (the vanilla
 HF-StaticCache / vLLM pattern, chosen so ONNX Runtime can capture a CUDA
@@ -168,6 +171,63 @@ def debug_meta_path_for(onnx_path: str) -> str:
     return base + ".debug.json"
 
 
+@dataclass(frozen=True)
+class OnnxArtifact:
+    """Paths + small build metadata for a finished ONNX export.
+
+    Returned by :func:`compile_to_onnx` and :func:`compile_headless_to_onnx`
+    so consumers stop reconstructing paths and build facts by convention
+    (re-reading layer counts out of the ONNX file, rebuilding
+    ``d_embed``/``vocab_size`` from globals).
+
+    HARD INVARIANT: built strictly from paths and scalars AFTER export
+    completes — holds no graph, no weights, no exporter state.  The
+    exporters' streaming memory bound (one dense layer's worth of weights
+    in RAM regardless of depth) is sacred; this handle must never grow a
+    field that would anchor the compiled model in memory.
+    """
+
+    path: str
+    meta_path: str
+    debug_path: Optional[str]  # None when debug_sidecar=False
+    kind: str  # "token" | "headless"
+    n_layers: int
+    per_layer_n_heads: Tuple[int, ...]  # tuple copy, not the exporter's live list
+    d: int
+    d_head: int
+    cache_stride: int
+    cache_window: Optional[int]
+    d_embed: Optional[int] = None  # token kind only
+    vocab_size: Optional[int] = None  # token kind only
+
+    def load(self, providers=None):
+        """Load the artifact via :func:`torchwright.compiler.onnx_load.load_onnx`.
+
+        Returns an ``OnnxTokenModule`` or ``OnnxHeadlessModule`` per
+        ``kind``.
+        """
+        # Function-level import: onnx_load imports from this module at
+        # module level, so importing it at the top would be a cycle.
+        from torchwright.compiler.onnx_load import load_onnx
+
+        return load_onnx(self.path, providers=providers)
+
+    def debug_session(self, output_node, pos_encoding, providers=None):
+        """Open a :class:`torchwright.debug.onnx_debug.OnnxDebugSession`.
+
+        ``output_node``/``pos_encoding`` must come from the same
+        deterministic graph-construction code the export used (the
+        session fingerprint-checks the rebuild).
+        """
+        # Function-level import: onnx_debug imports from this module at
+        # module level, so importing it at the top would be a cycle.
+        from torchwright.debug.onnx_debug import OnnxDebugSession
+
+        return OnnxDebugSession(
+            self.path, output_node, pos_encoding, providers=providers
+        )
+
+
 def _write_debug_sidecar(
     onnx_path: str,
     *,
@@ -943,18 +1003,23 @@ def compile_headless_to_onnx(
     d: int = 1024,
     d_head: int = 16,
     max_seq_len: int = 512,
-    max_layers: int = 200,
-    verbose: bool = True,
+    max_layers: int = 400,
+    verbose: bool = False,
     extra_metadata: Optional[dict] = None,
     d_hidden: Optional[int] = None,
     trim_heads: bool = True,
+    optimize: int = 0,
     assume_zero_init: bool = True,
     cache_stride: Optional[int] = None,
     cache_window: Optional[int] = None,
     debug_sidecar: bool = True,
-) -> None:
+) -> OnnxArtifact:
     """Compile a float-I/O graph to a KV-cached ONNX model.
 
+    Returns an :class:`OnnxArtifact` (paths + small build metadata;
+    ``artifact.load()`` for the runtime, ``artifact.debug_session(...)``
+    for the debug surface).
+
     Writes three files:
         ``<output_path>``     — the ONNX model
         ``<stem>.meta.json``  — ``{"format": "torchwright.headless.v1",
@@ -1042,6 +1107,7 @@ def compile_headless_to_onnx(
         on_layer_compiled=on_layer_compiled,
         d_hidden=d_hidden,
         trim_heads=trim_heads,
+        optimize=optimize,
         assume_zero_init=assume_zero_init,
     )
     t_compile = time.perf_counter() - t0
@@ -1240,6 +1306,19 @@ def add(op, ins, outs, **attrs):
         print(f"Wrote {output_path} ({model_size:,} bytes)")
         print(f"Wrote {meta_path}")
 
+    return OnnxArtifact(
+        path=output_path,
+        meta_path=meta_path,
+        debug_path=debug_meta_path_for(output_path) if debug_sidecar else None,
+        kind="headless",
+        n_layers=n_layers,
+        per_layer_n_heads=tuple(per_layer_n_heads),
+        d=d,
+        d_head=d_head,
+        cache_stride=cache_stride_resolved,
+        cache_window=cache_window,
+    )
+
 
 def compile_to_onnx(
     output_node: Node,
@@ -1249,18 +1328,27 @@ def compile_to_onnx(
     d: int = 1024,
     d_head: int = 16,
     max_seq_len: int = 512,
-    max_layers: int = 200,
-    verbose: bool = True,
+    max_layers: int = 400,
+    verbose: bool = False,
     trim_heads: bool = True,
     optimize: int = 0,
     assume_zero_init: bool = True,
     d_hidden: Optional[int] = None,
+    extra_metadata: Optional[dict] = None,
     cache_stride: Optional[int] = None,
     cache_window: Optional[int] = None,
     debug_sidecar: bool = True,
-) -> None:
+) -> OnnxArtifact:
     """Compile a token-I/O graph to a KV-cached ONNX model.
 
+    Returns an :class:`OnnxArtifact` (paths + small build metadata;
+    ``artifact.load()`` for the runtime, ``artifact.debug_session(...)``
+    for the debug surface).
+
+    ``extra_metadata`` is a free-form dict written under the sidecar's
+    ``"extra"`` key (surfaced via ``OnnxTokenModule.metadata``), mirroring
+    the headless exporter; top-level sidecar keys are unchanged.
+
     Writes three files:
         ``<output_path>``     — the ONNX model
         ``<stem>.meta.json``  — ``{"format": "torchwright.token.v1",
@@ -1526,6 +1614,8 @@ def add(op, ins, outs, **attrs):
         # Windowed-cache protocol discriminator: bindings must be exactly
         # C + n_new wide (see the module docstring).
         token_meta["cache_window"] = int(cache_window)
+    if extra_metadata:
+        token_meta["extra"] = dict(extra_metadata)
     meta_path = _write_meta(output_path, token_meta)
 
     if debug_sidecar:
@@ -1561,6 +1651,21 @@ def add(op, ins, outs, **attrs):
         print(f"Wrote {output_path} ({model_size:,} bytes)")
         print(f"Wrote {meta_path}")
 
+    return OnnxArtifact(
+        path=output_path,
+        meta_path=meta_path,
+        debug_path=debug_meta_path_for(output_path) if debug_sidecar else None,
+        kind="token",
+        n_layers=n_layers,
+        per_layer_n_heads=tuple(per_layer_n_heads),
+        d=d,
+        d_head=d_head,
+        cache_stride=cache_stride_resolved,
+        cache_window=cache_window,
+        d_embed=d_embed,
+        vocab_size=int(vocab_size),
+    )
+
 
 # ---------------------------------------------------------------------------
 # In-process headless callable: a thin adapter over HeadlessTransformer.compute()