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.
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/PKG-INFO +4 -2
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/__init__.py +43 -3
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/model_storage.py +225 -24
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/docs/cortexgrid/README.md +3 -1
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/pyproject.toml +1 -1
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/.gitignore +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/LICENSE +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/_bundle.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/_ray_job_driver.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/_serve_entry.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/checkpoint.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/experiment.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/infra.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/jobs.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/mlflow_util.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/model_serving.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/py.typed +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/ray_util.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/s3_util.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/secrets.py +0 -0
- {cortexgrid-0.3.1 → cortexgrid-0.3.3}/cortexgrid/serve.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: cortexgrid
|
|
3
|
-
Version: 0.3.
|
|
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
|
-
|
|
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
|
|
155
|
-
|
|
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(
|
|
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
|
-
|
|
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
|
|
17
|
-
|
|
18
|
-
|
|
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,
|
|
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`
|
|
209
|
-
only if the version has none yet, so values changed since
|
|
210
|
-
`set_model_requirements` are kept. To replace
|
|
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
|
-
|
|
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
|
-
|
|
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 =
|
|
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.
|
|
300
|
-
#
|
|
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
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
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
|
-
|
|
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
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|