cortexgrid-infer 0.1.14__tar.gz → 0.2.2__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 (29) hide show
  1. cortexgrid_infer-0.2.2/PKG-INFO +310 -0
  2. cortexgrid_infer-0.2.2/README.md +284 -0
  3. cortexgrid_infer-0.2.2/cortexgrid_infer/__init__.py +69 -0
  4. cortexgrid_infer-0.2.2/cortexgrid_infer/compiling.py +70 -0
  5. cortexgrid_infer-0.2.2/cortexgrid_infer/completion.py +192 -0
  6. cortexgrid_infer-0.2.2/cortexgrid_infer/core.py +117 -0
  7. cortexgrid_infer-0.2.2/cortexgrid_infer/importing.py +158 -0
  8. cortexgrid_infer-0.2.2/cortexgrid_infer/providers/anthropic.py +65 -0
  9. cortexgrid_infer-0.2.2/cortexgrid_infer/providers/anthropic_serve.py +159 -0
  10. cortexgrid_infer-0.2.2/cortexgrid_infer/providers/huggingface_complete.py +23 -0
  11. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/providers/huggingface_complete_serve.py +15 -4
  12. cortexgrid_infer-0.2.2/cortexgrid_infer/providers/huggingface_image.py +110 -0
  13. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/providers/huggingface_image_serve.py +5 -3
  14. cortexgrid_infer-0.2.2/cortexgrid_infer/requirements.py +117 -0
  15. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/utils.py +0 -47
  16. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/pyproject.toml +10 -3
  17. cortexgrid_infer-0.1.14/PKG-INFO +0 -213
  18. cortexgrid_infer-0.1.14/README.md +0 -187
  19. cortexgrid_infer-0.1.14/cortexgrid_infer/__init__.py +0 -71
  20. cortexgrid_infer-0.1.14/cortexgrid_infer/core.py +0 -250
  21. cortexgrid_infer-0.1.14/cortexgrid_infer/providers/anthropic.py +0 -189
  22. cortexgrid_infer-0.1.14/cortexgrid_infer/providers/huggingface_complete.py +0 -308
  23. cortexgrid_infer-0.1.14/cortexgrid_infer/providers/huggingface_image.py +0 -231
  24. cortexgrid_infer-0.1.14/cortexgrid_infer/providers/serving.py +0 -32
  25. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/.gitignore +0 -0
  26. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/LICENSE +0 -0
  27. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/device.py +0 -0
  28. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/providers/__init__.py +0 -0
  29. {cortexgrid_infer-0.1.14 → cortexgrid_infer-0.2.2}/cortexgrid_infer/py.typed +0 -0
@@ -0,0 +1,310 @@
1
+ Metadata-Version: 2.5
2
+ Name: cortexgrid-infer
3
+ Version: 0.2.2
4
+ Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)
5
+ Project-URL: Homepage, https://github.com/robodatalab/cortexgrid-infer
6
+ Project-URL: Repository, https://github.com/robodatalab/cortexgrid-infer
7
+ License-Expression: Apache-2.0
8
+ License-File: LICENSE
9
+ Requires-Python: <3.12,>=3.11
10
+ Requires-Dist: accelerate>=0.33
11
+ Requires-Dist: anthropic>=0.40.0
12
+ Requires-Dist: cortexgrid>=0.3.4
13
+ Requires-Dist: diffusers>=0.38
14
+ Requires-Dist: fastapi>=0.133
15
+ Requires-Dist: httpx>=0.27.0
16
+ Requires-Dist: huggingface-hub>=0.20.0
17
+ Requires-Dist: pillow>=10.0
18
+ Requires-Dist: protobuf>=4.0
19
+ Requires-Dist: python-dotenv>=1.0.0
20
+ Requires-Dist: sentencepiece>=0.2
21
+ Requires-Dist: torch>=2.0.0
22
+ Requires-Dist: transformers>=4.47.0
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest>=8.0; extra == 'dev'
25
+ Description-Content-Type: text/markdown
26
+
27
+ # cortexgrid-infer
28
+
29
+ The model-specific pieces that [cortexgrid](https://github.com/robodatalab/cortexgrid)
30
+ needs but cannot work out for itself, plus the clients that talk to what it deploys.
31
+
32
+ cortexgrid stores whatever weights it is handed under whatever key it is given, and
33
+ serves them with whatever class it was told to bundle. It deliberately knows nothing
34
+ about where a model came from or what shape it is. This library supplies exactly that
35
+ missing knowledge for one family of models at a time, and drives no lifecycle of its
36
+ own: importing, deploying and deleting stay your calls against the cortexgrid SDK.
37
+
38
+ ## Division of labour
39
+
40
+ | | Owns |
41
+ |---|---|
42
+ | **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
43
+ | **cortexgrid-infer** | Per family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
44
+ | **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
45
+
46
+ An **importer** is the whole of this library's cluster-side surface. It gives you what
47
+ `cortexgrid.import_model` and `cortexgrid.deploy_model` need:
48
+
49
+ | | |
50
+ |---|---|
51
+ | `.family` / `.suffix` | how this model id names itself in the registry |
52
+ | `.serve_app` | the class cortexgrid bundles and runs on the cluster |
53
+ | `.requirements()` | what one replica needs to be placed and to run |
54
+ | `.config()` | settings the serve app reads at construction, if it needs any |
55
+ | `.client(url)` | a client that speaks the deployed app's routes |
56
+ | `.source` | *(models with weights)* a callable that downloads them and returns the directory |
57
+
58
+ | Importer | Serves | Client |
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.
69
+
70
+ ## Quick start
71
+
72
+ ```python
73
+ import asyncio
74
+ import cortexgrid
75
+ import cortexgrid_infer as mg
76
+
77
+ cortexgrid.Experiment.init("img-gen")
78
+ imp = mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
79
+
80
+ # 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
81
+ # your machine, so this runs as a job rather than in-process.
82
+ def import_weights(importer, requirements):
83
+ with importer:
84
+ cortexgrid.import_model(
85
+ importer.source, importer.serve_app,
86
+ family=importer.family, suffix=importer.suffix,
87
+ requirements=requirements,
88
+ )
89
+
90
+ # remote() hands back a JobFuture; result() blocks and re-raises whatever the
91
+ # job raised, so a failed import surfaces here with its own traceback.
92
+ cortexgrid.remote(
93
+ import_weights, imp, imp.requirements(), num_gpus=0, num_cpus=2
94
+ ).result()
95
+
96
+ # 2. Deploy — blocks until the Ray Serve app is running.
97
+ deployment = cortexgrid.deploy_model(
98
+ imp.family, imp.suffix, cortexgrid.IMPORTED, wait=True, timeout=1800
99
+ )
100
+
101
+ # 3. Inference.
102
+ model = imp.client(deployment.url)
103
+
104
+ async def run():
105
+ result = await mg.generate(model, "a red bicycle on a beach at sunrise")
106
+ open("out.png", "wb").write(result.image) # PNG bytes; also .width/.height/.params
107
+
108
+ asyncio.run(run())
109
+
110
+ # 4. Teardown.
111
+ cortexgrid.undeploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED) # free the GPU
112
+ cortexgrid.delete_model(imp.family, imp.suffix, cortexgrid.IMPORTED) # and the weights
113
+ ```
114
+
115
+ See [`examples/deploy_text_model`](examples/deploy_text_model/main.py) for the whole
116
+ lifecycle end to end.
117
+
118
+ An importer carries no open resources until it is entered, so it can be handed to
119
+ `cortexgrid.remote` and rebuilt on the node that runs the import. Entering it opens
120
+ the scratch directory `source` downloads into; `import_model` reads that directory
121
+ after `source` returns, which is why the lifetime spans the import call rather than
122
+ the download.
123
+
124
+ Call the import on every run. The weights are imported once and shared by every run:
125
+ on a model that is already `ready`, `import_model` downloads nothing, re-bundles the
126
+ serve code if it changed, and tags the run with the model.
127
+
128
+ ## Hardware requirements
129
+
130
+ `cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
131
+ nothing in the registry can derive it — cortexgrid stores whatever weights it is
132
+ handed, and an importer imports whatever repo id it is handed. `requirements()`
133
+ estimates it from the repo's own metadata: safetensors headers where the repo has
134
+ them, weight-file sizes otherwise. Both are read over HTTP; no weights are
135
+ downloaded.
136
+
137
+ ```python
138
+ >>> mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
139
+ ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
140
+ ```
141
+
142
+ These are estimates. Pass a `ModelRequirements` of your own to `import_model` where
143
+ you know better, or edit the figures on the model card in the dashboard afterwards —
144
+ `import_model` only fills in requirements a version does not already have, so a
145
+ hand-set value is never overwritten.
146
+
147
+ The estimate always asks for exactly one GPU, because both HuggingFace serve apps move
148
+ the whole model onto a single device. A model too large for any one card needs a serve
149
+ app that shards it, not a larger `num_gpus`: the requirement would simply never be
150
+ placed. (`AnthropicImport` asks for none — it holds no weights.)
151
+
152
+ ## Compilation
153
+
154
+ Both HuggingFace serve apps run `torch.compile` over the model they load, because
155
+ eager PyTorch bills per *operation* and a transformer forward is mostly glue —
156
+ views, broadcasts, permutes, elementwise adds — with only a small fraction of the
157
+ dispatches being the matrix multiplies that do the arithmetic. While the tensors
158
+ are large the host stays ahead of the accelerator and none of that shows; once
159
+ they are small, the accelerator drains its queue faster than the host can fill it
160
+ and step time tracks the operation count rather than the work. Feeding the model
161
+ less then stops helping, because the op count does not shrink with the input.
162
+ Inductor fuses the glue away; `reduce-overhead` additionally captures a CUDA graph,
163
+ so a step replays one graph instead of thousands of launches.
164
+
165
+ `cortexgrid_infer.compiling` offers two modes, and a serve app picks the one that
166
+ fits the work it does:
167
+
168
+ | | | for |
169
+ |---|---|---|
170
+ | `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
171
+ | `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
172
+
173
+ Both serve apps are tuned rather than defaulted:
174
+
175
+ - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
176
+ and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
177
+ every step is the same shapes — so the capture is made once and replayed for the
178
+ whole loop. The VAE is left alone: smallest win of the three, and the one whose
179
+ shapes move most, especially with tiling on.
180
+ - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
181
+ no stable shape until the KV cache is static, so it sets
182
+ `cache_implementation="static"` and a `CompileConfig(mode=Mode.GRAPHED)`, and
183
+ transformers compiles `generate` itself — capturing prefill and decode
184
+ separately, which it is better placed to do. Compiling the module here as well
185
+ would only compile the same forward twice. The cost is VRAM: a static cache is
186
+ preallocated to the length asked for, where a dynamic one grows into it.
187
+
188
+ Neither asks to be configured, for the same reason the serve apps do not ask which
189
+ dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off
190
+ it, and a host outrunning its accelerator is a GPU problem to begin with.
191
+
192
+ Two things worth knowing:
193
+
194
+ - **Warm-up is per input shape, and lands on a request.** Compilation is lazy, so
195
+ the first request at each image size — or each prompt and generation length —
196
+ pays for it.
197
+ - **Compilation never becomes load-bearing.** A backend that cannot compile a model
198
+ falls back to eager rather than failing the deployment.
199
+
200
+ It does not fix everything. Every forward re-reads the model's whole weight tensor
201
+ from memory, and no amount of fusion makes that read smaller. Below some input size
202
+ it is the floor; only quantizing the weights or running fewer forwards moves it.
203
+
204
+ ## Inference
205
+
206
+ `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
207
+
208
+ ```python
209
+ imp = mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
210
+ model = imp.client(deployment.url)
211
+ messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
212
+
213
+ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=0.7):
214
+ if chunk.has_text:
215
+ print(chunk.content, end="", flush=True)
216
+ ```
217
+
218
+ `generate` returns one `GeneratedImage` (`.image` PNG bytes, `.width`, `.height`,
219
+ `.params`); pass `image=` bytes for img2img.
220
+
221
+ ### Anthropic
222
+
223
+ An Anthropic model is a registry entry like any other — the difference is that
224
+ there are no weights to stage, so it is imported with no source, and the serve
225
+ app forwards to the API instead of loading anything:
226
+
227
+ ```python
228
+ cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
229
+
230
+ imp = mg.AnthropicImport("claude-sonnet-5")
231
+ cortexgrid.register_model(
232
+ imp.serve_app,
233
+ family=imp.family, suffix=imp.suffix,
234
+ requirements=imp.requirements(), config=imp.config(),
235
+ )
236
+ ```
237
+
238
+ `register_model` is `import_model` for a model that stages nothing: only the
239
+ serve bundle is stored, and the entry is indistinguishable from an imported one
240
+ to `deploy_model`, `list_models` and the dashboard.
241
+
242
+ `config` carries the Anthropic model name and the **name of** the cortexgrid
243
+ secret holding the key — never the key itself, since a registry entry is readable
244
+ by anyone who can see the model. The serve app reads it with
245
+ `cortexgrid.model_config` at construction, and both values are editable on the
246
+ model card, so pointing a deployment at a different model or key is an edit plus
247
+ a re-deploy rather than a re-registration.
248
+
249
+ A replica asks for no hardware at all, so it is placed on any node, CPU-only
250
+ included. From there it deploys and streams exactly like a cluster-served model:
251
+ `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
252
+ reports tool calls as structured blocks rather than as generated text, so the
253
+ serve app re-encodes them into the text form the client parses.
254
+
255
+ ## API
256
+
257
+ | | |
258
+ |---|---|
259
+ | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
260
+ | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
261
+ | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
262
+ | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
263
+ | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
264
+ | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
265
+ | `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
266
+ | `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
267
+ | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
268
+ | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
269
+ | `detect_device()` | The torch device a serve app should load onto. |
270
+ | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
271
+ | `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. |
272
+ | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
273
+
274
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,
275
+ `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
276
+
277
+ Adding a provider means one importer and one serve app. A completion app only has
278
+ to stream text from `POST /complete`, inlining tool calls as
279
+ `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
280
+ `encode_tool_call` for upstreams that report them structurally) — then
281
+ `ServedCompletingModel` is its client, unchanged.
282
+
283
+ ## Notes
284
+
285
+ - **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
286
+ read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
287
+ that downloads.
288
+ - **Image weight pruning.** `HuggingFaceImageImport` skips example images, docs and
289
+ `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
290
+ skip a repo's consolidated single-file checkpoint, which is weights it cannot
291
+ tell apart from the component ones: name it per-model in
292
+ `HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
293
+ `flux-2-klein-base-4b.safetensors`) to save staging it, or pass
294
+ `ignore_patterns=` to replace the defaults outright. Whatever the globs, the
295
+ download and the hardware estimate use the same list, so the figures always
296
+ describe what actually gets staged:
297
+
298
+ ```python
299
+ >>> mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
300
+ ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
301
+ ```
302
+ - **Only model files are stored.** HuggingFace's download bookkeeping
303
+ (`.cache/huggingface/` inside the download folder) is dropped before the upload.
304
+ - **Shared across runs.** A model's identity is `(family, suffix, cortexgrid.IMPORTED)`,
305
+ derived from the repo id alone, so every run imports, deploys and deletes the same
306
+ copy — deleting removes it for all runs. Importing needs an active cortexgrid run,
307
+ which it tags with the model.
308
+
309
+ See the [cortexgrid model-serving docs](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md)
310
+ for the registry and serving state machines.
@@ -0,0 +1,284 @@
1
+ # cortexgrid-infer
2
+
3
+ The model-specific pieces that [cortexgrid](https://github.com/robodatalab/cortexgrid)
4
+ needs but cannot work out for itself, plus the clients that talk to what it deploys.
5
+
6
+ cortexgrid stores whatever weights it is handed under whatever key it is given, and
7
+ serves them with whatever class it was told to bundle. It deliberately knows nothing
8
+ about where a model came from or what shape it is. This library supplies exactly that
9
+ missing knowledge for one family of models at a time, and drives no lifecycle of its
10
+ own: importing, deploying and deleting stay your calls against the cortexgrid SDK.
11
+
12
+ ## Division of labour
13
+
14
+ | | Owns |
15
+ |---|---|
16
+ | **cortexgrid** | Registry identity, weights, serve bundle, hardware placement, deployment, jobs. Never knows what kind of model it holds. |
17
+ | **cortexgrid-infer** | Per family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
18
+ | **you** | Policy: which model, which requirements, when to import, deploy and tear down. |
19
+
20
+ An **importer** is the whole of this library's cluster-side surface. It gives you what
21
+ `cortexgrid.import_model` and `cortexgrid.deploy_model` need:
22
+
23
+ | | |
24
+ |---|---|
25
+ | `.family` / `.suffix` | how this model id names itself in the registry |
26
+ | `.serve_app` | the class cortexgrid bundles and runs on the cluster |
27
+ | `.requirements()` | what one replica needs to be placed and to run |
28
+ | `.config()` | settings the serve app reads at construction, if it needs any |
29
+ | `.client(url)` | a client that speaks the deployed app's routes |
30
+ | `.source` | *(models with weights)* a callable that downloads them and returns the directory |
31
+
32
+ | Importer | Serves | Client |
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.
43
+
44
+ ## Quick start
45
+
46
+ ```python
47
+ import asyncio
48
+ import cortexgrid
49
+ import cortexgrid_infer as mg
50
+
51
+ cortexgrid.Experiment.init("img-gen")
52
+ imp = mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
53
+
54
+ # 1. Import — weights go HuggingFace ─▶ cluster node ─▶ registry, never through
55
+ # your machine, so this runs as a job rather than in-process.
56
+ def import_weights(importer, requirements):
57
+ with importer:
58
+ cortexgrid.import_model(
59
+ importer.source, importer.serve_app,
60
+ family=importer.family, suffix=importer.suffix,
61
+ requirements=requirements,
62
+ )
63
+
64
+ # remote() hands back a JobFuture; result() blocks and re-raises whatever the
65
+ # job raised, so a failed import surfaces here with its own traceback.
66
+ cortexgrid.remote(
67
+ import_weights, imp, imp.requirements(), num_gpus=0, num_cpus=2
68
+ ).result()
69
+
70
+ # 2. Deploy — blocks until the Ray Serve app is running.
71
+ deployment = cortexgrid.deploy_model(
72
+ imp.family, imp.suffix, cortexgrid.IMPORTED, wait=True, timeout=1800
73
+ )
74
+
75
+ # 3. Inference.
76
+ model = imp.client(deployment.url)
77
+
78
+ async def run():
79
+ result = await mg.generate(model, "a red bicycle on a beach at sunrise")
80
+ open("out.png", "wb").write(result.image) # PNG bytes; also .width/.height/.params
81
+
82
+ asyncio.run(run())
83
+
84
+ # 4. Teardown.
85
+ cortexgrid.undeploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED) # free the GPU
86
+ cortexgrid.delete_model(imp.family, imp.suffix, cortexgrid.IMPORTED) # and the weights
87
+ ```
88
+
89
+ See [`examples/deploy_text_model`](examples/deploy_text_model/main.py) for the whole
90
+ lifecycle end to end.
91
+
92
+ An importer carries no open resources until it is entered, so it can be handed to
93
+ `cortexgrid.remote` and rebuilt on the node that runs the import. Entering it opens
94
+ the scratch directory `source` downloads into; `import_model` reads that directory
95
+ after `source` returns, which is why the lifetime spans the import call rather than
96
+ the download.
97
+
98
+ Call the import on every run. The weights are imported once and shared by every run:
99
+ on a model that is already `ready`, `import_model` downloads nothing, re-bundles the
100
+ serve code if it changed, and tags the run with the model.
101
+
102
+ ## Hardware requirements
103
+
104
+ `cortexgrid.ModelRequirements` decides which node Ray will place a replica on, and
105
+ nothing in the registry can derive it — cortexgrid stores whatever weights it is
106
+ handed, and an importer imports whatever repo id it is handed. `requirements()`
107
+ estimates it from the repo's own metadata: safetensors headers where the repo has
108
+ them, weight-file sizes otherwise. Both are read over HTTP; no weights are
109
+ downloaded.
110
+
111
+ ```python
112
+ >>> mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
113
+ ModelRequirements(num_gpus=1, ram_gb=3.0, vram_gb=3.2)
114
+ ```
115
+
116
+ These are estimates. Pass a `ModelRequirements` of your own to `import_model` where
117
+ you know better, or edit the figures on the model card in the dashboard afterwards —
118
+ `import_model` only fills in requirements a version does not already have, so a
119
+ hand-set value is never overwritten.
120
+
121
+ The estimate always asks for exactly one GPU, because both HuggingFace serve apps move
122
+ the whole model onto a single device. A model too large for any one card needs a serve
123
+ app that shards it, not a larger `num_gpus`: the requirement would simply never be
124
+ placed. (`AnthropicImport` asks for none — it holds no weights.)
125
+
126
+ ## Compilation
127
+
128
+ Both HuggingFace serve apps run `torch.compile` over the model they load, because
129
+ eager PyTorch bills per *operation* and a transformer forward is mostly glue —
130
+ views, broadcasts, permutes, elementwise adds — with only a small fraction of the
131
+ dispatches being the matrix multiplies that do the arithmetic. While the tensors
132
+ are large the host stays ahead of the accelerator and none of that shows; once
133
+ they are small, the accelerator drains its queue faster than the host can fill it
134
+ and step time tracks the operation count rather than the work. Feeding the model
135
+ less then stops helping, because the op count does not shrink with the input.
136
+ Inductor fuses the glue away; `reduce-overhead` additionally captures a CUDA graph,
137
+ so a step replays one graph instead of thousands of launches.
138
+
139
+ `cortexgrid_infer.compiling` offers two modes, and a serve app picks the one that
140
+ fits the work it does:
141
+
142
+ | | | for |
143
+ |---|---|---|
144
+ | `Mode.GRAPHED` (`reduce-overhead`) | fuses **and** captures a CUDA graph, so a forward replays one graph instead of issuing its launches | a loop of identically shaped forwards, where the capture is reused every iteration. A graph cannot go dynamic, so each distinct shape captures separately |
145
+ | `Mode.FUSED` (`default`) | fuses only | work whose shapes move under it, where there is nothing stable to capture and dynamo is free to compile one dynamic kernel set for all of them |
146
+
147
+ Both serve apps are tuned rather than defaulted:
148
+
149
+ - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
150
+ and every `text_encoder*` as `Mode.GRAPHED`. A denoising schedule is the ideal case —
151
+ every step is the same shapes — so the capture is made once and replayed for the
152
+ whole loop. The VAE is left alone: smallest win of the three, and the one whose
153
+ shapes move most, especially with tiling on.
154
+ - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
155
+ no stable shape until the KV cache is static, so it sets
156
+ `cache_implementation="static"` and a `CompileConfig(mode=Mode.GRAPHED)`, and
157
+ transformers compiles `generate` itself — capturing prefill and decode
158
+ separately, which it is better placed to do. Compiling the module here as well
159
+ would only compile the same forward twice. The cost is VRAM: a static cache is
160
+ preallocated to the length asked for, where a dynamic one grows into it.
161
+
162
+ Neither asks to be configured, for the same reason the serve apps do not ask which
163
+ dtype to load in. Both happen on CUDA and nowhere else — inductor is weakest off
164
+ it, and a host outrunning its accelerator is a GPU problem to begin with.
165
+
166
+ Two things worth knowing:
167
+
168
+ - **Warm-up is per input shape, and lands on a request.** Compilation is lazy, so
169
+ the first request at each image size — or each prompt and generation length —
170
+ pays for it.
171
+ - **Compilation never becomes load-bearing.** A backend that cannot compile a model
172
+ falls back to eager rather than failing the deployment.
173
+
174
+ It does not fix everything. Every forward re-reads the model's whole weight tensor
175
+ from memory, and no amount of fusion makes that read smaller. Below some input size
176
+ it is the floor; only quantizing the weights or running fewer forwards moves it.
177
+
178
+ ## Inference
179
+
180
+ `complete` streams `CompletionChunk`s (`.content`, `.tool_calls`, `.finish_reason`):
181
+
182
+ ```python
183
+ imp = mg.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
184
+ model = imp.client(deployment.url)
185
+ messages = [{"role": "user", "content": "Explain RAG in one sentence."}]
186
+
187
+ async for chunk in mg.complete(model, messages, max_new_tokens=512, temperature=0.7):
188
+ if chunk.has_text:
189
+ print(chunk.content, end="", flush=True)
190
+ ```
191
+
192
+ `generate` returns one `GeneratedImage` (`.image` PNG bytes, `.width`, `.height`,
193
+ `.params`); pass `image=` bytes for img2img.
194
+
195
+ ### Anthropic
196
+
197
+ An Anthropic model is a registry entry like any other — the difference is that
198
+ there are no weights to stage, so it is imported with no source, and the serve
199
+ app forwards to the API instead of loading anything:
200
+
201
+ ```python
202
+ cortexgrid.set_secret("ANTHROPIC_API_KEY", "sk-ant-...") # once
203
+
204
+ imp = mg.AnthropicImport("claude-sonnet-5")
205
+ cortexgrid.register_model(
206
+ imp.serve_app,
207
+ family=imp.family, suffix=imp.suffix,
208
+ requirements=imp.requirements(), config=imp.config(),
209
+ )
210
+ ```
211
+
212
+ `register_model` is `import_model` for a model that stages nothing: only the
213
+ serve bundle is stored, and the entry is indistinguishable from an imported one
214
+ to `deploy_model`, `list_models` and the dashboard.
215
+
216
+ `config` carries the Anthropic model name and the **name of** the cortexgrid
217
+ secret holding the key — never the key itself, since a registry entry is readable
218
+ by anyone who can see the model. The serve app reads it with
219
+ `cortexgrid.model_config` at construction, and both values are editable on the
220
+ model card, so pointing a deployment at a different model or key is an edit plus
221
+ a re-deploy rather than a re-registration.
222
+
223
+ A replica asks for no hardware at all, so it is placed on any node, CPU-only
224
+ included. From there it deploys and streams exactly like a cluster-served model:
225
+ `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
226
+ reports tool calls as structured blocks rather than as generated text, so the
227
+ serve app re-encodes them into the text form the client parses.
228
+
229
+ ## API
230
+
231
+ | | |
232
+ |---|---|
233
+ | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
234
+ | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
235
+ | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
236
+ | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
237
+ | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
238
+ | `complete(model, messages, tools=None, max_new_tokens=2048, temperature=0.7, **kw)` | Async stream of `CompletionChunk` for a `CompletingModel`. |
239
+ | `generate(model, prompt, *, image=None, **kw) -> GeneratedImage` | One image from a `GeneratingModel`. |
240
+ | `ServedCompletingModel(url, model_id)` | Client for any deployed completion app. |
241
+ | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
242
+ | `split_model_id(model_id) -> (family, suffix)` | The registry identity an importer derives. |
243
+ | `detect_device()` | The torch device a serve app should load onto. |
244
+ | `compiling.Mode.GRAPHED` / `compiling.Mode.FUSED` | The two modes a serve app chooses between. See [Compilation](#compilation). |
245
+ | `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. |
246
+ | `ModelDeployFailed` | Re-exported from cortexgrid; subclasses `RuntimeError`. |
247
+
248
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `CompletionChunk`,
249
+ `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
250
+
251
+ Adding a provider means one importer and one serve app. A completion app only has
252
+ to stream text from `POST /complete`, inlining tool calls as
253
+ `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
254
+ `encode_tool_call` for upstreams that report them structurally) — then
255
+ `ServedCompletingModel` is its client, unchanged.
256
+
257
+ ## Notes
258
+
259
+ - **HuggingFace auth.** For gated/private repos, pass `token=` to the importer (or
260
+ read `HF_TOKEN` yourself and pass it); it travels with the importer to the node
261
+ that downloads.
262
+ - **Image weight pruning.** `HuggingFaceImageImport` skips example images, docs and
263
+ `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
264
+ skip a repo's consolidated single-file checkpoint, which is weights it cannot
265
+ tell apart from the component ones: name it per-model in
266
+ `HF_IMAGE_SNAPSHOT_IGNORE` (comma-separated globs, e.g.
267
+ `flux-2-klein-base-4b.safetensors`) to save staging it, or pass
268
+ `ignore_patterns=` to replace the defaults outright. Whatever the globs, the
269
+ download and the hardware estimate use the same list, so the figures always
270
+ describe what actually gets staged:
271
+
272
+ ```python
273
+ >>> mg.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
274
+ ModelRequirements(num_gpus=1, ram_gb=26.3, vram_gb=29.6) # single-file checkpoint included
275
+ ```
276
+ - **Only model files are stored.** HuggingFace's download bookkeeping
277
+ (`.cache/huggingface/` inside the download folder) is dropped before the upload.
278
+ - **Shared across runs.** A model's identity is `(family, suffix, cortexgrid.IMPORTED)`,
279
+ derived from the repo id alone, so every run imports, deploys and deletes the same
280
+ copy — deleting removes it for all runs. Importing needs an active cortexgrid run,
281
+ which it tags with the model.
282
+
283
+ See the [cortexgrid model-serving docs](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md)
284
+ for the registry and serving state machines.
@@ -0,0 +1,69 @@
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 per-family knowledge those calls need
6
+ and nothing else: an importer that names, fetches, sizes and serves one family
7
+ of models, and a client that speaks the deployed app's routes.
8
+
9
+ imp = cortexgrid_infer.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct")
10
+ with imp:
11
+ cortexgrid.import_model(
12
+ imp.source, imp.serve_app,
13
+ family=imp.family, suffix=imp.suffix,
14
+ requirements=imp.requirements(),
15
+ )
16
+ deployment = cortexgrid.deploy_model(imp.family, imp.suffix, cortexgrid.IMPORTED, wait=True)
17
+ model = imp.client(deployment.url)
18
+
19
+ See `examples/deploy_text_model` for the whole lifecycle, import job included.
20
+ """
21
+
22
+ from cortexgrid import ModelDeployFailed
23
+
24
+ from cortexgrid_infer.core import (
25
+ CompletionChunk,
26
+ DeployedModel,
27
+ CompletingModel,
28
+ GeneratedImage,
29
+ GeneratingModel,
30
+ Message,
31
+ Tool,
32
+ ToolCall,
33
+ ToolSpec,
34
+ complete,
35
+ generate,
36
+ )
37
+ from cortexgrid_infer.completion import ServedCompletingModel
38
+ from cortexgrid_infer.device import detect_device
39
+ from cortexgrid_infer.importing import HuggingFaceImport, ModelImport, split_model_id
40
+ from cortexgrid_infer.providers.anthropic import AnthropicImport
41
+ from cortexgrid_infer.providers.huggingface_complete import HuggingFaceCompletingImport
42
+ from cortexgrid_infer.providers.huggingface_image import (
43
+ HuggingFaceImageImport,
44
+ HuggingFaceImageModel,
45
+ )
46
+
47
+ __all__ = [
48
+ "CompletionChunk",
49
+ "DeployedModel",
50
+ "CompletingModel",
51
+ "GeneratedImage",
52
+ "GeneratingModel",
53
+ "Message",
54
+ "Tool",
55
+ "ToolCall",
56
+ "ToolSpec",
57
+ "complete",
58
+ "generate",
59
+ "ModelDeployFailed",
60
+ "detect_device",
61
+ "split_model_id",
62
+ "ModelImport",
63
+ "HuggingFaceImport",
64
+ "HuggingFaceCompletingImport",
65
+ "HuggingFaceImageImport",
66
+ "AnthropicImport",
67
+ "ServedCompletingModel",
68
+ "HuggingFaceImageModel",
69
+ ]