downshift-server 0.2.0__py3-none-any.whl

downshift/loading.py

@@ -0,0 +1,166 @@
+"""Turn a CLI model argument into something the export layer can use.
+
+Accepted forms:
+  model.onnx                pre-built ONNX, served as-is (UNVERIFIED without --reference)
+  pkg.module:attr           import spec; attr is a module instance or a zero-arg factory.
+                            A sibling `make_inputs` in the same module is picked up automatically.
+  weights.pt                state dict; needs --model-class pkg.module:Class. Loaded with
+                            weights_only=True. A pickled full module needs --unsafe-load.
+  org/repo                  Hugging Face hub id (needs the [hf] extra)
+  path/to/repo/dir          Locally downloaded Hugging Face repo, i.e. a directory containing
+                            a config.json (needs the [hf] extra)
+"""
+
+import re
+from dataclasses import dataclass
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch import nn
+
+_IMPORT_SPEC = re.compile(r"^[A-Za-z_][\w.]*:[A-Za-z_]\w*$")
+_STATE_DICT_SUFFIXES = {".pt", ".pth", ".bin", ".ckpt"}
+
+
+class LoadError(ValueError):
+    pass
+
+
+@dataclass
+class LoadedModel:
+    source: str
+    model: nn.Module | None = None
+    onnx_path: Path | None = None
+    example_inputs: tuple | None = None
+    adapter_hint: str | None = None
+
+    @property
+    def source_path(self) -> Path | None:
+        path = Path(self.source)
+        return path if path.exists() else None
+
+
+def import_object(spec: str) -> Any:
+    if not _IMPORT_SPEC.match(spec):
+        raise LoadError(f"{spec!r} is not an import spec of the form package.module:attr")
+    module_name, _, attr = spec.partition(":")
+    try:
+        module = import_module(module_name)
+    except ImportError as exc:
+        raise LoadError(f"can't import {module_name!r}: {exc}") from exc
+    try:
+        return getattr(module, attr)
+    except AttributeError as exc:
+        raise LoadError(f"{module_name!r} has no attribute {attr!r}") from exc
+
+
+def _instantiate(obj: Any) -> nn.Module:
+    if isinstance(obj, nn.Module):
+        return obj
+    if callable(obj):
+        model = obj()
+        if isinstance(model, nn.Module):
+            return model
+        raise LoadError(f"{obj!r}() returned {type(model).__name__}, not an nn.Module")
+    raise LoadError(f"{obj!r} is neither an nn.Module nor a callable that builds one")
+
+
+def load_inputs(spec: str) -> tuple:
+    inputs = import_object(spec)
+    if callable(inputs):
+        inputs = inputs()
+    return inputs if isinstance(inputs, tuple) else (inputs,)
+
+
+def _load_from_import_spec(spec: str, inputs_spec: str | None) -> LoadedModel:
+    model = _instantiate(import_object(spec))
+    inputs: tuple | None = None
+    if inputs_spec:
+        inputs = load_inputs(inputs_spec)
+    else:
+        module_name = spec.partition(":")[0]
+        if callable(getattr(import_module(module_name), "make_inputs", None)):
+            inputs = load_inputs(f"{module_name}:make_inputs")
+    return LoadedModel(source=spec, model=model, example_inputs=inputs)
+
+
+def _load_checkpoint(path: Path, model_class: str | None, unsafe_load: bool) -> nn.Module:
+    try:
+        payload = torch.load(path, map_location="cpu", weights_only=not unsafe_load)
+    except Exception as exc:  # torch raises a few different types here
+        if unsafe_load:
+            raise LoadError(f"failed to load {path}: {exc}") from exc
+        raise LoadError(
+            f"{path} isn't loadable with weights_only=True ({type(exc).__name__}). If it "
+            "holds a pickled nn.Module from a source you trust, re-run with --unsafe-load."
+        ) from exc
+
+    if isinstance(payload, nn.Module):
+        return payload
+    if isinstance(payload, dict):
+        state = payload.get("state_dict", payload)
+        if model_class is None:
+            raise LoadError(
+                f"{path} is a state dict; pass --model-class package.module:Class so it can "
+                "be instantiated and the weights loaded into it."
+            )
+        model = _instantiate(import_object(model_class))
+        model.load_state_dict(state)
+        return model
+    raise LoadError(f"{path} contained {type(payload).__name__}, expected a state dict")
+
+
+def load_model(
+    spec: str,
+    inputs: str | None = None,
+    model_class: str | None = None,
+    unsafe_load: bool = False,
+) -> LoadedModel:
+    path = Path(spec)
+    if path.suffix == ".onnx":
+        if not path.exists():
+            raise LoadError(f"{path} does not exist")
+        return LoadedModel(source=spec, onnx_path=path)
+
+    if path.exists() and path.suffix in _STATE_DICT_SUFFIXES:
+        model = _load_checkpoint(path, model_class, unsafe_load)
+        return LoadedModel(
+            source=spec, model=model, example_inputs=load_inputs(inputs) if inputs else None
+        )
+
+    if _IMPORT_SPEC.match(spec):
+        return _load_from_import_spec(spec, inputs)
+
+    if path.is_dir():
+        if not (path / "config.json").exists():
+            raise LoadError(f"{path} has no config.json — not a Hugging Face repo")
+        return _load_hf(spec, inputs)
+
+    if path.exists():
+        raise LoadError(f"don't know how to load {path} (suffix {path.suffix!r})")
+    if path.suffix or path.is_absolute():
+        raise LoadError(f"{path} does not exist")
+
+    return _load_hf(spec, inputs)
+
+
+def _load_hf(spec: str, inputs: str | None) -> LoadedModel:
+    try:
+        from downshift.adapters import hf
+    except ImportError as exc:
+        raise LoadError(
+            f"{spec!r} isn't a file or an import spec; loading it as a Hugging Face repo "
+            "needs the [hf] extra: pip install 'downshift-server[hf]'"
+        ) from exc
+    try:
+        model = hf.load_pretrained(spec)
+    except (OSError, ValueError) as exc:  # hub errors and bad local repos surface as either
+        raise LoadError(f"can't load {spec!r} as a Hugging Face repo: {exc}") from exc
+    return LoadedModel(
+        source=spec,
+        model=model,
+        example_inputs=load_inputs(inputs) if inputs else None,
+        adapter_hint="hf",
+    )

downshift/serve/app.py

@@ -0,0 +1,93 @@
+"""FastAPI app over a ServingState. The same routes regardless of which backend is behind it."""
+
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+
+import downshift
+from downshift.serve.engine import ServingState
+from downshift.serve.middleware import load_middleware
+from downshift.serve.schemas import (
+    GraphPredictRequest,
+    HealthResponse,
+    MetadataResponse,
+    PredictRequest,
+    PredictResponse,
+    ReadyResponse,
+    to_numpy,
+)
+
+
+def run_predict(state: ServingState, inputs: dict[str, Any]) -> PredictResponse:
+    """Validate, convert, infer. Raises HTTPException(400) for anything the client got wrong."""
+    missing = [n for n in state.input_names if n not in inputs]
+    if missing:
+        raise HTTPException(400, f"missing inputs: {missing}")
+
+    declared = {spec.name: spec.dtype for spec in state.backend.metadata().inputs}
+    try:
+        feeds = {n: to_numpy(n, inputs[n], declared.get(n)) for n in state.input_names}
+        outputs = state.backend.infer(feeds)
+    except HTTPException:
+        raise
+    except Exception as exc:  # shape/dtype errors from ORT or torch are the client's problem
+        raise HTTPException(400, str(exc)) from exc
+
+    arrays = {name: np.asarray(arr) for name, arr in outputs.items()}
+    return PredictResponse(
+        outputs={name: arr.tolist() for name, arr in arrays.items()},
+        shapes={name: list(arr.shape) for name, arr in arrays.items()},
+        dtypes={name: arr.dtype.name for name, arr in arrays.items()},
+    )
+
+
+def build_app(state: ServingState, middleware: Sequence[str] = ()) -> FastAPI:
+    app = FastAPI(title="downshift", version=downshift.__version__)
+    app.state.serving = state
+    load_middleware(app, middleware)
+
+    @app.get("/health", response_model=HealthResponse)
+    def health() -> HealthResponse:
+        return HealthResponse()
+
+    @app.get("/ready", response_model=ReadyResponse)
+    def ready() -> JSONResponse:
+        status = 200 if state.ready else 503
+        return JSONResponse({"ready": state.ready}, status_code=status)
+
+    @app.get("/metadata", response_model=MetadataResponse)
+    def metadata() -> MetadataResponse:
+        return MetadataResponse(
+            model=state.source,
+            family=state.verdict.model_family,
+            verdict=state.verdict.to_dict(),
+            backend=state.backend.metadata().to_dict(),
+            input_names=list(state.input_names),
+            notes=list(state.notes),
+            version=downshift.__version__,
+        )
+
+    @app.post("/predict", response_model=PredictResponse)
+    def predict(req: PredictRequest) -> PredictResponse:
+        return run_predict(state, req.inputs)
+
+    @app.post("/predict/graph", response_model=PredictResponse)
+    def predict_graph(req: GraphPredictRequest) -> PredictResponse:
+        if not {"x", "edge_index"} <= set(state.input_names):
+            raise HTTPException(
+                400,
+                f"model is not graph-shaped: inputs are {list(state.input_names)}, "
+                "expected at least 'x' and 'edge_index'",
+            )
+        inputs: dict[str, Any] = {
+            "x": req.x,
+            "edge_index": {"data": req.edge_index, "dtype": "int64"},
+        }
+        if req.edge_attr is not None:
+            inputs["edge_attr"] = req.edge_attr
+        return run_predict(state, inputs)
+
+    return app

downshift/serve/backends.py

@@ -0,0 +1,181 @@
+"""Inference backends. Both take and return dicts of numpy arrays keyed by input name,
+so the HTTP layer doesn't care which one is behind it.
+"""
+
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Protocol
+
+import numpy as np
+import onnxruntime as ort
+import torch
+from torch import nn
+
+from downshift.export.verdict import BackendName
+
+_CUDA_EP = "CUDAExecutionProvider"
+_CPU_EP = "CPUExecutionProvider"
+
+
+@dataclass
+class IOSpec:
+    name: str
+    dtype: str | None
+    shape: list[int | str | None] | None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class BackendMeta:
+    name: BackendName
+    device: str
+    inputs: list[IOSpec]
+    outputs: list[IOSpec]
+
+    def to_dict(self) -> dict:
+        return {
+            "name": self.name,
+            "device": self.device,
+            "inputs": [i.to_dict() for i in self.inputs],
+            "outputs": [o.to_dict() for o in self.outputs],
+        }
+
+
+class Backend(Protocol):
+    name: BackendName
+    input_names: list[str]
+
+    def infer(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: ...
+
+    def metadata(self) -> BackendMeta: ...
+
+
+def output_names(count: int) -> list[str]:
+    return [f"output_{i}" for i in range(count)]
+
+
+def resolve_device(device: str) -> str:
+    if device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    return device
+
+
+def _ort_providers(device: str) -> list[str]:
+    available = ort.get_available_providers()
+    if device == "cuda" and _CUDA_EP in available:
+        return [_CUDA_EP, _CPU_EP]
+    return [_CPU_EP]
+
+
+def _session_options(intra_op_threads: int, inter_op_threads: int) -> ort.SessionOptions:
+    """Max graph optimization always on. Thread counts of 0 mean "let ONNX Runtime choose",
+    which is also its own default, so this is safe to set unconditionally."""
+    options = ort.SessionOptions()
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    options.intra_op_num_threads = intra_op_threads
+    options.inter_op_num_threads = inter_op_threads
+    return options
+
+
+class OnnxRuntimeBackend:
+    name = BackendName.onnxruntime
+
+    def __init__(
+        self,
+        model: bytes | str | Path,
+        device: str = "auto",
+        intra_op_threads: int = 0,
+        inter_op_threads: int = 0,
+    ) -> None:
+        source = model if isinstance(model, bytes) else str(model)
+        providers = _ort_providers(resolve_device(device))
+        options = _session_options(intra_op_threads, inter_op_threads)
+        self.session = ort.InferenceSession(source, sess_options=options, providers=providers)
+        self.provider = self.session.get_providers()[0]
+        self.input_names = [i.name for i in self.session.get_inputs()]
+        self.onnx_output_names = [o.name for o in self.session.get_outputs()]
+        # Outputs are keyed positionally so responses look the same from either backend.
+        self.output_names = output_names(len(self.onnx_output_names))
+
+    def infer(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+        outputs = self.session.run(self.onnx_output_names, inputs)
+        return dict(zip(self.output_names, outputs, strict=True))
+
+    def metadata(self) -> BackendMeta:
+        def spec(name: str, node) -> IOSpec:
+            return IOSpec(name, node.type, list(node.shape) if node.shape else None)
+
+        return BackendMeta(
+            name=self.name,
+            device=self.provider,
+            inputs=[spec(i.name, i) for i in self.session.get_inputs()],
+            outputs=[
+                spec(name, o)
+                for name, o in zip(self.output_names, self.session.get_outputs(), strict=True)
+            ],
+        )
+
+
+class TorchBackend:
+    """Eager PyTorch. The fallback path, and a first-class one: same contract as ORT."""
+
+    name = BackendName.torch
+
+    def __init__(
+        self,
+        module: nn.Module,
+        input_names: tuple[str, ...],
+        device: str = "auto",
+        example_inputs: tuple | None = None,
+    ) -> None:
+        self.device = resolve_device(device)
+        self.module = module.eval().to(self.device)
+        self.input_names = list(input_names)
+        self._input_specs = [IOSpec(n, None, None) for n in self.input_names]
+        self._output_specs: list[IOSpec] = []
+        if example_inputs is not None:
+            # One pass over the example fills in dtypes and shapes for /metadata.
+            self._input_specs = [
+                _spec_from_tensor(n, t)
+                for n, t in zip(self.input_names, example_inputs, strict=True)
+            ]
+            feeds = {n: t.numpy() for n, t in zip(self.input_names, example_inputs, strict=True)}
+            self._output_specs = [_spec_from_array(n, a) for n, a in self.infer(feeds).items()]
+
+    def infer(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+        missing = [n for n in self.input_names if n not in inputs]
+        if missing:
+            raise KeyError(f"missing inputs: {missing}")
+        args = [
+            torch.from_numpy(np.ascontiguousarray(inputs[n])).to(self.device)
+            for n in self.input_names
+        ]
+        with torch.inference_mode():
+            out = self.module(*args)
+        if isinstance(out, torch.Tensor):
+            tensors = [out]
+        else:
+            tensors = [t for t in out if isinstance(t, torch.Tensor)]
+        return {
+            name: t.detach().cpu().numpy()
+            for name, t in zip(output_names(len(tensors)), tensors, strict=True)
+        }
+
+    def metadata(self) -> BackendMeta:
+        return BackendMeta(self.name, self.device, self._input_specs, self._output_specs)
+
+
+def _dynamic_shape(shape: tuple[int, ...]) -> list[int | str | None]:
+    # Axis 0 is dynamic for anything we serve; report it the way ORT does.
+    return ["batch", *shape[1:]] if shape else []
+
+
+def _spec_from_tensor(name: str, t: torch.Tensor) -> IOSpec:
+    dtype = str(t.dtype).removeprefix("torch.")
+    return IOSpec(name, f"tensor({dtype})", _dynamic_shape(tuple(t.shape)))
+
+
+def _spec_from_array(name: str, a: np.ndarray) -> IOSpec:
+    return IOSpec(name, f"tensor({a.dtype.name})", _dynamic_shape(a.shape))

downshift/serve/engine.py

@@ -0,0 +1,154 @@
+"""From a loaded model to a warmed-up backend. The CLI's `serve` is render(prepare_serving())."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from downshift.export.prevalidated import intake
+from downshift.export.verdict import BackendName, ExportVerdict, build_verdict, prepare_model
+from downshift.loading import LoadedModel
+from downshift.serve.backends import Backend, OnnxRuntimeBackend, TorchBackend
+
+
+class BackendChoice(str, Enum):
+    """What the caller asked for; "auto" defers to the verdict's recommendation."""
+
+    auto = "auto"
+    onnxruntime = "onnxruntime"
+    torch = "torch"
+
+
+@dataclass
+class ServeOptions:
+    backend: BackendChoice = BackendChoice.auto
+    force_onnx: bool = False  # serve a DEGRADED graph via ORT anyway
+    device: str = "auto"
+    warmup: int = 3
+    k: int = 8
+    adapter: str | None = None
+    dynamic: dict[str, list[int]] | None = None
+    intra_op_threads: int = 0  # ORT SessionOptions; 0 = let ONNX Runtime choose
+    inter_op_threads: int = 0
+
+
+@dataclass
+class ServingState:
+    source: str
+    verdict: ExportVerdict
+    backend: Backend
+    input_names: tuple[str, ...]
+    options: ServeOptions
+    example_inputs: tuple | None = None
+    ready: bool = False
+    notes: list[str] = field(default_factory=list)  # things the banner should say
+
+    @property
+    def backend_auto_selected(self) -> bool:
+        return self.options.backend == BackendChoice.auto and not self.forced_onnx
+
+    @property
+    def forced_onnx(self) -> bool:
+        return self.options.force_onnx and self.verdict.status == "DEGRADED"
+
+
+def _verdict_for(
+    loaded: LoadedModel, reference: LoadedModel | None, opts: ServeOptions
+) -> ExportVerdict:
+    adapter = opts.adapter or loaded.adapter_hint
+    if loaded.onnx_path is not None:
+        ref_model = reference.model if reference else None
+        ref_inputs = reference.example_inputs if reference else None
+        return intake(
+            loaded.onnx_path, ref_model, ref_inputs, adapter, k=opts.k, dynamic=opts.dynamic
+        )
+    assert loaded.model is not None
+    prepared = prepare_model(loaded.model, loaded.example_inputs, adapter, opts.dynamic)
+    if opts.backend == BackendChoice.torch:
+        # Skip the export entirely; the user asked for eager.
+        return ExportVerdict(
+            status="UNVERIFIED",
+            model_family=prepared.family,
+            capture_strategy=None,
+            opset=None,
+            op_types=[],
+            numerics=None,
+            recommended_backend=BackendName.torch,
+            reason="--backend torch: export skipped",
+            input_names=prepared.input_names,
+            dynamic_dims=prepared.dynamic_dims,
+            prepared=prepared,
+        )
+    return build_verdict(prepared, k=opts.k)
+
+
+def choose_backend(verdict: ExportVerdict, opts: ServeOptions) -> tuple[BackendName, list[str]]:
+    """Return (backend name, notes for the banner)."""
+    notes: list[str] = []
+    wanted: BackendName = (
+        verdict.recommended_backend
+        if opts.backend == BackendChoice.auto
+        else BackendName(opts.backend)
+    )
+    if opts.force_onnx and verdict.status == "DEGRADED":
+        wanted = BackendName.onnxruntime
+        notes.append("--force-onnx: serving a DEGRADED graph; outputs may be wrong")
+
+    has_onnx = verdict.onnx_program is not None or verdict.onnx_path is not None
+    has_torch = verdict.prepared is not None
+    if wanted == BackendName.onnxruntime and not has_onnx:
+        notes.append("no ONNX graph available; falling back to torch")
+        wanted = BackendName.torch
+    if wanted == BackendName.torch and not has_torch:
+        raise ValueError("torch backend requested but there is no PyTorch model to run")
+    return wanted, notes
+
+
+def _build_backend(name: BackendName, verdict: ExportVerdict, opts: ServeOptions) -> Backend:
+    if name == BackendName.onnxruntime:
+        if verdict.onnx_path is not None:
+            source: bytes | Path = verdict.onnx_path
+        else:
+            program: Any = verdict.onnx_program
+            source = program.model_proto.SerializeToString()
+        return OnnxRuntimeBackend(source, opts.device, opts.intra_op_threads, opts.inter_op_threads)
+    prepared = verdict.prepared
+    assert prepared is not None
+    return TorchBackend(prepared.model, prepared.input_names, opts.device, prepared.inputs)
+
+
+def prepare_serving(
+    loaded: LoadedModel, opts: ServeOptions | None = None, reference: LoadedModel | None = None
+) -> ServingState:
+    opts = opts or ServeOptions()
+    verdict = _verdict_for(loaded, reference, opts)
+    name, notes = choose_backend(verdict, opts)
+    backend = _build_backend(name, verdict, opts)
+
+    if verdict.prepared is not None:
+        input_names = verdict.prepared.input_names
+        example_inputs = verdict.prepared.inputs
+    else:
+        input_names = tuple(backend.input_names)
+        example_inputs = None
+
+    state = ServingState(
+        loaded.source, verdict, backend, input_names, opts, example_inputs, notes=notes
+    )
+    warmup(state, opts.warmup)
+    return state
+
+
+def warmup(state: ServingState, n: int) -> None:
+    """Run a few inferences before /ready flips; first-call costs shouldn't hit users."""
+    if state.example_inputs is not None:
+        feeds = {
+            name: t.numpy() if isinstance(t, torch.Tensor) else np.asarray(t)
+            for name, t in zip(state.input_names, state.example_inputs, strict=True)
+        }
+        for _ in range(n):
+            state.backend.infer(feeds)
+    state.ready = True

downshift/serve/middleware.py

@@ -0,0 +1,24 @@
+"""Attach user-supplied middleware given as "pkg.module:attr" import specs."""
+
+import inspect
+from collections.abc import Sequence
+
+from fastapi import FastAPI
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from downshift.loading import import_object
+
+
+def load_middleware(app: FastAPI, specs: Sequence[str]) -> None:
+    """Each spec names a BaseHTTPMiddleware subclass or an async (request, call_next) function."""
+    for spec in specs:
+        obj = import_object(spec)
+        if inspect.isclass(obj) and issubclass(obj, BaseHTTPMiddleware):
+            app.add_middleware(obj)
+        elif inspect.iscoroutinefunction(obj):
+            app.middleware("http")(obj)
+        else:
+            raise ValueError(
+                f"{spec!r} is not middleware: expected a BaseHTTPMiddleware subclass or an "
+                f"async function taking (request, call_next), got {type(obj).__name__}"
+            )