cortexgrid 0.2.94__tar.gz → 0.2.96__tar.gz

{cortexgrid-0.2.94 → cortexgrid-0.2.96}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.2.94
+Version: 0.2.96
 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
@@ -178,7 +178,7 @@ deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True
 print(deployed.url)
 ```
 
-`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 is a no-op on later runs; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
+`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.2.94 → cortexgrid-0.2.96}/cortexgrid/__init__.py

@@ -77,11 +77,11 @@ from cortexgrid.model_storage import (
     IMPORTED,
     SavedModel,
     delete_model,
-    import_model,
     list_models,
     load_model,
     model_registry_status,
 )
+from cortexgrid.model_storage import import_model as _import_model_storage
 from cortexgrid.model_storage import save_model as _save_model_storage
 from cortexgrid.model_serving import (
     Deployment,
@@ -137,6 +137,27 @@ def save_model(
     )
 
 
+def import_model(
+    source: str | Path | Callable[[], str | Path],
+    serve_app: type,
+    family: str,
+    suffix: str,
+) -> 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.
+
+    The model belongs to no run (see `cortexgrid.model_storage.import_model`),
+    so the run keeps the link instead: the tag
+    `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)
+    get_mlflow_client().set_tag(
+        experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
+    )
+    return model
+
+
 __all__ = [
     "Experiment",
     "delete_experiment",

{cortexgrid-0.2.94 → cortexgrid-0.2.96}/cortexgrid/_bundle.py

@@ -12,6 +12,8 @@ Bundles of several seeds combine with `BundleDesc.merge`.
 
 `stage(files, dest)` lays a bundle out under `dest` at each file's import path,
 so `dest` on sys.path (e.g. a Ray working_dir) makes every module importable.
+`digest(files)` hashes that layout, so bundles that stage identically compare
+equal wherever their files live.
 
 `BundleDesc.pip_requirements(worker_provides())` pins the third-party
 distributions the Ray worker image does not already have, for a Ray `pip`
@@ -24,6 +26,7 @@ import ast
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 import functools
+import hashlib
 import importlib.machinery
 import importlib.metadata
 import importlib.util
@@ -105,11 +108,23 @@ def stage(files: set[Path], dest: Path) -> None:
     in installed distributions)."""
     dest.mkdir(parents=True, exist_ok=True)
     for file in files:
-        target = dest / file.relative_to(_sys_path_root(file))
+        target = dest / _import_path(file)
         target.parent.mkdir(parents=True, exist_ok=True)
         shutil.copy2(file, target)
 
 
+def digest(files: set[Path]) -> str:
+    """SHA-256 of `files` as `stage` lays them out: each file's import path and
+    contents, in import-path order. Files that stage identically digest
+    identically, wherever they live on disk."""
+    sha = hashlib.sha256()
+    for path, file in sorted((_import_path(file).as_posix(), file) for file in files):
+        content = file.read_bytes()
+        sha.update(f"{path}\0{len(content)}\0".encode())
+        sha.update(content)
+    return sha.hexdigest()
+
+
 # What the Ray worker image pip-installs, as the Dockerfile spells it. They and
 # their dependency trees are on the worker already, so they are never installed
 # there again -- a second copy in the job's virtualenv would shadow the image's.
@@ -335,6 +350,11 @@ def _package(file: Path) -> str:
     return ".".join(parts)
 
 
+def _import_path(file: Path) -> Path:
+    """`file` relative to the sys.path entry it is imported from."""
+    return file.relative_to(_sys_path_root(file))
+
+
 def _sys_path_root(file: Path) -> Path:
     """The sys.path entry `file` is imported from: its first ancestor directory
     without an __init__.py."""

{cortexgrid-0.2.94 → cortexgrid-0.2.96}/cortexgrid/model_serving.py

@@ -3,11 +3,12 @@
 Caller stays HTTP-only: deploy/undeploy/list talk to the Ray dashboard's
 declarative `/api/serve/applications/` endpoint via [cortexgrid.ray_util],
 never `ray.init`. The deployment class is bundled at `save_model` time, zipped,
-uploaded to MinIO under `serve-bundles/<run_name>/<family>__<suffix>.zip`, and
+uploaded to MinIO under
+`serve-bundles/<run_name>/<family>__<suffix>/<fingerprint>.zip`, and
 referenced via `runtime_env.working_dir` so Ray workers fetch it from there.
-The bundle URL, class import path, and pip list are persisted as MLflow tags
-on the ModelVersion so `deploy_model` can find them later without the caller
-holding the class object.
+The bundle URL, class import path, pip list, and fingerprint are persisted as
+MLflow tags on the ModelVersion so `deploy_model` can find them later without
+the caller holding the class object.
 
 Naming: the Ray Serve application is named "<family>__<suffix>__<run_name>".
 This relies on family/suffix/run_name not containing the literal "__".
@@ -26,13 +27,14 @@ import shutil
 import tempfile
 import time
 from dataclasses import dataclass, field
+import hashlib
 from pathlib import Path
 from typing import Any
 
 from mlflow.tracking import MlflowClient
 from ray.serve.schema import ApplicationStatus
 
-from cortexgrid._bundle import bundle, stage, worker_provides
+from cortexgrid._bundle import bundle, digest, stage, worker_provides
 from cortexgrid.infra import get_mlflow_tracking_uri, get_ray_serve_uri
 from cortexgrid.ray_util import (
     get_serve_details,
@@ -95,17 +97,28 @@ class BundleMetadata:
     # pinned third-party requirements the replica pip-installs (the bundle's
     # distributions the Ray image does not already provide)
     pip_requirements: list[str] = field(default_factory=list)
+    # ServeBundle.fingerprint of the uploaded bundle; empty for models saved
+    # before bundles were fingerprinted
+    fingerprint: str = ""
 
 
-def bundle_class(
-    cls: type, family: str, suffix: str, run_name: str
-) -> BundleMetadata:
-    """Bundle the serve-app class's code (and the serve entrypoint), zip it, and
-    upload to MinIO.
+@dataclass
+class ServeBundle:
+    """A serve-app's bundle, resolved locally but not uploaded yet: what
+    `build_bundle` finds and `upload_bundle` ships."""
 
-    Returns the metadata `deploy_model` needs later; callers (typically
-    `save_model`) persist it on the ModelVersion so the deploy step can run
-    without holding the class object.
+    files: set[Path]
+    class_import_path: str
+    pip_requirements: list[str]
+    # Hash of everything the replica runs: the staged files, the class it
+    # imports, and the requirements it installs. Equal fingerprints mean the
+    # same code, so an uploaded bundle can be reused.
+    fingerprint: str
+
+
+def build_bundle(cls: type) -> ServeBundle:
+    """Resolve the serve-app class's code (and the serve entrypoint) into a
+    bundle, without uploading it.
 
     Raises ValueError for a class wrapped by `ray.serve.ingress`: that wrapper
     is a subclass Ray defines in its own module, and on older Ray (e.g. 2.9) it
@@ -120,17 +133,51 @@ def bundle_class(
     entry_file = Path(inspect.getfile(cls)).resolve()
     serve_entry = Path(__file__).with_name("_serve_entry.py")
     desc = bundle(entry_file).merge(bundle(serve_entry))
+    class_import_path = f"{cls.__module__}:{cls.__name__}"
     pip_requirements = desc.pip_requirements(worker_provides())
+    fingerprint = hashlib.sha256(
+        json.dumps(
+            [digest(desc.local_files), class_import_path, pip_requirements]
+        ).encode()
+    ).hexdigest()
+    return ServeBundle(
+        files=desc.local_files,
+        class_import_path=class_import_path,
+        pip_requirements=pip_requirements,
+        fingerprint=fingerprint,
+    )
+
+
+def bundle_class(
+    cls: type, family: str, suffix: str, run_name: str
+) -> BundleMetadata:
+    """Bundle the serve-app class's code (and the serve entrypoint), zip it, and
+    upload to MinIO: `build_bundle` followed by `upload_bundle`.
+
+    Returns the metadata `deploy_model` needs later; callers (typically
+    `save_model`) persist it on the ModelVersion so the deploy step can run
+    without holding the class object."""
+    return upload_bundle(build_bundle(cls), family, suffix, run_name)
+
+
+def upload_bundle(
+    serve_bundle: ServeBundle, family: str, suffix: str, run_name: str
+) -> BundleMetadata:
+    """Zip a built bundle and upload it under its fingerprint.
+
+    The fingerprint is part of the URL because Ray keeps a remote working_dir
+    it has downloaded and reuses it for the same URL: new code at an old URL
+    would never reach a replica."""
     with tempfile.TemporaryDirectory() as tmp:
         code_root = Path(tmp) / "code"
-        stage(desc.local_files, code_root)
+        stage(serve_bundle.files, code_root)
         log.info(
             "Serve bundle for %s/%s/%s: %d files, pip: %s",
             family,
             suffix,
             run_name,
-            len(desc.local_files),
-            pip_requirements,
+            len(serve_bundle.files),
+            serve_bundle.pip_requirements,
         )
         # Ray unpacks a remote (s3://) working_dir zip by stripping its
         # top-level directory when there is exactly one, so a bundle of a single
@@ -146,12 +193,16 @@ def bundle_class(
         )
         bundle_url = upload(
             archive,
-            dest_path=f"serve-bundles/{run_name}/{family}__{suffix}.zip",
+            dest_path=(
+                f"serve-bundles/{run_name}/{family}__{suffix}/"
+                f"{serve_bundle.fingerprint}.zip"
+            ),
         )
     return BundleMetadata(
         bundle_url=bundle_url,
-        class_import_path=f"{cls.__module__}:{cls.__name__}",
-        pip_requirements=pip_requirements,
+        class_import_path=serve_bundle.class_import_path,
+        pip_requirements=serve_bundle.pip_requirements,
+        fingerprint=serve_bundle.fingerprint,
     )
 
 
@@ -189,19 +240,34 @@ def _build_application_spec(
 _CLASS_IMPORT_PATH_TAG = "class_import_path"
 _BUNDLE_URL_TAG = "serve_bundle_url"
 _PIP_REQUIREMENTS_TAG = "serve_pip_requirements"
+_BUNDLE_FINGERPRINT_TAG = "serve_bundle_fingerprint"
 
 
 def metadata_to_tags(meta: BundleMetadata) -> dict[str, str]:
     """Serialise BundleMetadata to MLflow tags. The inverse of
-    `_load_bundle_metadata`; lives here next to the consumer so the tag schema
+    `metadata_from_tags`; lives here next to the consumer so the tag schema
     stays in one place."""
     return {
         _CLASS_IMPORT_PATH_TAG: meta.class_import_path,
         _BUNDLE_URL_TAG: meta.bundle_url,
         _PIP_REQUIREMENTS_TAG: json.dumps(meta.pip_requirements),
+        _BUNDLE_FINGERPRINT_TAG: meta.fingerprint,
     }
 
 
+def metadata_from_tags(tags: dict[str, str]) -> BundleMetadata:
+    """Deserialise BundleMetadata from a ModelVersion's MLflow tags. Raises
+    KeyError for a missing bundle URL or class import path."""
+    return BundleMetadata(
+        bundle_url=tags[_BUNDLE_URL_TAG],
+        class_import_path=tags[_CLASS_IMPORT_PATH_TAG],
+        # Absent on models saved before dependencies were pip-installed.
+        pip_requirements=json.loads(tags.get(_PIP_REQUIREMENTS_TAG, "[]")),
+        # Absent on models saved before bundles were fingerprinted.
+        fingerprint=tags.get(_BUNDLE_FINGERPRINT_TAG, ""),
+    )
+
+
 def _load_bundle_metadata(
     family: str, suffix: str, run_name: str
 ) -> BundleMetadata:
@@ -215,14 +281,8 @@ def _load_bundle_metadata(
         raise ValueError(
             f"No saved model for {family}/{suffix}/{run_name}; cannot deploy."
         )
-    tags = versions[0].tags or {}
     try:
-        return BundleMetadata(
-            bundle_url=tags[_BUNDLE_URL_TAG],
-            class_import_path=tags[_CLASS_IMPORT_PATH_TAG],
-            # Absent on models saved before dependencies were pip-installed.
-            pip_requirements=json.loads(tags.get(_PIP_REQUIREMENTS_TAG, "[]")),
-        )
+        return metadata_from_tags(versions[0].tags or {})
     except KeyError as exc:
         raise ValueError(
             f"Saved model {family}/{suffix}/{run_name} is missing the deployment "

{cortexgrid-0.2.94 → cortexgrid-0.2.96}/cortexgrid/model_storage.py

@@ -11,8 +11,9 @@ Mapping cortexgrid taxonomy <-> MLflow Registry:
 
 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
-model produced elsewhere once, under the fixed run_name IMPORTED, and is a
-no-op after that. Both write the same layout, so every
+model produced elsewhere once, under the fixed run_name IMPORTED, and after
+that only re-bundles the serve-app when its code changed. Both write the same
+layout, so every
 (family, suffix, run_name) consumer - load_model, deploy_model - handles both.
 
 storage.py is pure: it takes run_id/run_name as explicit args and never reads
@@ -35,8 +36,11 @@ from mlflow.tracking import MlflowClient
 from cortexgrid import s3_util
 from cortexgrid.infra import get_mlflow_tracking_uri, get_s3_bucket
 from cortexgrid.model_serving import (
+    build_bundle,
     bundle_class,
+    metadata_from_tags,
     metadata_to_tags,
+    upload_bundle,
 )
 
 
@@ -50,9 +54,10 @@ _PHASE_UPLOAD_FAILED = "upload_failed"
 _PHASE_BROKEN = "broken"
 
 # An upload still marked "uploading" this long after the version was created is
-# treated as broken: save_model creates the version immediately before the
-# upload begins, so creation_timestamp is the upload start, and a process that
-# dies mid-upload never flips the tag to "ready"/"upload_failed". Expiry is
+# treated as broken: the version is created before the weights are resolved
+# (for `import_model`, before its download), so creation_timestamp is the
+# start of the whole upload, and a process that dies mid-way never flips the
+# tag to "ready"/"upload_failed". Expiry is
 # derived lazily on read (see `_phase_for`); nothing is written back.
 _UPLOAD_DEADLINE = timedelta(hours=3)
 
@@ -178,11 +183,15 @@ def import_model(
 
     `source` is the weights directory, or a callable returning it; the callable
     runs only when the upload actually happens, so an expensive download can be
-    skipped on every run after the first.
+    skipped on every run after the first. It runs after the version is
+    registered as "uploading", so a concurrent import sees this one in flight
+    while it downloads.
 
     If a version is already registered under the key:
-      - "ready": no-op, returns it. `source` and `serve_app` are ignored; to
-        replace the weights or the serve-app, `delete_model` it first.
+      - "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`). To replace the weights,
+        `delete_model` it first.
       - "uploading": raises RuntimeError - another process is importing it.
       - "upload_failed" / "broken": deleted and imported again.
 
@@ -192,6 +201,7 @@ def import_model(
     existing = model_registry_status(family, suffix, IMPORTED)
     if existing is not None:
         if existing.phase == _PHASE_READY:
+            _refresh_bundle(serve_app, family, suffix)
             return existing
         if existing.phase == _PHASE_UPLOADING:
             raise RuntimeError(
@@ -199,30 +209,53 @@ def import_model(
                 "another process"
             )
         delete_model(family, suffix, IMPORTED)
-    weights_dir = source() if callable(source) else source
-    return _upload_model(weights_dir, serve_app, suffix, family, None, IMPORTED)
+    return _upload_model(source, serve_app, suffix, family, None, IMPORTED)
+
+
+def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
+    """Re-bundle an imported model's serve-app when its fingerprint differs
+    from the bundle stored on the version, leaving the weights in place.
+
+    The new bundle is uploaded next to the old one and the version's tags are
+    pointed at it, so the next `deploy_model` runs the new code. An app that is
+    already running keeps the code it started with until it is deployed again;
+    the old bundle stays in storage so that app can still restart."""
+    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]
+    serve_bundle = build_bundle(serve_app)
+    if metadata_from_tags(version.tags).fingerprint == serve_bundle.fingerprint:
+        return
+    meta = upload_bundle(serve_bundle, family, suffix, IMPORTED)
+    for key, value in metadata_to_tags(meta).items():
+        client.set_model_version_tag(name, version.version, key, value)
 
 
 def _upload_model(
-    weights_dir: str | Path,
+    weights: str | Path | Callable[[], str | Path],
     serve_app: type,
     suffix: str,
     family: str,
     run_id: str | None,
     run_name: str,
 ) -> SavedModel:
-    """Register a ModelVersion in "uploading", upload the weights and the
-    serve-app bundle, and flip it to "ready" (or "upload_failed")."""
+    """Register a ModelVersion in "uploading", resolve `weights` to a directory
+    (calling it when it is a callable), upload the weights and the serve-app
+    bundle, and flip it to "ready" (or "upload_failed")."""
     bucket = get_s3_bucket()
     prefix = f"models/{run_name}/{family}/{suffix}"
-    size_bytes = _dir_size_bytes(weights_dir)
     source = f"s3://{bucket}/{prefix}/weights/"
     name = f"{family}__{suffix}"
     client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
     _ensure_registered_model(client, name)
-    # Register the version up front in the "uploading" phase so the dashboard
-    # can surface a model while its weights are still streaming to storage. The
-    # bundle tags and the flip to "ready" happen only after the upload lands.
+    # Register the version up front in the "uploading" phase, before `weights`
+    # is resolved: the dashboard surfaces the model while it is still being
+    # 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.
     version = client.create_model_version(
         name=name,
         source=source,
@@ -231,11 +264,14 @@ def _upload_model(
             "family": family,
             "suffix": suffix,
             "run_name": run_name,
-            "size_bytes": str(size_bytes),
             _LIFECYCLE_TAG: _PHASE_UPLOADING,
         },
     )
     try:
+        weights_dir = weights() if callable(weights) else weights
+        client.set_model_version_tag(
+            name, version.version, "size_bytes", str(_dir_size_bytes(weights_dir))
+        )
         s3_util.upload_dir(str(weights_dir), dest_path=f"{prefix}/weights")
         bundle_meta = bundle_class(serve_app, family, suffix, run_name)
         for key, value in metadata_to_tags(bundle_meta).items():

{cortexgrid-0.2.94 → cortexgrid-0.2.96}/docs/cortexgrid/README.md

@@ -150,7 +150,7 @@ deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True
 print(deployed.url)
 ```
 
-`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 is a no-op on later runs; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
+`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.2.94 → cortexgrid-0.2.96}/pyproject.toml

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