cortexgrid 0.3.14__tar.gz → 0.3.15__tar.gz

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.3.14
+Version: 0.3.15
 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.3.14 → cortexgrid-0.3.15}/cortexgrid/__init__.py

@@ -111,6 +111,7 @@ from cortexgrid.model_serving import (
     list_deployed_models,
     model_serving_status,
     redeploy_model,
+    required_models,
     undeploy_model,
     wait_for_model_serving,
 )
@@ -170,6 +171,7 @@ def save_model(
     run produced (e.g. a fine-tune). For a model produced elsewhere that should
     be uploaded once and reused across runs, use `import_model`."""
     experiment = Experiment.get_instance()
+    _register_required_models(requirements)
     return _save_model_storage(
         weights_dir,
         serve_app,
@@ -198,6 +200,7 @@ def import_model(
     `created_at` under (family, suffix), whether this call uploaded the model
     or reused it."""
     experiment = Experiment.get_instance()
+    _register_required_models(requirements)
     model = _import_model_storage(
         source, serve_app, family, suffix, requirements, config
     )
@@ -221,6 +224,7 @@ def register_model(
     (see `cortexgrid.model_storage.register_model`). The model belongs to no
     run, so the run keeps the link the same way."""
     experiment = Experiment.get_instance()
+    _register_required_models(requirements)
     model = _register_model_storage(
         serve_app, family, suffix, requirements, config
     )
@@ -235,7 +239,7 @@ def deploy_model(
     num_replicas: int = 1,
     wait: bool = False,
     timeout: float | None = 300.0,
-    config: DeploymentConfig | None = None,
+    config: dict[str, str] | None = None,
 ) -> Deployment:
     experiment = active_experiment()
     return _deploy_model_serving(
@@ -250,6 +254,16 @@ def deploy_model(
     )
 
 
+def _register_required_models(requirements: ModelRequirements | None) -> None:
+    for model in [] if requirements is None else requirements.models:
+        if model.serve_app is None:
+            continue
+        if model.source is None:
+            register_model(model.serve_app, model.family, model.suffix)
+        else:
+            import_model(model.source, model.serve_app, model.family, model.suffix)
+
+
 def _record_imported_model(
     run_id: str, family: str, suffix: str, model: SavedModel
 ) -> None:
@@ -342,6 +356,7 @@ __all__ = [
     "ServingStatus",
     "deploy_model",
     "redeploy_model",
+    "required_models",
     "wait_for_model_serving",
     "model_serving_status",
     "undeploy_model",

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/_model_scheduler.py

@@ -16,9 +16,10 @@ log = logging.getLogger(__name__)
 
 
 _RayResources = dict[str, float]
+_ModelGraph = dict[str, set[str]]
 
 
-_SCHEDULER_PROTOCOL_VERSION = 2
+_SCHEDULER_PROTOCOL_VERSION = 3
 _SCHEDULER_ACTOR_NAME = f"cortexgrid-model-scheduler-v{_SCHEDULER_PROTOCOL_VERSION}"
 _SCHEDULER_ACTOR_NAMESPACE = "cortexgrid"
 
@@ -33,7 +34,7 @@ _STATE_API_RESULT_LIMIT = 10_000
 _RESOURCE_FLOAT_TOLERANCE = 1e-6
 
 
-def model_autoscaling_config(max_replicas: int) -> dict[str, Any]:
+def model_autoscaling_config(max_replicas: int, required_apps: list[str]) -> dict[str, Any]:
     return {
         "min_replicas": 0,
         "initial_replicas": max_replicas,
@@ -48,12 +49,14 @@ def model_autoscaling_config(max_replicas: int) -> dict[str, Any]:
                 f"{ModelAutoscalingPolicy.__module__}:"
                 f"{ModelAutoscalingPolicy.__qualname__}"
             ),
+            "policy_kwargs": {"required_apps": required_apps},
         },
     }
 
 
 class ModelAutoscalingPolicy:
-    def __init__(self) -> None:
+    def __init__(self, required_apps: list[str]) -> None:
+        self._required_apps = required_apps
         self._scheduler: ActorProxy[_ModelScheduler] | None = None
         self._pending_pause_answer: Future[bool] | None = None
         self._pause_requested_by_scheduler = False
@@ -63,7 +66,9 @@ class ModelAutoscalingPolicy:
         if not self._awaiting_pause_answer():
             self._pause_requested_by_scheduler = self._collect_pause_answer()
             self._send_activity_report(
-                context.deployment_id.to_replica_actor_class_name(), has_requests
+                context.deployment_id.to_replica_actor_class_name(),
+                context.deployment_id.app_name,
+                has_requests,
             )
         return self._replica_count(context, has_requests), context.policy_state
 
@@ -90,12 +95,14 @@ class ModelAutoscalingPolicy:
             self._scheduler = None
             return False
 
-    def _send_activity_report(self, replica_class_name: str, has_requests: bool) -> None:
+    def _send_activity_report(
+        self, replica_class_name: str, app_name: str, has_requests: bool
+    ) -> None:
         if self._scheduler is None:
             self._scheduler = _get_or_create_scheduler_actor()
         self._pending_pause_answer = (
             self._scheduler.report_activity_and_check_pause.remote(
-                replica_class_name, has_requests
+                replica_class_name, app_name, self._required_apps, has_requests
             ).future()
         )
 
@@ -117,6 +124,8 @@ def _get_or_create_scheduler_actor() -> ActorProxy[_ModelScheduler]:
 
 @dataclass
 class _ScheduledModel:
+    app_name: str = ""
+    required_apps: list[str] = field(default_factory=list)
     has_requests: bool = False
     last_request_at: float = 0.0
     last_report_at: float = 0.0
@@ -131,6 +140,7 @@ class _NodeOccupancy:
 @dataclass
 class _ReplicaWaitingForRoom:
     actor_id: str
+    replica_class_name: str
     required_resources: _RayResources
 
 
@@ -149,16 +159,27 @@ class _ModelScheduler:
 
     @ray.method
     def report_activity_and_check_pause(
-        self, replica_class_name: str, has_requests: bool
+        self,
+        replica_class_name: str,
+        app_name: str,
+        required_apps: list[str],
+        has_requests: bool,
     ) -> bool:
         return self._record_activity_and_check_pause(
-            replica_class_name, has_requests, time.monotonic()
+            replica_class_name, app_name, required_apps, has_requests, time.monotonic()
         )
 
     def _record_activity_and_check_pause(
-        self, replica_class_name: str, has_requests: bool, now: float
+        self,
+        replica_class_name: str,
+        app_name: str,
+        required_apps: list[str],
+        has_requests: bool,
+        now: float,
     ) -> bool:
-        self._record_activity(replica_class_name, has_requests, now)
+        self._record_activity(
+            replica_class_name, app_name, required_apps, has_requests, now
+        )
         if now - self._last_pause_decision_at >= _PAUSE_DECISION_INTERVAL_S:
             self._last_pause_decision_at = now
             self._forget_models_silent_since(now - _FORGET_MODELS_SILENT_FOR_S)
@@ -166,9 +187,16 @@ class _ModelScheduler:
         return replica_class_name in self._models_to_pause
 
     def _record_activity(
-        self, replica_class_name: str, has_requests: bool, now: float
+        self,
+        replica_class_name: str,
+        app_name: str,
+        required_apps: list[str],
+        has_requests: bool,
+        now: float,
     ) -> None:
         model = self._models.setdefault(replica_class_name, _ScheduledModel())
+        model.app_name = app_name
+        model.required_apps = required_apps
         if has_requests:
             model.last_request_at = now
         model.has_requests = has_requests
@@ -194,25 +222,56 @@ class _ModelScheduler:
             occupancy.replicas_waiting_for_room,
             key=lambda replica: self._first_seen_waiting_at[replica.actor_id],
         )
+        running = {
+            name for node in occupancy.nodes for name in node.resources_held_by_model
+        }
+        required = self._models_required_directly()
         models_to_pause: set[str] = set()
         for replica in waiting_longest_first:
             models_to_pause |= self._fewest_idle_models_to_pause_for(
-                replica.required_resources, occupancy.nodes, models_to_pause
+                replica.required_resources,
+                occupancy.nodes,
+                self._pause_groups(
+                    running - models_to_pause,
+                    {replica.replica_class_name}
+                    | _reachable(replica.replica_class_name, required),
+                    required,
+                ),
             )
         return models_to_pause
 
+    def _models_required_directly(self) -> _ModelGraph:
+        model_of_app = {model.app_name: name for name, model in self._models.items()}
+        return {
+            name: {model_of_app[app] for app in model.required_apps if app in model_of_app}
+            for name, model in self._models.items()
+        }
+
+    def _pause_groups(
+        self, running: set[str], unpausable: set[str], required: _ModelGraph
+    ) -> _ModelGraph:
+        dependents = _reversed(required)
+        groups: _ModelGraph = {}
+        for name in running:
+            group = {name} | (_reachable(name, dependents) & running)
+            if unpausable.isdisjoint(group) and not any(
+                self._models[member].has_requests for member in group
+            ):
+                groups[name] = group
+        return groups
+
     def _fewest_idle_models_to_pause_for(
         self,
         required_resources: _RayResources,
         nodes: list[_NodeOccupancy],
-        already_pausing: set[str],
+        pause_groups: _ModelGraph,
     ) -> set[str]:
         fewest_models_to_pause: list[str] | None = None
         for node in nodes:
             if _has_room_for(required_resources, node.free_resources):
                 return set()
             models_to_pause = self._least_recently_used_idle_models_freeing(
-                required_resources, node, already_pausing
+                required_resources, node, pause_groups
             )
             if models_to_pause is not None and (
                 fewest_models_to_pause is None
@@ -225,23 +284,21 @@ class _ModelScheduler:
         self,
         required_resources: _RayResources,
         node: _NodeOccupancy,
-        already_pausing: set[str],
+        pause_groups: _ModelGraph,
     ) -> list[str] | None:
         idle_models_least_recently_used_first = sorted(
-            (
-                name
-                for name in node.resources_held_by_model
-                if name not in already_pausing and not self._models[name].has_requests
-            ),
+            (name for name in node.resources_held_by_model if name in pause_groups),
             key=lambda name: self._models[name].last_request_at,
         )
         resources_free_after_pause = dict(node.free_resources)
         models_to_pause: list[str] = []
         for name in idle_models_least_recently_used_first:
-            models_to_pause.append(name)
-            _add_resources(
-                resources_free_after_pause, node.resources_held_by_model[name]
-            )
+            for member in sorted(pause_groups[name] - set(models_to_pause)):
+                models_to_pause.append(member)
+                _add_resources(
+                    resources_free_after_pause,
+                    node.resources_held_by_model.get(member, {}),
+                )
             if _has_room_for(required_resources, resources_free_after_pause):
                 return models_to_pause
         return None
@@ -269,7 +326,11 @@ def _read_cluster_occupancy(scheduled_replica_class_names: set[str]) -> _Cluster
         reserved_resources = actor_fields["required_resources"] or {}
         if _is_serve_replica_waiting_for_room(actor_fields):
             replicas_waiting_for_room.append(
-                _ReplicaWaitingForRoom(actor_fields["actor_id"], reserved_resources)
+                _ReplicaWaitingForRoom(
+                    actor_fields["actor_id"],
+                    actor_fields["class_name"],
+                    reserved_resources,
+                )
             )
             continue
         occupancy = occupancy_by_node_id.get(actor_fields["node_id"])
@@ -305,6 +366,25 @@ def _is_serve_replica_waiting_for_room(actor_fields: dict[str, Any]) -> bool:
     )
 
 
+def _reachable(start: str, edges: _ModelGraph) -> set[str]:
+    reached: set[str] = set()
+    to_visit = list(edges.get(start, ()))
+    while to_visit:
+        name = to_visit.pop()
+        if name not in reached:
+            reached.add(name)
+            to_visit.extend(edges.get(name, ()))
+    return reached
+
+
+def _reversed(edges: _ModelGraph) -> _ModelGraph:
+    reversed_edges: _ModelGraph = {}
+    for source, targets in edges.items():
+        for target in targets:
+            reversed_edges.setdefault(target, set()).add(source)
+    return reversed_edges
+
+
 def _add_resources(target: _RayResources, amounts: _RayResources) -> None:
     for name, amount in amounts.items():
         target[name] = target.get(name, 0.0) + amount

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/_serve_entry.py

@@ -77,7 +77,9 @@ def build(args: dict[str, Any]) -> Application:
         # This builder ships in the bundle, frozen at save time, while `args`
         # come from the cortexgrid that deploys it; one older than the bundle
         # sends neither key.
-        autoscaling_config=model_autoscaling_config(args.get("num_replicas", 1)),
+        autoscaling_config=model_autoscaling_config(
+            args.get("num_replicas", 1), args.get("required_apps", [])
+        ),
         # Ray 2.32 lowered the default from 100 to 5; keep what serve-apps
         # had on Ray 2.9.
         max_ongoing_requests=_MAX_ONGOING_REQUESTS,

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/__init__.py

@@ -39,6 +39,7 @@ from cortexgrid.model_serving.lifecycle import (
     ModelNotDeployed,
     deploy_model,
     redeploy_model,
+    required_models,
     undeploy_model,
     wait_for_model_serving,
 )
@@ -48,6 +49,7 @@ from cortexgrid.model_serving.registry_tags import (
     has_requirement_tags,
     metadata_from_tags,
     metadata_to_tags,
+    required_models_to_tags,
     requirements_from_tags,
     requirements_to_tags,
 )

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/application_spec.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from cortexgrid.model_serving.deployment_key import DeploymentKey
+from cortexgrid.model_serving.deployment_key import DeploymentKey, deployment_key
 from cortexgrid.model_serving.placement import ModelRequirements, ray_actor_options
 from cortexgrid.model_serving.serve_bundle import (
     BundleMetadata,
@@ -41,6 +41,22 @@ def build_application_spec(
     runtime_env: dict[str, Any] = {"working_dir": meta.bundle_url}
     if meta.pip_requirements:
         runtime_env["pip"] = meta.pip_requirements
+    args: dict[str, Any] = {
+        "class_import_path": meta.class_import_path,
+        "family": key.family,
+        "suffix": key.suffix,
+        "run_name": key.run_name,
+        "config_fingerprint": key.config_fingerprint,
+        "num_replicas": num_replicas,
+        "ray_actor_options": ray_actor_options(requirements, tiers),
+    }
+    if requirements.models:
+        args["required_apps"] = [
+            app_name(
+                deployment_key(model.family, model.suffix, model.run_name, model.config)
+            )
+            for model in requirements.models
+        ]
     return {
         "name": app_name(key),
         "route_prefix": route_prefix(key),
@@ -49,15 +65,7 @@ def build_application_spec(
         # Deployment class is rejected, so cortexgrid.deploy_model goes through
         # a generic builder that re-imports the user's class and binds it.
         "import_path": "cortexgrid._serve_entry:build",
-        "args": {
-            "class_import_path": meta.class_import_path,
-            "family": key.family,
-            "suffix": key.suffix,
-            "run_name": key.run_name,
-            "config_fingerprint": key.config_fingerprint,
-            "num_replicas": num_replicas,
-            "ray_actor_options": ray_actor_options(requirements, tiers),
-        },
+        "args": args,
         "runtime_env": runtime_env,
     }
 

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/deployment_key.py

@@ -2,10 +2,27 @@ from __future__ import annotations
 
 import hashlib
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
 
 
-DeploymentConfig = dict[str, str]
+@dataclass
+class DeploymentConfig:
+    family: str
+    suffix: str
+    run_name: str
+    num_replicas: int = 1
+    config: dict[str, str] = field(default_factory=dict)
+    serve_app: type | None = None
+    source: Callable[[], str | Path] | None = None
+
+    def __post_init__(self) -> None:
+        if self.source is not None and self.serve_app is None:
+            raise ValueError(
+                f"{self.family}/{self.suffix} has a source but no serve_app to front it"
+            )
+
 
 _CONFIG_FINGERPRINT_LENGTH = 12
 
@@ -19,12 +36,12 @@ class DeploymentKey:
 
 
 def deployment_key(
-    family: str, suffix: str, run_name: str, config: DeploymentConfig
+    family: str, suffix: str, run_name: str, config: dict[str, str]
 ) -> DeploymentKey:
     return DeploymentKey(family, suffix, run_name, _config_fingerprint(config))
 
 
-def _config_fingerprint(config: DeploymentConfig) -> str:
+def _config_fingerprint(config: dict[str, str]) -> str:
     for name, value in config.items():
         if not isinstance(name, str) or not isinstance(value, str):
             raise ValueError(

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/lifecycle.py

@@ -77,6 +77,10 @@ def _past(deadline: float | None) -> bool:
     return deadline is not None and time.monotonic() >= deadline
 
 
+def _time_left(deadline: float | None) -> float | None:
+    return None if deadline is None else max(0.0, deadline - time.monotonic())
+
+
 def wait_for_model_serving(key: DeploymentKey, timeout: float | None = None) -> None:
     """Block until the model's Serve app is RUNNING.
 
@@ -219,7 +223,7 @@ def deploy_model(
     num_replicas: int = 1,
     wait: bool = False,
     timeout: float | None = 300.0,
-    config: DeploymentConfig | None = None,
+    config: dict[str, str] | None = None,
     experiment_name: str = "",
 ) -> Deployment:
     """Schedule a Ray Serve app for a previously-saved model and return a
@@ -259,9 +263,20 @@ def deploy_model(
     stuck in DEPLOYING) will hang forever.
     """
     deadline = _deadline(timeout)
+    meta, requirements = load_deploy_metadata(family, suffix, run_name)
+    for model in requirements.models:
+        deploy_model(
+            model.family,
+            model.suffix,
+            model.run_name,
+            num_replicas=model.num_replicas,
+            wait=True,
+            timeout=_time_left(deadline),
+            config=model.config,
+            experiment_name=experiment_name,
+        )
     deployment_config = config or {}
     key = deployment_key(family, suffix, run_name, deployment_config)
-    meta, requirements = load_deploy_metadata(family, suffix, run_name)
     record = get_deployment_record(key)
     if record is not None and _record_is_current(
         record, key, meta, requirements, num_replicas
@@ -334,6 +349,21 @@ def redeploy_model(key: DeploymentKey) -> Deployment:
     )
 
 
+def required_models(deployment: DeploymentKey) -> list[Deployment]:
+    _, requirements = load_deploy_metadata(
+        deployment.family, deployment.suffix, deployment.run_name
+    )
+    return [_deployment_of(model) for model in requirements.models]
+
+
+def _deployment_of(model: DeploymentConfig) -> Deployment:
+    key = deployment_key(model.family, model.suffix, model.run_name, model.config)
+    record = get_deployment_record(key)
+    if record is None:
+        raise ModelNotDeployed(f"{key} is not deployed")
+    return deployment_of_record(record)
+
+
 def undeploy_model(key: DeploymentKey) -> None:
     """Tear down the Ray Serve app of this deployment and drop its record."""
     name = app_name(key)

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/placement.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
+from cortexgrid.model_serving.deployment_key import DeploymentConfig
 from cortexgrid.ray_util import get_ray_nodes
 
 
@@ -24,6 +25,7 @@ class ModelRequirements:
     ram_gb: float = 0.0
     # GPU memory across the replica's num_gpus GPUs, so it needs a GPU share.
     vram_gb: float = 0.0
+    models: list[DeploymentConfig] = field(default_factory=list)
 
     def __post_init__(self) -> None:
         if self.num_gpus < 0 or self.ram_gb < 0 or self.vram_gb < 0:

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/registry_tags.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import json
 
 from cortexgrid import state
+from cortexgrid.model_serving.deployment_key import DeploymentConfig
 from cortexgrid.model_serving.placement import ModelRequirements
 from cortexgrid.model_serving.serve_bundle import BundleMetadata
 
@@ -48,6 +49,7 @@ def bundle_fingerprint_from_tags(tags: dict[str, str]) -> str:
 _NUM_GPUS_TAG = "num_gpus"
 _RAM_GB_TAG = "ram_gb"
 _VRAM_GB_TAG = "vram_gb"
+_MODELS_TAG = "models"
 
 
 def has_requirement_tags(tags: dict[str, str]) -> bool:
@@ -62,6 +64,24 @@ def requirements_to_tags(requirements: ModelRequirements) -> dict[str, str]:
         _NUM_GPUS_TAG: str(requirements.num_gpus),
         _RAM_GB_TAG: str(requirements.ram_gb),
         _VRAM_GB_TAG: str(requirements.vram_gb),
+        **required_models_to_tags(requirements.models),
+    }
+
+
+def required_models_to_tags(models: list[DeploymentConfig]) -> dict[str, str]:
+    return {
+        _MODELS_TAG: json.dumps(
+            [
+                {
+                    "family": model.family,
+                    "suffix": model.suffix,
+                    "run_name": model.run_name,
+                    "num_replicas": model.num_replicas,
+                    "config": model.config,
+                }
+                for model in models
+            ]
+        )
     }
 
 
@@ -74,6 +94,9 @@ def requirements_from_tags(tags: dict[str, str]) -> ModelRequirements:
         num_gpus=float(tags.get(_NUM_GPUS_TAG, "0")),
         ram_gb=float(tags.get(_RAM_GB_TAG, "0")),
         vram_gb=float(tags.get(_VRAM_GB_TAG, "0")),
+        models=[
+            DeploymentConfig(**model) for model in json.loads(tags.get(_MODELS_TAG, "[]"))
+        ],
     )
 
 

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_serving/status.py

@@ -10,7 +10,7 @@ from cortexgrid.model_serving.application_spec import (
     app_name,
     bundle_fingerprint_in_spec,
 )
-from cortexgrid.model_serving.deployment_key import DeploymentConfig, DeploymentKey
+from cortexgrid.model_serving.deployment_key import DeploymentKey
 from cortexgrid.model_serving.deployment_records import (
     DeploymentRecord,
     get_deployment_record,
@@ -51,7 +51,7 @@ class Deployment:
     listing always exists, so its phase is never "not_deployed"."""
 
     key: DeploymentKey
-    config: DeploymentConfig
+    config: dict[str, str]
     url: str
     phase: str
     bundle_fingerprint: str
@@ -87,7 +87,7 @@ def deployment_of_record(record: DeploymentRecord) -> Deployment:
     )
 
 
-def deployment_config(key: DeploymentKey) -> DeploymentConfig:
+def deployment_config(key: DeploymentKey) -> dict[str, str]:
     record = get_deployment_record(key)
     if record is None:
         raise ValueError(f"No deployment {key}")

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/cortexgrid/model_storage.py

@@ -8,7 +8,8 @@ One registry entry (a `models` row) per (family, suffix, run_name):
                            "s3://<bucket>/models/<run_name>/<family>/<suffix>/weights/"
                            (NO_WEIGHTS for a model registered without any)
     run linkage         -> run_id  (unset for imported models)
-    requirements        -> tags["num_gpus"], ["ram_gb"], ["vram_gb"]
+    requirements        -> tags["num_gpus"], ["ram_gb"], ["vram_gb"], ["models"]
+                           (models: JSON list)
     config              -> tags["config"]  (JSON object)
 
 Three ways in: `save_model` registers a fresh copy under the calling run's
@@ -28,9 +29,11 @@ cortexgrid/__init__.py.
 
 from __future__ import annotations
 
+import fcntl
 import json
 import os
-from dataclasses import dataclass, field
+import shutil
+from dataclasses import dataclass, field, replace
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 import tempfile
@@ -48,6 +51,7 @@ from cortexgrid.model_serving import (
     has_requirement_tags,
     metadata_from_tags,
     metadata_to_tags,
+    required_models_to_tags,
     requirements_from_tags,
     requirements_to_tags,
     upload_bundle,
@@ -93,6 +97,9 @@ IMPORTED = "imported"
 # why there is no path instead of an empty one.
 NO_WEIGHTS = "cortexgrid://no-weights"
 
+_MODEL_CACHE_DIR_ENV = "CORTEXGRID_MODEL_CACHE_DIR"
+_DEFAULT_MODEL_CACHE_DIR = Path.home() / ".cache" / "cortexgrid" / "models"
+
 
 @dataclass
 class SavedModel:
@@ -223,6 +230,23 @@ def _download_s3_uri(uri: str, dest_dir: str | Path | None) -> Path:
     return dest
 
 
+def _model_cache_dir() -> Path:
+    return Path(os.environ.get(_MODEL_CACHE_DIR_ENV, _DEFAULT_MODEL_CACHE_DIR))
+
+
+def _load_imported_weights(uri: str, family: str, suffix: str) -> Path:
+    cached = _model_cache_dir() / family / suffix
+    cached.parent.mkdir(parents=True, exist_ok=True)
+    with open(cached.with_name(f"{suffix}.lock"), "w") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        if not cached.is_dir():
+            partial = cached.with_name(f"{suffix}.partial")
+            shutil.rmtree(partial, ignore_errors=True)
+            _download_s3_uri(uri, partial)
+            partial.rename(cached)
+    return cached
+
+
 def save_model(
     weights_dir: str | Path,
     serve_app: type,
@@ -292,8 +316,9 @@ def import_model(
         longer matches the stored bundle, it is re-bundled first and the
         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.
+        with `set_model_requirements` / `set_model_config` are kept; the
+        required models in `requirements` replace the stored ones regardless.
+        To replace the weights, `delete_model` it first.
       - "uploading": raises RuntimeError - another process is importing it.
       - "upload_failed" / "broken": deleted and imported again.
 
@@ -386,11 +411,13 @@ def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
 def _set_missing_requirements(
     family: str, suffix: str, requirements: ModelRequirements
 ) -> ModelRequirements:
-    """Store `requirements` on an imported model whose version has none yet.
+    """Store `requirements` on an imported model whose version has none yet;
+    its required models are stored either way, since only code declares them.
     Returns the requirements the version holds afterwards."""
     version = _get_version(family, suffix, IMPORTED)
     if has_requirement_tags(version.tags):
-        return requirements_from_tags(version.tags)
+        _set_tags(family, suffix, IMPORTED, required_models_to_tags(requirements.models))
+        return replace(requirements_from_tags(version.tags), models=requirements.models)
     _set_tags(family, suffix, IMPORTED, requirements_to_tags(requirements))
     return requirements
 
@@ -499,6 +526,8 @@ def load_model(family: str, suffix: str, run_name: str) -> Path:
             f"Model {family}/{suffix}/{run_name} was registered without "
             "weights; there is nothing to load"
         )
+    if run_name == IMPORTED:
+        return _load_imported_weights(version.source, family, suffix)
     return _download_s3_uri(version.source, None)
 
 

{cortexgrid-0.3.14 → cortexgrid-0.3.15}/pyproject.toml

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