cortexgrid 0.2.94__tar.gz → 0.2.95__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.2.94
+Version: 0.2.95
 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

{cortexgrid-0.2.94 → cortexgrid-0.2.95}/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.95}/cortexgrid/model_storage.py

@@ -50,9 +50,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,7 +179,9 @@ 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
@@ -199,30 +202,32 @@ 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 _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 +236,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.95}/pyproject.toml

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