cortexgrid 0.3.1__tar.gz → 0.3.2__tar.gz

{cortexgrid-0.3.1 → cortexgrid-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.3.1
+Version: 0.3.2
 Summary: Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid
@@ -232,6 +232,8 @@ print(deployed.url)
 
 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.
 
+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.
+
 `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)`.
 
 `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.

{cortexgrid-0.3.1 → cortexgrid-0.3.2}/cortexgrid/__init__.py

@@ -89,7 +89,9 @@ from cortexgrid.model_storage import (
     delete_model,
     list_models,
     load_model,
+    model_config,
     model_registry_status,
+    set_model_config,
     set_model_requirements,
 )
 from cortexgrid.model_storage import import_model as _import_model_storage
@@ -149,10 +151,12 @@ def save_model(
     family: str,
     suffix: str,
     requirements: ModelRequirements | None = None,
+    config: dict[str, str] | None = None,
 ) -> SavedModel:
     """Persist a weights directory under the current Experiment's run, paired
-    with the serve-app class that will front it at deploy time and the
-    hardware one replica of it needs.
+    with the serve-app class that will front it at deploy time, the hardware
+    one replica of it needs, and any `config` the serve-app reads with
+    `model_config`.
 
     Every run saves a new copy under its own run_name - meant for weights the
     run produced (e.g. a fine-tune). For a model produced elsewhere that should
@@ -166,6 +170,7 @@ def save_model(
         run_id=experiment.run_id,
         run_name=experiment.run_name(),
         requirements=requirements,
+        config=config,
     )
 
 
@@ -175,6 +180,7 @@ def import_model(
     family: str,
     suffix: str,
     requirements: ModelRequirements | None = None,
+    config: dict[str, str] | None = None,
 ) -> SavedModel:
     """Register a model produced elsewhere once, reuse it on every later call,
     and record on the current Experiment's run which imported model it used.
@@ -184,7 +190,9 @@ def import_model(
     `imported_model/<family>/<suffix>` holds the version's `created_at`, set
     whether this call uploaded the model or reused it."""
     experiment = Experiment.get_instance()
-    model = _import_model_storage(source, serve_app, family, suffix, requirements)
+    model = _import_model_storage(
+        source, serve_app, family, suffix, requirements, config
+    )
     get_mlflow_client().set_tag(
         experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
     )
@@ -255,6 +263,8 @@ __all__ = [
     "model_registry_status",
     "ModelRequirements",
     "set_model_requirements",
+    "model_config",
+    "set_model_config",
     "delete_model",
     # Model serving
     "Deployment",

{cortexgrid-0.3.1 → cortexgrid-0.3.2}/cortexgrid/model_storage.py

@@ -9,6 +9,7 @@ Mapping cortexgrid taxonomy <-> MLflow Registry:
     run linkage         -> ModelVersion.run_id  (built-in MLflow field; unset
                            for imported models)
     requirements        -> ModelVersion.tags["num_gpus"], ["ram_gb"], ["vram_gb"]
+    config              -> ModelVersion.tags["config"]  (JSON object)
 
 Two ways in: `save_model` registers a fresh copy under the calling run's
 run_name every time it runs (fine-tuned output); `import_model` registers a
@@ -24,8 +25,9 @@ cortexgrid/__init__.py.
 
 from __future__ import annotations
 
+import json
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 import tempfile
@@ -58,6 +60,19 @@ _PHASE_READY = "ready"
 _PHASE_UPLOAD_FAILED = "upload_failed"
 _PHASE_BROKEN = "broken"
 
+# MLflow ModelVersion tag holding the model's config: whatever settings its
+# serve-app needs that are not the weights (a provider's model id, an endpoint,
+# the name of a secret to read). cortexgrid never interprets it - it is the
+# serve-app's own vocabulary, stored next to the model so the app reads it at
+# construction instead of being redeployed to change a setting.
+#
+# One JSON object in one tag, not a tag per key: the keys are the serve-app's
+# to choose, free of MLflow's tag-key charset, and the whole mapping is
+# replaced in a single write, so removing a key needs no tag deletion. MLflow
+# caps a tag value at 8000 characters, which bounds how big a config can get.
+_CONFIG_TAG = "config"
+
+
 # An upload still marked "uploading" this long after the version was created is
 # treated as broken: the version is created before the weights are resolved
 # (for `import_model`, before its download), so creation_timestamp is the
@@ -88,6 +103,9 @@ class SavedModel:
     phase: str
     # Hardware one replica needs; defaults for versions stored without it.
     requirements: ModelRequirements
+    # Free-form settings the serve-app reads at construction; empty for
+    # versions stored without any.
+    config: dict[str, str] = field(default_factory=dict)
 
 
 def _phase_for(version: Any) -> str:
@@ -107,6 +125,28 @@ def _phase_for(version: Any) -> str:
     return phase
 
 
+def _config_to_tag(config: dict[str, str]) -> str:
+    """Serialise a config mapping to its MLflow tag value.
+
+    Keys and values are strings: they round-trip through a tag, and the model
+    card edits them as text. A caller with a number or a flag spells it as a
+    string and the serve-app parses it back."""
+    for key, value in config.items():
+        if not isinstance(key, str) or not isinstance(value, str):
+            raise ValueError(
+                f"Model config must map strings to strings: {key!r}: {value!r}"
+            )
+        if not key.strip():
+            raise ValueError("Model config keys cannot be blank")
+    return json.dumps(config)
+
+
+def _config_from_tags(tags: dict[str, str]) -> dict[str, str]:
+    """Deserialise a config mapping from a ModelVersion's MLflow tags; a
+    missing tag reads as no config."""
+    return json.loads(tags.get(_CONFIG_TAG, "{}"))
+
+
 def _to_saved_model(version: Any) -> SavedModel:
     return SavedModel(
         family=version.tags["family"],
@@ -119,6 +159,7 @@ def _to_saved_model(version: Any) -> SavedModel:
         size_bytes=int(version.tags.get("size_bytes", "0")),
         phase=_phase_for(version),
         requirements=requirements_from_tags(version.tags),
+        config=_config_from_tags(version.tags),
     )
 
 
@@ -159,6 +200,7 @@ def save_model(
     run_id: str,
     run_name: str,
     requirements: ModelRequirements | None = None,
+    config: dict[str, str] | None = None,
 ) -> SavedModel:
     """Upload a weights directory to S3 and register a new MLflow ModelVersion
     paired with the serve-app that fronts it.
@@ -180,9 +222,20 @@ def save_model(
     without the caller holding the class object.
 
     `requirements` is the hardware one replica needs; None stores none, which
-    reads as no requirement. Change it later with `set_model_requirements`."""
+    reads as no requirement. Change it later with `set_model_requirements`.
+
+    `config` is whatever else the serve-app needs to know about this model,
+    as a string mapping it reads with `model_config` at construction; None
+    stores none. Change it later with `set_model_config`."""
     return _upload_model(
-        weights_dir, serve_app, suffix, family, run_id, run_name, requirements
+        weights_dir,
+        serve_app,
+        suffix,
+        family,
+        run_id,
+        run_name,
+        requirements,
+        config,
     )
 
 
@@ -192,6 +245,7 @@ def import_model(
     family: str,
     suffix: str,
     requirements: ModelRequirements | None = None,
+    config: dict[str, str] | None = None,
 ) -> SavedModel:
     """Register a model produced elsewhere (e.g. a pretrained base model) under
     the fixed key (family, suffix, IMPORTED), once.
@@ -205,10 +259,10 @@ def import_model(
     If a version is already registered under the key:
       - "ready": returns it without calling `source`. If `serve_app`'s code no
         longer matches the stored bundle, it is re-bundled first and the
-        weights are kept (see `_refresh_bundle`). `requirements` are stored
-        only if the version has none yet, so values changed since with
-        `set_model_requirements` are kept. To replace the weights,
-        `delete_model` it first.
+        weights are kept (see `_refresh_bundle`). `requirements` and `config`
+        are stored only if the version has none yet, so values changed since
+        with `set_model_requirements` / `set_model_config` are kept. To replace
+        the weights, `delete_model` it first.
       - "uploading": raises RuntimeError - another process is importing it.
       - "upload_failed" / "broken": deleted and imported again.
 
@@ -223,6 +277,8 @@ def import_model(
                 existing.requirements = _set_missing_requirements(
                     family, suffix, requirements
                 )
+            if config is not None:
+                existing.config = _set_missing_config(family, suffix, config)
             return existing
         if existing.phase == _PHASE_UPLOADING:
             raise RuntimeError(
@@ -231,7 +287,7 @@ def import_model(
             )
         delete_model(family, suffix, IMPORTED)
     return _upload_model(
-        source, serve_app, suffix, family, None, IMPORTED, requirements
+        source, serve_app, suffix, family, None, IMPORTED, requirements, config
     )
 
 
@@ -273,6 +329,24 @@ def _set_missing_requirements(
     return requirements
 
 
+def _set_missing_config(
+    family: str, suffix: str, config: dict[str, str]
+) -> dict[str, str]:
+    """Store `config` on an imported model whose version has none yet. Returns
+    the config the version holds afterwards."""
+    name = f"{family}__{suffix}"
+    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+    version = client.search_model_versions(
+        f"name='{name}' and tags.run_name='{IMPORTED}'"
+    )[0]
+    if _CONFIG_TAG in version.tags:
+        return _config_from_tags(version.tags)
+    client.set_model_version_tag(
+        name, version.version, _CONFIG_TAG, _config_to_tag(config)
+    )
+    return config
+
+
 def _upload_model(
     weights: str | Path | Callable[[], str | Path],
     serve_app: type,
@@ -281,6 +355,7 @@ def _upload_model(
     run_id: str | None,
     run_name: str,
     requirements: ModelRequirements | None,
+    config: dict[str, str] | None,
 ) -> SavedModel:
     """Register a ModelVersion in "uploading", resolve `weights` to a directory
     (calling it when it is a callable), upload the weights and the serve-app
@@ -296,8 +371,9 @@ def _upload_model(
     # fetched and uploaded, and a concurrent `import_model` sees the import in
     # flight for the whole download instead of starting one of its own. The
     # size is stamped once the directory exists; the bundle tags and the flip
-    # to "ready" happen only after the upload lands. No requirements leaves
-    # their tags unset, so a later `import_model` can still store them.
+    # to "ready" happen only after the upload lands. No requirements and no
+    # config leave their tags unset, so a later `import_model` can still store
+    # them.
     version = client.create_model_version(
         name=name,
         source=source,
@@ -308,6 +384,7 @@ def _upload_model(
             "run_name": run_name,
             _LIFECYCLE_TAG: _PHASE_UPLOADING,
             **(requirements_to_tags(requirements) if requirements is not None else {}),
+            **({_CONFIG_TAG: _config_to_tag(config)} if config is not None else {}),
         },
     )
     try:
@@ -393,6 +470,46 @@ def set_model_requirements(
         client.set_model_version_tag(name, versions[0].version, key, value)
 
 
+def model_config(family: str, suffix: str, run_name: str) -> dict[str, str]:
+    """The config mapping stored on a model, empty if it has none.
+
+    Meant for the serve-app to call in `__init__` with the
+    (family, suffix, run_name) it was constructed with: the settings that are
+    not the weights - a provider's model id, an endpoint, the name of a secret
+    to read - travel with the registry entry instead of the bundled code, so
+    changing one is an edit on the model card rather than a re-save.
+
+    Read at construction, so a replica keeps the values it started with until
+    it is deployed again. Raises ValueError if the model was never
+    registered."""
+    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+    versions = client.search_model_versions(
+        f"name='{family}__{suffix}' and tags.run_name='{run_name}'"
+    )
+    if not versions:
+        raise ValueError(f"No model {family}/{suffix}/{run_name}")
+    return _config_from_tags(versions[0].tags)
+
+
+def set_model_config(
+    family: str, suffix: str, run_name: str, config: dict[str, str]
+) -> None:
+    """Replace the config mapping stored on a model - the whole mapping, so a
+    key left out of `config` is gone. Takes effect on its next `deploy_model`;
+    a replica already running keeps the values it read at construction.
+    Raises ValueError if the model was never registered."""
+    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+    name = f"{family}__{suffix}"
+    versions = client.search_model_versions(
+        f"name='{name}' and tags.run_name='{run_name}'"
+    )
+    if not versions:
+        raise ValueError(f"No model {family}/{suffix}/{run_name}")
+    client.set_model_version_tag(
+        name, versions[0].version, _CONFIG_TAG, _config_to_tag(config)
+    )
+
+
 def delete_model(family: str, suffix: str, run_name: str) -> None:
     """Delete the ModelVersion in MLflow, its weights blob, and its serve bundle."""
     client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())

{cortexgrid-0.3.1 → cortexgrid-0.3.2}/docs/cortexgrid/README.md

@@ -204,6 +204,8 @@ print(deployed.url)
 
 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.
 
+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.
+
 `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)`.
 
 `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.

{cortexgrid-0.3.1 → cortexgrid-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid"
-version = "0.3.1"
+version = "0.3.2"
 description = "Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3"
 readme = "docs/cortexgrid/README.md"
 license = "Apache-2.0"