cortexgrid-infer 0.2.6__tar.gz → 0.3.1__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.
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/PKG-INFO +91 -55
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/README.md +88 -53
- cortexgrid_infer-0.3.1/cortexgrid_infer/__init__.py +95 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/core.py +27 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/imaging.py +74 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/importers/__init__.py +10 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/importers/base.py +99 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/importers/huggingface.py +116 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/meshing.py +72 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/models/__init__.py +18 -0
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/anthropic_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/anthropic.py +30 -10
- cortexgrid_infer-0.3.1/cortexgrid_infer/models/base.py +121 -0
- cortexgrid_infer-0.3.1/cortexgrid_infer/models/image2mesh.py +84 -0
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_image_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/text2image.py +27 -4
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_complete_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/text2text.py +15 -2
- cortexgrid_infer-0.3.1/cortexgrid_infer/registry.py +102 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/pyproject.toml +5 -3
- cortexgrid_infer-0.2.6/cortexgrid_infer/__init__.py +0 -69
- cortexgrid_infer-0.2.6/cortexgrid_infer/importing.py +0 -158
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/__init__.py +0 -0
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/anthropic.py +0 -65
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_complete.py +0 -23
- cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_image.py +0 -110
- cortexgrid_infer-0.2.6/cortexgrid_infer/requirements.py +0 -117
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/.gitignore +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/LICENSE +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/compiling.py +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/completion.py +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/device.py +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/py.typed +0 -0
- {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/utils.py +0 -0
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: cortexgrid-infer
|
|
3
|
-
Version: 0.
|
|
4
|
-
Summary:
|
|
3
|
+
Version: 0.3.1
|
|
4
|
+
Summary: Serve apps per task (text-to-text, text-to-image, image-to-mesh) and importers per source (HuggingFace) for running models on cortexgrid
|
|
5
5
|
Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
|
|
6
6
|
Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
|
|
7
7
|
License-Expression: Apache-2.0
|
|
@@ -14,6 +14,7 @@ Requires-Dist: diffusers>=0.38
|
|
|
14
14
|
Requires-Dist: fastapi>=0.133
|
|
15
15
|
Requires-Dist: httpx>=0.27.0
|
|
16
16
|
Requires-Dist: huggingface-hub>=0.20.0
|
|
17
|
+
Requires-Dist: numpy>=1.24
|
|
17
18
|
Requires-Dist: pillow>=10.0
|
|
18
19
|
Requires-Dist: protobuf>=4.0
|
|
19
20
|
Requires-Dist: python-dotenv>=1.0.0
|
|
@@ -32,19 +33,47 @@ needs but cannot work out for itself, plus the clients that talk to what it depl
|
|
|
32
33
|
cortexgrid stores whatever weights it is handed under whatever key it is given, and
|
|
33
34
|
serves them with whatever class it was told to bundle. It deliberately knows nothing
|
|
34
35
|
about where a model came from or what shape it is. This library supplies exactly that
|
|
35
|
-
missing knowledge
|
|
36
|
-
|
|
36
|
+
missing knowledge, and drives no lifecycle of its own: importing, deploying and
|
|
37
|
+
deleting stay your calls against the cortexgrid SDK.
|
|
38
|
+
|
|
39
|
+
Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgrid-infer/blob/main/MIGRATION.md).
|
|
37
40
|
|
|
38
41
|
## Division of labour
|
|
39
42
|
|
|
40
43
|
| | Owns |
|
|
41
44
|
|---|---|
|
|
42
45
|
| **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
|
|
43
|
-
| **cortexgrid-infer** | Per
|
|
46
|
+
| **cortexgrid-infer** | Per source: how to name a model, fetch it and size it. Per task: how to run it and talk to it. Never drives a lifecycle. |
|
|
44
47
|
| **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
|
|
45
48
|
|
|
46
|
-
|
|
47
|
-
|
|
49
|
+
A model is two things, and the library keeps them apart:
|
|
50
|
+
|
|
51
|
+
- **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
|
|
52
|
+
runs, not for where the model came from. It loads the weights, answers the task's
|
|
53
|
+
routes, and names the client that speaks them. One serve app runs every model of
|
|
54
|
+
its task, whichever source staged it.
|
|
55
|
+
- **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
|
|
56
|
+
per source. It names the model in the registry, downloads the weights, and sizes
|
|
57
|
+
them from the source's metadata. It is handed the serve app, so one importer takes
|
|
58
|
+
any model its source holds.
|
|
59
|
+
|
|
60
|
+
| Serve app | Task | Runs | Client |
|
|
61
|
+
|---|---|---|---|
|
|
62
|
+
| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
|
|
63
|
+
| `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
|
|
64
|
+
| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
|
|
65
|
+
| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
|
|
66
|
+
|
|
67
|
+
| Importer | Source |
|
|
68
|
+
|---|---|
|
|
69
|
+
| `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | the HuggingFace Hub |
|
|
70
|
+
|
|
71
|
+
A model hosted elsewhere has no weights, and so no importer: its serve app forwards
|
|
72
|
+
to it, and a `Hosted(model_id, serve_app, **settings)` entry registers it.
|
|
73
|
+
|
|
74
|
+
An importer and a `Hosted` entry are both a `ModelEntry`, which gives you what
|
|
75
|
+
`cortexgrid.import_model`, `cortexgrid.register_model` and `cortexgrid.deploy_model`
|
|
76
|
+
need:
|
|
48
77
|
|
|
49
78
|
| | |
|
|
50
79
|
|---|---|
|
|
@@ -53,19 +82,10 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
|
|
|
53
82
|
| `.requirements()` | what one replica needs to be placed and to run |
|
|
54
83
|
| `.config()` | settings the serve app reads at construction, if it needs any |
|
|
55
84
|
| `.client(url)` | a client that speaks the deployed app's routes |
|
|
56
|
-
| `.source` | *(
|
|
85
|
+
| `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
|
|
57
86
|
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
| `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
|
|
61
|
-
| `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
|
|
62
|
-
| `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
|
|
63
|
-
|
|
64
|
-
Not every model has weights, so `.source` belongs to the HuggingFace importers
|
|
65
|
-
rather than to all of them: one that has weights goes through
|
|
66
|
-
`cortexgrid.import_model`, one that has none through `cortexgrid.register_model`.
|
|
67
|
-
Everything else is common, which is why one client serves all three — every
|
|
68
|
-
completion app speaks the same `/complete` protocol.
|
|
87
|
+
One with weights goes through `cortexgrid.import_model`, one without through
|
|
88
|
+
`cortexgrid.register_model`.
|
|
69
89
|
|
|
70
90
|
## Quick start
|
|
71
91
|
|
|
@@ -75,7 +95,7 @@ import cortexgrid
|
|
|
75
95
|
import cortexgrid_infer as mg
|
|
76
96
|
|
|
77
97
|
cortexgrid.Experiment.init("img-gen")
|
|
78
|
-
imp = mg.
|
|
98
|
+
imp = mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image)
|
|
79
99
|
|
|
80
100
|
# 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
|
|
81
101
|
# your machine, so this runs as a job rather than in-process.
|
|
@@ -130,12 +150,14 @@ serve code if it changed, and tags the run with the model.
|
|
|
130
150
|
`cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
|
|
131
151
|
nothing in the registry can derive it — cortexgrid stores whatever weights it is
|
|
132
152
|
handed, and an importer imports whatever repo id it is handed. `requirements()`
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
153
|
+
splits the estimate along the same line as the rest: the importer reads what the
|
|
154
|
+
source records about the weights — for HuggingFace, safetensors headers where the
|
|
155
|
+
repo has them, weight-file sizes otherwise, both over HTTP with no weights
|
|
156
|
+
downloaded — and the serve app turns that into what a replica needs, since it knows
|
|
157
|
+
the dtype it loads in and what its work takes on top.
|
|
136
158
|
|
|
137
159
|
```python
|
|
138
|
-
>>> mg.
|
|
160
|
+
>>> mg.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", mg.Text2Text).requirements()
|
|
139
161
|
ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
|
|
140
162
|
```
|
|
141
163
|
|
|
@@ -144,14 +166,15 @@ you know better, or edit the figures on the model card in the dashboard afterwar
|
|
|
144
166
|
`import_model` only fills in requirements a version does not already have, so a
|
|
145
167
|
hand-set value is never overwritten.
|
|
146
168
|
|
|
147
|
-
The estimate always asks for exactly one GPU, because
|
|
148
|
-
the whole model onto a single device. A model too large for any one card needs a
|
|
149
|
-
app that shards it, not a larger `num_gpus`: the requirement would simply never
|
|
150
|
-
placed.
|
|
169
|
+
The estimate always asks for exactly one GPU, because every serve app that runs weights
|
|
170
|
+
moves the whole model onto a single device. A model too large for any one card needs a
|
|
171
|
+
serve app that shards it, not a larger `num_gpus`: the requirement would simply never
|
|
172
|
+
be placed. An `Image2Mesh` subclass can set `min_vram_gb` for the memory meshing takes
|
|
173
|
+
beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
|
|
151
174
|
|
|
152
175
|
## Compilation
|
|
153
176
|
|
|
154
|
-
|
|
177
|
+
`Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
|
|
155
178
|
eager PyTorch bills per *operation* and a transformer forward is mostly glue —
|
|
156
179
|
views, broadcasts, permutes, elementwise adds — with only a small fraction of the
|
|
157
180
|
dispatches being the matrix multiplies that do the arithmetic. While the tensors
|
|
@@ -172,12 +195,12 @@ fits the work it does:
|
|
|
172
195
|
|
|
173
196
|
Both serve apps are tuned rather than defaulted:
|
|
174
197
|
|
|
175
|
-
- **`
|
|
198
|
+
- **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
|
|
176
199
|
and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
|
|
177
200
|
every step is the same shapes — so the capture is made once and replayed for the
|
|
178
201
|
whole loop. The VAE is left alone: smallest win of the three, and the one whose
|
|
179
202
|
shapes move most, especially with tiling on.
|
|
180
|
-
- **`
|
|
203
|
+
- **`Text2Text`** does *not* compile the module. Decode has
|
|
181
204
|
no stable shape of its own — the sequence grows a token per forward — so it
|
|
182
205
|
pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
|
|
183
206
|
`CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
|
|
@@ -221,7 +244,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
|
|
|
221
244
|
`complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
|
|
222
245
|
|
|
223
246
|
```python
|
|
224
|
-
imp = mg.
|
|
247
|
+
imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)
|
|
225
248
|
model = imp.client(deployment.url)
|
|
226
249
|
messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
|
|
227
250
|
|
|
@@ -236,17 +259,17 @@ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=
|
|
|
236
259
|
### Anthropic
|
|
237
260
|
|
|
238
261
|
An Anthropic model is a registry entry like any other — the difference is that
|
|
239
|
-
there are no weights to stage, so
|
|
240
|
-
|
|
262
|
+
there are no weights to stage, so there is no importer: the serve app forwards to
|
|
263
|
+
the API instead of loading anything, and a `Hosted` entry registers it:
|
|
241
264
|
|
|
242
265
|
```python
|
|
243
266
|
cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
|
|
244
267
|
|
|
245
|
-
|
|
268
|
+
entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
|
|
246
269
|
cortexgrid.register_model(
|
|
247
|
-
|
|
248
|
-
family=
|
|
249
|
-
requirements=
|
|
270
|
+
entry.serve_app,
|
|
271
|
+
family=entry.family, suffix=entry.suffix,
|
|
272
|
+
requirements=entry.requirements(), config=entry.config(),
|
|
250
273
|
)
|
|
251
274
|
```
|
|
252
275
|
|
|
@@ -263,7 +286,7 @@ a re-deploy rather than a re-registration.
|
|
|
263
286
|
|
|
264
287
|
A replica asks for no hardware at all, so it is placed on any node, CPU-only
|
|
265
288
|
included. From there it deploys and streams exactly like a cluster-served model:
|
|
266
|
-
`
|
|
289
|
+
`entry.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
|
|
267
290
|
reports tool calls as structured blocks rather than as generated text, so the
|
|
268
291
|
serve app re-encodes them into the text form the client parses.
|
|
269
292
|
|
|
@@ -271,37 +294,50 @@ serve app re-encodes them into the text form the client parses.
|
|
|
271
294
|
|
|
272
295
|
| | |
|
|
273
296
|
|---|---|
|
|
274
|
-
| `
|
|
275
|
-
| `
|
|
276
|
-
| `
|
|
277
|
-
| `
|
|
278
|
-
| `
|
|
297
|
+
| `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | Importer for a model on the HuggingFace Hub, run by `serve_app`. |
|
|
298
|
+
| `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
|
|
299
|
+
| `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
|
|
300
|
+
| `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
|
|
301
|
+
| `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
|
|
302
|
+
| `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
|
|
303
|
+
| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
|
|
304
|
+
| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
|
|
305
|
+
| `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
|
|
279
306
|
| `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
|
|
280
307
|
| `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
|
|
281
308
|
| `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
|
|
282
|
-
| `
|
|
309
|
+
| `ServedGeneratingModel(url, model_id)` | Client for any deployed text-to-image app. |
|
|
310
|
+
| `ServedMeshingModel(url, model_id)` | Client for any deployed image-to-mesh app. |
|
|
283
311
|
| `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
|
|
284
312
|
| `detect_device()` | The torch device a serve app should load onto. |
|
|
285
313
|
| `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
|
|
286
314
|
| `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
|
|
287
315
|
| `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
|
|
288
316
|
|
|
289
|
-
Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `
|
|
290
|
-
`GeneratedImage`, `ToolCall`, `Message`, `Tool`,
|
|
317
|
+
Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
|
|
318
|
+
`CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
|
|
319
|
+
`ToolSpec`.
|
|
320
|
+
|
|
321
|
+
Extending it:
|
|
291
322
|
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
`
|
|
323
|
+
- **A new source** is one importer: subclass `Importer` with `download(local_dir)`
|
|
324
|
+
and `weights()`, and every serve app whose weight format it delivers runs its
|
|
325
|
+
models unchanged.
|
|
326
|
+
- **A new hosted API** is one serve app: subclass `HostedModel`, and register it
|
|
327
|
+
with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
|
|
328
|
+
inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
|
|
329
|
+
(there is an `encode_tool_call` for upstreams that report them structurally) —
|
|
330
|
+
then `ServedCompletingModel` is its client, unchanged.
|
|
331
|
+
- **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
|
|
332
|
+
subclass, say — and, if its source is unusual too, its own importer.
|
|
297
333
|
|
|
298
334
|
## Notes
|
|
299
335
|
|
|
300
336
|
- **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
|
|
301
337
|
read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
|
|
302
338
|
that downloads.
|
|
303
|
-
- **Image weight pruning.** `
|
|
304
|
-
`.gitattributes`
|
|
339
|
+
- **Image weight pruning.** `Text2Image` tells the importer to skip example images,
|
|
340
|
+
docs and `.gitattributes` — nothing `from_pretrained` reads. It does **not**
|
|
305
341
|
skip a repo's consolidated single-file checkpoint, which is weights it cannot
|
|
306
342
|
tell apart from the component ones: name it per-model in
|
|
307
343
|
`HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
|
|
@@ -311,7 +347,7 @@ to stream text from `POST /complete`, inlining tool calls as
|
|
|
311
347
|
describe what actually gets staged:
|
|
312
348
|
|
|
313
349
|
```python
|
|
314
|
-
>>> mg.
|
|
350
|
+
>>> mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image).requirements()
|
|
315
351
|
ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
|
|
316
352
|
```
|
|
317
353
|
- **Only model files are stored.** HuggingFace's download bookkeeping
|
|
@@ -6,19 +6,47 @@ needs but cannot work out for itself, plus the clients that talk to what it depl
|
|
|
6
6
|
cortexgrid stores whatever weights it is handed under whatever key it is given, and
|
|
7
7
|
serves them with whatever class it was told to bundle. It deliberately knows nothing
|
|
8
8
|
about where a model came from or what shape it is. This library supplies exactly that
|
|
9
|
-
missing knowledge
|
|
10
|
-
|
|
9
|
+
missing knowledge, and drives no lifecycle of its own: importing, deploying and
|
|
10
|
+
deleting stay your calls against the cortexgrid SDK.
|
|
11
|
+
|
|
12
|
+
Upgrading from 0.2.x? See [MIGRATION.md](https://github.com/robodatalab/cortexgrid-infer/blob/main/MIGRATION.md).
|
|
11
13
|
|
|
12
14
|
## Division of labour
|
|
13
15
|
|
|
14
16
|
| | Owns |
|
|
15
17
|
|---|---|
|
|
16
18
|
| **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
|
|
17
|
-
| **cortexgrid-infer** | Per
|
|
19
|
+
| **cortexgrid-infer** | Per source: how to name a model, fetch it and size it. Per task: how to run it and talk to it. Never drives a lifecycle. |
|
|
18
20
|
| **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
|
|
19
21
|
|
|
20
|
-
|
|
21
|
-
|
|
22
|
+
A model is two things, and the library keeps them apart:
|
|
23
|
+
|
|
24
|
+
- **What runs it** — a serve app in `cortexgrid_infer.models`, named for the task it
|
|
25
|
+
runs, not for where the model came from. It loads the weights, answers the task's
|
|
26
|
+
routes, and names the client that speaks them. One serve app runs every model of
|
|
27
|
+
its task, whichever source staged it.
|
|
28
|
+
- **Where its weights come from** — an importer in `cortexgrid_infer.importers`, one
|
|
29
|
+
per source. It names the model in the registry, downloads the weights, and sizes
|
|
30
|
+
them from the source's metadata. It is handed the serve app, so one importer takes
|
|
31
|
+
any model its source holds.
|
|
32
|
+
|
|
33
|
+
| Serve app | Task | Runs | Client |
|
|
34
|
+
|---|---|---|---|
|
|
35
|
+
| `Text2Text` | text to text | any causal LM in the transformers layout, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
|
|
36
|
+
| `Text2Image` | text (and image) to image | any diffusers pipeline | `ServedGeneratingModel` (`generate`) |
|
|
37
|
+
| `Image2Mesh` — a base: subclass it with the model's `load` and `make_mesh` | image to mesh | the model the subclass loads, behind the task's `POST /mesh` | `ServedMeshingModel` (`mesh`: one picture in, a `GeneratedMesh` out) |
|
|
38
|
+
| `AnthropicText2Text` | text to text | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
|
|
39
|
+
|
|
40
|
+
| Importer | Source |
|
|
41
|
+
|---|---|
|
|
42
|
+
| `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | the HuggingFace Hub |
|
|
43
|
+
|
|
44
|
+
A model hosted elsewhere has no weights, and so no importer: its serve app forwards
|
|
45
|
+
to it, and a `Hosted(model_id, serve_app, **settings)` entry registers it.
|
|
46
|
+
|
|
47
|
+
An importer and a `Hosted` entry are both a `ModelEntry`, which gives you what
|
|
48
|
+
`cortexgrid.import_model`, `cortexgrid.register_model` and `cortexgrid.deploy_model`
|
|
49
|
+
need:
|
|
22
50
|
|
|
23
51
|
| | |
|
|
24
52
|
|---|---|
|
|
@@ -27,19 +55,10 @@ An **importer** is the whole of this library's cluster-side surface. It gives yo
|
|
|
27
55
|
| `.requirements()` | what one replica needs to be placed and to run |
|
|
28
56
|
| `.config()` | settings the serve app reads at construction, if it needs any |
|
|
29
57
|
| `.client(url)` | a client that speaks the deployed app's routes |
|
|
30
|
-
| `.source` | *(
|
|
58
|
+
| `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
|
|
31
59
|
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
| `HuggingFaceCompletingImport(hf_id, token=None)` | HuggingFace causal LM, via `AutoModelForCausalLM` | `ServedCompletingModel` (`complete`) |
|
|
35
|
-
| `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | HuggingFace diffusers pipeline | `HuggingFaceImageModel` (`generate`) |
|
|
36
|
-
| `AnthropicImport(model_id, api_key_secret=...)` | the Anthropic API, forwarded from the cluster | `ServedCompletingModel` (`complete`) |
|
|
37
|
-
|
|
38
|
-
Not every model has weights, so `.source` belongs to the HuggingFace importers
|
|
39
|
-
rather than to all of them: one that has weights goes through
|
|
40
|
-
`cortexgrid.import_model`, one that has none through `cortexgrid.register_model`.
|
|
41
|
-
Everything else is common, which is why one client serves all three — every
|
|
42
|
-
completion app speaks the same `/complete` protocol.
|
|
60
|
+
One with weights goes through `cortexgrid.import_model`, one without through
|
|
61
|
+
`cortexgrid.register_model`.
|
|
43
62
|
|
|
44
63
|
## Quick start
|
|
45
64
|
|
|
@@ -49,7 +68,7 @@ import cortexgrid
|
|
|
49
68
|
import cortexgrid_infer as mg
|
|
50
69
|
|
|
51
70
|
cortexgrid.Experiment.init("img-gen")
|
|
52
|
-
imp = mg.
|
|
71
|
+
imp = mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image)
|
|
53
72
|
|
|
54
73
|
# 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
|
|
55
74
|
# your machine, so this runs as a job rather than in-process.
|
|
@@ -104,12 +123,14 @@ serve code if it changed, and tags the run with the model.
|
|
|
104
123
|
`cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
|
|
105
124
|
nothing in the registry can derive it — cortexgrid stores whatever weights it is
|
|
106
125
|
handed, and an importer imports whatever repo id it is handed. `requirements()`
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
126
|
+
splits the estimate along the same line as the rest: the importer reads what the
|
|
127
|
+
source records about the weights — for HuggingFace, safetensors headers where the
|
|
128
|
+
repo has them, weight-file sizes otherwise, both over HTTP with no weights
|
|
129
|
+
downloaded — and the serve app turns that into what a replica needs, since it knows
|
|
130
|
+
the dtype it loads in and what its work takes on top.
|
|
110
131
|
|
|
111
132
|
```python
|
|
112
|
-
>>> mg.
|
|
133
|
+
>>> mg.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", mg.Text2Text).requirements()
|
|
113
134
|
ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
|
|
114
135
|
```
|
|
115
136
|
|
|
@@ -118,14 +139,15 @@ you know better, or edit the figures on the model card in the dashboard afterwar
|
|
|
118
139
|
`import_model` only fills in requirements a version does not already have, so a
|
|
119
140
|
hand-set value is never overwritten.
|
|
120
141
|
|
|
121
|
-
The estimate always asks for exactly one GPU, because
|
|
122
|
-
the whole model onto a single device. A model too large for any one card needs a
|
|
123
|
-
app that shards it, not a larger `num_gpus`: the requirement would simply never
|
|
124
|
-
placed.
|
|
142
|
+
The estimate always asks for exactly one GPU, because every serve app that runs weights
|
|
143
|
+
moves the whole model onto a single device. A model too large for any one card needs a
|
|
144
|
+
serve app that shards it, not a larger `num_gpus`: the requirement would simply never
|
|
145
|
+
be placed. An `Image2Mesh` subclass can set `min_vram_gb` for the memory meshing takes
|
|
146
|
+
beyond the weights. (A `Hosted` entry asks for none — it holds no weights.)
|
|
125
147
|
|
|
126
148
|
## Compilation
|
|
127
149
|
|
|
128
|
-
|
|
150
|
+
`Text2Text` and `Text2Image` run `torch.compile` over the model they load, because
|
|
129
151
|
eager PyTorch bills per *operation* and a transformer forward is mostly glue —
|
|
130
152
|
views, broadcasts, permutes, elementwise adds — with only a small fraction of the
|
|
131
153
|
dispatches being the matrix multiplies that do the arithmetic. While the tensors
|
|
@@ -146,12 +168,12 @@ fits the work it does:
|
|
|
146
168
|
|
|
147
169
|
Both serve apps are tuned rather than defaulted:
|
|
148
170
|
|
|
149
|
-
- **`
|
|
171
|
+
- **`Text2Image`** compiles the denoiser (`transformer` or `unet`)
|
|
150
172
|
and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
|
|
151
173
|
every step is the same shapes — so the capture is made once and replayed for the
|
|
152
174
|
whole loop. The VAE is left alone: smallest win of the three, and the one whose
|
|
153
175
|
shapes move most, especially with tiling on.
|
|
154
|
-
- **`
|
|
176
|
+
- **`Text2Text`** does *not* compile the module. Decode has
|
|
155
177
|
no stable shape of its own — the sequence grows a token per forward — so it
|
|
156
178
|
pins one: `cache_implementation="static"` with a fixed `max_cache_len`, plus a
|
|
157
179
|
`CompileConfig(mode=Mode.GRAPHED)`. transformers then compiles `generate`
|
|
@@ -195,7 +217,7 @@ it is the floor; only quantizing the weights or running fewer forwards moves it.
|
|
|
195
217
|
`complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
|
|
196
218
|
|
|
197
219
|
```python
|
|
198
|
-
imp = mg.
|
|
220
|
+
imp = mg.HuggingFaceImporter("Qwen/Qwen2.5-7B-Instruct", mg.Text2Text)
|
|
199
221
|
model = imp.client(deployment.url)
|
|
200
222
|
messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
|
|
201
223
|
|
|
@@ -210,17 +232,17 @@ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=
|
|
|
210
232
|
### Anthropic
|
|
211
233
|
|
|
212
234
|
An Anthropic model is a registry entry like any other — the difference is that
|
|
213
|
-
there are no weights to stage, so
|
|
214
|
-
|
|
235
|
+
there are no weights to stage, so there is no importer: the serve app forwards to
|
|
236
|
+
the API instead of loading anything, and a `Hosted` entry registers it:
|
|
215
237
|
|
|
216
238
|
```python
|
|
217
239
|
cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
|
|
218
240
|
|
|
219
|
-
|
|
241
|
+
entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
|
|
220
242
|
cortexgrid.register_model(
|
|
221
|
-
|
|
222
|
-
family=
|
|
223
|
-
requirements=
|
|
243
|
+
entry.serve_app,
|
|
244
|
+
family=entry.family, suffix=entry.suffix,
|
|
245
|
+
requirements=entry.requirements(), config=entry.config(),
|
|
224
246
|
)
|
|
225
247
|
```
|
|
226
248
|
|
|
@@ -237,7 +259,7 @@ a re-deploy rather than a re-registration.
|
|
|
237
259
|
|
|
238
260
|
A replica asks for no hardware at all, so it is placed on any node, CPU-only
|
|
239
261
|
included. From there it deploys and streams exactly like a cluster-served model:
|
|
240
|
-
`
|
|
262
|
+
`entry.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
|
|
241
263
|
reports tool calls as structured blocks rather than as generated text, so the
|
|
242
264
|
serve app re-encodes them into the text form the client parses.
|
|
243
265
|
|
|
@@ -245,37 +267,50 @@ serve app re-encodes them into the text form the client parses.
|
|
|
245
267
|
|
|
246
268
|
| | |
|
|
247
269
|
|---|---|
|
|
248
|
-
| `
|
|
249
|
-
| `
|
|
250
|
-
| `
|
|
251
|
-
| `
|
|
252
|
-
| `
|
|
270
|
+
| `HuggingFaceImporter(hf_id, serve_app, token=None, ignore_patterns=None)` | Importer for a model on the HuggingFace Hub, run by `serve_app`. |
|
|
271
|
+
| `Importer` | Base of the importers: identity, scratch directory, `source`, and `requirements()` / `client(url)` taken from the serve app. A new source implements `download(local_dir)` and `weights()`. |
|
|
272
|
+
| `Hosted(model_id, serve_app, **settings)` | Entry for a model hosted elsewhere; no weights. `settings` are the serve app's `config` arguments. |
|
|
273
|
+
| `ModelEntry` | Base of both: `model_id`, `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
|
|
274
|
+
| `Text2Text`, `Text2Image`, `Image2Mesh` | Serve apps that run weights (`LocalModel`s). |
|
|
275
|
+
| `AnthropicText2Text` | Serve app that forwards to the Anthropic API (a `HostedModel`). |
|
|
276
|
+
| `LocalModel` | Base of the serve apps that run weights: `ignore_patterns()`, `bytes_per_param`, `min_vram_gb`, `requirements(weights)`, `client(url, name)`. |
|
|
277
|
+
| `HostedModel` | Base of the serve apps that forward: `config(model_id, **settings)`, `requirements()`, `client(url, name)`. |
|
|
278
|
+
| `Weights(params=None, file_bytes=0)` | What an importer reads about the weights, handed to the serve app to size a replica. |
|
|
253
279
|
| `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
|
|
254
280
|
| `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
|
|
255
281
|
| `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
|
|
256
|
-
| `
|
|
282
|
+
| `ServedGeneratingModel(url, model_id)` | Client for any deployed text-to-image app. |
|
|
283
|
+
| `ServedMeshingModel(url, model_id)` | Client for any deployed image-to-mesh app. |
|
|
257
284
|
| `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
|
|
258
285
|
| `detect_device()` | The torch device a serve app should load onto. |
|
|
259
286
|
| `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
|
|
260
287
|
| `compiling.supported(device)` / `compile_module(module, device, mode)` / `compile_pipeline(pipe, device, mode)` | Whether to compile here, and what a serve app calls to do it. |
|
|
261
288
|
| `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
|
|
262
289
|
|
|
263
|
-
Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `
|
|
264
|
-
`GeneratedImage`, `ToolCall`, `Message`, `Tool`,
|
|
290
|
+
Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
|
|
291
|
+
`CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
|
|
292
|
+
`ToolSpec`.
|
|
293
|
+
|
|
294
|
+
Extending it:
|
|
265
295
|
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
`
|
|
296
|
+
- **A new source** is one importer: subclass `Importer` with `download(local_dir)`
|
|
297
|
+
and `weights()`, and every serve app whose weight format it delivers runs its
|
|
298
|
+
models unchanged.
|
|
299
|
+
- **A new hosted API** is one serve app: subclass `HostedModel`, and register it
|
|
300
|
+
with `Hosted`. A text-to-text one only has to stream text from `POST /complete`,
|
|
301
|
+
inlining tool calls as `<tool_call>{"name": ..., "arguments": {...}}</tool_call>`
|
|
302
|
+
(there is an `encode_tool_call` for upstreams that report them structurally) —
|
|
303
|
+
then `ServedCompletingModel` is its client, unchanged.
|
|
304
|
+
- **A model no task's serve app can load** gets its own serve app — an `Image2Mesh`
|
|
305
|
+
subclass, say — and, if its source is unusual too, its own importer.
|
|
271
306
|
|
|
272
307
|
## Notes
|
|
273
308
|
|
|
274
309
|
- **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
|
|
275
310
|
read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
|
|
276
311
|
that downloads.
|
|
277
|
-
- **Image weight pruning.** `
|
|
278
|
-
`.gitattributes`
|
|
312
|
+
- **Image weight pruning.** `Text2Image` tells the importer to skip example images,
|
|
313
|
+
docs and `.gitattributes` — nothing `from_pretrained` reads. It does **not**
|
|
279
314
|
skip a repo's consolidated single-file checkpoint, which is weights it cannot
|
|
280
315
|
tell apart from the component ones: name it per-model in
|
|
281
316
|
`HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
|
|
@@ -285,7 +320,7 @@ to stream text from `POST /complete`, inlining tool calls as
|
|
|
285
320
|
describe what actually gets staged:
|
|
286
321
|
|
|
287
322
|
```python
|
|
288
|
-
>>> mg.
|
|
323
|
+
>>> mg.HuggingFaceImporter("black-forest-labs/FLUX.2-klein-base-4B", mg.Text2Image).requirements()
|
|
289
324
|
ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
|
|
290
325
|
```
|
|
291
326
|
- **Only model files are stored.** HuggingFace's download bookkeeping
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
"""Run inference on models the cortexgrid cluster serves.
|
|
2
|
+
|
|
3
|
+
cortexgrid owns a model's life - registry identity, weights, serve bundle,
|
|
4
|
+
hardware placement, deployment, jobs - and knows nothing about what kind of
|
|
5
|
+
model it holds. This library supplies the knowledge those calls need and
|
|
6
|
+
nothing else, split along the two things a model is:
|
|
7
|
+
|
|
8
|
+
- where its weights come from - an importer per source (`importers`), which
|
|
9
|
+
names, fetches and sizes them;
|
|
10
|
+
- what runs it - a serve app per task (`models`), which loads the weights,
|
|
11
|
+
answers the task's routes, and names the client that speaks them.
|
|
12
|
+
|
|
13
|
+
imp = cortexgrid_infer.HuggingFaceImporter("Qwen/Qwen2.5-0.5B-Instruct", cortexgrid_infer.Text2Text)
|
|
14
|
+
with imp:
|
|
15
|
+
cortexgrid.import_model(
|
|
16
|
+
imp.source, imp.serve_app,
|
|
17
|
+
family=imp.family, suffix=imp.suffix,
|
|
18
|
+
requirements=imp.requirements(),
|
|
19
|
+
)
|
|
20
|
+
deployment = cortexgrid.deploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED, wait=True)
|
|
21
|
+
model = imp.client(deployment.url)
|
|
22
|
+
|
|
23
|
+
A model hosted elsewhere has no weights and so no importer: its serve app
|
|
24
|
+
forwards to it, and a `Hosted` entry registers it.
|
|
25
|
+
|
|
26
|
+
See `examples/deploy_text_model` for the whole lifecycle, import job included.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
from cortexgrid import ModelDeployFailed
|
|
30
|
+
|
|
31
|
+
from cortexgrid_infer.core import (
|
|
32
|
+
CompletionChunk,
|
|
33
|
+
DeployedModel,
|
|
34
|
+
CompletingModel,
|
|
35
|
+
GeneratedImage,
|
|
36
|
+
GeneratedMesh,
|
|
37
|
+
GeneratingModel,
|
|
38
|
+
MeshingModel,
|
|
39
|
+
Message,
|
|
40
|
+
Tool,
|
|
41
|
+
ToolCall,
|
|
42
|
+
ToolSpec,
|
|
43
|
+
complete,
|
|
44
|
+
generate,
|
|
45
|
+
mesh,
|
|
46
|
+
)
|
|
47
|
+
from cortexgrid_infer.completion import ServedCompletingModel
|
|
48
|
+
from cortexgrid_infer.imaging import ServedGeneratingModel
|
|
49
|
+
from cortexgrid_infer.meshing import ServedMeshingModel
|
|
50
|
+
from cortexgrid_infer.device import detect_device
|
|
51
|
+
from cortexgrid_infer.models import (
|
|
52
|
+
AnthropicText2Text,
|
|
53
|
+
HostedModel,
|
|
54
|
+
Image2Mesh,
|
|
55
|
+
LocalModel,
|
|
56
|
+
Text2Image,
|
|
57
|
+
Text2Text,
|
|
58
|
+
Weights,
|
|
59
|
+
)
|
|
60
|
+
from cortexgrid_infer.registry import Hosted, ModelEntry, split_model_id
|
|
61
|
+
from cortexgrid_infer.importers import HuggingFaceImporter, Importer
|
|
62
|
+
|
|
63
|
+
__all__ = [
|
|
64
|
+
"CompletionChunk",
|
|
65
|
+
"DeployedModel",
|
|
66
|
+
"CompletingModel",
|
|
67
|
+
"GeneratedImage",
|
|
68
|
+
"GeneratedMesh",
|
|
69
|
+
"GeneratingModel",
|
|
70
|
+
"MeshingModel",
|
|
71
|
+
"Message",
|
|
72
|
+
"Tool",
|
|
73
|
+
"ToolCall",
|
|
74
|
+
"ToolSpec",
|
|
75
|
+
"complete",
|
|
76
|
+
"generate",
|
|
77
|
+
"mesh",
|
|
78
|
+
"ModelDeployFailed",
|
|
79
|
+
"detect_device",
|
|
80
|
+
"ServedCompletingModel",
|
|
81
|
+
"ServedGeneratingModel",
|
|
82
|
+
"ServedMeshingModel",
|
|
83
|
+
"LocalModel",
|
|
84
|
+
"HostedModel",
|
|
85
|
+
"Weights",
|
|
86
|
+
"Text2Text",
|
|
87
|
+
"Text2Image",
|
|
88
|
+
"Image2Mesh",
|
|
89
|
+
"AnthropicText2Text",
|
|
90
|
+
"ModelEntry",
|
|
91
|
+
"Hosted",
|
|
92
|
+
"split_model_id",
|
|
93
|
+
"Importer",
|
|
94
|
+
"HuggingFaceImporter",
|
|
95
|
+
]
|