downshift-server 0.2.0__py3-none-any.whl

downshift/cli/main.py

@@ -0,0 +1,428 @@
+"""downshift CLI. Each command loads, calls the library, and hands the result to render."""
+
+import json
+import logging
+import os
+import re
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass, replace
+from enum import Enum
+from pathlib import Path
+from typing import Annotated
+
+import typer
+import uvicorn
+from fastapi import FastAPI
+
+import downshift
+from downshift import __version__, settings
+from downshift.cli import render
+from downshift.export.manifest import manifest_path_for
+from downshift.export.shapes import parse_dynamic_spec
+from downshift.export.verdict import ExportVerdict
+from downshift.loading import LoadedModel, LoadError, load_model
+from downshift.serve.app import build_app
+from downshift.serve.engine import BackendChoice, ServeOptions, ServingState, prepare_serving
+
+EXIT_USAGE = 4  # bad model spec, bad option, unloadable file
+EXIT_CRASH = 5
+
+app = typer.Typer(
+    help="Check whether a PyTorch model survives ONNX export, then serve it.",
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+
+class LogLevel(str, Enum):
+    debug = "debug"
+    info = "info"
+    warning = "warning"
+    error = "error"
+
+
+class LogFormat(str, Enum):
+    text = "text"
+    json = "json"
+
+
+ModelArg = Annotated[
+    str,
+    typer.Argument(metavar="MODEL", help="model.onnx | pkg.module:attr | weights.pt | org/repo | hf-repo-dir/"),
+]
+InputsOpt = Annotated[
+    str | None,
+    typer.Option(
+        "--inputs", metavar="pkg.module:fn", help="Example inputs: a tuple, or a factory for one"
+    ),
+]
+ModelClassOpt = Annotated[
+    str | None,
+    typer.Option(
+        "--model-class", metavar="pkg.module:Class", help="Class to load a state dict into"
+    ),
+]
+UnsafeLoadOpt = Annotated[
+    bool,
+    typer.Option(
+        "--unsafe-load", help="Allow torch.load(weights_only=False); runs code from the file"
+    ),
+]
+AdapterOpt = Annotated[
+    str | None,
+    typer.Option(
+        "--adapter",
+        metavar="NAME|path/to/adapter.py[:attr]",
+        help="Model-family adapter: generic, pyg, hf, or your own adapter.py; default: detect",
+    ),
+]
+SamplesOpt = Annotated[
+    int, typer.Option("-k", "--samples", min=1, help="Number of verification samples")
+]
+DynamicOpt = Annotated[
+    str | None,
+    typer.Option(
+        "--dynamic",
+        metavar="NAME:AXIS[,...]",
+        help='Dynamic axes, e.g. "x:0,edge_index:1". Default: axis 0 of every input.',
+    ),
+]
+ReferenceOpt = Annotated[
+    str | None,
+    typer.Option("--reference", metavar="MODEL", help="PyTorch model to verify a .onnx against"),
+]
+IntraOpThreadsOpt = Annotated[
+    int,
+    typer.Option(
+        "--intra-op-threads", min=0, help="ORT threads within one op; 0 = let ONNX Runtime choose"
+    ),
+]
+InterOpThreadsOpt = Annotated[
+    int,
+    typer.Option(
+        "--inter-op-threads", min=0, help="ORT threads across ops; 0 = let ONNX Runtime choose"
+    ),
+]
+WorkersOpt = Annotated[
+    int,
+    typer.Option(
+        "--workers",
+        min=1,
+        help="Uvicorn worker processes; each independently loads/exports/warms the model",
+    ),
+]
+JsonOpt = Annotated[
+    bool, typer.Option("--json", help="Print the verdict as JSON and nothing else")
+]
+LogLevelOpt = Annotated[LogLevel, typer.Option("--log-level")]
+LogFormatOpt = Annotated[LogFormat, typer.Option("--log-format")]
+
+
+class _JsonFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        payload = {
+            "time": self.formatTime(record, "%Y-%m-%dT%H:%M:%S"),
+            "level": record.levelname,
+            "logger": record.name,
+            "message": record.getMessage(),
+        }
+        if record.exc_info:
+            payload["exc_info"] = self.formatException(record.exc_info)
+        return json.dumps(payload)
+
+
+def _setup_logging(level: LogLevel, fmt: LogFormat) -> None:
+    handler = logging.StreamHandler(sys.stderr)  # keep stdout clean for --json
+    if fmt is LogFormat.json:
+        handler.setFormatter(_JsonFormatter())
+    else:
+        handler.setFormatter(logging.Formatter("%(levelname)s %(name)s: %(message)s"))
+    logging.basicConfig(level=level.value.upper(), handlers=[handler], force=True)
+    # The ONNX optimizer passes log every rewrite at INFO. Only show them when debugging.
+    if level is not LogLevel.debug:
+        for name in ("onnxscript", "onnx_ir"):
+            logging.getLogger(name).setLevel(logging.WARNING)
+
+
+@contextmanager
+def _exit_on_error(debug: bool) -> Iterator[None]:
+    """User errors exit 4, anything else exits 5. typer.Exit passes through untouched."""
+    try:
+        yield
+    except typer.Exit:
+        raise
+    except ValueError as exc:  # LoadError, bad --dynamic, bad backend combination
+        render.error(str(exc))
+        raise typer.Exit(EXIT_USAGE) from exc
+    except Exception as exc:
+        if debug:
+            render.print_traceback()
+        render.error(f"{type(exc).__name__}: {exc}")
+        raise typer.Exit(EXIT_CRASH) from exc
+
+
+def _load(spec: str, inputs: str | None, model_class: str | None, unsafe_load: bool) -> LoadedModel:
+    if unsafe_load:
+        render.warn(
+            f"--unsafe-load: torch.load(weights_only=False) on {spec}; arbitrary code may run"
+        )
+    return load_model(spec, inputs=inputs, model_class=model_class, unsafe_load=unsafe_load)
+
+
+_IMPORT_SPEC = re.compile(r"^[A-Za-z_][\w.]*:[A-Za-z_]\w*$")
+
+
+def slug(spec: str) -> str:
+    """tests.models.clean_mlp:make_model -> clean_mlp; ./gat_v3.pt -> gat_v3; org/repo -> repo."""
+    if _IMPORT_SPEC.match(spec):
+        return spec.partition(":")[0].rsplit(".", 1)[-1]
+    return Path(spec).stem
+
+
+def _emit(
+    verdict: ExportVerdict, model_name: str, as_json: bool, extra: dict | None = None
+) -> None:
+    if as_json:
+        typer.echo(json.dumps(verdict.to_dict() | (extra or {}), indent=2))
+    else:
+        render.print_verdict(verdict, model_name)
+
+
+@app.command("check")
+def check_cmd(
+    model: ModelArg,
+    json_out: JsonOpt = False,
+    reference: ReferenceOpt = None,
+    inputs: InputsOpt = None,
+    model_class: ModelClassOpt = None,
+    unsafe_load: UnsafeLoadOpt = False,
+    adapter: AdapterOpt = None,
+    k: SamplesOpt = settings.SAMPLES,
+    dynamic: DynamicOpt = None,
+    log_level: LogLevelOpt = LogLevel.warning,
+    log_format: LogFormatOpt = LogFormat.text,
+) -> None:
+    """Export in memory and verify numerics. Exit 0 CLEAN, 1 FAILED, 2 DEGRADED, 3 UNVERIFIED."""
+    _setup_logging(log_level, log_format)
+    with _exit_on_error(log_level is LogLevel.debug):
+        loaded = _load(model, inputs, model_class, unsafe_load)
+        dynamic_spec = parse_dynamic_spec(dynamic) if dynamic else None
+        if loaded.onnx_path is not None:
+            ref = _load(reference, inputs, model_class, unsafe_load) if reference else None
+            verdict = downshift.intake(
+                loaded.onnx_path,
+                ref.model if ref else None,
+                ref.example_inputs if ref else None,
+                adapter or (ref.adapter_hint if ref else None),
+                k=k,
+                dynamic=dynamic_spec,
+            )
+        else:
+            assert loaded.model is not None
+            verdict = downshift.check(
+                loaded.model,
+                loaded.example_inputs,
+                k=k,
+                adapter=adapter or loaded.adapter_hint,
+                dynamic=dynamic_spec,
+            )
+        _emit(verdict, model, json_out)
+        raise typer.Exit(verdict.exit_code)
+
+
+@app.command("export")
+def export_cmd(
+    model: ModelArg,
+    output: Annotated[Path, typer.Option("-o", "--output", help="Output directory")],
+    name: Annotated[
+        str | None, typer.Option("--name", help="Artifact stem; default: model slug")
+    ] = None,
+    fp16: Annotated[
+        bool, typer.Option("--fp16", help="Cast the model to fp16 before export")
+    ] = False,
+    no_verify: Annotated[
+        bool, typer.Option("--no-verify", help="Skip numerics; the verdict is UNVERIFIED")
+    ] = False,
+    json_out: JsonOpt = False,
+    inputs: InputsOpt = None,
+    model_class: ModelClassOpt = None,
+    unsafe_load: UnsafeLoadOpt = False,
+    adapter: AdapterOpt = None,
+    k: SamplesOpt = settings.SAMPLES,
+    dynamic: DynamicOpt = None,
+    log_level: LogLevelOpt = LogLevel.warning,
+    log_format: LogFormatOpt = LogFormat.text,
+) -> None:
+    """Export to DIR/NAME.onnx with a NAME.manifest.json sidecar. Nothing is written if FAILED."""
+    _setup_logging(log_level, log_format)
+    with _exit_on_error(log_level is LogLevel.debug):
+        loaded = _load(model, inputs, model_class, unsafe_load)
+        if loaded.model is None:
+            raise LoadError(f"{model} is already ONNX; export needs a PyTorch model")
+        onnx_path = output / f"{name or slug(model)}.onnx"
+        if no_verify:
+            render.warn("--no-verify: the graph is saved without checking its numerics")
+        verdict = downshift.export(
+            loaded.model,
+            onnx_path,
+            loaded.example_inputs,
+            k=k,
+            adapter=adapter or loaded.adapter_hint,
+            dynamic=parse_dynamic_spec(dynamic) if dynamic else None,
+            fp16=fp16,
+            source_path=loaded.source_path,
+            verify_numerics=not no_verify,
+        )
+        manifest = manifest_path_for(onnx_path) if verdict.onnx_path else None
+        _emit(verdict, model, json_out, {"manifest_path": str(manifest) if manifest else None})
+        if not json_out:
+            render.print_artifacts(verdict.onnx_path, manifest)
+        raise typer.Exit(verdict.exit_code)
+
+
+@dataclass
+class ServeArgs:
+    """Everything needed to rebuild a ServingState + FastAPI app from scratch. Plain JSON-able
+    types only: a multi-worker run ships this to each worker process via an env var."""
+
+    model: str
+    inputs: str | None
+    model_class: str | None
+    unsafe_load: bool
+    adapter: str | None
+    k: int
+    dynamic: str | None
+    reference: str | None
+    middleware: list[str] | None
+    backend: str
+    force_onnx: bool
+    device: str
+    warmup: int
+    intra_op_threads: int
+    inter_op_threads: int
+    log_level: str
+    log_format: str
+
+
+_SERVE_ARGS_ENV = "_DOWNSHIFT_SERVE_ARGS"
+
+
+def _build_serving_app(args: ServeArgs) -> tuple[ServingState, FastAPI]:
+    loaded = _load(args.model, args.inputs, args.model_class, args.unsafe_load)
+    ref = (
+        _load(args.reference, args.inputs, args.model_class, args.unsafe_load)
+        if args.reference
+        else None
+    )
+    opts = ServeOptions(
+        backend=BackendChoice(args.backend),
+        force_onnx=args.force_onnx,
+        device=args.device,
+        warmup=args.warmup,
+        k=args.k,
+        adapter=args.adapter,
+        dynamic=parse_dynamic_spec(args.dynamic) if args.dynamic else None,
+        intra_op_threads=args.intra_op_threads,
+        inter_op_threads=args.inter_op_threads,
+    )
+    state = prepare_serving(loaded, opts, ref)
+    api = build_app(state, tuple(args.middleware or ()))
+    return state, api
+
+
+def _serve_app_factory() -> FastAPI:
+    """Import-string target for uvicorn's multi-worker mode (`downshift.cli.main:_serve_app_factory`).
+    Each worker process calls this on its own, independently reloading/re-exporting/re-warming
+    the model from the args the parent process serialized into _SERVE_ARGS_ENV."""
+    args = ServeArgs(**json.loads(os.environ[_SERVE_ARGS_ENV]))
+    _setup_logging(LogLevel(args.log_level), LogFormat(args.log_format))
+    _, api = _build_serving_app(args)
+    return api
+
+
+@app.command("serve")
+def serve_cmd(
+    model: ModelArg,
+    host: Annotated[str, typer.Option("--host")] = settings.HOST,
+    port: Annotated[int, typer.Option("--port")] = settings.PORT,
+    backend: Annotated[BackendChoice, typer.Option("--backend")] = BackendChoice(settings.BACKEND),
+    force_onnx: Annotated[
+        bool, typer.Option("--force-onnx", help="Serve a DEGRADED graph via ONNX Runtime anyway")
+    ] = False,
+    device: Annotated[str, typer.Option("--device", help="auto | cpu | cuda")] = settings.DEVICE,
+    warmup: Annotated[
+        int, typer.Option("--warmup", min=0, help="Warm-up inferences before /ready flips")
+    ] = settings.WARMUP,
+    reference: ReferenceOpt = None,
+    middleware: Annotated[
+        list[str] | None,
+        typer.Option("--middleware", metavar="pkg.module:Attr", help="Middleware to attach; repeatable"),
+    ] = None,
+    inputs: InputsOpt = None,
+    model_class: ModelClassOpt = None,
+    unsafe_load: UnsafeLoadOpt = False,
+    adapter: AdapterOpt = None,
+    k: SamplesOpt = settings.SAMPLES,
+    dynamic: DynamicOpt = None,
+    intra_op_threads: IntraOpThreadsOpt = settings.INTRA_OP_THREADS,
+    inter_op_threads: InterOpThreadsOpt = settings.INTER_OP_THREADS,
+    workers: WorkersOpt = settings.WORKERS,
+    log_level: LogLevelOpt = LogLevel.info,
+    log_format: LogFormatOpt = LogFormat.text,
+) -> None:
+    """Check the model, pick a backend from the verdict, and serve it over HTTP."""
+    _setup_logging(log_level, log_format)
+    with _exit_on_error(log_level is LogLevel.debug):
+        args = ServeArgs(
+            model=model,
+            inputs=inputs,
+            model_class=model_class,
+            unsafe_load=unsafe_load,
+            adapter=adapter,
+            k=k,
+            dynamic=dynamic,
+            reference=reference,
+            middleware=list(middleware) if middleware else None,
+            backend=backend.value,
+            force_onnx=force_onnx,
+            device=device,
+            warmup=warmup,
+            intra_op_threads=intra_op_threads,
+            inter_op_threads=inter_op_threads,
+            log_level=log_level.value,
+            log_format=log_format.value,
+        )
+        if workers <= 1:
+            state, api = _build_serving_app(args)
+            render.print_banner(state, host, port)
+            uvicorn.run(api, host=host, port=port, log_level=log_level.value)
+        else:
+            # Only for the banner/fail-fast check: each of the N workers rebuilds its own
+            # backend anyway, so this throwaway copy skips warmup, it'll never serve traffic.
+            state, _ = _build_serving_app(replace(args, warmup=0))
+            render.print_banner(state, host, port)
+            render.warn(
+                f"--workers {workers}: each worker independently reloads, re-exports, and "
+                "re-warms the model (memory and startup time scale with this number)"
+            )
+            os.environ[_SERVE_ARGS_ENV] = json.dumps(asdict(args))
+            uvicorn.run(
+                "downshift.cli.main:_serve_app_factory",
+                host=host,
+                port=port,
+                workers=workers,
+                log_level=log_level.value,
+                factory=True,
+            )
+
+
+@app.command()
+def version() -> None:
+    """Print the version."""
+    typer.echo(f"downshift v{__version__}")
+
+
+if __name__ == "__main__":  # pragma: no cover
+    app()

downshift/cli/render.py

@@ -0,0 +1,174 @@
+"""All rich output for the CLI lives here. Commands hand over objects; this module prints."""
+
+from pathlib import Path
+
+from rich import box
+from rich.console import Console
+from rich.markup import escape
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+from downshift import __version__
+from downshift.export.verdict import ExportVerdict
+from downshift.serve.engine import ServingState
+
+console = Console()
+err_console = Console(stderr=True)
+
+STATUS_STYLE = {
+    "CLEAN": "bold green",
+    "DEGRADED": "bold yellow",
+    "FAILED": "bold red",
+    "UNVERIFIED": "bold magenta",
+}
+
+
+def _sym(utf: str, ascii_: str) -> str:
+    """Unicode glyph unless the console can't encode it (cp1252 Windows pipes)."""
+    return ascii_ if console.options.ascii_only else utf
+
+
+def _status_text(verdict: ExportVerdict) -> Text:
+    text = Text(verdict.status, style=STATUS_STYLE[verdict.status])
+    details = [verdict.capture_strategy] if verdict.capture_strategy else []
+    if verdict.opset is not None:
+        details.append(f"opset {verdict.opset}")
+    if details:
+        text.append(f"  ({', '.join(details)})", style="dim")
+    return text
+
+
+def _numerics_text(verdict: ExportVerdict) -> Text:
+    n = verdict.numerics
+    if n is None:
+        return Text("not checked", style="dim")
+    text = Text(f"max abs err {n.max_abs_err:.2e} over {n.samples_tested} samples  ")
+    if n.passed:
+        text.append(_sym("✓", "OK"), style="bold green")
+    else:
+        text.append(f"{_sym('✗', 'X')} {n.failures}/{n.samples_tested} failed", style="bold red")
+    return text
+
+
+def _dynamic_text(verdict: ExportVerdict) -> str:
+    if not verdict.dynamic_dims:
+        return _sym("—", "-")
+    return ", ".join(
+        f"{name}[{axis}]" for name, axes in verdict.dynamic_dims.items() for axis in axes
+    )
+
+
+def _shape_text(verdict: ExportVerdict) -> str:
+    if verdict.shape_generalization is None:
+        return _sym("—", "-")
+    return "yes" if verdict.shape_generalization else "no"
+
+
+def print_verdict(verdict: ExportVerdict, model_name: str) -> None:
+    table = Table(show_header=False, box=box.ROUNDED, border_style=STATUS_STYLE[verdict.status])
+    table.add_column(style="bold", no_wrap=True)
+    table.add_column()
+    table.add_row("Model", escape(model_name))
+    table.add_row("Family", verdict.model_family)
+    table.add_row("Export", _status_text(verdict))
+    table.add_row("Numerics", _numerics_text(verdict))
+    table.add_row("Shape-general", _shape_text(verdict))
+    table.add_row("Dynamic dims", _dynamic_text(verdict))
+    if verdict.unsupported_ops:
+        table.add_row("Unsupported ops", Text(", ".join(verdict.unsupported_ops), style="red"))
+    if verdict.warnings:
+        table.add_row("Warnings", Text("\n".join(verdict.warnings), style="yellow"))
+    table.add_row("Backend", verdict.recommended_backend)
+    table.add_row("Reason", escape(verdict.reason))
+    console.print(table)
+
+
+def print_artifacts(onnx_path: Path | None, manifest_path: Path | None) -> None:
+    if onnx_path is None:
+        console.print("[red]nothing written[/]: the export failed")
+        return
+    console.print(f"[bold]Wrote[/]     {escape(str(onnx_path))}")
+    if manifest_path is not None:
+        console.print(f"[bold]Manifest[/]  {escape(str(manifest_path))}")
+
+
+def _backend_text(state: ServingState) -> Text:
+    meta = state.backend.metadata()
+    label = "torch (eager)" if meta.name == "torch" else meta.name
+    text = Text(f"{label} {_sym('·', '|')} {meta.device}")
+    arrow = _sym("←", "<-")
+    if state.forced_onnx:
+        text.append(f"  {arrow} --force-onnx", style="yellow")
+    elif state.backend_auto_selected:
+        text.append(f"  {arrow} auto-selected", style="dim")
+    return text
+
+
+def print_banner(state: ServingState, host: str, port: int) -> None:
+    """Boot banner for `serve`. Says what is served, how it was judged, and where it listens."""
+    verdict = state.verdict
+    sub = _sym("└", "\\")
+    grid = Table.grid(padding=(0, 3))
+    grid.add_column(style="bold cyan", no_wrap=True)
+    grid.add_column()
+
+    grid.add_row("Model", escape(state.source))
+    grid.add_row("Family", verdict.model_family)
+
+    unverified_onnx = verdict.status == "UNVERIFIED" and verdict.prepared is None
+    if unverified_onnx:
+        grid.add_row("Verdict", _status_text(verdict) + Text(" - no reference model supplied"))
+        grid.add_row("", Text(f"{sub} served as-is; numerics were never checked", style="dim"))
+        grid.add_row("Tip", "pass --reference <model> to verify")
+    else:
+        grid.add_row("Verdict", _status_text(verdict))
+        if verdict.status in ("FAILED", "UNVERIFIED"):
+            grid.add_row("", Text(f"{sub} {verdict.reason}", style="dim"))
+
+    grid.add_row("Numerics", _numerics_text(verdict))
+
+    if verdict.status == "DEGRADED":
+        n = verdict.numerics
+        if n is None:
+            detail = verdict.reason
+        else:
+            detail = (
+                f"numerics diverge on {n.failures}/{n.samples_tested} samples "
+                f"(max abs err {n.max_abs_err:.2e})"
+            )
+        grid.add_row("", Text(f"{_sym('⚠', '!')} {detail}", style="yellow"))
+        if state.backend.name != "onnxruntime":
+            grid.add_row("Override", "--force-onnx to serve the ONNX graph anyway")
+
+    grid.add_row("Backend", _backend_text(state))
+    for note in state.notes:  # backend-selection notes from the engine
+        style = "bold red" if "outputs may be wrong" in note else "yellow"
+        grid.add_row("", Text(f"{_sym('⚠', '!')} {note}", style=style))
+    grid.add_row("Dynamic dims", _dynamic_text(verdict))
+    for warning in verdict.warnings:
+        grid.add_row("", Text(f"{_sym('⚠', '!')} {warning}", style="yellow"))
+    grid.add_row("Endpoint", f"http://{host}:{port}")
+
+    console.print(
+        Panel(
+            grid,
+            title=f"downshift v{__version__}",
+            title_align="left",
+            border_style=STATUS_STYLE[verdict.status],
+            expand=False,
+            padding=(1, 2),
+        )
+    )
+
+
+def warn(msg: str) -> None:
+    err_console.print(f"[bold yellow]warning:[/] {escape(msg)}")
+
+
+def error(msg: str) -> None:
+    err_console.print(f"[bold red]error:[/] {escape(msg)}")
+
+
+def print_traceback() -> None:
+    err_console.print_exception()

downshift/export/capture.py

@@ -0,0 +1,93 @@
+"""Drive torch.export + torch.onnx.export and say which strategy worked.
+
+We run torch.export ourselves (strict=False, then strict=True) so the winning strategy is
+a value we return, not something scraped from console output. The ONNX translation is
+still entirely torch.onnx.export's; we only hand it the ExportedProgram.
+
+torch 2.14 note: the dynamo exporter documents only the two strict modes. There's no
+draft_export step or TorchScript fallback any more.
+"""
+
+import contextlib
+import io
+import logging
+from dataclasses import dataclass, field
+
+import torch
+
+_STRATEGIES: tuple[tuple[str, bool], ...] = (("strict=False", False), ("strict=True", True))
+
+# torch.onnx logs a warning per missing torchvision op on every export. Not actionable.
+logging.getLogger("torch.onnx._internal.exporter._registration").setLevel(logging.ERROR)
+
+
+@dataclass
+class CaptureResult:
+    success: bool
+    capture_strategy: str | None  # a _STRATEGIES name; None when nothing traced
+    onnx_program: "torch.onnx.ONNXProgram | None" = None
+    opset: int | None = None
+    op_types: list[str] = field(default_factory=list)
+    exception: Exception | None = None
+    stderr: str = ""  # whatever torch printed while we tried; useful at debug level
+
+
+def capture(
+    model: torch.nn.Module,
+    example_inputs: tuple,
+    dynamic_shapes: tuple | None = None,
+) -> CaptureResult:
+    exported_program = None
+    strategy_used: str | None = None
+    last_exception: Exception | None = None
+    # torch prints whole FX graphs straight to stderr when a data-dependent guard fails.
+    # Keep that out of the user's terminal; the exception message is what matters.
+    captured = io.StringIO()
+
+    with contextlib.redirect_stderr(captured):
+        for name, strict in _STRATEGIES:
+            try:
+                exported_program = torch.export.export(
+                    model, example_inputs, dynamic_shapes=dynamic_shapes, strict=strict
+                )
+            except Exception as exc:  # noqa: BLE001 - a failed strategy means try the next
+                last_exception = exc
+            else:
+                strategy_used = name
+                break
+
+        if exported_program is None:
+            return CaptureResult(
+                success=False,
+                capture_strategy=None,
+                exception=last_exception,
+                stderr=captured.getvalue(),
+            )
+
+        try:
+            onnx_program = torch.onnx.export(exported_program, verbose=False, report=False)
+        except Exception as exc:  # noqa: BLE001 - a failed translation is a failed capture
+            return CaptureResult(
+                success=False,
+                capture_strategy=strategy_used,
+                exception=exc,
+                stderr=captured.getvalue(),
+            )
+
+    if onnx_program is None:
+        # The type stub allows None for the legacy path; with an ExportedProgram it never is.
+        return CaptureResult(
+            success=False,
+            capture_strategy=strategy_used,
+            exception=RuntimeError("torch.onnx.export returned None"),
+        )
+
+    proto = onnx_program.model_proto
+    return CaptureResult(
+        success=True,
+        capture_strategy=strategy_used,
+        onnx_program=onnx_program,
+        opset=proto.opset_import[0].version if proto.opset_import else None,
+        op_types=[node.op_type for node in proto.graph.node],
+        stderr=captured.getvalue(),
+    )