cortexgrid 0.2.96__tar.gz → 0.2.97__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.2.96
3
+ Version: 0.2.97
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
@@ -165,19 +165,23 @@ app = FastAPI()
165
165
 
166
166
  @serve.ingress(app)
167
167
  class MyServeApp:
168
- num_gpus = 1
169
-
170
168
  def __init__(self, family: str, suffix: str, run_name: str) -> None:
171
169
  self._weights_dir = cortexgrid.load_model(family, suffix, run_name)
172
170
 
173
171
  @app.post("/complete")
174
172
  async def complete(self, body: dict): ...
175
173
 
176
- saved = cortexgrid.save_model(weights_dir, MyServeApp, family="qwen", suffix="instruct")
174
+ saved = cortexgrid.save_model(
175
+ weights_dir, MyServeApp, family="qwen", suffix="instruct",
176
+ # What one replica needs; the model is deployed only on a host that has it.
177
+ requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
178
+ )
177
179
  deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
178
180
  print(deployed.url)
179
181
  ```
180
182
 
183
+ 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.
184
+
181
185
  `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)`.
182
186
 
183
187
  `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.
@@ -80,12 +80,14 @@ from cortexgrid.model_storage import (
80
80
  list_models,
81
81
  load_model,
82
82
  model_registry_status,
83
+ set_model_requirements,
83
84
  )
84
85
  from cortexgrid.model_storage import import_model as _import_model_storage
85
86
  from cortexgrid.model_storage import save_model as _save_model_storage
86
87
  from cortexgrid.model_serving import (
87
88
  Deployment,
88
89
  ModelDeployFailed,
90
+ ModelRequirements,
89
91
  ServingStatus,
90
92
  deploy_model,
91
93
  list_deployed_models,
@@ -118,10 +120,15 @@ def remote(
118
120
 
119
121
 
120
122
  def save_model(
121
- weights_dir: str | Path, serve_app: type, family: str, suffix: str
123
+ weights_dir: str | Path,
124
+ serve_app: type,
125
+ family: str,
126
+ suffix: str,
127
+ requirements: ModelRequirements | None = None,
122
128
  ) -> SavedModel:
123
129
  """Persist a weights directory under the current Experiment's run, paired
124
- with the serve-app class that will front it at deploy time.
130
+ with the serve-app class that will front it at deploy time and the
131
+ hardware one replica of it needs.
125
132
 
126
133
  Every run saves a new copy under its own run_name - meant for weights the
127
134
  run produced (e.g. a fine-tune). For a model produced elsewhere that should
@@ -134,6 +141,7 @@ def save_model(
134
141
  family,
135
142
  run_id=experiment.run_id,
136
143
  run_name=experiment.run_name(),
144
+ requirements=requirements,
137
145
  )
138
146
 
139
147
 
@@ -142,6 +150,7 @@ def import_model(
142
150
  serve_app: type,
143
151
  family: str,
144
152
  suffix: str,
153
+ requirements: ModelRequirements | None = None,
145
154
  ) -> SavedModel:
146
155
  """Register a model produced elsewhere once, reuse it on every later call,
147
156
  and record on the current Experiment's run which imported model it used.
@@ -151,7 +160,7 @@ def import_model(
151
160
  `imported_model/<family>/<suffix>` holds the version's `created_at`, set
152
161
  whether this call uploaded the model or reused it."""
153
162
  experiment = Experiment.get_instance()
154
- model = _import_model_storage(source, serve_app, family, suffix)
163
+ model = _import_model_storage(source, serve_app, family, suffix, requirements)
155
164
  get_mlflow_client().set_tag(
156
165
  experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
157
166
  )
@@ -213,6 +222,8 @@ __all__ = [
213
222
  "load_model",
214
223
  "list_models",
215
224
  "model_registry_status",
225
+ "ModelRequirements",
226
+ "set_model_requirements",
216
227
  "delete_model",
217
228
  # Model serving
218
229
  "Deployment",
@@ -4,9 +4,9 @@ Ray Serve's REST `import_path` resolves to `cortexgrid._serve_entry:build`.
4
4
  On the cluster replica, `build` imports the serve-app class bundled at
5
5
  `save_model` time (its import path was stored as an MLflow tag), applies Ray's
6
6
  ingress with the app it was marked with by `cortexgrid.serve.ingress` (again on
7
- each replica, see `_IngressOnReplica`), reads its
8
- `num_gpus`/`num_replicas` class attributes for actor placement, wraps it as a
9
- Ray Serve deployment, and binds it with the (family, suffix, run_name)
7
+ each replica, see `_IngressOnReplica`), wraps it as a Ray Serve deployment
8
+ with the replica count and Ray resource requests `deploy_model` derived from
9
+ the model's requirements, and binds it with the (family, suffix, run_name)
10
10
  identifiers.
11
11
 
12
12
  The serve-app owns everything about traffic: its own routes, request schemas,
@@ -14,11 +14,9 @@ streaming, and timeouts. cortexgrid does not interpose a request/response
14
14
  contract - it only schedules the app and hands it the identifiers it needs to
15
15
  fetch its own weights via `cortexgrid.load_model`.
16
16
 
17
- Design note: resource needs (`num_gpus`/`num_replicas`) are read from plain
18
- class attributes rather than a cortexgrid decorator or base class. This is a
19
- deliberate, provisional choice - kept minimal until we see how serve-apps
20
- declare resources in practice; revisit if plain class attributes prove too
21
- limited.
17
+ The serve-app declares no resources: the hardware a replica needs belongs to
18
+ the model and is stored in the registry (`cortexgrid.ModelRequirements`), and
19
+ the replica count is chosen per `deploy_model`.
22
20
  """
23
21
 
24
22
  from __future__ import annotations
@@ -72,12 +70,13 @@ def build(args: dict[str, Any]) -> Application:
72
70
  {"_serve_app": serve_app},
73
71
  )
74
72
  serve_app = serve.ingress(app)(on_replica)
75
- num_gpus = getattr(serve_app, "num_gpus", 0)
76
- num_replicas = getattr(serve_app, "num_replicas", 1)
77
73
  return serve.deployment(serve_app).options(
78
- num_replicas=num_replicas,
74
+ # This builder ships in the bundle, frozen at save time, while `args`
75
+ # come from the cortexgrid that deploys it; one older than the bundle
76
+ # sends neither key.
77
+ num_replicas=args.get("num_replicas", 1),
79
78
  # Ray 2.32 lowered the default from 100 to 5; keep what serve-apps
80
79
  # had on Ray 2.9.
81
80
  max_ongoing_requests=_MAX_ONGOING_REQUESTS,
82
- ray_actor_options={"num_gpus": num_gpus},
81
+ ray_actor_options=args.get("ray_actor_options", {}),
83
82
  ).bind(args["family"], args["suffix"], args["run_name"])
@@ -8,7 +8,8 @@ uploaded to MinIO under
8
8
  referenced via `runtime_env.working_dir` so Ray workers fetch it from there.
9
9
  The bundle URL, class import path, pip list, and fingerprint are persisted as
10
10
  MLflow tags on the ModelVersion so `deploy_model` can find them later without
11
- the caller holding the class object.
11
+ the caller holding the class object. So are the model's `ModelRequirements`,
12
+ which `deploy_model` turns into the replica's Ray resource requests.
12
13
 
13
14
  Naming: the Ray Serve application is named "<family>__<suffix>__<run_name>".
14
15
  This relies on family/suffix/run_name not containing the literal "__".
@@ -206,10 +207,41 @@ def upload_bundle(
206
207
  )
207
208
 
208
209
 
210
+ # Custom Ray resource each GPU worker advertises: the MiB of memory its GPUs
211
+ # have (see the ray-worker DaemonSet). A replica requests its vram_gb of it in
212
+ # MiB, so Ray places it only on a node with that much VRAM left. MiB because
213
+ # nvidia-smi reports MiB and a GPU's memory is not a whole number of GiB.
214
+ _VRAM_RESOURCE = "vram_mib"
215
+
216
+ _MIB_PER_GIB = 1024
217
+
218
+ _GIB = 1024**3
219
+
220
+
221
+ def _ray_actor_options(requirements: ModelRequirements) -> dict[str, Any]:
222
+ """Translate ModelRequirements into a replica's Ray actor resource requests.
223
+ Ray places the replica only on a node with that much free and reserves it
224
+ there; a zero requirement requests nothing."""
225
+ options: dict[str, Any] = {"num_gpus": requirements.num_gpus}
226
+ if requirements.ram_gb > 0:
227
+ options["memory"] = int(requirements.ram_gb * _GIB)
228
+ if requirements.vram_gb > 0:
229
+ options["resources"] = {
230
+ _VRAM_RESOURCE: round(requirements.vram_gb * _MIB_PER_GIB)
231
+ }
232
+ return options
233
+
234
+
209
235
  def _build_application_spec(
210
- family: str, suffix: str, run_name: str, meta: BundleMetadata
236
+ family: str,
237
+ suffix: str,
238
+ run_name: str,
239
+ meta: BundleMetadata,
240
+ requirements: ModelRequirements,
241
+ num_replicas: int,
211
242
  ) -> dict[str, Any]:
212
- """Assemble a Ray Serve application schema from pre-bundled metadata."""
243
+ """Assemble a Ray Serve application schema from pre-bundled metadata and
244
+ the model's requirements."""
213
245
  # working_dir carries the serve-app's own source; Ray pip-installs the
214
246
  # third-party distributions the image lacks into a per-node cached
215
247
  # virtualenv layered on the image. No pip key when there are none, so Ray
@@ -230,6 +262,8 @@ def _build_application_spec(
230
262
  "family": family,
231
263
  "suffix": suffix,
232
264
  "run_name": run_name,
265
+ "num_replicas": num_replicas,
266
+ "ray_actor_options": _ray_actor_options(requirements),
233
267
  },
234
268
  "runtime_env": runtime_env,
235
269
  }
@@ -268,10 +302,66 @@ def metadata_from_tags(tags: dict[str, str]) -> BundleMetadata:
268
302
  )
269
303
 
270
304
 
271
- def _load_bundle_metadata(
305
+ @dataclass
306
+ class ModelRequirements:
307
+ """Hardware one replica of a model needs to be served, in GiB. Persisted as
308
+ tags on the ModelVersion next to the bundle metadata, so it is read without
309
+ touching the weights or importing the serve-app class.
310
+
311
+ Zero means no requirement: a model with no requirements is served on any
312
+ node, CPU-only included. Models saved before requirements existed carry
313
+ no tags and read as the defaults."""
314
+
315
+ num_gpus: int = 0
316
+ ram_gb: float = 0.0
317
+ # GPU memory across the replica's num_gpus GPUs, so it needs num_gpus >= 1.
318
+ vram_gb: float = 0.0
319
+
320
+ def __post_init__(self) -> None:
321
+ if self.num_gpus < 0 or self.ram_gb < 0 or self.vram_gb < 0:
322
+ raise ValueError(f"Model requirements cannot be negative: {self}")
323
+ if self.vram_gb > 0 and self.num_gpus == 0:
324
+ raise ValueError(
325
+ f"vram_gb={self.vram_gb} needs a GPU; set num_gpus >= 1"
326
+ )
327
+
328
+
329
+ # MLflow tag keys for the ModelRequirements.
330
+ _NUM_GPUS_TAG = "num_gpus"
331
+ _RAM_GB_TAG = "ram_gb"
332
+ _VRAM_GB_TAG = "vram_gb"
333
+
334
+
335
+ def has_requirement_tags(tags: dict[str, str]) -> bool:
336
+ """Whether requirements were ever stored on the ModelVersion."""
337
+ return any(key in tags for key in (_NUM_GPUS_TAG, _RAM_GB_TAG, _VRAM_GB_TAG))
338
+
339
+
340
+ def requirements_to_tags(requirements: ModelRequirements) -> dict[str, str]:
341
+ """Serialise ModelRequirements to MLflow tags. The inverse of
342
+ `requirements_from_tags`."""
343
+ return {
344
+ _NUM_GPUS_TAG: str(requirements.num_gpus),
345
+ _RAM_GB_TAG: str(requirements.ram_gb),
346
+ _VRAM_GB_TAG: str(requirements.vram_gb),
347
+ }
348
+
349
+
350
+ def requirements_from_tags(tags: dict[str, str]) -> ModelRequirements:
351
+ """Deserialise ModelRequirements from a ModelVersion's MLflow tags; a
352
+ missing tag reads as no requirement."""
353
+ return ModelRequirements(
354
+ num_gpus=int(tags.get(_NUM_GPUS_TAG, "0")),
355
+ ram_gb=float(tags.get(_RAM_GB_TAG, "0")),
356
+ vram_gb=float(tags.get(_VRAM_GB_TAG, "0")),
357
+ )
358
+
359
+
360
+ def _load_deploy_metadata(
272
361
  family: str, suffix: str, run_name: str
273
- ) -> BundleMetadata:
274
- """Read the bundle metadata `save_model` persisted on the ModelVersion."""
362
+ ) -> tuple[BundleMetadata, ModelRequirements]:
363
+ """Read the bundle metadata and requirements `save_model` persisted on the
364
+ ModelVersion."""
275
365
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
276
366
  name = f"{family}__{suffix}"
277
367
  versions = client.search_model_versions(
@@ -281,8 +371,9 @@ def _load_bundle_metadata(
281
371
  raise ValueError(
282
372
  f"No saved model for {family}/{suffix}/{run_name}; cannot deploy."
283
373
  )
374
+ tags = versions[0].tags or {}
284
375
  try:
285
- return metadata_from_tags(versions[0].tags or {})
376
+ return metadata_from_tags(tags), requirements_from_tags(tags)
286
377
  except KeyError as exc:
287
378
  raise ValueError(
288
379
  f"Saved model {family}/{suffix}/{run_name} is missing the deployment "
@@ -398,6 +489,7 @@ def deploy_model(
398
489
  family: str,
399
490
  suffix: str,
400
491
  run_name: str,
492
+ num_replicas: int = 1,
401
493
  wait: bool = False,
402
494
  timeout: float | None = 300.0,
403
495
  ) -> Deployment:
@@ -408,6 +500,8 @@ def deploy_model(
408
500
 
409
501
  The serve-app class is pulled from the MLflow ModelVersion tags `save_model`
410
502
  wrote at save time; the caller does not need to hold the class object.
503
+ Each of the `num_replicas` replicas requests the model's `ModelRequirements`
504
+ from Ray, so it is placed only on a node that has them free.
411
505
 
412
506
  A DEPLOY_FAILED app left by an earlier attempt is undeployed first, and it,
413
507
  or an app still DELETING, is waited out before the new spec is PUT, so the
@@ -421,8 +515,10 @@ def deploy_model(
421
515
  stuck in DEPLOYING) will hang forever.
422
516
  """
423
517
  deadline = _deadline(timeout)
424
- meta = _load_bundle_metadata(family, suffix, run_name)
425
- spec = _build_application_spec(family, suffix, run_name, meta)
518
+ meta, requirements = _load_deploy_metadata(family, suffix, run_name)
519
+ spec = _build_application_spec(
520
+ family, suffix, run_name, meta, requirements, num_replicas
521
+ )
426
522
  _clear_failed_application(family, suffix, run_name, timeout, deadline)
427
523
  existing = [a for a in _current_application_specs() if a["name"] != spec["name"]]
428
524
  # The controller registers the app, sets it DEPLOYING and stamps
@@ -8,6 +8,7 @@ Mapping cortexgrid taxonomy <-> MLflow Registry:
8
8
  "s3://<bucket>/models/<run_name>/<family>/<suffix>/weights/"
9
9
  run linkage -> ModelVersion.run_id (built-in MLflow field; unset
10
10
  for imported models)
11
+ requirements -> ModelVersion.tags["num_gpus"], ["ram_gb"], ["vram_gb"]
11
12
 
12
13
  Two ways in: `save_model` registers a fresh copy under the calling run's
13
14
  run_name every time it runs (fine-tuned output); `import_model` registers a
@@ -36,10 +37,14 @@ from mlflow.tracking import MlflowClient
36
37
  from cortexgrid import s3_util
37
38
  from cortexgrid.infra import get_mlflow_tracking_uri, get_s3_bucket
38
39
  from cortexgrid.model_serving import (
40
+ ModelRequirements,
39
41
  build_bundle,
40
42
  bundle_class,
43
+ has_requirement_tags,
41
44
  metadata_from_tags,
42
45
  metadata_to_tags,
46
+ requirements_from_tags,
47
+ requirements_to_tags,
43
48
  upload_bundle,
44
49
  )
45
50
 
@@ -81,6 +86,8 @@ class SavedModel:
81
86
  # _UPLOAD_DEADLINE (writer presumed dead). Versions written before this tag
82
87
  # existed report "ready".
83
88
  phase: str
89
+ # Hardware one replica needs; defaults for versions stored without it.
90
+ requirements: ModelRequirements
84
91
 
85
92
 
86
93
  def _phase_for(version: Any) -> str:
@@ -111,6 +118,7 @@ def _to_saved_model(version: Any) -> SavedModel:
111
118
  data_blob_path=version.source,
112
119
  size_bytes=int(version.tags.get("size_bytes", "0")),
113
120
  phase=_phase_for(version),
121
+ requirements=requirements_from_tags(version.tags),
114
122
  )
115
123
 
116
124
 
@@ -150,6 +158,7 @@ def save_model(
150
158
  family: str,
151
159
  run_id: str,
152
160
  run_name: str,
161
+ requirements: ModelRequirements | None = None,
153
162
  ) -> SavedModel:
154
163
  """Upload a weights directory to S3 and register a new MLflow ModelVersion
155
164
  paired with the serve-app that fronts it.
@@ -168,8 +177,13 @@ def save_model(
168
177
  weights.
169
178
  Its code is bundled and its import path, bundle URL, and pip list are
170
179
  stored as tags on the ModelVersion so `deploy_model` can bind it later
171
- without the caller holding the class object."""
172
- return _upload_model(weights_dir, serve_app, suffix, family, run_id, run_name)
180
+ without the caller holding the class object.
181
+
182
+ `requirements` is the hardware one replica needs; None stores none, which
183
+ reads as no requirement. Change it later with `set_model_requirements`."""
184
+ return _upload_model(
185
+ weights_dir, serve_app, suffix, family, run_id, run_name, requirements
186
+ )
173
187
 
174
188
 
175
189
  def import_model(
@@ -177,6 +191,7 @@ def import_model(
177
191
  serve_app: type,
178
192
  family: str,
179
193
  suffix: str,
194
+ requirements: ModelRequirements | None = None,
180
195
  ) -> SavedModel:
181
196
  """Register a model produced elsewhere (e.g. a pretrained base model) under
182
197
  the fixed key (family, suffix, IMPORTED), once.
@@ -190,7 +205,9 @@ def import_model(
190
205
  If a version is already registered under the key:
191
206
  - "ready": returns it without calling `source`. If `serve_app`'s code no
192
207
  longer matches the stored bundle, it is re-bundled first and the
193
- weights are kept (see `_refresh_bundle`). To replace the weights,
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,
194
211
  `delete_model` it first.
195
212
  - "uploading": raises RuntimeError - another process is importing it.
196
213
  - "upload_failed" / "broken": deleted and imported again.
@@ -202,6 +219,10 @@ def import_model(
202
219
  if existing is not None:
203
220
  if existing.phase == _PHASE_READY:
204
221
  _refresh_bundle(serve_app, family, suffix)
222
+ if requirements is not None:
223
+ existing.requirements = _set_missing_requirements(
224
+ family, suffix, requirements
225
+ )
205
226
  return existing
206
227
  if existing.phase == _PHASE_UPLOADING:
207
228
  raise RuntimeError(
@@ -209,7 +230,9 @@ def import_model(
209
230
  "another process"
210
231
  )
211
232
  delete_model(family, suffix, IMPORTED)
212
- return _upload_model(source, serve_app, suffix, family, None, IMPORTED)
233
+ return _upload_model(
234
+ source, serve_app, suffix, family, None, IMPORTED, requirements
235
+ )
213
236
 
214
237
 
215
238
  def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
@@ -233,6 +256,23 @@ def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
233
256
  client.set_model_version_tag(name, version.version, key, value)
234
257
 
235
258
 
259
+ def _set_missing_requirements(
260
+ family: str, suffix: str, requirements: ModelRequirements
261
+ ) -> ModelRequirements:
262
+ """Store `requirements` on an imported model whose version has none yet.
263
+ Returns the requirements the version holds afterwards."""
264
+ name = f"{family}__{suffix}"
265
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
266
+ version = client.search_model_versions(
267
+ f"name='{name}' and tags.run_name='{IMPORTED}'"
268
+ )[0]
269
+ if has_requirement_tags(version.tags):
270
+ return requirements_from_tags(version.tags)
271
+ for key, value in requirements_to_tags(requirements).items():
272
+ client.set_model_version_tag(name, version.version, key, value)
273
+ return requirements
274
+
275
+
236
276
  def _upload_model(
237
277
  weights: str | Path | Callable[[], str | Path],
238
278
  serve_app: type,
@@ -240,6 +280,7 @@ def _upload_model(
240
280
  family: str,
241
281
  run_id: str | None,
242
282
  run_name: str,
283
+ requirements: ModelRequirements | None,
243
284
  ) -> SavedModel:
244
285
  """Register a ModelVersion in "uploading", resolve `weights` to a directory
245
286
  (calling it when it is a callable), upload the weights and the serve-app
@@ -255,7 +296,8 @@ def _upload_model(
255
296
  # fetched and uploaded, and a concurrent `import_model` sees the import in
256
297
  # flight for the whole download instead of starting one of its own. The
257
298
  # size is stamped once the directory exists; the bundle tags and the flip
258
- # to "ready" happen only after the upload lands.
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.
259
301
  version = client.create_model_version(
260
302
  name=name,
261
303
  source=source,
@@ -265,6 +307,7 @@ def _upload_model(
265
307
  "suffix": suffix,
266
308
  "run_name": run_name,
267
309
  _LIFECYCLE_TAG: _PHASE_UPLOADING,
310
+ **(requirements_to_tags(requirements) if requirements is not None else {}),
268
311
  },
269
312
  )
270
313
  try:
@@ -333,6 +376,23 @@ def model_registry_status(
333
376
  return _to_saved_model(versions[0]) if versions else None
334
377
 
335
378
 
379
+ def set_model_requirements(
380
+ family: str, suffix: str, run_name: str, requirements: ModelRequirements
381
+ ) -> None:
382
+ """Replace the hardware requirements stored on a model. Takes effect on
383
+ its next `deploy_model`; a replica already running keeps its placement.
384
+ Raises ValueError if the model was never registered."""
385
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
386
+ name = f"{family}__{suffix}"
387
+ versions = client.search_model_versions(
388
+ f"name='{name}' and tags.run_name='{run_name}'"
389
+ )
390
+ if not versions:
391
+ raise ValueError(f"No model {family}/{suffix}/{run_name}")
392
+ for key, value in requirements_to_tags(requirements).items():
393
+ client.set_model_version_tag(name, versions[0].version, key, value)
394
+
395
+
336
396
  def delete_model(family: str, suffix: str, run_name: str) -> None:
337
397
  """Delete the ModelVersion in MLflow, its weights blob, and its serve bundle."""
338
398
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
@@ -137,19 +137,23 @@ app = FastAPI()
137
137
 
138
138
  @serve.ingress(app)
139
139
  class MyServeApp:
140
- num_gpus = 1
141
-
142
140
  def __init__(self, family: str, suffix: str, run_name: str) -> None:
143
141
  self._weights_dir = cortexgrid.load_model(family, suffix, run_name)
144
142
 
145
143
  @app.post("/complete")
146
144
  async def complete(self, body: dict): ...
147
145
 
148
- saved = cortexgrid.save_model(weights_dir, MyServeApp, family="qwen", suffix="instruct")
146
+ saved = cortexgrid.save_model(
147
+ weights_dir, MyServeApp, family="qwen", suffix="instruct",
148
+ # What one replica needs; the model is deployed only on a host that has it.
149
+ requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
150
+ )
149
151
  deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
150
152
  print(deployed.url)
151
153
  ```
152
154
 
155
+ 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.
156
+
153
157
  `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)`.
154
158
 
155
159
  `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.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "cortexgrid"
3
- version = "0.2.96"
3
+ version = "0.2.97"
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