downshift-server 0.2.0__py3-none-any.whl

downshift/export/inputs.py

@@ -0,0 +1,25 @@
+"""Example-input synthesis ladder.
+
+1. user-supplied      always wins
+2. adapter-derived    the adapter knows its family (HF config, PyG in_channels, ...)
+3. signature guess    the generic adapter's first-Linear/Conv heuristic
+4. fail loudly        say exactly what to pass
+"""
+
+from torch import nn
+
+from downshift.adapters.base import Adapter
+
+
+def synthesize(model: nn.Module, adapter: Adapter, user_inputs: tuple | None) -> tuple:
+    if user_inputs is not None:
+        return user_inputs
+    guessed = adapter.example_inputs(model)
+    if guessed is not None:
+        return guessed
+    raise ValueError(
+        f"Couldn't work out example inputs for {type(model).__name__} with the "
+        f"{adapter.name!r} adapter. Pass them explicitly: from Python, "
+        "check(model, example_inputs=(tensor, ...)); from the CLI, --inputs module:function "
+        "where the function returns a tuple of forward() arguments."
+    )

downshift/export/manifest.py

@@ -0,0 +1,89 @@
+"""Provenance sidecar written next to every exported .onnx."""
+
+import hashlib
+import json
+from datetime import UTC, datetime
+from enum import IntEnum
+from pathlib import Path
+
+import onnx
+import onnxruntime
+import torch
+
+from downshift.export.verdict import ExportVerdict
+
+
+class _DTYPE_NAMES(IntEnum):
+    """ONNX TensorProto dtype codes we can label, keyed by their short name."""
+
+    fp32 = onnx.TensorProto.FLOAT
+    fp16 = onnx.TensorProto.FLOAT16
+    fp64 = onnx.TensorProto.DOUBLE
+    bf16 = onnx.TensorProto.BFLOAT16
+
+
+# A plain value->name dict, so a miss is a dict lookup rather than an IntEnum ValueError;
+# most initializers (int64 indices, bools, ...) are misses, and this runs in a loop.
+_DTYPE_LOOKUP: dict[int, str] = {member.value: member.name for member in _DTYPE_NAMES}
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _dtype_name(code: int) -> str | None:
+    return _DTYPE_LOOKUP.get(code)
+
+
+def observed_dtype(onnx_path: Path) -> str | None:
+    """Float dtype of the graph's weights, i.e. what the export actually produced."""
+    proto = onnx.load(str(onnx_path), load_external_data=False)
+    for init in proto.graph.initializer:
+        name = _dtype_name(init.data_type)
+        if name is not None:
+            return name
+    for inp in proto.graph.input:
+        name = _dtype_name(inp.type.tensor_type.elem_type)
+        if name is not None:
+            return name
+    return None
+
+
+def build_manifest(
+    onnx_path: Path, verdict: ExportVerdict, source_path: Path | None, package_version: str
+) -> dict:
+    return {
+        "downshift_version": package_version,
+        "created_utc": datetime.now(UTC).isoformat(timespec="seconds"),
+        "onnx_file": onnx_path.name,
+        "onnx_sha256": _sha256(onnx_path),
+        "source_model": str(source_path) if source_path else None,
+        "source_sha256": _sha256(source_path) if source_path and source_path.is_file() else None,
+        "versions": {
+            "torch": torch.__version__,
+            "onnx": onnx.__version__,
+            "onnxruntime": onnxruntime.__version__,
+        },
+        "opset": verdict.opset,
+        "observed_dtype": observed_dtype(onnx_path),
+        "execution_providers_available": onnxruntime.get_available_providers(),
+        "verdict": verdict.to_dict(),
+    }
+
+
+def manifest_path_for(onnx_path: Path) -> Path:
+    return onnx_path.with_suffix(".manifest.json")
+
+
+def write_manifest(
+    onnx_path: Path, verdict: ExportVerdict, source_path: Path | None, package_version: str
+) -> Path:
+    path = manifest_path_for(onnx_path)
+    path.write_text(
+        json.dumps(build_manifest(onnx_path, verdict, source_path, package_version), indent=2)
+    )
+    return path

downshift/export/prevalidated.py

@@ -0,0 +1,70 @@
+"""Intake for a .onnx file someone else produced (Olive, a notebook, whatever).
+
+No reference model -> UNVERIFIED. We serve it, we just say we never checked it.
+With --reference -> the normal verify path, exactly as for a fresh export.
+"""
+
+from pathlib import Path
+
+import onnx
+import torch
+
+from downshift.adapters.base import Adapter
+from downshift.export.verdict import BackendName, ExportVerdict, numerics_outcome, prepare_model
+from downshift.export.verify import verify
+
+
+def _graph_summary(onnx_path: Path) -> tuple[int | None, list[str], tuple[str, ...]]:
+    proto = onnx.load(str(onnx_path), load_external_data=False)
+    opset = next((imp.version for imp in proto.opset_import if imp.domain in ("", "ai.onnx")), None)
+    initializers = {init.name for init in proto.graph.initializer}
+    input_names = tuple(i.name for i in proto.graph.input if i.name not in initializers)
+    return opset, [n.op_type for n in proto.graph.node], input_names
+
+
+def intake(
+    onnx_path: str | Path,
+    reference: torch.nn.Module | None = None,
+    example_inputs: tuple | None = None,
+    adapter: Adapter | str | None = None,
+    k: int = 8,
+    dynamic: dict[str, list[int]] | None = None,
+) -> ExportVerdict:
+    onnx_path = Path(onnx_path)
+    opset, op_types, input_names = _graph_summary(onnx_path)
+
+    if reference is None:
+        return ExportVerdict(
+            status="UNVERIFIED",
+            model_family="onnx",
+            capture_strategy=None,
+            opset=opset,
+            op_types=op_types,
+            numerics=None,
+            recommended_backend=BackendName.onnxruntime,
+            reason="no reference model supplied; served as-is, numerics never checked",
+            input_names=input_names,
+            onnx_path=onnx_path,
+        )
+
+    prepared = prepare_model(reference, example_inputs, adapter, dynamic)
+    numerics = verify(
+        prepared.model, onnx_path, prepared.inputs, prepared.dynamic_shapes, prepared.vary_fn, k=k
+    )
+    status, backend, reason = numerics_outcome(
+        numerics, "pre-built ONNX matches reference", "pre-built ONNX diverges from reference"
+    )
+    return ExportVerdict(
+        status=status,
+        model_family=prepared.family,
+        capture_strategy=None,
+        opset=opset,
+        op_types=op_types,
+        numerics=numerics,
+        recommended_backend=backend,
+        reason=reason,
+        input_names=prepared.input_names,
+        dynamic_dims=prepared.dynamic_dims,
+        onnx_path=onnx_path,
+        prepared=prepared,
+    )

downshift/export/shapes.py

@@ -0,0 +1,61 @@
+"""Dynamic-shape inference and the `--dynamic` override.
+
+Default heuristic: axis 0 of every tensor input is dynamic and they all share one Dim
+(the batch case). Adapters override this where it's wrong, e.g. PyG's independent N/E.
+"""
+
+import torch
+
+
+def alternative_sizes(base_size: int) -> list[int]:
+    """Sizes to exercise a dynamic axis with, excluding the export-time size."""
+    return sorted({1, 2, 3, base_size + 1, base_size * 2} - {base_size})
+
+
+def infer_dynamic_shapes(inputs: tuple) -> tuple:
+    dim0 = torch.export.Dim("dim0", min=1, max=1 << 16)
+    return tuple({0: dim0} if isinstance(t, torch.Tensor) and t.ndim > 0 else None for t in inputs)
+
+
+def parse_dynamic_spec(spec: str) -> dict[str, list[int]]:
+    """Parse "x:0,edge_index:1" or "x:0:1" into {name: [axes]}."""
+    result: dict[str, list[int]] = {}
+    for item in filter(None, (s.strip() for s in spec.split(","))):
+        name, _, axes = item.partition(":")
+        if not name or not axes:
+            raise ValueError(f"bad --dynamic entry {item!r}; expected name:axis[:axis...]")
+        result.setdefault(name, []).extend(int(a) for a in axes.split(":"))
+    return result
+
+
+def apply_dynamic_override(
+    input_names: tuple[str, ...], inputs: tuple, override: dict[str, list[int]]
+) -> tuple:
+    """Build a dynamic_shapes tuple from an explicit {name: [axes]} spec. Every
+    (name, axis) pair gets its own independent Dim."""
+    unknown = set(override) - set(input_names)
+    if unknown:
+        raise ValueError(f"--dynamic names {sorted(unknown)} not in inputs {list(input_names)}")
+    shapes: list[dict[int, torch.export.Dim] | None] = []
+    for name, tensor in zip(input_names, inputs, strict=True):
+        axes = override.get(name)
+        if not axes or not isinstance(tensor, torch.Tensor):
+            shapes.append(None)
+            continue
+        shapes.append(
+            {axis: torch.export.Dim(f"{name}_{axis}", min=1, max=1 << 16) for axis in axes}
+        )
+    return tuple(shapes)
+
+
+def safe_capture_inputs(inputs: tuple, dynamic_shapes: tuple) -> tuple:
+    """torch.export specialises a size-1 dim to a constant even when it's marked dynamic.
+    Double any such axis for the trace only; verification still uses the real sizes."""
+    safe = []
+    for t, spec in zip(inputs, dynamic_shapes, strict=True):
+        if isinstance(t, torch.Tensor) and spec:
+            for axis in spec:
+                if t.shape[axis] == 1:
+                    t = torch.cat([t, t], dim=axis)
+        safe.append(t)
+    return tuple(safe)

downshift/export/verdict.py

@@ -0,0 +1,211 @@
+"""ExportVerdict: the one object everything else reads.
+
+CLEAN      exports, numerics match, survives shapes it wasn't traced on -> serve via ORT
+DEGRADED   exports but numerics drift past tolerance             -> serve via torch
+FAILED     won't export                                          -> serve via torch
+UNVERIFIED a .onnx handed to us with no reference model         -> serve via ORT, say so
+"""
+
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Literal
+
+import torch
+
+from downshift.adapters import registry
+from downshift.adapters.base import Adapter, Prepared
+from downshift.export.capture import capture
+from downshift.export.inputs import synthesize
+from downshift.export.shapes import apply_dynamic_override, safe_capture_inputs
+from downshift.export.verify import NumericsReport, verify
+
+Status = Literal["CLEAN", "DEGRADED", "FAILED", "UNVERIFIED"]
+
+
+class BackendName(str, Enum):
+    """The concrete backends a verdict can recommend/serve; never "auto" (that's a CLI-only
+    selection sentinel, not a real backend) - see engine.BackendChoice."""
+
+    onnxruntime = "onnxruntime"
+    torch = "torch"
+
+    def __str__(self) -> str:
+        # Python 3.11 made str(Enum)/format(Enum) print "BackendName.onnxruntime" instead of
+        # the plain value for any (str, Enum) mixin that isn't ReprEnum; banners embed this
+        # in f-strings, so pin it back to the value.
+        return self.value
+
+EXIT_CODES: dict[str, int] = {"CLEAN": 0, "FAILED": 1, "DEGRADED": 2, "UNVERIFIED": 3}
+
+_ATEN_OP = re.compile(r"(?:torch\.ops\.)?aten\.(\w+)(?:\.\w+)?")
+
+
+@dataclass
+class ExportVerdict:
+    status: Status
+    model_family: str
+    capture_strategy: str | None
+    opset: int | None
+    op_types: list[str]
+    numerics: NumericsReport | None
+    recommended_backend: BackendName
+    reason: str
+    input_names: tuple[str, ...] = ()
+    dynamic_dims: dict[str, list[int]] = field(default_factory=dict)
+    unsupported_ops: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    onnx_path: Path | None = None
+    onnx_program: object | None = field(default=None, repr=False)  # torch.onnx.ONNXProgram
+    prepared: Prepared | None = field(default=None, repr=False)
+
+    @property
+    def shape_generalization(self) -> bool | None:
+        return self.numerics.shape_generalization if self.numerics else None
+
+    @property
+    def exit_code(self) -> int:
+        return EXIT_CODES[self.status]
+
+    def to_dict(self) -> dict:
+        return {
+            "status": self.status,
+            "model_family": self.model_family,
+            "capture_strategy": self.capture_strategy,
+            "opset": self.opset,
+            "op_types": self.op_types,
+            "numerics": self.numerics.to_dict() if self.numerics else None,
+            "shape_generalization": self.shape_generalization,
+            "recommended_backend": self.recommended_backend,
+            "reason": self.reason,
+            "input_names": list(self.input_names),
+            "dynamic_dims": self.dynamic_dims,
+            "unsupported_ops": self.unsupported_ops,
+            "warnings": self.warnings,
+            "onnx_path": str(self.onnx_path) if self.onnx_path else None,
+        }
+
+
+def numerics_outcome(
+    numerics: NumericsReport, passed_prefix: str, failed_prefix: str
+) -> tuple[Status, BackendName, str]:
+    """Numerics decide the verdict: pass -> CLEAN via ORT, fail -> DEGRADED via torch.
+
+    The prefixes open the reason string; the sample counts and error are appended.
+    """
+    err = f"(max abs err {numerics.max_abs_err:.2e})"
+    if numerics.passed:
+        reason = f"{passed_prefix} across {numerics.samples_tested} samples {err}"
+        return "CLEAN", BackendName.onnxruntime, reason
+    reason = f"{failed_prefix} on {numerics.failures}/{numerics.samples_tested} samples {err}"
+    return "DEGRADED", BackendName.torch, reason
+
+
+def _tied_weight_warnings(model: torch.nn.Module) -> list[str]:
+    seen: dict[int, str] = {}
+    tied: list[str] = []
+    for name, param in model.named_parameters(remove_duplicate=False):
+        first = seen.setdefault(id(param), name)
+        if first != name:
+            tied.append(f"tied weights: {name} shares storage with {first}")
+    return tied
+
+
+def prepare_model(
+    model: torch.nn.Module,
+    example_inputs: tuple | None = None,
+    adapter: Adapter | str | None = None,
+    dynamic: dict[str, list[int]] | None = None,
+) -> Prepared:
+    """Pick an adapter, synthesise inputs if needed, and flatten into export form."""
+    if isinstance(adapter, str):
+        adapter = registry.get(adapter)
+    if adapter is None:
+        adapter = registry.detect(model, example_inputs)
+    example_inputs = synthesize(model, adapter, example_inputs)
+    prepared = adapter.prepare(model, example_inputs)
+    if dynamic:
+        prepared.dynamic_shapes = apply_dynamic_override(
+            prepared.input_names, prepared.inputs, dynamic
+        )
+    return prepared
+
+
+def build_verdict(prepared: Prepared, k: int = 8, verify_numerics: bool = True) -> ExportVerdict:
+    """Capture, then verify. verify_numerics=False is the --no-verify escape hatch: the
+    graph is still produced but the verdict is UNVERIFIED, never CLEAN."""
+    warnings = _tied_weight_warnings(prepared.model)
+    if prepared.model.training:
+        warnings.append("model was in training mode; switched to eval() for export")
+        prepared.model.eval()
+
+    result = capture(
+        prepared.model,
+        safe_capture_inputs(prepared.inputs, prepared.dynamic_shapes),
+        prepared.dynamic_shapes,
+    )
+    verdict = ExportVerdict(
+        status="FAILED",
+        model_family=prepared.family,
+        capture_strategy=result.capture_strategy,
+        opset=result.opset,
+        op_types=result.op_types,
+        numerics=None,
+        recommended_backend=BackendName.torch,
+        reason="",
+        input_names=prepared.input_names,
+        dynamic_dims=prepared.dynamic_dims,
+        warnings=warnings,
+        onnx_program=result.onnx_program,
+        prepared=prepared,
+    )
+
+    if not result.success:
+        exc = result.exception
+        message = f"{type(exc).__name__}: {exc}" if exc is not None else "export failed"
+        verdict.reason = message.splitlines()[0]
+        verdict.unsupported_ops = sorted(set(_ATEN_OP.findall(message)))
+        return verdict
+
+    if not verify_numerics:
+        verdict.status, verdict.recommended_backend = "UNVERIFIED", BackendName.onnxruntime
+        verdict.reason = f"exported via {result.capture_strategy}; numerics never checked"
+        return verdict
+
+    numerics = verify(
+        prepared.model,
+        result.onnx_program,
+        prepared.inputs,
+        prepared.dynamic_shapes,
+        vary_fn=prepared.vary_fn,
+        k=k,
+    )
+    verdict.numerics = numerics
+    verdict.status, verdict.recommended_backend, verdict.reason = numerics_outcome(
+        numerics,
+        f"exported via {result.capture_strategy}; numerics ok",
+        f"exported via {result.capture_strategy} but numerics diverge",
+    )
+    return verdict
+
+
+def check(
+    model: torch.nn.Module,
+    example_inputs: tuple | None = None,
+    k: int = 8,
+    adapter: Adapter | str | None = None,
+    dynamic: dict[str, list[int]] | None = None,
+    fp16: bool = False,
+    verify_numerics: bool = True,
+) -> ExportVerdict:
+    """Export in memory, verify, and return the verdict. Writes nothing to disk."""
+    if fp16:
+        model = model.half()
+        if example_inputs is not None:
+            example_inputs = tuple(
+                t.half() if isinstance(t, torch.Tensor) and t.is_floating_point() else t
+                for t in example_inputs
+            )
+    prepared = prepare_model(model, example_inputs, adapter, dynamic)
+    return build_verdict(prepared, k=k, verify_numerics=verify_numerics)

downshift/export/verify.py

@@ -0,0 +1,162 @@
+"""Numerical verification. Mandatory: an export isn't a success until this passes.
+
+Runs K samples through the torch model and the ONNX graph, varying dynamic dims so at
+least some samples have shapes the exporter never saw. That's what catches a graph that
+traced fine but froze a shape or specialised a data-dependent branch.
+"""
+
+import random
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import torch
+
+from downshift.adapters.base import VaryFn
+from downshift.export.shapes import alternative_sizes
+from downshift.settings import TOLERANCES
+
+
+@dataclass
+class NumericsReport:
+    samples_tested: int
+    max_abs_err: float
+    max_rel_err: float
+    failures: int
+    shape_generalization: bool  # did every non-baseline-shape sample also pass?
+    tolerance_abs: float
+    tolerance_rel: float
+
+    @property
+    def passed(self) -> bool:
+        return self.failures == 0
+
+    def to_dict(self) -> dict:
+        return asdict(self) | {"passed": self.passed}
+
+
+def default_tolerances(model: torch.nn.Module) -> tuple[float, float]:
+    dtypes = {p.dtype for p in model.parameters() if p.is_floating_point()}
+    for dtype in (torch.bfloat16, torch.float16):
+        if dtype in dtypes:
+            return TOLERANCES[dtype]
+    return TOLERANCES[torch.float32]
+
+
+def _resize_dim0(tensor: torch.Tensor, new_size: int) -> torch.Tensor:
+    if tensor.ndim == 0 or tensor.shape[0] == new_size:
+        return tensor
+    shape = list(tensor.shape)
+    shape[0] = new_size
+    if tensor.is_floating_point():
+        return torch.randn(*shape, dtype=tensor.dtype)
+    # Integer inputs are usually indices; stay inside the observed range.
+    lo = int(tensor.min().item())
+    hi = max(int(tensor.max().item()) + 1, lo + 1)
+    return torch.randint(lo, hi, shape, dtype=tensor.dtype)
+
+
+def make_shared_axis0_vary_fn(base_inputs: tuple, dynamic_shapes: tuple, seed: int = 0) -> VaryFn:
+    """Default sampler: every dynamic tensor shares one axis-0 size (the batch case)."""
+    base_size = next(
+        (t.shape[0] for t, spec in zip(base_inputs, dynamic_shapes, strict=True) if spec),
+        None,
+    )
+    rng = random.Random(seed)
+    candidates = alternative_sizes(base_size) if base_size is not None else []
+
+    def vary(i: int) -> tuple:
+        if i == 0 or not candidates:
+            return base_inputs
+        size = rng.choice(candidates)
+        return tuple(
+            _resize_dim0(t, size) if isinstance(t, torch.Tensor) and spec else t
+            for t, spec in zip(base_inputs, dynamic_shapes, strict=True)
+        )
+
+    return vary
+
+
+def _to_session(onnx_model) -> ort.InferenceSession:
+    if isinstance(onnx_model, (str, Path)):
+        source: str | bytes = str(onnx_model)
+    elif isinstance(onnx_model, bytes):
+        source = onnx_model
+    else:  # torch.onnx.ONNXProgram
+        source = onnx_model.model_proto.SerializeToString()
+    return ort.InferenceSession(source, providers=["CPUExecutionProvider"])
+
+
+def _as_tensor_list(output) -> list[torch.Tensor]:
+    if isinstance(output, torch.Tensor):
+        return [output]
+    if isinstance(output, (tuple, list)):
+        return [t for t in output if isinstance(t, torch.Tensor)]
+    raise TypeError(f"can't compare model output of type {type(output).__name__}")
+
+
+def verify(
+    model: torch.nn.Module,
+    onnx_model,
+    base_inputs: tuple,
+    dynamic_shapes: tuple | None = None,
+    vary_fn: VaryFn | None = None,
+    k: int = 8,
+    atol: float | None = None,
+    rtol: float | None = None,
+    seed: int = 0,
+) -> NumericsReport:
+    if vary_fn is None:
+        if dynamic_shapes is None:
+            raise ValueError("verify() needs either dynamic_shapes or an explicit vary_fn")
+        vary_fn = make_shared_axis0_vary_fn(base_inputs, dynamic_shapes, seed=seed)
+
+    default_atol, default_rtol = default_tolerances(model)
+    atol = default_atol if atol is None else atol
+    rtol = default_rtol if rtol is None else rtol
+
+    session = _to_session(onnx_model)
+    input_names = [inp.name for inp in session.get_inputs()]
+
+    model.eval()
+    max_abs_err = 0.0
+    max_rel_err = 0.0
+    failures = 0
+    non_baseline_failures = 0
+
+    # Seed inside a forked RNG so callers' global random state is untouched afterwards.
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(seed)
+        for i in range(k):
+            sample = vary_fn(i)
+            with torch.inference_mode():
+                torch_outs = _as_tensor_list(model(*sample))
+            feeds = {name: t.numpy() for name, t in zip(input_names, sample, strict=True)}
+            ort_outs = session.run(None, feeds)
+
+            sample_abs = 0.0
+            sample_rel = 0.0
+            for expected, got in zip(torch_outs, ort_outs, strict=True):
+                expected_np = expected.detach().numpy().astype(np.float64)
+                abs_err = np.abs(expected_np - np.asarray(got).astype(np.float64))
+                rel_err = abs_err / (np.abs(expected_np) + 1e-8)
+                sample_abs = max(sample_abs, float(abs_err.max(initial=0.0)))
+                sample_rel = max(sample_rel, float(rel_err.max(initial=0.0)))
+
+            max_abs_err = max(max_abs_err, sample_abs)
+            max_rel_err = max(max_rel_err, sample_rel)
+            if sample_abs > atol and sample_rel > rtol:
+                failures += 1
+                if i > 0:
+                    non_baseline_failures += 1
+
+    return NumericsReport(
+        samples_tested=k,
+        max_abs_err=max_abs_err,
+        max_rel_err=max_rel_err,
+        failures=failures,
+        shape_generalization=non_baseline_failures == 0,
+        tolerance_abs=atol,
+        tolerance_rel=rtol,
+    )