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