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.
Files changed (31) hide show
  1. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/PKG-INFO +91 -55
  2. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/README.md +88 -53
  3. cortexgrid_infer-0.3.1/cortexgrid_infer/__init__.py +95 -0
  4. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/core.py +27 -0
  5. cortexgrid_infer-0.3.1/cortexgrid_infer/imaging.py +74 -0
  6. cortexgrid_infer-0.3.1/cortexgrid_infer/importers/__init__.py +10 -0
  7. cortexgrid_infer-0.3.1/cortexgrid_infer/importers/base.py +99 -0
  8. cortexgrid_infer-0.3.1/cortexgrid_infer/importers/huggingface.py +116 -0
  9. cortexgrid_infer-0.3.1/cortexgrid_infer/meshing.py +72 -0
  10. cortexgrid_infer-0.3.1/cortexgrid_infer/models/__init__.py +18 -0
  11. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/anthropic_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/anthropic.py +30 -10
  12. cortexgrid_infer-0.3.1/cortexgrid_infer/models/base.py +121 -0
  13. cortexgrid_infer-0.3.1/cortexgrid_infer/models/image2mesh.py +84 -0
  14. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_image_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/text2image.py +27 -4
  15. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_complete_serve.py → cortexgrid_infer-0.3.1/cortexgrid_infer/models/text2text.py +15 -2
  16. cortexgrid_infer-0.3.1/cortexgrid_infer/registry.py +102 -0
  17. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/pyproject.toml +5 -3
  18. cortexgrid_infer-0.2.6/cortexgrid_infer/__init__.py +0 -69
  19. cortexgrid_infer-0.2.6/cortexgrid_infer/importing.py +0 -158
  20. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/__init__.py +0 -0
  21. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/anthropic.py +0 -65
  22. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_complete.py +0 -23
  23. cortexgrid_infer-0.2.6/cortexgrid_infer/providers/huggingface_image.py +0 -110
  24. cortexgrid_infer-0.2.6/cortexgrid_infer/requirements.py +0 -117
  25. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/.gitignore +0 -0
  26. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/LICENSE +0 -0
  27. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/compiling.py +0 -0
  28. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/completion.py +0 -0
  29. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/device.py +0 -0
  30. {cortexgrid_infer-0.2.6 → cortexgrid_infer-0.3.1}/cortexgrid_infer/py.typed +0 -0
  31. {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.2.6
4
- Summary: Deploy and interact with models across providers (Anthropic, HuggingFace LLM + diffusion)
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 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.
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 family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
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
- 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:
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` | *(models with weights)* a callable that downloads them and returns the directory |
85
+ | `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
57
86
 
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.
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.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
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
- 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.
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.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
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 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.)
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
- Both HuggingFace serve apps run `torch.compile` over the model they load, because
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
- - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
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
- - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
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.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
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 it is imported with no source, and the serve
240
- app forwards to the API instead of loading anything:
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
- imp = mg.AnthropicImport("claude-sonnet-5")
268
+ entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
246
269
  cortexgrid.register_model(
247
- imp.serve_app,
248
- family=imp.family, suffix=imp.suffix,
249
- requirements=imp.requirements(), config=imp.config(),
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
- `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
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
- | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
275
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
276
- | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
277
- | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
278
- | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
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
- | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
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`, `CompletionChunk`,
290
- `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
317
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
318
+ `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
319
+ `ToolSpec`.
320
+
321
+ Extending it:
291
322
 
292
- Adding a provider means one importer and one serve app. A completion app only has
293
- to stream text from `POST /complete`, inlining tool calls as
294
- `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
295
- `encode_tool_call` for upstreams that report them structurally) — then
296
- `ServedCompletingModel` is its client, unchanged.
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.** `HuggingFaceImageImport` skips example images, docs and
304
- `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
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.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
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 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.
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 family: how to name it, fetch it, run it, size it, talk to it. Never drives a lifecycle. |
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
- 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
+ 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` | *(models with weights)* a callable that downloads them and returns the directory |
58
+ | `.source` | *(importers only)* a callable that downloads the weights and returns the directory |
31
59
 
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.
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.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B")
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
- 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.
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.HuggingFaceCompletingImport("Qwen/Qwen2.5-0.5B-Instruct").requirements()
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 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.)
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
- Both HuggingFace serve apps run `torch.compile` over the model they load, because
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
- - **`HuggingFaceImageDeployment`** compiles the denoiser (`transformer` or `unet`)
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
- - **`HuggingFaceCompletingDeployment`** does *not* compile the module. Decode has
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.HuggingFaceCompletingImport("Qwen/Qwen2.5-7B-Instruct")
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 it is imported with no source, and the serve
214
- app forwards to the API instead of loading anything:
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
- imp = mg.AnthropicImport("claude-sonnet-5")
241
+ entry = mg.Hosted("claude-sonnet-5", mg.AnthropicText2Text) # api_key_secret="..." to use another secret
220
242
  cortexgrid.register_model(
221
- imp.serve_app,
222
- family=imp.family, suffix=imp.suffix,
223
- requirements=imp.requirements(), config=imp.config(),
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
- `imp.client(deployment.url)` returns the same `ServedCompletingModel`. Anthropic
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
- | `HuggingFaceCompletingImport(hf_id, token=None)` | Importer for a causal LM. |
249
- | `HuggingFaceImageImport(hf_id, token=None, ignore_patterns=None)` | Importer for a diffusers pipeline. |
250
- | `AnthropicImport(model_id, api_key_secret="ANTHROPIC_API_KEY")` | Importer for an Anthropic model; no weights. |
251
- | `ModelImport` | Base of all three: `family`, `suffix`, `serve_app`, `requirements()`, `config()`, `client(url)`. |
252
- | `HuggingFaceImport` | Adds the download half — `source` and the scratch directory. Subclass it for another HuggingFace family. |
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
- | `HuggingFaceImageModel(url, model_id)` | Client for a deployed diffusers app. |
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`, `CompletionChunk`,
264
- `GeneratedImage`, `ToolCall`, `Message`, `Tool`, `ToolSpec`.
290
+ Types: `DeployedModel`, `CompletingModel`, `GeneratingModel`, `MeshingModel`,
291
+ `CompletionChunk`, `GeneratedImage`, `GeneratedMesh`, `ToolCall`, `Message`, `Tool`,
292
+ `ToolSpec`.
293
+
294
+ Extending it:
265
295
 
266
- Adding a provider means one importer and one serve app. A completion app only has
267
- to stream text from `POST /complete`, inlining tool calls as
268
- `<tool_call>{"name": ..., "arguments": {...}}</tool_call>` (there is an
269
- `encode_tool_call` for upstreams that report them structurally) — then
270
- `ServedCompletingModel` is its client, unchanged.
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.** `HuggingFaceImageImport` skips example images, docs and
278
- `.gitattributes` by default — nothing `from_pretrained` reads. It does **not**
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.HuggingFaceImageImport("black-forest-labs/FLUX.2-klein-base-4B").requirements()
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
+ ]