downshift-server 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (86) hide show
  1. downshift_server-0.2.0/.github/workflows/ci.yml +25 -0
  2. downshift_server-0.2.0/.github/workflows/compatibility-matrix.yml +32 -0
  3. downshift_server-0.2.0/.gitignore +48 -0
  4. downshift_server-0.2.0/LICENSE +21 -0
  5. downshift_server-0.2.0/PKG-INFO +265 -0
  6. downshift_server-0.2.0/README.md +224 -0
  7. downshift_server-0.2.0/docs/blog/your-gnn-exports-cleanly.md +74 -0
  8. downshift_server-0.2.0/docs/compatibility.md +25 -0
  9. downshift_server-0.2.0/examples/01_check_a_model.ipynb +243 -0
  10. downshift_server-0.2.0/examples/01_check_a_model.py +153 -0
  11. downshift_server-0.2.0/examples/02_export_and_manifest.ipynb +193 -0
  12. downshift_server-0.2.0/examples/02_export_and_manifest.py +117 -0
  13. downshift_server-0.2.0/examples/03_serve_and_query.ipynb +219 -0
  14. downshift_server-0.2.0/examples/03_serve_and_query.py +133 -0
  15. downshift_server-0.2.0/examples/04_pyg_graph_neural_networks.ipynb +155 -0
  16. downshift_server-0.2.0/examples/04_pyg_graph_neural_networks.py +89 -0
  17. downshift_server-0.2.0/examples/05_huggingface_encoder.ipynb +149 -0
  18. downshift_server-0.2.0/examples/05_huggingface_encoder.py +83 -0
  19. downshift_server-0.2.0/examples/06_custom_adapter.ipynb +207 -0
  20. downshift_server-0.2.0/examples/06_custom_adapter.py +146 -0
  21. downshift_server-0.2.0/examples/07_cli_walkthrough.ipynb +178 -0
  22. downshift_server-0.2.0/examples/07_cli_walkthrough.py +103 -0
  23. downshift_server-0.2.0/examples/README.md +43 -0
  24. downshift_server-0.2.0/examples/_build_notebooks.py +97 -0
  25. downshift_server-0.2.0/pyproject.toml +93 -0
  26. downshift_server-0.2.0/scripts/gen_matrix.py +135 -0
  27. downshift_server-0.2.0/src/downshift/__init__.py +60 -0
  28. downshift_server-0.2.0/src/downshift/adapters/__init__.py +0 -0
  29. downshift_server-0.2.0/src/downshift/adapters/_flatten.py +40 -0
  30. downshift_server-0.2.0/src/downshift/adapters/base.py +47 -0
  31. downshift_server-0.2.0/src/downshift/adapters/generic.py +99 -0
  32. downshift_server-0.2.0/src/downshift/adapters/hf.py +95 -0
  33. downshift_server-0.2.0/src/downshift/adapters/pyg.py +120 -0
  34. downshift_server-0.2.0/src/downshift/adapters/registry.py +116 -0
  35. downshift_server-0.2.0/src/downshift/cli/__init__.py +0 -0
  36. downshift_server-0.2.0/src/downshift/cli/main.py +428 -0
  37. downshift_server-0.2.0/src/downshift/cli/render.py +174 -0
  38. downshift_server-0.2.0/src/downshift/export/__init__.py +0 -0
  39. downshift_server-0.2.0/src/downshift/export/capture.py +93 -0
  40. downshift_server-0.2.0/src/downshift/export/inputs.py +25 -0
  41. downshift_server-0.2.0/src/downshift/export/manifest.py +89 -0
  42. downshift_server-0.2.0/src/downshift/export/prevalidated.py +70 -0
  43. downshift_server-0.2.0/src/downshift/export/shapes.py +61 -0
  44. downshift_server-0.2.0/src/downshift/export/verdict.py +211 -0
  45. downshift_server-0.2.0/src/downshift/export/verify.py +162 -0
  46. downshift_server-0.2.0/src/downshift/loading.py +166 -0
  47. downshift_server-0.2.0/src/downshift/serve/__init__.py +0 -0
  48. downshift_server-0.2.0/src/downshift/serve/app.py +93 -0
  49. downshift_server-0.2.0/src/downshift/serve/backends.py +181 -0
  50. downshift_server-0.2.0/src/downshift/serve/engine.py +154 -0
  51. downshift_server-0.2.0/src/downshift/serve/middleware.py +24 -0
  52. downshift_server-0.2.0/src/downshift/serve/schemas.py +124 -0
  53. downshift_server-0.2.0/src/downshift/settings.py +69 -0
  54. downshift_server-0.2.0/tests/__init__.py +0 -0
  55. downshift_server-0.2.0/tests/conftest.py +47 -0
  56. downshift_server-0.2.0/tests/models/__init__.py +0 -0
  57. downshift_server-0.2.0/tests/models/clean_mlp.py +27 -0
  58. downshift_server-0.2.0/tests/models/custom_autograd.py +40 -0
  59. downshift_server-0.2.0/tests/models/data_dependent_branch.py +27 -0
  60. downshift_server-0.2.0/tests/models/dict_input.py +32 -0
  61. downshift_server-0.2.0/tests/models/dropout_model.py +25 -0
  62. downshift_server-0.2.0/tests/models/dynamic_batch_cnn.py +27 -0
  63. downshift_server-0.2.0/tests/models/gnn_gat.py +37 -0
  64. downshift_server-0.2.0/tests/models/gnn_gcn.py +30 -0
  65. downshift_server-0.2.0/tests/models/gnn_sage.py +30 -0
  66. downshift_server-0.2.0/tests/models/scatter_include_self_false.py +34 -0
  67. downshift_server-0.2.0/tests/models/tied_weights.py +27 -0
  68. downshift_server-0.2.0/tests/models/tiny_bert.py +23 -0
  69. downshift_server-0.2.0/tests/test_adapters.py +406 -0
  70. downshift_server-0.2.0/tests/test_capture.py +28 -0
  71. downshift_server-0.2.0/tests/test_cli.py +255 -0
  72. downshift_server-0.2.0/tests/test_engine.py +204 -0
  73. downshift_server-0.2.0/tests/test_export.py +87 -0
  74. downshift_server-0.2.0/tests/test_export_pyg.py +40 -0
  75. downshift_server-0.2.0/tests/test_fixtures.py +43 -0
  76. downshift_server-0.2.0/tests/test_hf.py +37 -0
  77. downshift_server-0.2.0/tests/test_inputs.py +13 -0
  78. downshift_server-0.2.0/tests/test_loading.py +210 -0
  79. downshift_server-0.2.0/tests/test_manifest.py +95 -0
  80. downshift_server-0.2.0/tests/test_prevalidated.py +47 -0
  81. downshift_server-0.2.0/tests/test_render.py +188 -0
  82. downshift_server-0.2.0/tests/test_schemas.py +38 -0
  83. downshift_server-0.2.0/tests/test_serve.py +184 -0
  84. downshift_server-0.2.0/tests/test_settings.py +61 -0
  85. downshift_server-0.2.0/tests/test_shapes.py +78 -0
  86. downshift_server-0.2.0/tests/test_verify_internals.py +35 -0
@@ -0,0 +1,25 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [ main, develop ]
6
+ pull_request:
7
+ branches: [ main, develop ]
8
+
9
+ jobs:
10
+ test:
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ matrix:
14
+ python-version: ["3.11", "3.12"]
15
+ steps:
16
+ - uses: actions/checkout@v4
17
+ - uses: actions/setup-python@v5
18
+ with:
19
+ python-version: ${{ matrix.python-version }}
20
+ cache: pip
21
+ - run: pip install torch --index-url https://download.pytorch.org/whl/cpu
22
+ - run: pip install -e ".[dev,gnn,hf]"
23
+ - run: ruff check .
24
+ - run: mypy src
25
+ - run: pytest -q
@@ -0,0 +1,32 @@
1
+ name: Compatibility matrix
2
+
3
+ # The matrix is the project's core claim. A stale one is worse than none, so it is
4
+ # regenerated weekly against current torch/onnxruntime and opened as a PR when it changes.
5
+
6
+ on:
7
+ schedule:
8
+ - cron: "0 6 * * 1" # Mondays 06:00 UTC
9
+ workflow_dispatch:
10
+
11
+ jobs:
12
+ regenerate:
13
+ runs-on: ubuntu-latest
14
+ permissions:
15
+ contents: write
16
+ pull-requests: write
17
+ steps:
18
+ - uses: actions/checkout@v4
19
+ - uses: actions/setup-python@v5
20
+ with:
21
+ python-version: "3.12"
22
+ cache: pip
23
+ - run: pip install torch --index-url https://download.pytorch.org/whl/cpu
24
+ - run: pip install -e ".[dev,gnn,hf]"
25
+ - run: python scripts/gen_matrix.py
26
+ - uses: peter-evans/create-pull-request@v7
27
+ with:
28
+ commit-message: "docs: regenerate compatibility matrix"
29
+ title: "Weekly compatibility matrix refresh"
30
+ body: "Automated run of `scripts/gen_matrix.py` against current torch / onnxruntime."
31
+ branch: ci/compatibility-matrix
32
+ add-paths: docs/compatibility.md
@@ -0,0 +1,48 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ *.so
6
+ .Python
7
+ *.egg-info/
8
+ dist/
9
+ build/
10
+
11
+ # Virtual environments
12
+ .venv/
13
+ venv/
14
+ ENV/
15
+ env/
16
+
17
+ # IDEs
18
+ .vscode/
19
+ .idea/
20
+ *.swp
21
+ *.swo
22
+ *~
23
+
24
+ # Testing
25
+ .pytest_cache/
26
+ .mypy_cache/
27
+ .ruff_cache/
28
+ htmlcov/
29
+ .coverage
30
+
31
+ # OS
32
+ .DS_Store
33
+ Thumbs.db
34
+
35
+ # Project-specific
36
+ *.onnx
37
+ *.onnx.data
38
+ artifacts/
39
+ tmp/
40
+
41
+ # Local tooling config
42
+ .claude/
43
+ .docs/
44
+
45
+ # Local pre-publish validation: hits test.pypi.org/pypi.org and builds a throwaway venv;
46
+ # not part of the package or CI, kept out of version control on purpose.
47
+ scripts/validate_pypi_release.py
48
+ scripts/.validation/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Nikhil Ranjan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,265 @@
1
+ Metadata-Version: 2.5
2
+ Name: downshift-server
3
+ Version: 0.2.0
4
+ Summary: Universal PyTorch → ONNX verification and serving, with graceful fallback.
5
+ Author-email: Nikhil Ranjan <nrnikhilranjan1997@gmail.com>
6
+ License: MIT
7
+ License-File: LICENSE
8
+ Keywords: export,model-serving,onnx,pytorch
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Classifier: Intended Audience :: Developers
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Requires-Python: <3.15,>=3.11
17
+ Requires-Dist: fastapi>=0.100
18
+ Requires-Dist: numpy>=1.24
19
+ Requires-Dist: onnx>=1.14
20
+ Requires-Dist: onnxruntime>=1.17
21
+ Requires-Dist: onnxscript>=0.1
22
+ Requires-Dist: pydantic>=2.0
23
+ Requires-Dist: rich>=13.0
24
+ Requires-Dist: torch>=2.0
25
+ Requires-Dist: typer>=0.12
26
+ Requires-Dist: uvicorn[standard]>=0.23
27
+ Provides-Extra: all
28
+ Requires-Dist: torch-geometric>=2.5; extra == 'all'
29
+ Requires-Dist: transformers>=4.40; extra == 'all'
30
+ Provides-Extra: dev
31
+ Requires-Dist: httpx>=0.24; extra == 'dev'
32
+ Requires-Dist: mypy>=1.0; extra == 'dev'
33
+ Requires-Dist: pytest-cov>=5.0; extra == 'dev'
34
+ Requires-Dist: pytest>=7.0; extra == 'dev'
35
+ Requires-Dist: ruff>=0.6; extra == 'dev'
36
+ Provides-Extra: gnn
37
+ Requires-Dist: torch-geometric>=2.5; extra == 'gnn'
38
+ Provides-Extra: hf
39
+ Requires-Dist: transformers>=4.40; extra == 'hf'
40
+ Description-Content-Type: text/markdown
41
+
42
+ # downshift
43
+
44
+ Check whether your PyTorch model survives ONNX export. Then serve it, falling back to eager PyTorch when ONNX would be lying to you.
45
+
46
+ ```
47
+ $ downshift serve tests.models.scatter_include_self_false:make_model
48
+
49
+ ┌─ downshift v0.1.0 ────────────────────────────────────────────────────────┐
50
+ │ │
51
+ │ Model tests.models.scatter_include_self_false:make_model │
52
+ │ Family generic-torch │
53
+ │ Verdict DEGRADED (strict=False, opset 20) │
54
+ │ Numerics max abs err 1.16e+00 over 8 samples ✗ 7/8 failed │
55
+ │ ⚠ numerics diverge on 7/8 samples (max abs err 1.16e+00) │
56
+ │ Override --force-onnx to serve the ONNX graph anyway │
57
+ │ Backend torch (eager) · cpu ← auto-selected │
58
+ │ Dynamic dims x[0], segment_ids[0] │
59
+ │ Endpoint http://127.0.0.1:8000 │
60
+ │ │
61
+ └───────────────────────────────────────────────────────────────────────────┘
62
+ ```
63
+
64
+ This graph exported without a single error and produces wrong numbers on 7 of 8 inputs. The tool caught it and served PyTorch instead.
65
+
66
+ ## Install
67
+
68
+ ```bash
69
+ pip install downshift-server # core: any nn.Module, any .onnx
70
+ pip install "downshift-server[gnn]" # + PyTorch Geometric adapter
71
+ pip install "downshift-server[hf]" # + Hugging Face encoder adapter
72
+ pip install "downshift-server[all]"
73
+ ```
74
+
75
+ Until the PyPI release lands, install from a checkout with `pip install -e ".[all]"`.
76
+
77
+ Python 3.11 to 3.13. CPU-only is what this release was tested on. CUDA execution-provider selection exists (`--device cuda`) but is untested in this release.
78
+
79
+ ## Quick start
80
+
81
+ ### `check`: is the export trustworthy?
82
+
83
+ ```
84
+ $ downshift check tests.models.scatter_include_self_false:make_model
85
+
86
+ ┌───────────────┬─────────────────────────────────────────────────────────────┐
87
+ │ Model │ tests.models.scatter_include_self_false:make_model │
88
+ │ Family │ generic-torch │
89
+ │ Export │ DEGRADED (strict=False, opset 20) │
90
+ │ Numerics │ max abs err 1.34e+00 over 8 samples ✗ 7/8 failed │
91
+ │ Shape-general │ no │
92
+ │ Dynamic dims │ x[0], segment_ids[0] │
93
+ │ Backend │ torch │
94
+ │ Reason │ exported via strict=False but numerics diverge on 7/8 │
95
+ │ │ samples (max abs err 1.34e+00) │
96
+ └───────────────┴─────────────────────────────────────────────────────────────┘
97
+ ```
98
+
99
+ `check` exports in memory, runs `k` random samples (default 8) through both PyTorch and ONNX Runtime, and varies the dynamic axes so some samples have shapes the exporter never saw. Nothing is written to disk.
100
+
101
+ The exit code is the verdict, so it can gate CI: `0` CLEAN, `1` FAILED, `2` DEGRADED, `3` UNVERIFIED. (`4` is a usage error such as an unloadable model; `5` is a crash.) `--json` prints the full verdict as JSON and nothing else:
102
+
103
+ ```bash
104
+ downshift check my_pkg.models:build --json -k 16 > verdict.json
105
+ ```
106
+
107
+ Useful options: `-k/--samples`, `--dynamic "x:0,edge_index:1"` to override which axes are dynamic (default: axis 0 of every input), `--adapter generic|pyg|hf` to skip detection, `--inputs pkg.module:fn` to supply example inputs.
108
+
109
+ ### `export`: write the artifact
110
+
111
+ ```bash
112
+ downshift export my_pkg.models:build -o artifacts/ --name classifier
113
+ ```
114
+
115
+ Writes `artifacts/classifier.onnx` and `artifacts/classifier.manifest.json`. The manifest records the SHA-256 of the artifact and of the source checkpoint (when the model came from a file), torch/onnx/onnxruntime versions, opset, the observed weight dtype, and the full verdict including the numerics report. A DEGRADED export is still written, because the manifest records exactly how far off it is; a FAILED export writes nothing.
116
+
117
+ `--fp16` casts the model to half before export (a plain `.half()`, not quantization). `--no-verify` skips the numerics check and marks the verdict UNVERIFIED, with a warning.
118
+
119
+ ### `serve`: one endpoint, backend chosen by the verdict
120
+
121
+ ```bash
122
+ downshift serve my_pkg.models:build --port 8000
123
+ ```
124
+
125
+ Prints the banner above, then starts uvicorn. Routes:
126
+
127
+ | Route | What it does |
128
+ |---|---|
129
+ | `POST /predict` | Named tensor inputs, any model |
130
+ | `POST /predict/graph` | One graph: `x`, `edge_index`, optional `edge_attr` |
131
+ | `GET /health` | Liveness |
132
+ | `GET /ready` | `200` once warmup is done, `503` before |
133
+ | `GET /metadata` | Family, backend, full verdict, input names |
134
+
135
+ ```bash
136
+ curl -s localhost:8000/predict -H 'content-type: application/json' \
137
+ -d '{"inputs": {"x": [[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5]]}}'
138
+ ```
139
+
140
+ ```json
141
+ {"outputs": {"output_0": [[0.199, -0.206, 0.561, 0.405]]},
142
+ "shapes": {"output_0": [1, 4]},
143
+ "dtypes": {"output_0": "float32"}}
144
+ ```
145
+
146
+ Integer lists become `int64`, everything else `float32`. To be explicit, pass `{"data": [...], "dtype": "float16", "shape": [1, 16]}` instead of a bare list. For graph models:
147
+
148
+ ```bash
149
+ curl -s localhost:8000/predict/graph -H 'content-type: application/json' \
150
+ -d '{"x": [[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
151
+ [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
152
+ [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]],
153
+ "edge_index": [[0, 1, 2], [1, 2, 0]]}'
154
+ ```
155
+
156
+ The response has the same shape as `/predict`, with one output row per node.
157
+
158
+ Options that change what gets served:
159
+
160
+ - `--backend auto|onnxruntime|torch`. `auto` follows the verdict. `torch` skips the export entirely.
161
+ - `--force-onnx` serves a DEGRADED graph through ONNX Runtime anyway. The banner says so in red.
162
+ - `--reference model` verifies a pre-built `.onnx` against a PyTorch model; without it the verdict is UNVERIFIED.
163
+ - `--middleware pkg.module:Attr` (repeatable) attaches a `BaseHTTPMiddleware` subclass or an `async (request, call_next)` function. No middleware means no overhead.
164
+ - `--device auto|cpu|cuda`, `--warmup N` (inferences before `/ready` flips), `--host`, `--port`, `--log-level`, `--log-format json`.
165
+
166
+ ## Accepted model forms
167
+
168
+ | Argument | Meaning |
169
+ |---|---|
170
+ | `model.onnx` | Pre-built ONNX, served as-is. UNVERIFIED unless `--reference` is given. |
171
+ | `pkg.module:attr` | Import spec. `attr` is an `nn.Module` instance or a zero-argument factory. A sibling `make_inputs` in the same module is picked up automatically; otherwise pass `--inputs pkg.module:fn`. |
172
+ | `weights.pt` | State dict. Needs `--model-class pkg.module:Class`. Also `.pth`, `.bin`, `.ckpt`. |
173
+ | `org/repo` | Hugging Face hub id. Needs the `[hf]` extra. |
174
+ | `path/to/repo/dir/` | Locally downloaded Hugging Face repo: a directory containing `config.json`. Needs the `[hf]` extra. |
175
+
176
+ Checkpoints are loaded with `torch.load(weights_only=True)`. A file that holds a pickled full module will not load that way; `--unsafe-load` switches to `weights_only=False`, which means running arbitrary code from the file. Only use it on files you would run as a script.
177
+
178
+ ## The four verdicts
179
+
180
+ - **CLEAN**: exports, matches PyTorch on every sample, survives shapes it was not traced on. Served via ONNX Runtime.
181
+ - **DEGRADED**: exports without error, but numerics drift past tolerance on at least one sample. Served via eager PyTorch; `--force-onnx` overrides.
182
+ - **FAILED**: does not export. Served via eager PyTorch. Not an error, a supported path.
183
+ - **UNVERIFIED**: a `.onnx` with no reference model, or `--no-verify`. Served via ONNX Runtime and labelled as never checked.
184
+
185
+ ## Compatibility matrix
186
+
187
+ Generated by `scripts/gen_matrix.py` from the fixture corpus in `tests/models/`, each fixture isolating one export hazard. CI regenerates it weekly against current torch and onnxruntime and opens a PR when it changes. Full file, with versions and legend: [docs/compatibility.md](docs/compatibility.md).
188
+
189
+ | Model | Hazard | Family | Export | Capture | Numerics | Shape-general | Backend |
190
+ |---|---|---|---|---|---|---|---|
191
+ | `clean_mlp` | Control fixture: no export hazards | generic-torch | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
192
+ | `custom_autograd` | custom autograd.Function with no symbolic override | generic-torch | CLEAN | strict=False | 4.2e-07 | ✓ | onnxruntime |
193
+ | `data_dependent_branch` | data-dependent control flow | generic-torch | FAILED | — | — | — | torch |
194
+ | `dict_input` | dataclass container input | generic-torch | CLEAN | strict=False | 4.8e-07 | ✓ | onnxruntime |
195
+ | `dropout_model` | stochastic layer | generic-torch | CLEAN | strict=False | 3.6e-07 | ✓ | onnxruntime |
196
+ | `dynamic_batch_cnn` | batch-dim generalization | generic-torch | CLEAN | strict=False | 6.0e-08 | ✓ | onnxruntime |
197
+ | `gnn_gat` | GNN fixture: 3-layer GAT node classifier | pyg | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
198
+ | `gnn_gcn` | GNN fixture: 2-layer GCN node classifier | pyg | CLEAN | strict=False | 2.4e-07 | ✓ | onnxruntime |
199
+ | `gnn_sage` | GNN fixture: 2-layer GraphSAGE node classifier | pyg | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
200
+ | `scatter_include_self_false` | scatter_reduce(include_self=False) has no faithful ONNX translation | generic-torch | DEGRADED | strict=False | 1.3e+00 | ✗ | torch |
201
+ | `tied_weights` | tied embedding/output weight (GPT-2/OPT-style) | generic-torch | CLEAN | strict=False | 1.9e-06 | ✓ | onnxruntime |
202
+ | `tiny_bert` | HF fixture: a randomly initialised two-layer BERT encoder | hf-transformers | CLEAN | strict=False | 6.0e-07 | ✓ | onnxruntime |
203
+
204
+ Two rows worth reading twice. `custom_autograd` was expected to fail and is CLEAN, because `torch.export` traces straight through a `Function.forward` made of ordinary ops. `scatter_include_self_false` was expected to fail loudly and instead exports with zero errors and returns the wrong numbers; the only thing standing between that graph and production is the numerics check.
205
+
206
+ ## What this is not
207
+
208
+ - **No quantization or graph optimization, ever.** Not deferred, cut. Run Olive, `onnxruntime.quantization`, or your own script, then hand the result to `downshift serve model.onnx --reference model.pt` and it gets verified against the original weights like any other export. `--fp16` is a cast before tracing, nothing lower exists here.
209
+ - **No LLM path.** No `onnxruntime-genai` backend, no OpenAI-compatible endpoints, no KV cache, no sampling loop. Encoder-only Hugging Face models work; causal LMs are not a target yet.
210
+ - **No continuous batching, no PagedAttention.** The boot banner is a visual homage to vLLM. That is the full extent of the resemblance.
211
+ - **No dynamic request batching yet.** One request, one inference.
212
+ - **No graph batching yet.** `/predict/graph` takes one graph. Concatenate graphs client-side with offset edge indices if you need more.
213
+ - **No Prometheus metrics, no Docker image.** `--middleware` is the hook for the former; pip plus version pins is the path for the latter.
214
+ - **No DGL adapter yet.** PyG only.
215
+
216
+ **vs. anydeploy.** `anydeploy` also does export, validate, and serve, with a pass/fail validation step and an edge/mobile focus. downshift differs in three places: the verdict is tiered, with DEGRADED as a real middle state between "works" and "crashes"; the eager PyTorch fallback sits behind the same endpoint so a FAILED or DEGRADED model still serves; and GNNs (PyTorch Geometric) are a supported family with independent node and edge dynamic dims.
217
+
218
+ ## Writing your own adapter
219
+
220
+ An adapter knows one model family well enough to build example inputs when the user gave none, and to turn the model plus inputs into something `torch.export` can trace: a module with a flat tensor signature. Implement the `Adapter` protocol from `downshift.adapters.base`:
221
+
222
+ ```python
223
+ from downshift.adapters.base import Prepared
224
+
225
+ class MyAdapter:
226
+ name = "myfamily"
227
+ family = "myfamily"
228
+
229
+ def matches(self, model, example_inputs) -> bool: ...
230
+ def example_inputs(self, model) -> tuple | None: ... # None if you can't guess
231
+ def prepare(self, model, example_inputs) -> Prepared: ...
232
+
233
+ ADAPTER = MyAdapter()
234
+ ```
235
+
236
+ `Prepared` carries the export-ready module, the flat example inputs, their names, the per-input `dynamic_shapes` spec, an optional `vary_fn(i) -> inputs` that generates verification samples, and the family string. Register it under the `downshift.adapters` entry-point group in your own package:
237
+
238
+ ```toml
239
+ [project.entry-points."downshift.adapters"]
240
+ myfamily = "my_pkg.adapter:ADAPTER"
241
+ ```
242
+
243
+ Adapters are tried most-specific first; `generic` always goes last. An adapter whose optional dependency is missing is skipped silently.
244
+
245
+ For a one-off adapter that isn't worth packaging, `--adapter` (and `check()`'s `adapter=`) also accepts a bare `.py` file directly, no install or entry point required:
246
+
247
+ ```bash
248
+ downshift check my_model.py:model --adapter path/to/pointcloud_adapter.py
249
+ ```
250
+
251
+ The file needs a module-level `ADAPTER = MyAdapter()`, or point at the class directly with `--adapter path/to/pointcloud_adapter.py:MyAdapter` and it's instantiated with no arguments.
252
+
253
+ ## Development
254
+
255
+ ```bash
256
+ pip install -e ".[dev,all]"
257
+ ruff check .
258
+ mypy src
259
+ pytest
260
+ python scripts/gen_matrix.py # regenerates docs/compatibility.md
261
+ ```
262
+
263
+ ## License
264
+
265
+ MIT.
@@ -0,0 +1,224 @@
1
+ # downshift
2
+
3
+ Check whether your PyTorch model survives ONNX export. Then serve it, falling back to eager PyTorch when ONNX would be lying to you.
4
+
5
+ ```
6
+ $ downshift serve tests.models.scatter_include_self_false:make_model
7
+
8
+ ┌─ downshift v0.1.0 ────────────────────────────────────────────────────────┐
9
+ │ │
10
+ │ Model tests.models.scatter_include_self_false:make_model │
11
+ │ Family generic-torch │
12
+ │ Verdict DEGRADED (strict=False, opset 20) │
13
+ │ Numerics max abs err 1.16e+00 over 8 samples ✗ 7/8 failed │
14
+ │ ⚠ numerics diverge on 7/8 samples (max abs err 1.16e+00) │
15
+ │ Override --force-onnx to serve the ONNX graph anyway │
16
+ │ Backend torch (eager) · cpu ← auto-selected │
17
+ │ Dynamic dims x[0], segment_ids[0] │
18
+ │ Endpoint http://127.0.0.1:8000 │
19
+ │ │
20
+ └───────────────────────────────────────────────────────────────────────────┘
21
+ ```
22
+
23
+ This graph exported without a single error and produces wrong numbers on 7 of 8 inputs. The tool caught it and served PyTorch instead.
24
+
25
+ ## Install
26
+
27
+ ```bash
28
+ pip install downshift-server # core: any nn.Module, any .onnx
29
+ pip install "downshift-server[gnn]" # + PyTorch Geometric adapter
30
+ pip install "downshift-server[hf]" # + Hugging Face encoder adapter
31
+ pip install "downshift-server[all]"
32
+ ```
33
+
34
+ Until the PyPI release lands, install from a checkout with `pip install -e ".[all]"`.
35
+
36
+ Python 3.11 to 3.13. CPU-only is what this release was tested on. CUDA execution-provider selection exists (`--device cuda`) but is untested in this release.
37
+
38
+ ## Quick start
39
+
40
+ ### `check`: is the export trustworthy?
41
+
42
+ ```
43
+ $ downshift check tests.models.scatter_include_self_false:make_model
44
+
45
+ ┌───────────────┬─────────────────────────────────────────────────────────────┐
46
+ │ Model │ tests.models.scatter_include_self_false:make_model │
47
+ │ Family │ generic-torch │
48
+ │ Export │ DEGRADED (strict=False, opset 20) │
49
+ │ Numerics │ max abs err 1.34e+00 over 8 samples ✗ 7/8 failed │
50
+ │ Shape-general │ no │
51
+ │ Dynamic dims │ x[0], segment_ids[0] │
52
+ │ Backend │ torch │
53
+ │ Reason │ exported via strict=False but numerics diverge on 7/8 │
54
+ │ │ samples (max abs err 1.34e+00) │
55
+ └───────────────┴─────────────────────────────────────────────────────────────┘
56
+ ```
57
+
58
+ `check` exports in memory, runs `k` random samples (default 8) through both PyTorch and ONNX Runtime, and varies the dynamic axes so some samples have shapes the exporter never saw. Nothing is written to disk.
59
+
60
+ The exit code is the verdict, so it can gate CI: `0` CLEAN, `1` FAILED, `2` DEGRADED, `3` UNVERIFIED. (`4` is a usage error such as an unloadable model; `5` is a crash.) `--json` prints the full verdict as JSON and nothing else:
61
+
62
+ ```bash
63
+ downshift check my_pkg.models:build --json -k 16 > verdict.json
64
+ ```
65
+
66
+ Useful options: `-k/--samples`, `--dynamic "x:0,edge_index:1"` to override which axes are dynamic (default: axis 0 of every input), `--adapter generic|pyg|hf` to skip detection, `--inputs pkg.module:fn` to supply example inputs.
67
+
68
+ ### `export`: write the artifact
69
+
70
+ ```bash
71
+ downshift export my_pkg.models:build -o artifacts/ --name classifier
72
+ ```
73
+
74
+ Writes `artifacts/classifier.onnx` and `artifacts/classifier.manifest.json`. The manifest records the SHA-256 of the artifact and of the source checkpoint (when the model came from a file), torch/onnx/onnxruntime versions, opset, the observed weight dtype, and the full verdict including the numerics report. A DEGRADED export is still written, because the manifest records exactly how far off it is; a FAILED export writes nothing.
75
+
76
+ `--fp16` casts the model to half before export (a plain `.half()`, not quantization). `--no-verify` skips the numerics check and marks the verdict UNVERIFIED, with a warning.
77
+
78
+ ### `serve`: one endpoint, backend chosen by the verdict
79
+
80
+ ```bash
81
+ downshift serve my_pkg.models:build --port 8000
82
+ ```
83
+
84
+ Prints the banner above, then starts uvicorn. Routes:
85
+
86
+ | Route | What it does |
87
+ |---|---|
88
+ | `POST /predict` | Named tensor inputs, any model |
89
+ | `POST /predict/graph` | One graph: `x`, `edge_index`, optional `edge_attr` |
90
+ | `GET /health` | Liveness |
91
+ | `GET /ready` | `200` once warmup is done, `503` before |
92
+ | `GET /metadata` | Family, backend, full verdict, input names |
93
+
94
+ ```bash
95
+ curl -s localhost:8000/predict -H 'content-type: application/json' \
96
+ -d '{"inputs": {"x": [[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5]]}}'
97
+ ```
98
+
99
+ ```json
100
+ {"outputs": {"output_0": [[0.199, -0.206, 0.561, 0.405]]},
101
+ "shapes": {"output_0": [1, 4]},
102
+ "dtypes": {"output_0": "float32"}}
103
+ ```
104
+
105
+ Integer lists become `int64`, everything else `float32`. To be explicit, pass `{"data": [...], "dtype": "float16", "shape": [1, 16]}` instead of a bare list. For graph models:
106
+
107
+ ```bash
108
+ curl -s localhost:8000/predict/graph -H 'content-type: application/json' \
109
+ -d '{"x": [[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
110
+ [0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2],
111
+ [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]],
112
+ "edge_index": [[0, 1, 2], [1, 2, 0]]}'
113
+ ```
114
+
115
+ The response has the same shape as `/predict`, with one output row per node.
116
+
117
+ Options that change what gets served:
118
+
119
+ - `--backend auto|onnxruntime|torch`. `auto` follows the verdict. `torch` skips the export entirely.
120
+ - `--force-onnx` serves a DEGRADED graph through ONNX Runtime anyway. The banner says so in red.
121
+ - `--reference model` verifies a pre-built `.onnx` against a PyTorch model; without it the verdict is UNVERIFIED.
122
+ - `--middleware pkg.module:Attr` (repeatable) attaches a `BaseHTTPMiddleware` subclass or an `async (request, call_next)` function. No middleware means no overhead.
123
+ - `--device auto|cpu|cuda`, `--warmup N` (inferences before `/ready` flips), `--host`, `--port`, `--log-level`, `--log-format json`.
124
+
125
+ ## Accepted model forms
126
+
127
+ | Argument | Meaning |
128
+ |---|---|
129
+ | `model.onnx` | Pre-built ONNX, served as-is. UNVERIFIED unless `--reference` is given. |
130
+ | `pkg.module:attr` | Import spec. `attr` is an `nn.Module` instance or a zero-argument factory. A sibling `make_inputs` in the same module is picked up automatically; otherwise pass `--inputs pkg.module:fn`. |
131
+ | `weights.pt` | State dict. Needs `--model-class pkg.module:Class`. Also `.pth`, `.bin`, `.ckpt`. |
132
+ | `org/repo` | Hugging Face hub id. Needs the `[hf]` extra. |
133
+ | `path/to/repo/dir/` | Locally downloaded Hugging Face repo: a directory containing `config.json`. Needs the `[hf]` extra. |
134
+
135
+ Checkpoints are loaded with `torch.load(weights_only=True)`. A file that holds a pickled full module will not load that way; `--unsafe-load` switches to `weights_only=False`, which means running arbitrary code from the file. Only use it on files you would run as a script.
136
+
137
+ ## The four verdicts
138
+
139
+ - **CLEAN**: exports, matches PyTorch on every sample, survives shapes it was not traced on. Served via ONNX Runtime.
140
+ - **DEGRADED**: exports without error, but numerics drift past tolerance on at least one sample. Served via eager PyTorch; `--force-onnx` overrides.
141
+ - **FAILED**: does not export. Served via eager PyTorch. Not an error, a supported path.
142
+ - **UNVERIFIED**: a `.onnx` with no reference model, or `--no-verify`. Served via ONNX Runtime and labelled as never checked.
143
+
144
+ ## Compatibility matrix
145
+
146
+ Generated by `scripts/gen_matrix.py` from the fixture corpus in `tests/models/`, each fixture isolating one export hazard. CI regenerates it weekly against current torch and onnxruntime and opens a PR when it changes. Full file, with versions and legend: [docs/compatibility.md](docs/compatibility.md).
147
+
148
+ | Model | Hazard | Family | Export | Capture | Numerics | Shape-general | Backend |
149
+ |---|---|---|---|---|---|---|---|
150
+ | `clean_mlp` | Control fixture: no export hazards | generic-torch | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
151
+ | `custom_autograd` | custom autograd.Function with no symbolic override | generic-torch | CLEAN | strict=False | 4.2e-07 | ✓ | onnxruntime |
152
+ | `data_dependent_branch` | data-dependent control flow | generic-torch | FAILED | — | — | — | torch |
153
+ | `dict_input` | dataclass container input | generic-torch | CLEAN | strict=False | 4.8e-07 | ✓ | onnxruntime |
154
+ | `dropout_model` | stochastic layer | generic-torch | CLEAN | strict=False | 3.6e-07 | ✓ | onnxruntime |
155
+ | `dynamic_batch_cnn` | batch-dim generalization | generic-torch | CLEAN | strict=False | 6.0e-08 | ✓ | onnxruntime |
156
+ | `gnn_gat` | GNN fixture: 3-layer GAT node classifier | pyg | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
157
+ | `gnn_gcn` | GNN fixture: 2-layer GCN node classifier | pyg | CLEAN | strict=False | 2.4e-07 | ✓ | onnxruntime |
158
+ | `gnn_sage` | GNN fixture: 2-layer GraphSAGE node classifier | pyg | CLEAN | strict=False | 1.2e-07 | ✓ | onnxruntime |
159
+ | `scatter_include_self_false` | scatter_reduce(include_self=False) has no faithful ONNX translation | generic-torch | DEGRADED | strict=False | 1.3e+00 | ✗ | torch |
160
+ | `tied_weights` | tied embedding/output weight (GPT-2/OPT-style) | generic-torch | CLEAN | strict=False | 1.9e-06 | ✓ | onnxruntime |
161
+ | `tiny_bert` | HF fixture: a randomly initialised two-layer BERT encoder | hf-transformers | CLEAN | strict=False | 6.0e-07 | ✓ | onnxruntime |
162
+
163
+ Two rows worth reading twice. `custom_autograd` was expected to fail and is CLEAN, because `torch.export` traces straight through a `Function.forward` made of ordinary ops. `scatter_include_self_false` was expected to fail loudly and instead exports with zero errors and returns the wrong numbers; the only thing standing between that graph and production is the numerics check.
164
+
165
+ ## What this is not
166
+
167
+ - **No quantization or graph optimization, ever.** Not deferred, cut. Run Olive, `onnxruntime.quantization`, or your own script, then hand the result to `downshift serve model.onnx --reference model.pt` and it gets verified against the original weights like any other export. `--fp16` is a cast before tracing, nothing lower exists here.
168
+ - **No LLM path.** No `onnxruntime-genai` backend, no OpenAI-compatible endpoints, no KV cache, no sampling loop. Encoder-only Hugging Face models work; causal LMs are not a target yet.
169
+ - **No continuous batching, no PagedAttention.** The boot banner is a visual homage to vLLM. That is the full extent of the resemblance.
170
+ - **No dynamic request batching yet.** One request, one inference.
171
+ - **No graph batching yet.** `/predict/graph` takes one graph. Concatenate graphs client-side with offset edge indices if you need more.
172
+ - **No Prometheus metrics, no Docker image.** `--middleware` is the hook for the former; pip plus version pins is the path for the latter.
173
+ - **No DGL adapter yet.** PyG only.
174
+
175
+ **vs. anydeploy.** `anydeploy` also does export, validate, and serve, with a pass/fail validation step and an edge/mobile focus. downshift differs in three places: the verdict is tiered, with DEGRADED as a real middle state between "works" and "crashes"; the eager PyTorch fallback sits behind the same endpoint so a FAILED or DEGRADED model still serves; and GNNs (PyTorch Geometric) are a supported family with independent node and edge dynamic dims.
176
+
177
+ ## Writing your own adapter
178
+
179
+ An adapter knows one model family well enough to build example inputs when the user gave none, and to turn the model plus inputs into something `torch.export` can trace: a module with a flat tensor signature. Implement the `Adapter` protocol from `downshift.adapters.base`:
180
+
181
+ ```python
182
+ from downshift.adapters.base import Prepared
183
+
184
+ class MyAdapter:
185
+ name = "myfamily"
186
+ family = "myfamily"
187
+
188
+ def matches(self, model, example_inputs) -> bool: ...
189
+ def example_inputs(self, model) -> tuple | None: ... # None if you can't guess
190
+ def prepare(self, model, example_inputs) -> Prepared: ...
191
+
192
+ ADAPTER = MyAdapter()
193
+ ```
194
+
195
+ `Prepared` carries the export-ready module, the flat example inputs, their names, the per-input `dynamic_shapes` spec, an optional `vary_fn(i) -> inputs` that generates verification samples, and the family string. Register it under the `downshift.adapters` entry-point group in your own package:
196
+
197
+ ```toml
198
+ [project.entry-points."downshift.adapters"]
199
+ myfamily = "my_pkg.adapter:ADAPTER"
200
+ ```
201
+
202
+ Adapters are tried most-specific first; `generic` always goes last. An adapter whose optional dependency is missing is skipped silently.
203
+
204
+ For a one-off adapter that isn't worth packaging, `--adapter` (and `check()`'s `adapter=`) also accepts a bare `.py` file directly, no install or entry point required:
205
+
206
+ ```bash
207
+ downshift check my_model.py:model --adapter path/to/pointcloud_adapter.py
208
+ ```
209
+
210
+ The file needs a module-level `ADAPTER = MyAdapter()`, or point at the class directly with `--adapter path/to/pointcloud_adapter.py:MyAdapter` and it's instantiated with no arguments.
211
+
212
+ ## Development
213
+
214
+ ```bash
215
+ pip install -e ".[dev,all]"
216
+ ruff check .
217
+ mypy src
218
+ pytest
219
+ python scripts/gen_matrix.py # regenerates docs/compatibility.md
220
+ ```
221
+
222
+ ## License
223
+
224
+ MIT.