snoozeml 0.1.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 (41) hide show
  1. snoozeml-0.1.0/PKG-INFO +222 -0
  2. snoozeml-0.1.0/README.md +157 -0
  3. snoozeml-0.1.0/pyproject.toml +114 -0
  4. snoozeml-0.1.0/setup.cfg +4 -0
  5. snoozeml-0.1.0/snoozeml/__init__.py +140 -0
  6. snoozeml-0.1.0/snoozeml/_proto/__init__.py +25 -0
  7. snoozeml-0.1.0/snoozeml/_proto/snoozeml/__init__.py +1 -0
  8. snoozeml-0.1.0/snoozeml/_proto/snoozeml/v1/__init__.py +1 -0
  9. snoozeml-0.1.0/snoozeml/_proto/snoozeml/v1/inference_pb2.py +51 -0
  10. snoozeml-0.1.0/snoozeml/_proto/snoozeml/v1/inference_pb2.pyi +68 -0
  11. snoozeml-0.1.0/snoozeml/_proto/snoozeml/v1/inference_pb2_grpc.py +85 -0
  12. snoozeml-0.1.0/snoozeml/auth.py +313 -0
  13. snoozeml-0.1.0/snoozeml/client.py +1468 -0
  14. snoozeml-0.1.0/snoozeml/config.py +873 -0
  15. snoozeml-0.1.0/snoozeml/converters/__init__.py +11 -0
  16. snoozeml-0.1.0/snoozeml/converters/_onnx_builder.py +289 -0
  17. snoozeml-0.1.0/snoozeml/converters/sklearn.py +542 -0
  18. snoozeml-0.1.0/snoozeml/converters/torch.py +374 -0
  19. snoozeml-0.1.0/snoozeml/converters/xgboost.py +70 -0
  20. snoozeml-0.1.0/snoozeml/exceptions.py +268 -0
  21. snoozeml-0.1.0/snoozeml/inference.py +740 -0
  22. snoozeml-0.1.0/snoozeml/logging.py +306 -0
  23. snoozeml-0.1.0/snoozeml/metrics.py +182 -0
  24. snoozeml-0.1.0/snoozeml/models.py +871 -0
  25. snoozeml-0.1.0/snoozeml/py.typed +1 -0
  26. snoozeml-0.1.0/snoozeml/streams.py +771 -0
  27. snoozeml-0.1.0/snoozeml/worker.py +389 -0
  28. snoozeml-0.1.0/snoozeml.egg-info/PKG-INFO +222 -0
  29. snoozeml-0.1.0/snoozeml.egg-info/SOURCES.txt +39 -0
  30. snoozeml-0.1.0/snoozeml.egg-info/dependency_links.txt +1 -0
  31. snoozeml-0.1.0/snoozeml.egg-info/requires.txt +66 -0
  32. snoozeml-0.1.0/snoozeml.egg-info/top_level.txt +1 -0
  33. snoozeml-0.1.0/tests/test_client.py +94 -0
  34. snoozeml-0.1.0/tests/test_converters.py +217 -0
  35. snoozeml-0.1.0/tests/test_exceptions.py +113 -0
  36. snoozeml-0.1.0/tests/test_inference.py +149 -0
  37. snoozeml-0.1.0/tests/test_integration_e2e.py +455 -0
  38. snoozeml-0.1.0/tests/test_logging.py +88 -0
  39. snoozeml-0.1.0/tests/test_models.py +44 -0
  40. snoozeml-0.1.0/tests/test_sdk_extensions.py +549 -0
  41. snoozeml-0.1.0/tests/test_worker.py +208 -0
@@ -0,0 +1,222 @@
1
+ Metadata-Version: 2.4
2
+ Name: snoozeml
3
+ Version: 0.1.0
4
+ Summary: Official Python client and ML model conversion toolkit for the SnoozeML inference daemon
5
+ Author: SnoozeML Contributors
6
+ License: Apache-2.0
7
+ Keywords: machine-learning,inference,onnx,grpc,xgboost,rust,snoozeml
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: License :: OSI Approved :: Apache Software License
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.9
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
18
+ Requires-Python: >=3.9
19
+ Description-Content-Type: text/markdown
20
+ Requires-Dist: requests>=2.28.0
21
+ Provides-Extra: client
22
+ Requires-Dist: requests>=2.28.0; extra == "client"
23
+ Provides-Extra: grpc
24
+ Requires-Dist: grpcio>=1.60.0; extra == "grpc"
25
+ Requires-Dist: grpcio-health-checking>=1.60.0; extra == "grpc"
26
+ Requires-Dist: protobuf>=4.25.0; extra == "grpc"
27
+ Requires-Dist: numpy>=1.24.0; extra == "grpc"
28
+ Provides-Extra: worker
29
+ Requires-Dist: numpy>=1.24.0; extra == "worker"
30
+ Provides-Extra: onnx
31
+ Requires-Dist: numpy>=1.24.0; extra == "onnx"
32
+ Requires-Dist: scikit-learn>=1.2.0; extra == "onnx"
33
+ Requires-Dist: onnx>=1.15.0; extra == "onnx"
34
+ Requires-Dist: skl2onnx>=1.16.0; extra == "onnx"
35
+ Provides-Extra: xgboost
36
+ Requires-Dist: numpy>=1.24.0; extra == "xgboost"
37
+ Requires-Dist: xgboost-cpu>=2.0.0; sys_platform == "linux" and extra == "xgboost"
38
+ Requires-Dist: xgboost>=2.0.0; sys_platform != "linux" and extra == "xgboost"
39
+ Provides-Extra: torch
40
+ Requires-Dist: torch>=2.0.0; extra == "torch"
41
+ Provides-Extra: converters
42
+ Requires-Dist: numpy>=1.24.0; extra == "converters"
43
+ Requires-Dist: scikit-learn>=1.2.0; extra == "converters"
44
+ Requires-Dist: xgboost-cpu>=2.0.0; sys_platform == "linux" and extra == "converters"
45
+ Requires-Dist: xgboost>=2.0.0; sys_platform != "linux" and extra == "converters"
46
+ Requires-Dist: onnx>=1.15.0; extra == "converters"
47
+ Requires-Dist: skl2onnx>=1.16.0; extra == "converters"
48
+ Provides-Extra: full
49
+ Requires-Dist: grpcio>=1.60.0; extra == "full"
50
+ Requires-Dist: grpcio-health-checking>=1.60.0; extra == "full"
51
+ Requires-Dist: protobuf>=4.25.0; extra == "full"
52
+ Requires-Dist: numpy>=1.24.0; extra == "full"
53
+ Requires-Dist: scikit-learn>=1.2.0; extra == "full"
54
+ Requires-Dist: xgboost-cpu>=2.0.0; sys_platform == "linux" and extra == "full"
55
+ Requires-Dist: xgboost>=2.0.0; sys_platform != "linux" and extra == "full"
56
+ Requires-Dist: onnx>=1.15.0; extra == "full"
57
+ Requires-Dist: skl2onnx>=1.16.0; extra == "full"
58
+ Provides-Extra: dev
59
+ Requires-Dist: pytest>=8.0.0; extra == "dev"
60
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == "dev"
61
+ Requires-Dist: mypy>=1.9.0; extra == "dev"
62
+ Requires-Dist: ruff>=0.3.0; extra == "dev"
63
+ Requires-Dist: sphinx>=7.2.0; extra == "dev"
64
+ Requires-Dist: shibuya>=2024.8.0; extra == "dev"
65
+
66
+ # `snoozeml` — Official Python SDK for SnoozeML
67
+
68
+ `snoozeml` is the official Python client library and model export toolkit for the **SnoozeML** bare-metal Rust inference daemon.
69
+
70
+ It provides:
71
+ 1. **Dual-Transport Inference (`SnoozeClient` & `AsyncSnoozeClient`)**: Execute unary tensor inference over HTTP REST (JSON) or high-throughput native gRPC (raw little-endian `FP32` byte buffers) with automatic NumPy `ndarray` marshaling.
72
+ 2. **Model Lifecycle & Scale-to-Zero Control (`client.models`)**: List registered models, inspect live states (`idle`, `loading`, `ready`, `busy`), eagerly preload weights (`load`), pin critical models in memory (`pin` / `unpin`), and explicitly evict/unload models (`unload` / `evict`).
73
+ 3. **Built-in Model Converters (`snoozeml.converters`)**: Export scikit-learn estimators & `Pipeline`s (`export_sklearn`), PyTorch `nn.Module`s (`export_torch`), and XGBoost boosters (`export_xgboost`) directly into SnoozeML-ready `.onnx` and `.xgb.json` artifacts.
74
+ 4. **Structured Logging & Credential Redaction (`snoozeml.logging`)**: Built-in `RedactingFilter` and `StructuredFormatter` ensure API keys and JWT Bearer tokens are never leaked in logs or `__repr__` outputs.
75
+
76
+ ---
77
+
78
+ ## Installation
79
+
80
+ Install `snoozeml` directly from PyPI:
81
+
82
+ ```bash
83
+ # Lightweight HTTP REST client (zero compiled C-extensions, pure Python)
84
+ pip install snoozeml
85
+
86
+ # High-performance native gRPC inference + NumPy array support
87
+ pip install "snoozeml[grpc]"
88
+
89
+ # Complete toolkit: gRPC client, custom worker, and all model converters
90
+ pip install "snoozeml[full]"
91
+ ```
92
+
93
+ ### Modular Dependency Extras
94
+
95
+ | Extra | Command | Included Packages | Purpose |
96
+ | :--- | :--- | :--- | :--- |
97
+ | *(None / Base)* | `pip install snoozeml` | `requests` | Ultra-lightweight HTTP client for model lifecycle management (`load`, `unload`, `pin`, `unpin`), status polling, health probes, and metrics. Zero binary C-extensions. |
98
+ | `[grpc]` | `pip install "snoozeml[grpc]"` | `grpcio`, `protobuf`, `numpy` | Native binary gRPC transport with automatic NumPy array serialization and low-latency inference. |
99
+ | `[converters]` | `pip install "snoozeml[converters]"` | `numpy`, `scikit-learn`, `onnx`, `skl2onnx`, `xgboost[-cpu]` | Offline model export toolkit (`export_sklearn`, `export_xgboost`) targeting ONNX and pure-Rust tree formats. |
100
+ | `[full]` | `pip install "snoozeml[full]"` | `grpc`, `converters`, `worker` | Full-featured installation for client, gRPC, custom workers, and model converters with zero CUDA bloat. |
101
+ | `[onnx]` | `pip install "snoozeml[onnx]"` | `numpy`, `scikit-learn`, `onnx`, `skl2onnx` | Scikit-learn estimator and pipeline export to ONNX runtime format. |
102
+ | `[xgboost]` | `pip install "snoozeml[xgboost]"` | `numpy`, `xgboost-cpu` (Linux) / `xgboost` (macOS/Win) | XGBoost model export to `.xgb.json` for pure-Rust `hessboost` tree execution. Uses `xgboost-cpu` on Linux to avoid pulling heavy NVIDIA CUDA packages. |
103
+ | `[torch]` | `pip install "snoozeml[torch]"` | `torch` | PyTorch `nn.Module` export to ONNX. *(Note: If PyTorch is already installed in your training environment, `snoozeml` automatically duck-types PyTorch tensors without needing this extra).* |
104
+ | `[worker]` | `pip install "snoozeml[worker]"` | `numpy` | Minimal dependency profile for process-isolated custom Python worker models (`snooze_model.py`). |
105
+ | `[dev]` | `pip install "snoozeml[dev]"` | `pytest`, `mypy`, `ruff`, `sphinx` | Development, linting, testing, and documentation building. |
106
+
107
+ > [!NOTE]
108
+ > For editable local development from the repository root:
109
+ > ```bash
110
+ > pip install -e "./python[full,dev]"
111
+ > ```
112
+
113
+ ---
114
+
115
+ ## Quickstart
116
+
117
+ ### 1. Model Lifecycle Management
118
+
119
+ ```python
120
+ from snoozeml import SnoozeClient
121
+
122
+ client = SnoozeClient(endpoint="http://localhost:8080", api_key="sk-admin-key")
123
+
124
+ # List all registered models
125
+ models = client.models.list()
126
+ for m in models:
127
+ print(f"{m.id}: state={m.state.value}, pinned={m.pinned}, loaded={m.is_loaded}")
128
+
129
+ # Inspect a specific model
130
+ info = client.models.get("demo-model")
131
+ print(f"Path: {info.onnx_path}, Idle Timeout: {info.idle_timeout_secs}s")
132
+
133
+ # Pin model in memory (exempts from scale-to-zero and LRU eviction)
134
+ client.models.pin("demo-model")
135
+
136
+ # Unpin and explicitly evict (scale-to-zero)
137
+ client.models.unpin("demo-model")
138
+ client.models.evict("demo-model")
139
+ ```
140
+
141
+ ### 2. Synchronous & Asynchronous Inference (HTTP & gRPC)
142
+
143
+ ```python
144
+ import numpy as np
145
+ from snoozeml import SnoozeClient, AsyncSnoozeClient
146
+
147
+ # High-throughput gRPC inference
148
+ with SnoozeClient(endpoint="localhost:50051", use_grpc=True) as client:
149
+ x = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
150
+ response = client.infer(model_id="demo-model", inputs={"input": x})
151
+
152
+ output = response.get_numpy("output")
153
+ print(f"Output shape: {output.shape}, Server latency: {response.duration_ms:.2f} ms")
154
+ ```
155
+
156
+ Asynchronous usage with `AsyncSnoozeClient`:
157
+
158
+ ```python
159
+ import asyncio
160
+ import numpy as np
161
+ from snoozeml import AsyncSnoozeClient
162
+
163
+ async def main():
164
+ async with AsyncSnoozeClient(endpoint="localhost:50051", use_grpc=True) as client:
165
+ x = np.ones((2, 4), dtype=np.float32)
166
+ resp = await client.infer("demo-model", inputs={"input": x})
167
+ print(resp["output"])
168
+
169
+ asyncio.run(main())
170
+ ```
171
+
172
+ ### 3. Exporting scikit-learn, PyTorch & XGBoost Models
173
+
174
+ ```python
175
+ import numpy as np
176
+ from sklearn.ensemble import RandomForestClassifier
177
+ from sklearn.pipeline import make_pipeline
178
+ from sklearn.preprocessing import StandardScaler
179
+ from snoozeml.converters import export_sklearn, export_torch, export_xgboost
180
+
181
+ # 1. Export a scikit-learn Pipeline to ONNX
182
+ X = np.random.randn(100, 4).astype(np.float32)
183
+ y = (X[:, 0] + X[:, 1] > 0).astype(int)
184
+ pipe = make_pipeline(StandardScaler(), RandomForestClassifier(n_estimators=10, random_state=42))
185
+ pipe.fit(X, y)
186
+
187
+ onnx_path = export_sklearn(
188
+ model=pipe,
189
+ initial_types=[("float_input", (None, 4))],
190
+ output_path="models/rf_pipeline.onnx",
191
+ )
192
+
193
+ # 2. Export a PyTorch nn.Module to ONNX
194
+ import torch
195
+ import torch.nn as nn
196
+
197
+ net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2)).eval()
198
+ export_torch(
199
+ model=net,
200
+ sample_input=torch.randn(1, 4),
201
+ output_path="models/mlp_torch.onnx",
202
+ )
203
+ ```
204
+
205
+ ### 4. Structured Logging
206
+
207
+ ```python
208
+ from snoozeml import configure_logging
209
+
210
+ # Enable human-readable or JSON structured logs (with automatic secret redaction)
211
+ configure_logging(level="DEBUG", format_json=False)
212
+ ```
213
+
214
+ ---
215
+
216
+ ## Interactive Examples & Jupyter Notebooks
217
+
218
+ Check out the interactive tutorial notebooks under [`examples/`](../examples):
219
+ * [**`01_sklearn_onnx.ipynb`**](../examples/01_sklearn_onnx.ipynb): Train a Scikit-learn Random Forest classifier on Iris, export to ONNX, deploy to SnoozeML Docker, and run batched inference.
220
+ * [**`02_xgboost_tree.ipynb`**](../examples/02_xgboost_tree.ipynb): Train an XGBoost classifier on Breast Cancer, export to native `.xgb.json`, and run inference via SnoozeML's pure-Rust tree backend.
221
+ * [**`03_pipeline_regression_onnx.ipynb`**](../examples/03_pipeline_regression_onnx.ipynb): Train a Scikit-learn `Pipeline` (StandardScaler + Ridge) on Diabetes and infer continuous targets.
222
+
@@ -0,0 +1,157 @@
1
+ # `snoozeml` — Official Python SDK for SnoozeML
2
+
3
+ `snoozeml` is the official Python client library and model export toolkit for the **SnoozeML** bare-metal Rust inference daemon.
4
+
5
+ It provides:
6
+ 1. **Dual-Transport Inference (`SnoozeClient` & `AsyncSnoozeClient`)**: Execute unary tensor inference over HTTP REST (JSON) or high-throughput native gRPC (raw little-endian `FP32` byte buffers) with automatic NumPy `ndarray` marshaling.
7
+ 2. **Model Lifecycle & Scale-to-Zero Control (`client.models`)**: List registered models, inspect live states (`idle`, `loading`, `ready`, `busy`), eagerly preload weights (`load`), pin critical models in memory (`pin` / `unpin`), and explicitly evict/unload models (`unload` / `evict`).
8
+ 3. **Built-in Model Converters (`snoozeml.converters`)**: Export scikit-learn estimators & `Pipeline`s (`export_sklearn`), PyTorch `nn.Module`s (`export_torch`), and XGBoost boosters (`export_xgboost`) directly into SnoozeML-ready `.onnx` and `.xgb.json` artifacts.
9
+ 4. **Structured Logging & Credential Redaction (`snoozeml.logging`)**: Built-in `RedactingFilter` and `StructuredFormatter` ensure API keys and JWT Bearer tokens are never leaked in logs or `__repr__` outputs.
10
+
11
+ ---
12
+
13
+ ## Installation
14
+
15
+ Install `snoozeml` directly from PyPI:
16
+
17
+ ```bash
18
+ # Lightweight HTTP REST client (zero compiled C-extensions, pure Python)
19
+ pip install snoozeml
20
+
21
+ # High-performance native gRPC inference + NumPy array support
22
+ pip install "snoozeml[grpc]"
23
+
24
+ # Complete toolkit: gRPC client, custom worker, and all model converters
25
+ pip install "snoozeml[full]"
26
+ ```
27
+
28
+ ### Modular Dependency Extras
29
+
30
+ | Extra | Command | Included Packages | Purpose |
31
+ | :--- | :--- | :--- | :--- |
32
+ | *(None / Base)* | `pip install snoozeml` | `requests` | Ultra-lightweight HTTP client for model lifecycle management (`load`, `unload`, `pin`, `unpin`), status polling, health probes, and metrics. Zero binary C-extensions. |
33
+ | `[grpc]` | `pip install "snoozeml[grpc]"` | `grpcio`, `protobuf`, `numpy` | Native binary gRPC transport with automatic NumPy array serialization and low-latency inference. |
34
+ | `[converters]` | `pip install "snoozeml[converters]"` | `numpy`, `scikit-learn`, `onnx`, `skl2onnx`, `xgboost[-cpu]` | Offline model export toolkit (`export_sklearn`, `export_xgboost`) targeting ONNX and pure-Rust tree formats. |
35
+ | `[full]` | `pip install "snoozeml[full]"` | `grpc`, `converters`, `worker` | Full-featured installation for client, gRPC, custom workers, and model converters with zero CUDA bloat. |
36
+ | `[onnx]` | `pip install "snoozeml[onnx]"` | `numpy`, `scikit-learn`, `onnx`, `skl2onnx` | Scikit-learn estimator and pipeline export to ONNX runtime format. |
37
+ | `[xgboost]` | `pip install "snoozeml[xgboost]"` | `numpy`, `xgboost-cpu` (Linux) / `xgboost` (macOS/Win) | XGBoost model export to `.xgb.json` for pure-Rust `hessboost` tree execution. Uses `xgboost-cpu` on Linux to avoid pulling heavy NVIDIA CUDA packages. |
38
+ | `[torch]` | `pip install "snoozeml[torch]"` | `torch` | PyTorch `nn.Module` export to ONNX. *(Note: If PyTorch is already installed in your training environment, `snoozeml` automatically duck-types PyTorch tensors without needing this extra).* |
39
+ | `[worker]` | `pip install "snoozeml[worker]"` | `numpy` | Minimal dependency profile for process-isolated custom Python worker models (`snooze_model.py`). |
40
+ | `[dev]` | `pip install "snoozeml[dev]"` | `pytest`, `mypy`, `ruff`, `sphinx` | Development, linting, testing, and documentation building. |
41
+
42
+ > [!NOTE]
43
+ > For editable local development from the repository root:
44
+ > ```bash
45
+ > pip install -e "./python[full,dev]"
46
+ > ```
47
+
48
+ ---
49
+
50
+ ## Quickstart
51
+
52
+ ### 1. Model Lifecycle Management
53
+
54
+ ```python
55
+ from snoozeml import SnoozeClient
56
+
57
+ client = SnoozeClient(endpoint="http://localhost:8080", api_key="sk-admin-key")
58
+
59
+ # List all registered models
60
+ models = client.models.list()
61
+ for m in models:
62
+ print(f"{m.id}: state={m.state.value}, pinned={m.pinned}, loaded={m.is_loaded}")
63
+
64
+ # Inspect a specific model
65
+ info = client.models.get("demo-model")
66
+ print(f"Path: {info.onnx_path}, Idle Timeout: {info.idle_timeout_secs}s")
67
+
68
+ # Pin model in memory (exempts from scale-to-zero and LRU eviction)
69
+ client.models.pin("demo-model")
70
+
71
+ # Unpin and explicitly evict (scale-to-zero)
72
+ client.models.unpin("demo-model")
73
+ client.models.evict("demo-model")
74
+ ```
75
+
76
+ ### 2. Synchronous & Asynchronous Inference (HTTP & gRPC)
77
+
78
+ ```python
79
+ import numpy as np
80
+ from snoozeml import SnoozeClient, AsyncSnoozeClient
81
+
82
+ # High-throughput gRPC inference
83
+ with SnoozeClient(endpoint="localhost:50051", use_grpc=True) as client:
84
+ x = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
85
+ response = client.infer(model_id="demo-model", inputs={"input": x})
86
+
87
+ output = response.get_numpy("output")
88
+ print(f"Output shape: {output.shape}, Server latency: {response.duration_ms:.2f} ms")
89
+ ```
90
+
91
+ Asynchronous usage with `AsyncSnoozeClient`:
92
+
93
+ ```python
94
+ import asyncio
95
+ import numpy as np
96
+ from snoozeml import AsyncSnoozeClient
97
+
98
+ async def main():
99
+ async with AsyncSnoozeClient(endpoint="localhost:50051", use_grpc=True) as client:
100
+ x = np.ones((2, 4), dtype=np.float32)
101
+ resp = await client.infer("demo-model", inputs={"input": x})
102
+ print(resp["output"])
103
+
104
+ asyncio.run(main())
105
+ ```
106
+
107
+ ### 3. Exporting scikit-learn, PyTorch & XGBoost Models
108
+
109
+ ```python
110
+ import numpy as np
111
+ from sklearn.ensemble import RandomForestClassifier
112
+ from sklearn.pipeline import make_pipeline
113
+ from sklearn.preprocessing import StandardScaler
114
+ from snoozeml.converters import export_sklearn, export_torch, export_xgboost
115
+
116
+ # 1. Export a scikit-learn Pipeline to ONNX
117
+ X = np.random.randn(100, 4).astype(np.float32)
118
+ y = (X[:, 0] + X[:, 1] > 0).astype(int)
119
+ pipe = make_pipeline(StandardScaler(), RandomForestClassifier(n_estimators=10, random_state=42))
120
+ pipe.fit(X, y)
121
+
122
+ onnx_path = export_sklearn(
123
+ model=pipe,
124
+ initial_types=[("float_input", (None, 4))],
125
+ output_path="models/rf_pipeline.onnx",
126
+ )
127
+
128
+ # 2. Export a PyTorch nn.Module to ONNX
129
+ import torch
130
+ import torch.nn as nn
131
+
132
+ net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2)).eval()
133
+ export_torch(
134
+ model=net,
135
+ sample_input=torch.randn(1, 4),
136
+ output_path="models/mlp_torch.onnx",
137
+ )
138
+ ```
139
+
140
+ ### 4. Structured Logging
141
+
142
+ ```python
143
+ from snoozeml import configure_logging
144
+
145
+ # Enable human-readable or JSON structured logs (with automatic secret redaction)
146
+ configure_logging(level="DEBUG", format_json=False)
147
+ ```
148
+
149
+ ---
150
+
151
+ ## Interactive Examples & Jupyter Notebooks
152
+
153
+ Check out the interactive tutorial notebooks under [`examples/`](../examples):
154
+ * [**`01_sklearn_onnx.ipynb`**](../examples/01_sklearn_onnx.ipynb): Train a Scikit-learn Random Forest classifier on Iris, export to ONNX, deploy to SnoozeML Docker, and run batched inference.
155
+ * [**`02_xgboost_tree.ipynb`**](../examples/02_xgboost_tree.ipynb): Train an XGBoost classifier on Breast Cancer, export to native `.xgb.json`, and run inference via SnoozeML's pure-Rust tree backend.
156
+ * [**`03_pipeline_regression_onnx.ipynb`**](../examples/03_pipeline_regression_onnx.ipynb): Train a Scikit-learn `Pipeline` (StandardScaler + Ridge) on Diabetes and infer continuous targets.
157
+
@@ -0,0 +1,114 @@
1
+ [build-system]
2
+ requires = ["setuptools>=65.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "snoozeml"
7
+ version = "0.1.0"
8
+ description = "Official Python client and ML model conversion toolkit for the SnoozeML inference daemon"
9
+ readme = "README.md"
10
+ license = { text = "Apache-2.0" }
11
+ requires-python = ">=3.9"
12
+ authors = [
13
+ { name = "SnoozeML Contributors" }
14
+ ]
15
+ keywords = ["machine-learning", "inference", "onnx", "grpc", "xgboost", "rust", "snoozeml"]
16
+ classifiers = [
17
+ "Development Status :: 4 - Beta",
18
+ "Intended Audience :: Developers",
19
+ "Intended Audience :: Science/Research",
20
+ "License :: OSI Approved :: Apache Software License",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.9",
23
+ "Programming Language :: Python :: 3.10",
24
+ "Programming Language :: Python :: 3.11",
25
+ "Programming Language :: Python :: 3.12",
26
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
27
+ ]
28
+ dependencies = [
29
+ # Lightweight client core: used to upload models and call admin endpoints via HTTP
30
+ "requests>=2.28.0",
31
+ ]
32
+
33
+ [project.optional-dependencies]
34
+ # Explicit alias for the lightweight management client
35
+ client = [
36
+ "requests>=2.28.0",
37
+ ]
38
+
39
+ # Native high-performance gRPC inference
40
+ grpc = [
41
+ "grpcio>=1.60.0",
42
+ "grpcio-health-checking>=1.60.0",
43
+ "protobuf>=4.25.0",
44
+ "numpy>=1.24.0",
45
+ ]
46
+
47
+ # Process-isolated custom model worker runtime (numpy tensor formatting)
48
+ worker = [
49
+ "numpy>=1.24.0",
50
+ ]
51
+
52
+ # Scikit-Learn & ONNX conversion utilities
53
+ onnx = [
54
+ "numpy>=1.24.0",
55
+ "scikit-learn>=1.2.0",
56
+ "onnx>=1.15.0",
57
+ "skl2onnx>=1.16.0",
58
+ ]
59
+
60
+ # CPU-only XGBoost conversion utilities (avoids pulling nvidia-nccl on Linux)
61
+ xgboost = [
62
+ "numpy>=1.24.0",
63
+ "xgboost-cpu>=2.0.0; sys_platform == 'linux'",
64
+ "xgboost>=2.0.0; sys_platform != 'linux'",
65
+ ]
66
+
67
+ # Optional PyTorch exporter support (standard PyPI wheel includes CUDA; for CPU-only, use --index-url https://download.pytorch.org/whl/cpu)
68
+ torch = [
69
+ "torch>=2.0.0",
70
+ ]
71
+
72
+ # Offline model conversion toolkit (Scikit-Learn, XGBoost -> ONNX/JSON; PyTorch supported via duck-typing/lazy import)
73
+ converters = [
74
+ "numpy>=1.24.0",
75
+ "scikit-learn>=1.2.0",
76
+ "xgboost-cpu>=2.0.0; sys_platform == 'linux'",
77
+ "xgboost>=2.0.0; sys_platform != 'linux'",
78
+ "onnx>=1.15.0",
79
+ "skl2onnx>=1.16.0",
80
+ ]
81
+
82
+ # Full alias providing all client, gRPC, worker, and converter capabilities (zero CUDA bloat)
83
+ full = [
84
+ "grpcio>=1.60.0",
85
+ "grpcio-health-checking>=1.60.0",
86
+ "protobuf>=4.25.0",
87
+ "numpy>=1.24.0",
88
+ "scikit-learn>=1.2.0",
89
+ "xgboost-cpu>=2.0.0; sys_platform == 'linux'",
90
+ "xgboost>=2.0.0; sys_platform != 'linux'",
91
+ "onnx>=1.15.0",
92
+ "skl2onnx>=1.16.0",
93
+ ]
94
+
95
+ # Development, linting, testing, and documentation dependencies
96
+ dev = [
97
+ "pytest>=8.0.0",
98
+ "pytest-asyncio>=0.23.0",
99
+ "mypy>=1.9.0",
100
+ "ruff>=0.3.0",
101
+ "sphinx>=7.2.0",
102
+ "shibuya>=2024.8.0",
103
+ ]
104
+
105
+ [tool.setuptools.packages.find]
106
+ where = ["."]
107
+ include = ["snoozeml*"]
108
+
109
+ [tool.setuptools.package-data]
110
+ snoozeml = ["py.typed", "_proto/snoozeml/v1/*.pyi"]
111
+
112
+ [tool.pytest.ini_options]
113
+ testpaths = ["tests"]
114
+ addopts = "-v"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,140 @@
1
+ """Official Python SDK and ML conversion toolkit for the SnoozeML inference daemon."""
2
+
3
+ from .auth import (
4
+ JwtAlgorithm,
5
+ Role,
6
+ create_jwt_token,
7
+ decode_jwt_token,
8
+ mint_jwt_token,
9
+ )
10
+ from .config import (
11
+ ApiKeyEntry,
12
+ AuthSettings,
13
+ GlobalDefaultsConfig,
14
+ JwtSettings,
15
+ ModelConfigEntry,
16
+ SnoozeConfig,
17
+ SnoozeDaemon,
18
+ StreamBinding,
19
+ find_snooze_binary,
20
+ scaffold_worker,
21
+ )
22
+ from .exceptions import (
23
+ AuthenticationError,
24
+ ConversionError,
25
+ InferenceTimeoutError,
26
+ InvalidInputError,
27
+ ModelLoadingError,
28
+ ModelNotFoundError,
29
+ OutOfMemoryError,
30
+ PermissionDeniedError,
31
+ ResourceExhaustedError,
32
+ RuntimeInferenceError,
33
+ ShapeMismatchError,
34
+ SnoozeConnectionError,
35
+ SnoozeError,
36
+ UnsupportedArchitectureError,
37
+ )
38
+ from .inference import (
39
+ InferenceResponse,
40
+ OutputList,
41
+ TensorDataType,
42
+ TensorInput,
43
+ TensorOutput,
44
+ concat_inference_responses,
45
+ normalize_inputs,
46
+ slice_inputs_batch,
47
+ )
48
+ from .logging import (
49
+ RedactingFilter,
50
+ StructuredFormatter,
51
+ configure_logging,
52
+ get_logger,
53
+ mask_secret,
54
+ )
55
+ from .metrics import DaemonMetrics, parse_prometheus_metrics
56
+ from .models import AsyncModelsManager, ModelInfo, ModelsManager, ModelState
57
+ from .streams import (
58
+ MockStreamClient,
59
+ RedisStreamClient,
60
+ StreamInferRequest,
61
+ StreamInferResponse,
62
+ StreamWireFormat,
63
+ )
64
+ from .worker import RawTensor, SnoozeModel, WorkerError
65
+ from .client import AsyncSnoozeClient, SnoozeClient
66
+
67
+ # Alias for convenience and compatibility with design documentation
68
+ Client = SnoozeClient
69
+ AsyncClient = AsyncSnoozeClient
70
+
71
+ try:
72
+ from .converters import export_sklearn, export_torch, export_xgboost
73
+ except ImportError:
74
+ pass
75
+
76
+ __version__ = "0.1.0"
77
+
78
+ __all__ = [
79
+ "__version__",
80
+ "SnoozeClient",
81
+ "AsyncSnoozeClient",
82
+ "Client",
83
+ "AsyncClient",
84
+ "Role",
85
+ "JwtAlgorithm",
86
+ "create_jwt_token",
87
+ "mint_jwt_token",
88
+ "decode_jwt_token",
89
+ "SnoozeConfig",
90
+ "GlobalDefaultsConfig",
91
+ "ModelConfigEntry",
92
+ "StreamBinding",
93
+ "ApiKeyEntry",
94
+ "JwtSettings",
95
+ "AuthSettings",
96
+ "SnoozeDaemon",
97
+ "find_snooze_binary",
98
+ "scaffold_worker",
99
+ "DaemonMetrics",
100
+ "parse_prometheus_metrics",
101
+ "StreamWireFormat",
102
+ "StreamInferRequest",
103
+ "StreamInferResponse",
104
+ "MockStreamClient",
105
+ "RedisStreamClient",
106
+ "ModelInfo",
107
+ "ModelState",
108
+ "ModelsManager",
109
+ "AsyncModelsManager",
110
+ "InferenceResponse",
111
+ "OutputList",
112
+ "TensorInput",
113
+ "TensorOutput",
114
+ "TensorDataType",
115
+ "normalize_inputs",
116
+ "slice_inputs_batch",
117
+ "concat_inference_responses",
118
+ "SnoozeError",
119
+ "ModelNotFoundError",
120
+ "InvalidInputError",
121
+ "ShapeMismatchError",
122
+ "InferenceTimeoutError",
123
+ "ModelLoadingError",
124
+ "ResourceExhaustedError",
125
+ "OutOfMemoryError",
126
+ "UnsupportedArchitectureError",
127
+ "AuthenticationError",
128
+ "PermissionDeniedError",
129
+ "SnoozeConnectionError",
130
+ "RuntimeInferenceError",
131
+ "ConversionError",
132
+ "export_sklearn",
133
+ "export_torch",
134
+ "export_xgboost",
135
+ "configure_logging",
136
+ "get_logger",
137
+ "mask_secret",
138
+ "RedactingFilter",
139
+ "StructuredFormatter",
140
+ ]
@@ -0,0 +1,25 @@
1
+ """Internal generated Protocol Buffer and gRPC stubs for SnoozeML."""
2
+
3
+ from .snoozeml.v1.inference_pb2 import (
4
+ DataType,
5
+ ModelInferRequest,
6
+ ModelInferResponse,
7
+ TensorData,
8
+ )
9
+ from .snoozeml.v1.inference_pb2_grpc import (
10
+ InferenceService,
11
+ InferenceServiceServicer,
12
+ InferenceServiceStub,
13
+ add_InferenceServiceServicer_to_server,
14
+ )
15
+
16
+ __all__ = [
17
+ "DataType",
18
+ "ModelInferRequest",
19
+ "ModelInferResponse",
20
+ "TensorData",
21
+ "InferenceService",
22
+ "InferenceServiceServicer",
23
+ "InferenceServiceStub",
24
+ "add_InferenceServiceServicer_to_server",
25
+ ]
@@ -0,0 +1 @@
1
+ """Generated protobuf namespace package."""
@@ -0,0 +1 @@
1
+ """Generated snoozeml.v1 protobuf and gRPC modules."""