downshift-server 0.2.0__py3-none-any.whl

downshift/__init__.py

@@ -0,0 +1,60 @@
+"""downshift: check whether a PyTorch model survives ONNX export, then serve it."""
+
+from pathlib import Path
+
+from downshift.adapters.base import Adapter, Prepared
+from downshift.export.manifest import write_manifest
+from downshift.export.prevalidated import intake
+from downshift.export.verdict import ExportVerdict, build_verdict, check, prepare_model
+from downshift.export.verify import NumericsReport
+
+__version__ = "0.2.0"
+
+
+def export(
+    model,
+    output: str | Path,
+    example_inputs: tuple | None = None,
+    k: int = 8,
+    adapter: Adapter | str | None = None,
+    dynamic: dict[str, list[int]] | None = None,
+    fp16: bool = False,
+    source_path: Path | None = None,
+    verify_numerics: bool = True,
+) -> ExportVerdict:
+    """check() plus writing the .onnx and its manifest. `output` is the .onnx path.
+
+    A FAILED verdict writes nothing; a DEGRADED one still writes the artifact because
+    the manifest records exactly how far off it is.
+    """
+    verdict = check(
+        model,
+        example_inputs,
+        k=k,
+        adapter=adapter,
+        dynamic=dynamic,
+        fp16=fp16,
+        verify_numerics=verify_numerics,
+    )
+    if verdict.onnx_program is None:
+        return verdict
+    output = Path(output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    verdict.onnx_program.save(str(output))  # type: ignore[attr-defined]
+    verdict.onnx_path = output
+    write_manifest(output, verdict, source_path, __version__)
+    return verdict
+
+
+__all__ = [
+    "Adapter",
+    "ExportVerdict",
+    "NumericsReport",
+    "Prepared",
+    "build_verdict",
+    "check",
+    "export",
+    "intake",
+    "prepare_model",
+    "__version__",
+]

downshift/adapters/_flatten.py

@@ -0,0 +1,40 @@
+"""Shim that wraps a model taking a container argument (a dataclass, a PyG Data, ...) so
+torch.export sees a plain fixed-arity tensor signature. Export the shim, not the model.
+"""
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class FlattenShimBase(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        rebuild: Callable[[Mapping[str, torch.Tensor]], Any],
+        field_names: Sequence[str],
+    ) -> None:
+        super().__init__()
+        self.model = model
+        self._rebuild = rebuild
+        self._field_names = field_names
+        self.train(model.training)
+
+    def _call(self, tensors: tuple[torch.Tensor, ...]) -> Any:
+        return self.model(self._rebuild(dict(zip(self._field_names, tensors, strict=True))))
+
+
+def build_shim_class(field_names: Sequence[str]) -> type[FlattenShimBase]:
+    """Generate a subclass whose forward() has one named positional parameter per field.
+
+    A `def forward(self, *tensors)` would bind everything into one VAR_POSITIONAL arg and
+    torch.export would see a single tuple input, which doesn't line up with a per-input
+    dynamic_shapes tuple. The parameter names also become the ONNX graph's input names.
+    """
+    params = ", ".join(n if n.isidentifier() else f"t{i}" for i, n in enumerate(field_names))
+    src = f"def forward(self, {params}):\n    return self._call(({params},))\n"
+    namespace: dict[str, Any] = {}
+    exec(src, namespace)  # noqa: S102 - field names only; no user-controlled text
+    return type("FlattenShim", (FlattenShimBase,), {"forward": namespace["forward"]})

downshift/adapters/base.py

@@ -0,0 +1,47 @@
+"""What an adapter is.
+
+An adapter knows one model family well enough to (a) build example inputs when the user
+didn't give any, and (b) turn the model + inputs into something torch.export can trace:
+a module with a flat, fixed-arity tensor signature, plus the dynamic-shape spec and a way
+to generate more samples for verification.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from torch import nn
+
+VaryFn = Callable[[int], tuple]
+
+
+@dataclass
+class Prepared:
+    model: nn.Module  # export-ready module; forward takes flat tensors
+    inputs: tuple  # flat example inputs, one per input_names entry
+    input_names: tuple[str, ...]
+    dynamic_shapes: tuple  # per input: {axis: torch.export.Dim} or None
+    vary_fn: VaryFn | None  # sample i -> inputs; None means use the shared-axis-0 default
+    family: str
+
+    @property
+    def dynamic_dims(self) -> dict[str, list[int]]:
+        return {
+            name: sorted(spec)
+            for name, spec in zip(self.input_names, self.dynamic_shapes, strict=True)
+            if spec
+        }
+
+
+@runtime_checkable
+class Adapter(Protocol):
+    name: str
+    family: str
+
+    def matches(self, model: nn.Module, example_inputs: tuple | None) -> bool: ...
+
+    def example_inputs(self, model: nn.Module) -> tuple | None:
+        """Adapter-derived example inputs, used when the user gave none. None means no guess."""
+        ...
+
+    def prepare(self, model: nn.Module, example_inputs: tuple) -> Prepared: ...

downshift/adapters/generic.py

@@ -0,0 +1,99 @@
+"""Generic adapter: any nn.Module.
+
+Handles two things. A single dataclass positional argument gets flattened into plain
+tensors (torch.export rejects unregistered dataclasses outright). And when no example
+inputs are given, it guesses a shape from the first Linear/Conv layer, which is enough
+for the torchvision-style single-tensor case.
+"""
+
+import dataclasses
+import inspect
+
+import torch
+from torch import nn
+
+from downshift.adapters._flatten import build_shim_class
+from downshift.adapters.base import Prepared
+from downshift.export.shapes import infer_dynamic_shapes
+
+_GUESS_SPATIAL = 32
+
+
+def _forward_param_names(model: nn.Module) -> tuple[str, ...]:
+    params = inspect.signature(model.forward).parameters.values()
+    return tuple(
+        p.name
+        for p in params
+        if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) and p.name != "self"
+    )
+
+
+def _guess_single_tensor_input(model: nn.Module) -> torch.Tensor | None:
+    for layer in model.modules():
+        if isinstance(layer, nn.Linear):
+            return torch.randn(1, layer.in_features)
+        if isinstance(layer, nn.Conv1d):
+            return torch.randn(1, layer.in_channels, _GUESS_SPATIAL)
+        if isinstance(layer, nn.Conv2d):
+            return torch.randn(1, layer.in_channels, _GUESS_SPATIAL, _GUESS_SPATIAL)
+        if isinstance(layer, nn.Conv3d):
+            return torch.randn(1, layer.in_channels, 8, 8, 8)
+        if isinstance(layer, nn.Embedding):
+            return torch.randint(0, layer.num_embeddings, (1, 8))
+    return None
+
+
+class GenericAdapter:
+    name = "generic"
+    family = "generic-torch"
+
+    def matches(self, model: nn.Module, example_inputs: tuple | None) -> bool:
+        return True
+
+    def example_inputs(self, model: nn.Module) -> tuple | None:
+        if len(_forward_param_names(model)) != 1:
+            return None
+        guess = _guess_single_tensor_input(model)
+        return (guess,) if guess is not None else None
+
+    def prepare(self, model: nn.Module, example_inputs: tuple) -> Prepared:
+        flattened = _flatten_dataclass(model, example_inputs)
+        if flattened is not None:
+            model, inputs, names = flattened
+        else:
+            inputs = example_inputs
+            param_names = _forward_param_names(model)
+            names = tuple(
+                param_names[i] if i < len(param_names) else f"input_{i}"
+                for i in range(len(inputs))
+            )
+        return Prepared(
+            model=model,
+            inputs=inputs,
+            input_names=names,
+            dynamic_shapes=infer_dynamic_shapes(inputs),
+            vary_fn=None,
+            family=self.family,
+        )
+
+
+def _flatten_dataclass(
+    model: nn.Module, example_inputs: tuple
+) -> tuple[nn.Module, tuple, tuple[str, ...]] | None:
+    if len(example_inputs) != 1:
+        return None
+    (arg,) = example_inputs
+    if not dataclasses.is_dataclass(arg) or isinstance(arg, type):
+        return None
+
+    names = tuple(f.name for f in dataclasses.fields(arg))
+    tensors = tuple(getattr(arg, n) for n in names)
+    if not all(isinstance(t, torch.Tensor) for t in tensors):
+        return None  # can't flatten a non-tensor field; let export produce the real error
+
+    dataclass_type = type(arg)
+    shim = build_shim_class(names)(model, lambda fields: dataclass_type(**fields), names)
+    return shim, tensors, names
+
+
+ADAPTER = GenericAdapter()

downshift/adapters/hf.py

@@ -0,0 +1,95 @@
+"""Hugging Face adapter, encoder-only models.
+
+Builds input_ids / attention_mask straight from the model config rather than pulling in
+optimum. The export shim unwraps the ModelOutput so torch.export sees a plain tensor
+(last_hidden_state for base models, logits for heads).
+
+Only imported when transformers is installed.
+"""
+
+import random
+
+import torch
+from torch import nn
+from transformers import AutoModel, PreTrainedModel
+
+from downshift.adapters.base import Prepared, VaryFn
+from downshift.export.shapes import alternative_sizes
+
+INPUT_NAMES = ("input_ids", "attention_mask")
+_GUESS_BATCH = 2
+_GUESS_SEQ = 8
+
+
+class _FirstOutputShim(nn.Module):
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        self.model = model
+        self.train(model.training)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        out = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        first: torch.Tensor = out[0]
+        return first
+
+
+def load_pretrained(repo_id_or_path: str) -> PreTrainedModel:
+    """repo_id_or_path is either a Hugging Face hub id or a local directory containing a
+    previously downloaded repo (config.json, weights, etc.) -- from_pretrained handles both."""
+    model: PreTrainedModel = AutoModel.from_pretrained(repo_id_or_path)
+    return model
+
+
+class HFAdapter:
+    name = "hf"
+    family = "hf-transformers"
+
+    def matches(self, model: nn.Module, example_inputs: tuple | None) -> bool:
+        return isinstance(model, PreTrainedModel)
+
+    def example_inputs(self, model: nn.Module) -> tuple | None:
+        vocab = getattr(getattr(model, "config", None), "vocab_size", None)
+        if vocab is None:
+            return None
+        input_ids = torch.randint(0, vocab, (_GUESS_BATCH, _GUESS_SEQ))
+        return input_ids, torch.ones_like(input_ids)
+
+    def prepare(self, model: nn.Module, example_inputs: tuple) -> Prepared:
+        input_ids, attention_mask = example_inputs
+        config = getattr(model, "config")
+        # Position embeddings cap the sequence length; a looser bound trips export's guards.
+        max_seq = int(getattr(config, "max_position_embeddings", 1 << 12))
+        batch = torch.export.Dim("batch", min=1, max=1 << 12)
+        seq = torch.export.Dim("seq", min=1, max=max_seq)
+        spec = {0: batch, 1: seq}
+        vocab = int(config.vocab_size)
+        inputs = (input_ids, attention_mask)
+        return Prepared(
+            model=_FirstOutputShim(model),
+            inputs=inputs,
+            input_names=INPUT_NAMES,
+            dynamic_shapes=(spec, spec),
+            vary_fn=make_vary_fn(inputs, vocab),
+            family=self.family,
+        )
+
+
+def make_vary_fn(base_inputs: tuple, vocab_size: int, seed: int = 0) -> VaryFn:
+    input_ids, _ = base_inputs
+    base_batch, base_seq = input_ids.shape
+    rng = random.Random(seed)
+    batch_candidates = alternative_sizes(base_batch)
+    seq_candidates = alternative_sizes(base_seq)
+
+    def vary(i: int) -> tuple:
+        if i == 0:
+            return base_inputs
+        b = rng.choice(batch_candidates) if batch_candidates else base_batch
+        s = rng.choice(seq_candidates) if seq_candidates else base_seq
+        ids = torch.randint(0, vocab_size, (b, s), dtype=input_ids.dtype)
+        return ids, torch.ones_like(ids)
+
+    return vary
+
+
+ADAPTER = HFAdapter()

downshift/adapters/pyg.py

@@ -0,0 +1,120 @@
+"""PyG adapter: flattens torch_geometric.data.Data into (x, edge_index[, edge_attr]).
+
+Node count N and edge count E are independent dynamic dims. Tying them to one Dim is the
+classic way to get a GNN export that works on the example graph and throws
+INVALID_ARGUMENT on the next one.
+
+Only imported when a PyG Data input actually shows up, so torch_geometric stays optional.
+"""
+
+import random
+
+import torch
+from torch import nn
+from torch_geometric.data import Data
+from torch_geometric.nn import MessagePassing
+
+from downshift.adapters._flatten import build_shim_class
+from downshift.adapters.base import Prepared, VaryFn
+from downshift.export.shapes import alternative_sizes
+
+BASE_FIELD_NAMES = ("x", "edge_index")  # edge_attr appended when present on the input Data
+_GUESS_NODES = 8
+_GUESS_EDGES = 16
+
+
+def is_pyg_data(example_inputs: tuple | None) -> bool:
+    if example_inputs is None or len(example_inputs) != 1:
+        return False
+    return isinstance(example_inputs[0], Data)
+
+
+def _first_in_channels(model: nn.Module) -> int | None:
+    for layer in model.modules():
+        if isinstance(layer, MessagePassing):
+            in_channels = getattr(layer, "in_channels", None)
+            if isinstance(in_channels, tuple):
+                in_channels = in_channels[0]
+            if isinstance(in_channels, int):
+                return in_channels
+    return None
+
+
+class PyGAdapter:
+    name = "pyg"
+    family = "pyg"
+
+    def matches(self, model: nn.Module, example_inputs: tuple | None) -> bool:
+        if is_pyg_data(example_inputs):
+            return True
+        if example_inputs is not None:
+            return False
+        return any(isinstance(m, MessagePassing) for m in model.modules())
+
+    def example_inputs(self, model: nn.Module) -> tuple | None:
+        in_channels = _first_in_channels(model)
+        if in_channels is None:
+            return None
+        x = torch.randn(_GUESS_NODES, in_channels)
+        edge_index = torch.randint(0, _GUESS_NODES, (2, _GUESS_EDGES))
+        return (Data(x=x, edge_index=edge_index),)
+
+    def prepare(self, model: nn.Module, example_inputs: tuple) -> Prepared:
+        (data,) = example_inputs
+        names: tuple[str, ...] = BASE_FIELD_NAMES
+        if getattr(data, "edge_attr", None) is not None:
+            names += ("edge_attr",)
+
+        inputs = tuple(getattr(data, n) for n in names)
+        shim = build_shim_class(names)(model, lambda fields: Data(**fields), names)
+
+        n_dim = torch.export.Dim("num_nodes", min=1, max=1 << 16)
+        e_dim = torch.export.Dim("num_edges", min=1, max=1 << 16)
+        axis_by_field = {"x": {0: n_dim}, "edge_index": {1: e_dim}, "edge_attr": {0: e_dim}}
+
+        return Prepared(
+            model=shim,
+            inputs=inputs,
+            input_names=names,
+            dynamic_shapes=tuple(axis_by_field[n] for n in names),
+            vary_fn=make_vary_fn(inputs, names),
+            family=self.family,
+        )
+
+
+def make_vary_fn(base_inputs: tuple, field_names: tuple[str, ...], seed: int = 0) -> VaryFn:
+    """Regenerate (x, edge_index[, edge_attr]) with independently varied N and E.
+
+    edge_index is redrawn against the sample's own node count, not the original tensor's
+    value range, so a shrunken graph never references nodes it doesn't have.
+    """
+    x_idx = field_names.index("x")
+    ei_idx = field_names.index("edge_index")
+    ea_idx = field_names.index("edge_attr") if "edge_attr" in field_names else None
+
+    base_x, base_ei = base_inputs[x_idx], base_inputs[ei_idx]
+    base_n, in_channels = base_x.shape
+    base_e = base_ei.shape[1]
+
+    rng = random.Random(seed)
+    n_candidates = alternative_sizes(base_n)
+    e_candidates = alternative_sizes(base_e)
+
+    def vary(i: int) -> tuple:
+        if i == 0:
+            return base_inputs
+        n = rng.choice(n_candidates) if n_candidates else base_n
+        e = rng.choice(e_candidates) if e_candidates else base_e
+
+        sample: list = [None] * len(field_names)
+        sample[x_idx] = torch.randn(n, in_channels, dtype=base_x.dtype)
+        sample[ei_idx] = torch.randint(0, n, (2, e), dtype=base_ei.dtype)
+        if ea_idx is not None:
+            base_ea = base_inputs[ea_idx]
+            sample[ea_idx] = torch.randn(e, base_ea.shape[1], dtype=base_ea.dtype)
+        return tuple(sample)
+
+    return vary
+
+
+ADAPTER = PyGAdapter()

downshift/adapters/registry.py

@@ -0,0 +1,116 @@
+"""Adapter lookup: built-ins, anything registered under the `downshift.adapters` entry-point
+group, and one-off adapters loaded straight from a user's .py file. Adapters whose optional
+dependency is missing are skipped silently.
+"""
+
+import importlib.util
+from importlib import import_module
+from importlib.metadata import entry_points
+from pathlib import Path
+
+from torch import nn
+
+from downshift.adapters.base import Adapter
+from downshift.loading import LoadError
+
+ENTRY_POINT_GROUP = "downshift.adapters"
+
+# Most specific first; generic last so it only wins when nothing else matches.
+_BUILTIN_SPECS = (
+    "downshift.adapters.hf:ADAPTER",
+    "downshift.adapters.pyg:ADAPTER",
+    "downshift.adapters.generic:ADAPTER",
+)
+
+
+def _load_spec(spec: str) -> Adapter | None:
+    module_name, _, attr = spec.partition(":")
+    try:
+        adapter: Adapter = getattr(import_module(module_name), attr)
+    except ImportError:
+        return None
+    return adapter
+
+
+def _split_file_spec(name: str) -> tuple[str, str] | None:
+    """Split path/to/adapter.py[:attr] into (path, attr); None if `name` isn't a .py spec.
+
+    Splits on the literal ".py:" rather than the last ":", so a Windows drive letter's
+    colon (`C:\\...`) is never mistaken for the path:attr separator.
+    """
+    path, sep, attr = name.partition(".py:")
+    if sep:
+        return path + ".py", attr
+    if name.endswith(".py"):
+        return name, "ADAPTER"
+    return None
+
+
+def load_from_file(path_str: str, attr: str = "ADAPTER") -> Adapter:
+    """Load a user's adapter from a standalone .py file, outside any installed package.
+
+    `attr` names either an `Adapter`-shaped instance — the `generic`/`pyg`/`hf` convention of
+    a module-level `ADAPTER = MyAdapter()` — or the class itself, instantiated with no args.
+    """
+    path = Path(path_str)
+    if not path.is_file():
+        raise LoadError(f"{path} does not exist")
+    spec = importlib.util.spec_from_file_location(f"downshift._custom_adapter_{path.stem}", path)
+    if spec is None or spec.loader is None:
+        raise LoadError(f"can't import {path} as a Python module")
+    module = importlib.util.module_from_spec(spec)
+    try:
+        spec.loader.exec_module(module)
+    except Exception as exc:
+        raise LoadError(f"error while loading {path}: {exc}") from exc
+
+    try:
+        obj = getattr(module, attr)
+    except AttributeError as exc:
+        raise LoadError(f"{path} has no attribute {attr!r}") from exc
+
+    if isinstance(obj, type):
+        obj = obj()
+    if not isinstance(obj, Adapter):
+        raise LoadError(
+            f"{path}:{attr} is a {type(obj).__name__}, not an Adapter — it needs `name`, "
+            "`family`, matches(), example_inputs(), and prepare(); see GenericAdapter for "
+            "the shape to implement."
+        )
+    return obj
+
+
+def available() -> dict[str, Adapter]:
+    adapters: dict[str, Adapter] = {}
+    for ep in entry_points(group=ENTRY_POINT_GROUP):
+        try:
+            adapter = ep.load()
+        except ImportError:
+            continue
+        adapters[adapter.name] = adapter
+    for spec in _BUILTIN_SPECS:
+        adapter = _load_spec(spec)
+        if adapter is not None:
+            adapters.setdefault(adapter.name, adapter)
+    # Generic must be tried last regardless of registration order.
+    generic = adapters.pop("generic", None)
+    if generic is not None:
+        adapters["generic"] = generic
+    return adapters
+
+
+def get(name: str) -> Adapter:
+    file_spec = _split_file_spec(name)
+    if file_spec is not None:
+        return load_from_file(*file_spec)
+    adapters = available()
+    if name not in adapters:
+        raise KeyError(f"unknown adapter {name!r}; available: {', '.join(adapters)}")
+    return adapters[name]
+
+
+def detect(model: nn.Module, example_inputs: tuple | None) -> Adapter:
+    for adapter in available().values():
+        if adapter.matches(model, example_inputs):
+            return adapter
+    raise RuntimeError("no adapter matched and the generic adapter is unavailable")