cortexgrid 0.3.14__tar.gz → 0.3.16__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.3.14
+Version: 0.3.16
 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
@@ -213,8 +213,15 @@ from fastapi import FastAPI
 
 app = FastAPI()
 
+class MyClient(cortexgrid.DeploymentClient):
+    def complete(self, prompt: str) -> str: ...
+
 @serve.ingress(app)
 class MyServeApp:
+    @classmethod
+    def client(cls, deployment: cortexgrid.Deployment[MyClient]) -> MyClient:
+        return MyClient(key=deployment.key, url=deployment.url)
+
     def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
         self._weights_dir = cortexgrid.load_model(
             deployment.family, deployment.suffix, deployment.run_name
@@ -228,8 +235,8 @@ saved = cortexgrid.save_model(
     # What one replica needs; the model is deployed only on a host that has it.
     requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
 )
-deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
-print(deployed.url)
+deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name)
+model = deployed.client()   # MyServeApp's client, once the app serves
 ```
 
 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. Among the hosts that fit, the model goes to the **smallest GPU** that does, so a 4 GiB model does not occupy a 128 GiB card a bigger one needs; it moves up only once the smaller cards are full. `num_gpus` may be a fraction (`0.25`) to share one card between models, in which case `vram_gb` is what keeps them from overcommitting it. 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.

{cortexgrid-0.3.14 → cortexgrid-0.3.16}/cortexgrid/__init__.py

@@ -102,6 +102,7 @@ from cortexgrid.model_storage import register_model as _register_model_storage
 from cortexgrid.model_storage import save_model as _save_model_storage
 from cortexgrid.model_serving import (
     Deployment,
+    DeploymentClient,
     DeploymentConfig,
     DeploymentKey,
     ModelDeployFailed,
@@ -111,6 +112,7 @@ from cortexgrid.model_serving import (
     list_deployed_models,
     model_serving_status,
     redeploy_model,
+    required_models,
     undeploy_model,
     wait_for_model_serving,
 )
@@ -170,6 +172,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 +201,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 +225,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,8 +240,8 @@ def deploy_model(
     num_replicas: int = 1,
     wait: bool = False,
     timeout: float | None = 300.0,
-    config: DeploymentConfig | None = None,
-) -> Deployment:
+    config: dict[str, str] | None = None,
+) -> Deployment[Any]:
     experiment = active_experiment()
     return _deploy_model_serving(
         family,
@@ -250,6 +255,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:
@@ -335,6 +350,7 @@ __all__ = [
     "delete_model",
     # Model serving
     "Deployment",
+    "DeploymentClient",
     "DeploymentConfig",
     "DeploymentKey",
     "ModelDeployFailed",
@@ -342,6 +358,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.16}/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.16}/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.16}/cortexgrid/model_serving/__init__.py

@@ -35,10 +35,14 @@ from cortexgrid.model_serving.deployment_key import (
     deployment_key,
 )
 from cortexgrid.model_serving.lifecycle import (
+    Deployment,
+    DeploymentClient,
     ModelDeployFailed,
     ModelNotDeployed,
     deploy_model,
+    list_deployed_models,
     redeploy_model,
+    required_models,
     undeploy_model,
     wait_for_model_serving,
 )
@@ -48,6 +52,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,
 )
@@ -60,12 +65,10 @@ from cortexgrid.model_serving.serve_bundle import (
     upload_bundle,
 )
 from cortexgrid.model_serving.status import (
-    Deployment,
     ReplicaPlacement,
     ServingMessage,
     ServingStatus,
     deployment_config,
-    list_deployed_models,
     model_replica_placements,
     model_serving_messages,
     model_serving_status,

{cortexgrid-0.3.14 → cortexgrid-0.3.16}/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.16}/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.16}/cortexgrid/model_serving/lifecycle.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
+import asyncio
+import importlib
 import time
-from typing import Any
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
 
 from ray.serve.schema import ApplicationStatus
 
@@ -22,6 +25,8 @@ from cortexgrid.model_serving.deployment_records import (
     DeploymentRecord,
     delete_deployment_record,
     get_deployment_record,
+    key_of_record,
+    list_deployment_records,
     patch_deployment_record,
     put_deployment_record,
 )
@@ -29,9 +34,8 @@ from cortexgrid.model_serving.placement import ModelRequirements, vram_tiers
 from cortexgrid.model_serving.registry_tags import load_deploy_metadata
 from cortexgrid.model_serving.serve_bundle import BundleMetadata
 from cortexgrid.model_serving.status import (
+    PHASE_NOT_DEPLOYED,
     PHASE_PAUSED,
-    Deployment,
-    deployment_of_record,
     observed,
     replaced_bundle_fingerprint_until_rolled_out,
 )
@@ -77,6 +81,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.
 
@@ -97,16 +105,12 @@ def _wait_for_application_running(
     """`wait_for_model_serving` against a deadline already running. `timeout`
     only labels the TimeoutError."""
     while True:
-        app = get_serve_details().get("applications", {}).get(name)
-        if app is None:
-            raise ModelDeployFailed(f"Serve app {name!r} does not exist")
+        app = _application_still_able_to_serve(name)
         status = str(app.get("status", "(missing)"))
-        message = str(app.get("message", ""))
         if status == ApplicationStatus.RUNNING.value:
             return
-        if status == ApplicationStatus.DEPLOY_FAILED.value:
-            raise ModelDeployFailed(f"Serve app {name!r} DEPLOY_FAILED: {message}")
         if _past(deadline):
+            message = str(app.get("message", ""))
             raise TimeoutError(
                 f"Serve app {name!r} did not reach RUNNING within {timeout}s "
                 f"(last status={status!r}, message={message!r})"
@@ -114,6 +118,90 @@ def _wait_for_application_running(
         time.sleep(_SERVING_POLL_INTERVAL_S)
 
 
+def _application_still_able_to_serve(name: str) -> dict[str, Any]:
+    serve_details = get_serve_details()
+    applications = serve_details.get("applications", {})
+    app: dict[str, Any] | None = applications.get(name)
+    if app is None:
+        raise ModelDeployFailed(f"Serve app {name!r} does not exist")
+    status = app.get("status")
+    if status == ApplicationStatus.DEPLOY_FAILED.value:
+        message = app.get("message", "")
+        raise ModelDeployFailed(f"Serve app {name!r} DEPLOY_FAILED: {message}")
+    return app
+
+
+def _model_is_serving(key: DeploymentKey) -> bool:
+    name = app_name(key)
+    app = _application_still_able_to_serve(name)
+    status = app.get("status")
+    return status == ApplicationStatus.RUNNING.value
+
+
+@dataclass
+class DeploymentClient:
+    key: DeploymentKey
+    url: str
+
+    async def is_ready(self) -> bool:
+        serving = await asyncio.to_thread(_model_is_serving, self.key)
+        return serving
+
+
+DeploymentClientT = TypeVar("DeploymentClientT", bound=DeploymentClient)
+
+
+@dataclass
+class Deployment(Generic[DeploymentClientT]):
+    """A scheduled Ray Serve app fronting a model. `phase` is the normalized
+    serving lifecycle phase (see `ServingStatus`); an app that appears in a
+    listing always exists, so its phase is never "not_deployed"."""
+
+    key: DeploymentKey
+    config: dict[str, str]
+    url: str
+    phase: str
+    bundle_fingerprint: str
+    replaced_bundle_fingerprint: str
+    experiment_name: str
+    class_import_path: str
+
+    def client(self) -> DeploymentClientT:
+        wait_for_model_serving(self.key)
+        deployment_client = self.client_async()
+        return deployment_client
+
+    def client_async(self) -> DeploymentClientT:
+        module_name, class_name = self.class_import_path.split(":")
+        serve_app_module = importlib.import_module(module_name)
+        serve_app = getattr(serve_app_module, class_name)
+        deployment_client: DeploymentClientT = serve_app.client(self)
+        return deployment_client
+
+
+def list_deployed_models() -> list[Deployment[Any]]:
+    """Return a Deployment for every model `deploy_model` put on Ray Serve
+    whose app the control plane last saw existing."""
+    return [
+        deployment_of_record(record)
+        for record in list_deployment_records()
+        if record["phase"] != PHASE_NOT_DEPLOYED
+    ]
+
+
+def deployment_of_record(record: DeploymentRecord) -> Deployment[Any]:
+    return Deployment(
+        key=key_of_record(record),
+        config=record["config"],
+        url=record["url"],
+        phase=record["phase"],
+        bundle_fingerprint=bundle_fingerprint_in_spec(record["spec"]),
+        replaced_bundle_fingerprint=record["replaced_bundle_fingerprint"],
+        experiment_name=record["experiment_name"],
+        class_import_path=record["spec"]["args"]["class_import_path"],
+    )
+
+
 def _clear_failed_application(
     key: DeploymentKey, timeout: float | None, deadline: float | None
 ) -> None:
@@ -219,13 +307,12 @@ 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:
+) -> Deployment[Any]:
     """Schedule a Ray Serve app for a previously-saved model and return a
-    handle carrying its base URL. The caller (e.g. model-gateway) builds
-    whatever client the app's routes need - streaming, long timeouts, custom
-    request schemas - against that URL; cortexgrid imposes no traffic contract.
+    handle on it. The handle's `client()` is the client the serve-app's own
+    `client` factory makes for its routes; cortexgrid imposes no traffic contract.
 
     The serve-app class is pulled from the registry entry's tags `save_model`
     wrote at save time; the caller does not need to hold the class object.
@@ -259,9 +346,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
@@ -317,10 +415,11 @@ def deploy_model(
         bundle_fingerprint=meta.fingerprint,
         replaced_bundle_fingerprint=observation["replaced_bundle_fingerprint"],
         experiment_name=experiment_name,
+        class_import_path=meta.class_import_path,
     )
 
 
-def redeploy_model(key: DeploymentKey) -> Deployment:
+def redeploy_model(key: DeploymentKey) -> Deployment[Any]:
     record = get_deployment_record(key)
     if record is None:
         raise ModelNotDeployed(f"{key} is not deployed")
@@ -334,6 +433,21 @@ def redeploy_model(key: DeploymentKey) -> Deployment:
     )
 
 
+def required_models(deployment: DeploymentKey) -> list[Deployment[Any]]:
+    _, 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[Any]:
+    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.16}/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.16}/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.16}/cortexgrid/model_serving/serve_bundle.py

@@ -55,12 +55,19 @@ def build_bundle(cls: type) -> ServeBundle:
     is a subclass Ray defines in its own module, and on older Ray (e.g. 2.9) it
     reports that module as its own, so the class's source and import path would
     resolve to Ray instead of the serve-app. `cortexgrid.serve.ingress` leaves
-    the class unwrapped."""
+    the class unwrapped. Raises ValueError, too, for a class with no `client`
+    factory, which every `Deployment` of it makes its client with."""
     if any(klass.__module__.startswith("ray.serve") for klass in cls.__mro__):
         raise ValueError(
             f"{cls.__name__} is wrapped by ray.serve.ingress; decorate it with "
             "cortexgrid.serve.ingress instead (from cortexgrid import serve)"
         )
+    if not callable(getattr(cls, "client", None)):
+        raise ValueError(
+            f"{cls.__name__} has no client factory; give it a classmethod "
+            "client(cls, deployment: cortexgrid.Deployment) returning its "
+            "cortexgrid.DeploymentClient"
+        )
     entry_file = Path(inspect.getfile(cls)).resolve()
     serve_entry = Path(__file__).parent.with_name("_serve_entry.py")
     desc = bundle(entry_file).merge(bundle(serve_entry))

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

@@ -6,13 +6,9 @@ from typing import Any
 
 from ray.serve.schema import ApplicationStatus, ReplicaState
 
-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.application_spec import app_name
+from cortexgrid.model_serving.deployment_key import DeploymentKey
 from cortexgrid.model_serving.deployment_records import (
-    DeploymentRecord,
     get_deployment_record,
     key_of_record,
     list_deployment_records,
@@ -21,7 +17,7 @@ from cortexgrid.model_serving.deployment_records import (
 from cortexgrid.ray_util import get_serve_details
 
 
-_PHASE_NOT_DEPLOYED = "not_deployed"
+PHASE_NOT_DEPLOYED = "not_deployed"
 PHASE_PAUSED = "paused"
 _ROLLED_OUT_PHASES = ("running", PHASE_PAUSED)
 
@@ -44,50 +40,13 @@ def _serve_phase(raw_status: str) -> str:
     return _PHASE_BY_SERVE_STATUS.get(raw_status, "deploying")
 
 
-@dataclass
-class Deployment:
-    """A scheduled Ray Serve app fronting a model. `phase` is the normalized
-    serving lifecycle phase (see `ServingStatus`); an app that appears in a
-    listing always exists, so its phase is never "not_deployed"."""
-
-    key: DeploymentKey
-    config: DeploymentConfig
-    url: str
-    phase: str
-    bundle_fingerprint: str
-    replaced_bundle_fingerprint: str
-    experiment_name: str
-
-
 def replaced_bundle_fingerprint_until_rolled_out(
     replaced_bundle_fingerprint: str, phase: str
 ) -> str:
     return "" if phase in _ROLLED_OUT_PHASES else replaced_bundle_fingerprint
 
 
-def list_deployed_models() -> list[Deployment]:
-    """Return a Deployment for every model `deploy_model` put on Ray Serve
-    whose app the control plane last saw existing."""
-    return [
-        deployment_of_record(record)
-        for record in list_deployment_records()
-        if record["phase"] != _PHASE_NOT_DEPLOYED
-    ]
-
-
-def deployment_of_record(record: DeploymentRecord) -> Deployment:
-    return Deployment(
-        key=key_of_record(record),
-        config=record["config"],
-        url=record["url"],
-        phase=record["phase"],
-        bundle_fingerprint=bundle_fingerprint_in_spec(record["spec"]),
-        replaced_bundle_fingerprint=record["replaced_bundle_fingerprint"],
-        experiment_name=record["experiment_name"],
-    )
-
-
-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}")
@@ -134,8 +93,8 @@ def model_serving_status(key: DeploymentKey) -> ServingStatus:
     `cortexgrid.model_storage.model_registry_status`.
     """
     record = get_deployment_record(key)
-    if record is None or record["phase"] == _PHASE_NOT_DEPLOYED:
-        return ServingStatus(key, _PHASE_NOT_DEPLOYED, "", None)
+    if record is None or record["phase"] == PHASE_NOT_DEPLOYED:
+        return ServingStatus(key, PHASE_NOT_DEPLOYED, "", None)
     return ServingStatus(
         key=key,
         phase=record["phase"],
@@ -193,7 +152,7 @@ def observed(app: dict[str, Any] | None) -> dict[str, Any]:
     deployment record keeps; an app that does not exist reads as
     "not_deployed"."""
     if app is None:
-        return {"phase": _PHASE_NOT_DEPLOYED, "message": "", "replicas": []}
+        return {"phase": PHASE_NOT_DEPLOYED, "message": "", "replicas": []}
     raw = str(app.get("status", ""))
     return {
         "phase": _phase(app),

{cortexgrid-0.3.14 → cortexgrid-0.3.16}/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.16}/docs/cortexgrid/README.md

@@ -185,8 +185,15 @@ from fastapi import FastAPI
 
 app = FastAPI()
 
+class MyClient(cortexgrid.DeploymentClient):
+    def complete(self, prompt: str) -> str: ...
+
 @serve.ingress(app)
 class MyServeApp:
+    @classmethod
+    def client(cls, deployment: cortexgrid.Deployment[MyClient]) -> MyClient:
+        return MyClient(key=deployment.key, url=deployment.url)
+
     def __init__(self, deployment: cortexgrid.DeploymentKey) -> None:
         self._weights_dir = cortexgrid.load_model(
             deployment.family, deployment.suffix, deployment.run_name
@@ -200,8 +207,8 @@ saved = cortexgrid.save_model(
     # What one replica needs; the model is deployed only on a host that has it.
     requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
 )
-deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
-print(deployed.url)
+deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name)
+model = deployed.client()   # MyServeApp's client, once the app serves
 ```
 
 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. Among the hosts that fit, the model goes to the **smallest GPU** that does, so a 4 GiB model does not occupy a 128 GiB card a bigger one needs; it moves up only once the smaller cards are full. `num_gpus` may be a fraction (`0.25`) to share one card between models, in which case `vram_gb` is what keeps them from overcommitting it. 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.

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

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