cortexgrid 0.3.1__tar.gz → 0.3.3__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: cortexgrid
3
- Version: 0.3.1
3
+ Version: 0.3.3
4
4
  Summary: Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3
5
5
  Project-URL: Homepage, https://github.com/robodatalab/cortexgrid
6
6
  Project-URL: Repository, https://github.com/robodatalab/cortexgrid
@@ -232,7 +232,9 @@ print(deployed.url)
232
232
 
233
233
  The requirements are part of the model, not of the serve-app class: GPUs, RAM and VRAM (GiB, 0 meaning no requirement) are stored with it and matched against what the cluster's hosts have free. Correct them later with `cortexgrid.set_model_requirements(family, suffix, run_name, requirements)` or on the model card in the dashboard; `cortexgrid.deploy_model(..., num_replicas=2)` chooses how many copies to run.
234
234
 
235
- `save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
235
+ Anything else the serve-app has to know about the model - which model a provider should be asked for, an endpoint, the name of a secret to read - goes in a free-form string mapping on the same entry: `save_model(..., config={"model": "claude-opus-5"})`. The serve-app reads it in `__init__` with `cortexgrid.model_config(family, suffix, run_name)`; `cortexgrid.set_model_config(family, suffix, run_name, config)` or the model card replaces it. cortexgrid stores the mapping without interpreting it, and a tag is readable by anyone with registry access, so a credential belongs in `set_secret` with only its name in the config.
236
+
237
+ `save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`. A model with no weights to stage - one behind a provider's API, e.g. Gemini or OpenAI - is registered the same way by `cortexgrid.register_model(MyServeApp, family, suffix, config=...)`: same key, same reuse, only the bundle stored. [Serving a hosted-API model](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md#serving-a-hosted-api-model) walks through one end to end - serve-app, API key, registration, deploy.
236
238
 
237
239
  `save_model` is synchronous (registry lifecycle: `uploading` -> `ready`); `deploy_model` schedules the serving lifecycle (`deploying` -> `running`). With `wait=True` a failed deploy raises `cortexgrid.ModelDeployFailed`; `cortexgrid.wait_for_model_serving(family, suffix, run_name, timeout=...)` waits on a deploy started elsewhere, and re-deploying a failed model retries it from scratch. See [model-serving.md](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md) for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (`model_registry_status`, `model_serving_status`), and error handling.
238
240
 
@@ -85,14 +85,18 @@ from cortexgrid.ray_util import (
85
85
  from cortexgrid.s3_util import delete_prefix, download, get_s3_client, upload, upload_dir
86
86
  from cortexgrid.model_storage import (
87
87
  IMPORTED,
88
+ NO_WEIGHTS,
88
89
  SavedModel,
89
90
  delete_model,
90
91
  list_models,
91
92
  load_model,
93
+ model_config,
92
94
  model_registry_status,
95
+ set_model_config,
93
96
  set_model_requirements,
94
97
  )
95
98
  from cortexgrid.model_storage import import_model as _import_model_storage
99
+ from cortexgrid.model_storage import register_model as _register_model_storage
96
100
  from cortexgrid.model_storage import save_model as _save_model_storage
97
101
  from cortexgrid.model_serving import (
98
102
  Deployment,
@@ -149,10 +153,12 @@ def save_model(
149
153
  family: str,
150
154
  suffix: str,
151
155
  requirements: ModelRequirements | None = None,
156
+ config: dict[str, str] | None = None,
152
157
  ) -> SavedModel:
153
158
  """Persist a weights directory under the current Experiment's run, paired
154
- with the serve-app class that will front it at deploy time and the
155
- hardware one replica of it needs.
159
+ with the serve-app class that will front it at deploy time, the hardware
160
+ one replica of it needs, and any `config` the serve-app reads with
161
+ `model_config`.
156
162
 
157
163
  Every run saves a new copy under its own run_name - meant for weights the
158
164
  run produced (e.g. a fine-tune). For a model produced elsewhere that should
@@ -166,6 +172,7 @@ def save_model(
166
172
  run_id=experiment.run_id,
167
173
  run_name=experiment.run_name(),
168
174
  requirements=requirements,
175
+ config=config,
169
176
  )
170
177
 
171
178
 
@@ -175,6 +182,7 @@ def import_model(
175
182
  family: str,
176
183
  suffix: str,
177
184
  requirements: ModelRequirements | None = None,
185
+ config: dict[str, str] | None = None,
178
186
  ) -> SavedModel:
179
187
  """Register a model produced elsewhere once, reuse it on every later call,
180
188
  and record on the current Experiment's run which imported model it used.
@@ -184,7 +192,35 @@ def import_model(
184
192
  `imported_model/<family>/<suffix>` holds the version's `created_at`, set
185
193
  whether this call uploaded the model or reused it."""
186
194
  experiment = Experiment.get_instance()
187
- model = _import_model_storage(source, serve_app, family, suffix, requirements)
195
+ model = _import_model_storage(
196
+ source, serve_app, family, suffix, requirements, config
197
+ )
198
+ get_mlflow_client().set_tag(
199
+ experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
200
+ )
201
+ return model
202
+
203
+
204
+ def register_model(
205
+ serve_app: type,
206
+ family: str,
207
+ suffix: str,
208
+ requirements: ModelRequirements | None = None,
209
+ config: dict[str, str] | None = None,
210
+ ) -> SavedModel:
211
+ """Register a model with no weights of its own once - a serve-app that
212
+ forwards to a hosted API stages nothing - reuse it on every later call, and
213
+ record on the current Experiment's run which one it used.
214
+
215
+ `import_model` without the import: everything it needs beyond its code goes
216
+ in `config`, which the serve-app reads with `model_config` at construction
217
+ (see `cortexgrid.model_storage.register_model`). The model belongs to no
218
+ run, so the run keeps the link the same way, under the same tag
219
+ `imported_model/<family>/<suffix>`."""
220
+ experiment = Experiment.get_instance()
221
+ model = _register_model_storage(
222
+ serve_app, family, suffix, requirements, config
223
+ )
188
224
  get_mlflow_client().set_tag(
189
225
  experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
190
226
  )
@@ -247,14 +283,18 @@ __all__ = [
247
283
  "delete_secret",
248
284
  # Model registry
249
285
  "IMPORTED",
286
+ "NO_WEIGHTS",
250
287
  "SavedModel",
251
288
  "save_model",
252
289
  "import_model",
290
+ "register_model",
253
291
  "load_model",
254
292
  "list_models",
255
293
  "model_registry_status",
256
294
  "ModelRequirements",
257
295
  "set_model_requirements",
296
+ "model_config",
297
+ "set_model_config",
258
298
  "delete_model",
259
299
  # Model serving
260
300
  "Deployment",
@@ -6,16 +6,21 @@ Mapping cortexgrid taxonomy <-> MLflow Registry:
6
6
  family, suffix -> ModelVersion.tags["family"], ["suffix"] (denormalized)
7
7
  weights blob path -> ModelVersion.source =
8
8
  "s3://<bucket>/models/<run_name>/<family>/<suffix>/weights/"
9
+ (NO_WEIGHTS for a model registered without any)
9
10
  run linkage -> ModelVersion.run_id (built-in MLflow field; unset
10
11
  for imported models)
11
12
  requirements -> ModelVersion.tags["num_gpus"], ["ram_gb"], ["vram_gb"]
13
+ config -> ModelVersion.tags["config"] (JSON object)
12
14
 
13
- Two ways in: `save_model` registers a fresh copy under the calling run's
15
+ Three ways in: `save_model` registers a fresh copy under the calling run's
14
16
  run_name every time it runs (fine-tuned output); `import_model` registers a
15
17
  model produced elsewhere once, under the fixed run_name IMPORTED, and after
16
- that only re-bundles the serve-app when its code changed. Both write the same
17
- layout, so every
18
- (family, suffix, run_name) consumer - load_model, deploy_model - handles both.
18
+ that only re-bundles the serve-app when its code changed; `register_model`
19
+ registers a model that stages no weights at all - a serve-app that forwards to
20
+ a hosted API holds none - under that same fixed key. All three write the same
21
+ layout, so every (family, suffix, run_name) consumer - deploy_model, the
22
+ dashboard - handles them alike; only `load_model` parts them, having nothing to
23
+ hand back for a model registered without weights.
19
24
 
20
25
  storage.py is pure: it takes run_id/run_name as explicit args and never reads
21
26
  the active Experiment singleton. The facade that fills those in lives in
@@ -24,8 +29,9 @@ cortexgrid/__init__.py.
24
29
 
25
30
  from __future__ import annotations
26
31
 
32
+ import json
27
33
  import os
28
- from dataclasses import dataclass
34
+ from dataclasses import dataclass, field
29
35
  from datetime import datetime, timedelta, timezone
30
36
  from pathlib import Path
31
37
  import tempfile
@@ -58,6 +64,19 @@ _PHASE_READY = "ready"
58
64
  _PHASE_UPLOAD_FAILED = "upload_failed"
59
65
  _PHASE_BROKEN = "broken"
60
66
 
67
+ # MLflow ModelVersion tag holding the model's config: whatever settings its
68
+ # serve-app needs that are not the weights (a provider's model id, an endpoint,
69
+ # the name of a secret to read). cortexgrid never interprets it - it is the
70
+ # serve-app's own vocabulary, stored next to the model so the app reads it at
71
+ # construction instead of being redeployed to change a setting.
72
+ #
73
+ # One JSON object in one tag, not a tag per key: the keys are the serve-app's
74
+ # to choose, free of MLflow's tag-key charset, and the whole mapping is
75
+ # replaced in a single write, so removing a key needs no tag deletion. MLflow
76
+ # caps a tag value at 8000 characters, which bounds how big a config can get.
77
+ _CONFIG_TAG = "config"
78
+
79
+
61
80
  # An upload still marked "uploading" this long after the version was created is
62
81
  # treated as broken: the version is created before the weights are resolved
63
82
  # (for `import_model`, before its download), so creation_timestamp is the
@@ -71,6 +90,13 @@ _UPLOAD_DEADLINE = timedelta(hours=3)
71
90
  # them. Run names are haikunator "word-word-NN", so no run can take this name.
72
91
  IMPORTED = "imported"
73
92
 
93
+ # ModelVersion.source of a model registered with no weights of its own: nothing
94
+ # was staged, so there is no blob to point at. Spelled as a URI rather than left
95
+ # blank because MLflow rejects a source that is empty or a local path, and
96
+ # because every reader - `load_model`, the dashboard's storage field - then sees
97
+ # why there is no path instead of an empty one.
98
+ NO_WEIGHTS = "cortexgrid://no-weights"
99
+
74
100
 
75
101
  @dataclass
76
102
  class SavedModel:
@@ -78,7 +104,9 @@ class SavedModel:
78
104
  suffix: str
79
105
  run_name: str
80
106
  created_at: str
107
+ # Where the weights live, or NO_WEIGHTS for a model registered without any.
81
108
  data_blob_path: str
109
+ # Size of the weights; 0 for a model registered without any.
82
110
  size_bytes: int
83
111
  # Registry lifecycle phase: "uploading" while save_model streams the weights
84
112
  # and serve bundle to storage, "ready" once that finishes, "upload_failed"
@@ -88,6 +116,16 @@ class SavedModel:
88
116
  phase: str
89
117
  # Hardware one replica needs; defaults for versions stored without it.
90
118
  requirements: ModelRequirements
119
+ # Free-form settings the serve-app reads at construction; empty for
120
+ # versions stored without any.
121
+ config: dict[str, str] = field(default_factory=dict)
122
+
123
+ @property
124
+ def has_weights(self) -> bool:
125
+ """Whether this model staged weights of its own. False for one
126
+ registered with `register_model`, whose serve-app holds no bytes to
127
+ store; `load_model` on it raises."""
128
+ return self.data_blob_path != NO_WEIGHTS
91
129
 
92
130
 
93
131
  def _phase_for(version: Any) -> str:
@@ -107,6 +145,28 @@ def _phase_for(version: Any) -> str:
107
145
  return phase
108
146
 
109
147
 
148
+ def _config_to_tag(config: dict[str, str]) -> str:
149
+ """Serialise a config mapping to its MLflow tag value.
150
+
151
+ Keys and values are strings: they round-trip through a tag, and the model
152
+ card edits them as text. A caller with a number or a flag spells it as a
153
+ string and the serve-app parses it back."""
154
+ for key, value in config.items():
155
+ if not isinstance(key, str) or not isinstance(value, str):
156
+ raise ValueError(
157
+ f"Model config must map strings to strings: {key!r}: {value!r}"
158
+ )
159
+ if not key.strip():
160
+ raise ValueError("Model config keys cannot be blank")
161
+ return json.dumps(config)
162
+
163
+
164
+ def _config_from_tags(tags: dict[str, str]) -> dict[str, str]:
165
+ """Deserialise a config mapping from a ModelVersion's MLflow tags; a
166
+ missing tag reads as no config."""
167
+ return json.loads(tags.get(_CONFIG_TAG, "{}"))
168
+
169
+
110
170
  def _to_saved_model(version: Any) -> SavedModel:
111
171
  return SavedModel(
112
172
  family=version.tags["family"],
@@ -119,6 +179,7 @@ def _to_saved_model(version: Any) -> SavedModel:
119
179
  size_bytes=int(version.tags.get("size_bytes", "0")),
120
180
  phase=_phase_for(version),
121
181
  requirements=requirements_from_tags(version.tags),
182
+ config=_config_from_tags(version.tags),
122
183
  )
123
184
 
124
185
 
@@ -159,6 +220,7 @@ def save_model(
159
220
  run_id: str,
160
221
  run_name: str,
161
222
  requirements: ModelRequirements | None = None,
223
+ config: dict[str, str] | None = None,
162
224
  ) -> SavedModel:
163
225
  """Upload a weights directory to S3 and register a new MLflow ModelVersion
164
226
  paired with the serve-app that fronts it.
@@ -180,9 +242,20 @@ def save_model(
180
242
  without the caller holding the class object.
181
243
 
182
244
  `requirements` is the hardware one replica needs; None stores none, which
183
- reads as no requirement. Change it later with `set_model_requirements`."""
245
+ reads as no requirement. Change it later with `set_model_requirements`.
246
+
247
+ `config` is whatever else the serve-app needs to know about this model,
248
+ as a string mapping it reads with `model_config` at construction; None
249
+ stores none. Change it later with `set_model_config`."""
184
250
  return _upload_model(
185
- weights_dir, serve_app, suffix, family, run_id, run_name, requirements
251
+ weights_dir,
252
+ serve_app,
253
+ suffix,
254
+ family,
255
+ run_id,
256
+ run_name,
257
+ requirements,
258
+ config,
186
259
  )
187
260
 
188
261
 
@@ -192,6 +265,7 @@ def import_model(
192
265
  family: str,
193
266
  suffix: str,
194
267
  requirements: ModelRequirements | None = None,
268
+ config: dict[str, str] | None = None,
195
269
  ) -> SavedModel:
196
270
  """Register a model produced elsewhere (e.g. a pretrained base model) under
197
271
  the fixed key (family, suffix, IMPORTED), once.
@@ -205,16 +279,62 @@ def import_model(
205
279
  If a version is already registered under the key:
206
280
  - "ready": returns it without calling `source`. If `serve_app`'s code no
207
281
  longer matches the stored bundle, it is re-bundled first and the
208
- weights are kept (see `_refresh_bundle`). `requirements` are stored
209
- only if the version has none yet, so values changed since with
210
- `set_model_requirements` are kept. To replace the weights,
211
- `delete_model` it first.
282
+ weights are kept (see `_refresh_bundle`). `requirements` and `config`
283
+ are stored only if the version has none yet, so values changed since
284
+ with `set_model_requirements` / `set_model_config` are kept. To replace
285
+ the weights, `delete_model` it first.
212
286
  - "uploading": raises RuntimeError - another process is importing it.
213
287
  - "upload_failed" / "broken": deleted and imported again.
214
288
 
215
289
  The version is linked to no MLflow run, so deleting a run leaves it in
216
290
  place. Deploy it like any saved model:
217
291
  `deploy_model(family, suffix, IMPORTED)`."""
292
+ return _register_imported(
293
+ source, serve_app, family, suffix, requirements, config
294
+ )
295
+
296
+
297
+ def register_model(
298
+ serve_app: type,
299
+ family: str,
300
+ suffix: str,
301
+ requirements: ModelRequirements | None = None,
302
+ config: dict[str, str] | None = None,
303
+ ) -> SavedModel:
304
+ """Register a model that stages no weights under the fixed key
305
+ (family, suffix, IMPORTED), once.
306
+
307
+ For a model whose bytes are not ours to hold: a serve-app that forwards
308
+ requests to a hosted API, or one that reaches for the weights itself at
309
+ startup. There is nothing to upload, so only `serve_app`'s bundle is
310
+ stored and the version's source reads NO_WEIGHTS - `load_model` on such a
311
+ model raises, and `SavedModel.has_weights` is False. What it needs instead
312
+ of weights - which model a provider should be asked for, the name of the
313
+ secret holding the key - belongs in `config`, which the serve-app reads
314
+ with `model_config` at construction.
315
+
316
+ Registration is otherwise `import_model`'s, down to the phase an already
317
+ registered version leaves it in ("ready" is reused and its bundle
318
+ refreshed, "uploading" raises, "upload_failed"/"broken" is replaced), so
319
+ the entry is indistinguishable from an imported one to `deploy_model`,
320
+ `list_models` and the dashboard."""
321
+ return _register_imported(
322
+ None, serve_app, family, suffix, requirements, config
323
+ )
324
+
325
+
326
+ def _register_imported(
327
+ weights: str | Path | Callable[[], str | Path] | None,
328
+ serve_app: type,
329
+ family: str,
330
+ suffix: str,
331
+ requirements: ModelRequirements | None,
332
+ config: dict[str, str] | None,
333
+ ) -> SavedModel:
334
+ """Register `weights` under (family, suffix, IMPORTED) unless a version is
335
+ already there - the once-only registration `import_model` and
336
+ `register_model` share; `weights` is None for the model that stages
337
+ none."""
218
338
  existing = model_registry_status(family, suffix, IMPORTED)
219
339
  if existing is not None:
220
340
  if existing.phase == _PHASE_READY:
@@ -223,6 +343,8 @@ def import_model(
223
343
  existing.requirements = _set_missing_requirements(
224
344
  family, suffix, requirements
225
345
  )
346
+ if config is not None:
347
+ existing.config = _set_missing_config(family, suffix, config)
226
348
  return existing
227
349
  if existing.phase == _PHASE_UPLOADING:
228
350
  raise RuntimeError(
@@ -231,7 +353,7 @@ def import_model(
231
353
  )
232
354
  delete_model(family, suffix, IMPORTED)
233
355
  return _upload_model(
234
- source, serve_app, suffix, family, None, IMPORTED, requirements
356
+ weights, serve_app, suffix, family, None, IMPORTED, requirements, config
235
357
  )
236
358
 
237
359
 
@@ -273,21 +395,47 @@ def _set_missing_requirements(
273
395
  return requirements
274
396
 
275
397
 
398
+ def _set_missing_config(
399
+ family: str, suffix: str, config: dict[str, str]
400
+ ) -> dict[str, str]:
401
+ """Store `config` on an imported model whose version has none yet. Returns
402
+ the config the version holds afterwards."""
403
+ name = f"{family}__{suffix}"
404
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
405
+ version = client.search_model_versions(
406
+ f"name='{name}' and tags.run_name='{IMPORTED}'"
407
+ )[0]
408
+ if _CONFIG_TAG in version.tags:
409
+ return _config_from_tags(version.tags)
410
+ client.set_model_version_tag(
411
+ name, version.version, _CONFIG_TAG, _config_to_tag(config)
412
+ )
413
+ return config
414
+
415
+
276
416
  def _upload_model(
277
- weights: str | Path | Callable[[], str | Path],
417
+ weights: str | Path | Callable[[], str | Path] | None,
278
418
  serve_app: type,
279
419
  suffix: str,
280
420
  family: str,
281
421
  run_id: str | None,
282
422
  run_name: str,
283
423
  requirements: ModelRequirements | None,
424
+ config: dict[str, str] | None,
284
425
  ) -> SavedModel:
285
426
  """Register a ModelVersion in "uploading", resolve `weights` to a directory
286
427
  (calling it when it is a callable), upload the weights and the serve-app
287
- bundle, and flip it to "ready" (or "upload_failed")."""
288
- bucket = get_s3_bucket()
428
+ bundle, and flip it to "ready" (or "upload_failed").
429
+
430
+ `weights` is None for a model that stages none: its source reads
431
+ NO_WEIGHTS and the bundle is the only thing uploaded. The phases are the
432
+ same either way, so one registration path covers both."""
289
433
  prefix = f"models/{run_name}/{family}/{suffix}"
290
- source = f"s3://{bucket}/{prefix}/weights/"
434
+ source = (
435
+ NO_WEIGHTS
436
+ if weights is None
437
+ else f"s3://{get_s3_bucket()}/{prefix}/weights/"
438
+ )
291
439
  name = f"{family}__{suffix}"
292
440
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
293
441
  _ensure_registered_model(client, name)
@@ -296,8 +444,11 @@ def _upload_model(
296
444
  # fetched and uploaded, and a concurrent `import_model` sees the import in
297
445
  # flight for the whole download instead of starting one of its own. The
298
446
  # size is stamped once the directory exists; the bundle tags and the flip
299
- # to "ready" happen only after the upload lands. No requirements leaves
300
- # their tags unset, so a later `import_model` can still store them.
447
+ # to "ready" happen only after the upload lands. A weights-less
448
+ # registration has only its bundle to upload, and passes through the same
449
+ # phases. No requirements and no
450
+ # config leave their tags unset, so a later `import_model` can still store
451
+ # them.
301
452
  version = client.create_model_version(
302
453
  name=name,
303
454
  source=source,
@@ -308,14 +459,16 @@ def _upload_model(
308
459
  "run_name": run_name,
309
460
  _LIFECYCLE_TAG: _PHASE_UPLOADING,
310
461
  **(requirements_to_tags(requirements) if requirements is not None else {}),
462
+ **({_CONFIG_TAG: _config_to_tag(config)} if config is not None else {}),
311
463
  },
312
464
  )
313
465
  try:
314
- weights_dir = weights() if callable(weights) else weights
315
- client.set_model_version_tag(
316
- name, version.version, "size_bytes", str(_dir_size_bytes(weights_dir))
317
- )
318
- s3_util.upload_dir(str(weights_dir), dest_path=f"{prefix}/weights")
466
+ if weights is not None:
467
+ weights_dir = weights() if callable(weights) else weights
468
+ client.set_model_version_tag(
469
+ name, version.version, "size_bytes", str(_dir_size_bytes(weights_dir))
470
+ )
471
+ s3_util.upload_dir(str(weights_dir), dest_path=f"{prefix}/weights")
319
472
  bundle_meta = bundle_class(serve_app, family, suffix, run_name)
320
473
  for key, value in metadata_to_tags(bundle_meta).items():
321
474
  client.set_model_version_tag(name, version.version, key, value)
@@ -337,7 +490,10 @@ def load_model(family: str, suffix: str, run_name: str) -> Path:
337
490
  contents; the serve-app reconstructs the model from it however it likes
338
491
  (`from_pretrained`, `torch.load`, ...). The returned directory persists
339
492
  after this call - the caller (typically a serve-app loading weights at
340
- startup) owns its lifetime."""
493
+ startup) owns its lifetime.
494
+
495
+ Raises ValueError for a model registered with `register_model`: it stages
496
+ no weights, so there is none to hand back."""
341
497
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
342
498
  name = f"{family}__{suffix}"
343
499
  versions = client.search_model_versions(
@@ -345,6 +501,11 @@ def load_model(family: str, suffix: str, run_name: str) -> Path:
345
501
  )
346
502
  if not versions or not versions[0].source:
347
503
  raise ValueError(f"No model {family}/{suffix}/{run_name}")
504
+ if versions[0].source == NO_WEIGHTS:
505
+ raise ValueError(
506
+ f"Model {family}/{suffix}/{run_name} was registered without "
507
+ "weights; there is nothing to load"
508
+ )
348
509
  return _download_s3_uri(versions[0].source, None)
349
510
 
350
511
 
@@ -393,6 +554,46 @@ def set_model_requirements(
393
554
  client.set_model_version_tag(name, versions[0].version, key, value)
394
555
 
395
556
 
557
+ def model_config(family: str, suffix: str, run_name: str) -> dict[str, str]:
558
+ """The config mapping stored on a model, empty if it has none.
559
+
560
+ Meant for the serve-app to call in `__init__` with the
561
+ (family, suffix, run_name) it was constructed with: the settings that are
562
+ not the weights - a provider's model id, an endpoint, the name of a secret
563
+ to read - travel with the registry entry instead of the bundled code, so
564
+ changing one is an edit on the model card rather than a re-save.
565
+
566
+ Read at construction, so a replica keeps the values it started with until
567
+ it is deployed again. Raises ValueError if the model was never
568
+ registered."""
569
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
570
+ versions = client.search_model_versions(
571
+ f"name='{family}__{suffix}' and tags.run_name='{run_name}'"
572
+ )
573
+ if not versions:
574
+ raise ValueError(f"No model {family}/{suffix}/{run_name}")
575
+ return _config_from_tags(versions[0].tags)
576
+
577
+
578
+ def set_model_config(
579
+ family: str, suffix: str, run_name: str, config: dict[str, str]
580
+ ) -> None:
581
+ """Replace the config mapping stored on a model - the whole mapping, so a
582
+ key left out of `config` is gone. Takes effect on its next `deploy_model`;
583
+ a replica already running keeps the values it read at construction.
584
+ Raises ValueError if the model was never registered."""
585
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
586
+ name = f"{family}__{suffix}"
587
+ versions = client.search_model_versions(
588
+ f"name='{name}' and tags.run_name='{run_name}'"
589
+ )
590
+ if not versions:
591
+ raise ValueError(f"No model {family}/{suffix}/{run_name}")
592
+ client.set_model_version_tag(
593
+ name, versions[0].version, _CONFIG_TAG, _config_to_tag(config)
594
+ )
595
+
596
+
396
597
  def delete_model(family: str, suffix: str, run_name: str) -> None:
397
598
  """Delete the ModelVersion in MLflow, its weights blob, and its serve bundle."""
398
599
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
@@ -204,7 +204,9 @@ print(deployed.url)
204
204
 
205
205
  The requirements are part of the model, not of the serve-app class: GPUs, RAM and VRAM (GiB, 0 meaning no requirement) are stored with it and matched against what the cluster's hosts have free. Correct them later with `cortexgrid.set_model_requirements(family, suffix, run_name, requirements)` or on the model card in the dashboard; `cortexgrid.deploy_model(..., num_replicas=2)` chooses how many copies to run.
206
206
 
207
- `save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
207
+ Anything else the serve-app has to know about the model - which model a provider should be asked for, an endpoint, the name of a secret to read - goes in a free-form string mapping on the same entry: `save_model(..., config={"model": "claude-opus-5"})`. The serve-app reads it in `__init__` with `cortexgrid.model_config(family, suffix, run_name)`; `cortexgrid.set_model_config(family, suffix, run_name, config)` or the model card replaces it. cortexgrid stores the mapping without interpreting it, and a tag is readable by anyone with registry access, so a credential belongs in `set_secret` with only its name in the config.
208
+
209
+ `save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`. A model with no weights to stage - one behind a provider's API, e.g. Gemini or OpenAI - is registered the same way by `cortexgrid.register_model(MyServeApp, family, suffix, config=...)`: same key, same reuse, only the bundle stored. [Serving a hosted-API model](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md#serving-a-hosted-api-model) walks through one end to end - serve-app, API key, registration, deploy.
208
210
 
209
211
  `save_model` is synchronous (registry lifecycle: `uploading` -> `ready`); `deploy_model` schedules the serving lifecycle (`deploying` -> `running`). With `wait=True` a failed deploy raises `cortexgrid.ModelDeployFailed`; `cortexgrid.wait_for_model_serving(family, suffix, run_name, timeout=...)` waits on a deploy started elsewhere, and re-deploying a failed model retries it from scratch. See [model-serving.md](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md) for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (`model_registry_status`, `model_serving_status`), and error handling.
210
212
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "cortexgrid"
3
- version = "0.3.1"
3
+ version = "0.3.3"
4
4
  description = "Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3"
5
5
  readme = "docs/cortexgrid/README.md"
6
6
  license = "Apache-2.0"
File without changes
File without changes