cortexgrid 0.3.8__tar.gz → 0.3.10__tar.gz

{cortexgrid-0.3.8 → cortexgrid-0.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.3.8
+Version: 0.3.10
 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
@@ -55,7 +55,7 @@ import cortexgrid
 cortexgrid.init(experiment="weather-forecast")
 ```
 
-That single call reads the service URLs from the head's secrets server at `$CORTEXGRID_HEAD_URL` and connects to all services through them. It also creates (or finds) the named MLflow experiment and starts a new run inside it. If an experiment of that name was deleted (e.g. from the UI), a new experiment is created under the name: MLflow keeps a deleted experiment's name reserved, so the deleted one is renamed to `<name>__deleted__<id>` first (`delete_experiment` does that rename at deletion). Omit `experiment=` to auto-generate a unique name like `funky-koval-12`.
+That single call reads the service URLs from the head's secrets server at `$CORTEXGRID_HEAD_URL` and connects to all services through them. It also creates (or finds) the named experiment and starts a new run inside it. Experiments and runs are cortexgrid's own records, kept by the jobs control plane; each maps onto an MLflow experiment and run, which exist so MLflow can display the run's metrics. MLflow never frees the name of a deleted experiment, so an experiment re-created under the name of a deleted one gets its MLflow experiment under `<name>__<8 hex>`. Omit `experiment=` to auto-generate a unique name like `funky-koval-12`.
 
 **One experiment per binary run.** `cortexgrid.init()` may only be called once per process. Every subsequent `cortexgrid.log_metric`, `cortexgrid.log_artifact`, checkpoint, and `cortexgrid.remote()` submission is scoped to that experiment+run. Remote jobs dispatched by the control plane inherit the experiment+run via the pickled payload, so their logging flows into the same MLflow run as the parent binary.
 
@@ -80,7 +80,7 @@ No run-scoping context manager — `init()` starts the run, and every subsequent
 
 #### Checkpointing and resuming
 
-Inside a cortexgrid job, `cortexgrid.checkpoint()` returns an attribute-based checkpoint object that persists to MLflow artifacts when its `with` block exits. On job restart (either manual retry or `retry=True`), `cortexgrid.resume()` returns the last checkpoint for the same job ID, or `None` if there isn't one.
+Inside a cortexgrid job, `cortexgrid.checkpoint()` returns an attribute-based checkpoint object that persists to S3, with its manifest recorded by the jobs control plane, when its `with` block exits. On job restart (either manual retry or `retry=True`), `cortexgrid.resume()` returns the last checkpoint for the same job ID, or `None` if there isn't one.
 
 ```python
 ckpt = cortexgrid.resume()
@@ -111,7 +111,7 @@ job = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
 print(f"Submitted: {job.job_id}")
 ```
 
-`cortexgrid.remote` submits a job *request* (a pickled payload plus a `JobLifecycle` record) to MLflow and returns a `JobFuture` immediately. It does not wait for the job to run or finish — use the UI at `http://<DGX_IP>:8000`, `job.status()`, or poll `cortexgrid.list_experiment_run_jobs(run_id)`, to observe status.
+`cortexgrid.remote` submits a job *request* (a pickled payload plus a `JobLifecycle` record) to the jobs control plane and returns a `JobFuture` immediately. It does not wait for the job to run or finish — use the UI at `http://<DGX_IP>:8000`, `job.status()`, or poll `cortexgrid.list_experiment_run_jobs(run_id)`, to observe status.
 
 ##### Blocking on the result
 
@@ -163,7 +163,7 @@ The driver cloudpickles the outcome to `job/{job_id}/result.pkl`, beside the pay
 
 Waiting on a `retry=True` job raises `ValueError`: retries are unbounded by design (see below), so the wait would have no end. Fire-and-forget submission is the form training uses — submit, then watch the UI.
 
-A separate service — the **jobs control plane** — polls MLflow for pending job requests, matches them against the set of Ray submissions the cluster already has, and submits anything missing. It is also responsible for retrying failed jobs and honouring user-requested stops.
+A separate service — the **jobs control plane** — keeps cortexgrid's records (experiments, runs, jobs, the model registry, deployments) in Postgres and serves them over HTTP. Its poll loop reads the jobs it still has to act on, matches them against the set of Ray submissions the cluster already has, and submits anything missing. It is also responsible for retrying failed jobs and honouring user-requested stops.
 
 Each submission captures the code and dependencies the entry function needs automatically ([_bundle.py](https://github.com/robodatalab/cortexgrid/blob/main/cortexgrid/_bundle.py)):
 - `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`). The standard library is excluded (it ships with the interpreter)
@@ -182,7 +182,7 @@ Pass `retry=True` and the control plane will resubmit the job whenever Ray repor
 cortexgrid.stop_experiment_run_jobs(run_id)   # stops every job in the run
 ```
 
-`stop_experiment_run_jobs` never touches Ray directly. It only flips `stop_requested` on each job's lifecycle record in MLflow. The control plane observes the flag on its next poll and calls `ray.stop_job` for any attempt that has reached Ray. For jobs that have not yet been submitted, the same flag short-circuits the submission path inside the worker.
+`stop_experiment_run_jobs` never touches Ray directly. It only flips `stop_requested` on each job's lifecycle record, which the control plane keeps. The control plane observes the flag on its next poll and calls `ray.stop_job` for any attempt that has reached Ray. For jobs that have not yet been submitted, the same flag short-circuits the submission path inside the worker.
 
 #### Object storage (S3/MinIO)
 
@@ -242,12 +242,12 @@ Anything else the serve-app has to know about the model - which model a provider
 
 | Function | Description |
 |----------|-------------|
-| `cortexgrid.init(experiment=None)` | Configure connections + start a new MLflow run inside the named experiment. One call per binary. |
+| `cortexgrid.init(experiment=None)` | Configure connections + start a new run inside the named experiment. One call per binary. |
 | `cortexgrid.log_metric(key, value, step)` | Log a metric |
 | `cortexgrid.log_metrics(metrics, step)` | Log multiple metrics |
 | `cortexgrid.log_params(params)` | Log parameters |
 | `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
-| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
+| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved on exit |
 | `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
 | `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a `JobFuture` |
 | `JobFuture.status()` / `.done()` / `.result(timeout=None)` | Live status of a submitted job, and its function's return value (blocking) |
@@ -269,9 +269,10 @@ The DGX Spark runs the following services as k8s workloads managed by Argo CD (s
 | Service | Port | Purpose |
 |---------|------|---------|
 | Ray | 8265 | Dashboard + job submission (NodePort 30265) |
-| MLflow | 5000 | Experiment tracking, model registry |
+| Jobs control plane | 8000 | cortexgrid's records (experiments, runs, jobs, model registry, deployments); schedules jobs on Ray (NodePort 30700) |
+| MLflow | 5000 | Metrics and params of each run |
 | MinIO | 9000/9001 | S3-compatible artifact storage |
-| PostgreSQL | 5432 | MLflow metadata backend |
+| PostgreSQL | 5432 | `cortexgrid` database (the control plane's records), MLflow metadata, UI notes |
 | Prometheus | 9090 | Metrics collection |
 | Grafana | 3000 | Dashboards (GPU, jobs, system) |
 

{cortexgrid-0.3.8 → cortexgrid-0.3.10}/cortexgrid/__init__.py

@@ -31,6 +31,7 @@ from __future__ import annotations
 from pathlib import Path
 from typing import Any, Callable
 
+from cortexgrid import state
 from cortexgrid.checkpoint import checkpoint, resume
 from cortexgrid.experiment import (
     Experiment,
@@ -188,16 +189,14 @@ def import_model(
     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."""
+    so the run keeps the link instead: the run's record notes the model's
+    `created_at` under (family, suffix), whether this call uploaded the model
+    or reused it."""
     experiment = Experiment.get_instance()
     model = _import_model_storage(
         source, serve_app, family, suffix, requirements, config
     )
-    get_mlflow_client().set_tag(
-        experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
-    )
+    _record_imported_model(experiment.run_id, family, suffix, model)
     return model
 
 
@@ -215,18 +214,29 @@ def register_model(
     `import_model` without the import: everything it needs beyond its code goes
     in `config`, which the serve-app reads with `model_config` at construction
     (see `cortexgrid.model_storage.register_model`). The model belongs to no
-    run, so the run keeps the link the same way, under the same tag
-    `imported_model/<family>/<suffix>`."""
+    run, so the run keeps the link the same way."""
     experiment = Experiment.get_instance()
     model = _register_model_storage(
         serve_app, family, suffix, requirements, config
     )
-    get_mlflow_client().set_tag(
-        experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
-    )
+    _record_imported_model(experiment.run_id, family, suffix, model)
     return model
 
 
+def _record_imported_model(
+    run_id: str, family: str, suffix: str, model: SavedModel
+) -> None:
+    """Note on the run's record which imported model it used."""
+    state.put(
+        "runs",
+        run_id,
+        "imported-models",
+        family,
+        suffix,
+        body={"created_at": model.created_at},
+    )
+
+
 __all__ = [
     "Experiment",
     "delete_experiment",

cortexgrid-0.3.10/cortexgrid/_model_scheduler.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+
+import logging
+import time
+from concurrent.futures import Future
+from dataclasses import dataclass, field
+from typing import Any
+
+import ray
+from ray.actor import ActorProxy
+from ray.serve.config import AutoscalingContext
+from ray.util.state import list_actors
+
+
+log = logging.getLogger(__name__)
+
+
+_RayResources = dict[str, float]
+
+
+_SCHEDULER_ACTOR_NAME = "cortexgrid-model-scheduler"
+_SCHEDULER_ACTOR_NAMESPACE = "cortexgrid"
+
+_PAUSE_DECISION_INTERVAL_S = 1.0
+_FORGET_MODELS_SILENT_FOR_S = 30.0
+
+_REQUEST_METRICS_PUSH_INTERVAL_S = 0.5
+_REQUEST_METRICS_AVERAGING_WINDOW_S = 1.0
+
+_STATE_API_RESULT_LIMIT = 10_000
+_RESOURCE_FLOAT_TOLERANCE = 1e-6
+
+
+def model_autoscaling_config(
+    max_replicas: int, ray_actor_options: dict[str, Any]
+) -> dict[str, Any]:
+    return {
+        "min_replicas": 0,
+        "initial_replicas": max_replicas,
+        "max_replicas": max_replicas,
+        "upscale_delay_s": 0,
+        "downscale_delay_s": 0,
+        "downscale_to_zero_delay_s": 0,
+        "metrics_interval_s": _REQUEST_METRICS_PUSH_INTERVAL_S,
+        "look_back_period_s": _REQUEST_METRICS_AVERAGING_WINDOW_S,
+        "policy": {
+            "policy_function": (
+                f"{ModelAutoscalingPolicy.__module__}:"
+                f"{ModelAutoscalingPolicy.__qualname__}"
+            ),
+            "policy_kwargs": {
+                "replica_resources": _resources_requested_by_replica(
+                    ray_actor_options
+                )
+            },
+        },
+    }
+
+
+def _resources_requested_by_replica(
+    ray_actor_options: dict[str, Any],
+) -> _RayResources:
+    resources = {
+        name: float(amount)
+        for name, amount in ray_actor_options.get("resources", {}).items()
+    }
+    if ray_actor_options.get("num_gpus"):
+        resources["GPU"] = float(ray_actor_options["num_gpus"])
+    if ray_actor_options.get("memory"):
+        resources["memory"] = float(ray_actor_options["memory"])
+    return resources
+
+
+class ModelAutoscalingPolicy:
+    def __init__(self, replica_resources: _RayResources) -> None:
+        self._replica_resources = replica_resources
+        self._scheduler: ActorProxy[_ModelScheduler] | None = None
+        self._pending_pause_answer: Future[bool] | None = None
+        self._pause_requested_by_scheduler = False
+
+    def __call__(self, context: AutoscalingContext) -> tuple[int, dict[str, Any]]:
+        has_requests = context.total_num_requests > 0
+        waiting_for_replica = (
+            has_requests or context.target_num_replicas > 0
+        ) and not context.running_replicas
+        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,
+                waiting_for_replica,
+            )
+        return self._replica_count(context, has_requests), context.policy_state
+
+    def _replica_count(self, context: AutoscalingContext, has_requests: bool) -> int:
+        if has_requests:
+            return context.capacity_adjusted_max_replicas
+        if self._pause_requested_by_scheduler:
+            return 0
+        return context.target_num_replicas
+
+    def _awaiting_pause_answer(self) -> bool:
+        return (
+            self._pending_pause_answer is not None
+            and not self._pending_pause_answer.done()
+        )
+
+    def _collect_pause_answer(self) -> bool:
+        if self._pending_pause_answer is None:
+            return False
+        try:
+            return self._pending_pause_answer.result()
+        except Exception:
+            log.exception("Model scheduler unreachable")
+            self._scheduler = None
+            return False
+
+    def _send_activity_report(
+        self, replica_class_name: str, has_requests: bool, waiting_for_replica: 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,
+                self._replica_resources,
+                has_requests,
+                waiting_for_replica,
+            ).future()
+        )
+
+
+def _get_or_create_scheduler_actor() -> ActorProxy[_ModelScheduler]:
+    return (
+        ray.remote(_ModelScheduler)
+        .options(
+            name=_SCHEDULER_ACTOR_NAME,
+            namespace=_SCHEDULER_ACTOR_NAMESPACE,
+            get_if_exists=True,
+            lifetime="detached",
+            num_cpus=0,
+            max_restarts=-1,
+        )
+        .remote()
+    )
+
+
+@dataclass
+class _ScheduledModel:
+    replica_resources: _RayResources
+    has_requests: bool = False
+    waiting_for_replica: bool = False
+    waiting_since: float = 0.0
+    last_request_at: float = 0.0
+    last_report_at: float = 0.0
+
+
+@dataclass
+class _NodeOccupancy:
+    free_resources: _RayResources
+    resources_held_by_model: dict[str, _RayResources] = field(default_factory=dict)
+
+
+class _ModelScheduler:
+    def __init__(self) -> None:
+        self._models: dict[str, _ScheduledModel] = {}
+        self._models_to_pause: set[str] = set()
+        self._last_pause_decision_at = float("-inf")
+
+    @ray.method
+    def report_activity_and_check_pause(
+        self,
+        replica_class_name: str,
+        replica_resources: _RayResources,
+        has_requests: bool,
+        waiting_for_replica: bool,
+    ) -> bool:
+        now = time.monotonic()
+        self._record_activity(
+            replica_class_name, replica_resources, has_requests, waiting_for_replica, 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)
+            self._models_to_pause = self._select_models_to_pause()
+        return replica_class_name in self._models_to_pause
+
+    def _record_activity(
+        self,
+        replica_class_name: str,
+        replica_resources: _RayResources,
+        has_requests: bool,
+        waiting_for_replica: bool,
+        now: float,
+    ) -> None:
+        model = self._models.setdefault(
+            replica_class_name, _ScheduledModel(replica_resources)
+        )
+        model.replica_resources = replica_resources
+        if waiting_for_replica and not model.waiting_for_replica:
+            model.waiting_since = now
+        if has_requests:
+            model.last_request_at = now
+        model.has_requests = has_requests
+        model.waiting_for_replica = waiting_for_replica
+        model.last_report_at = now
+
+    def _forget_models_silent_since(self, cutoff: float) -> None:
+        self._models = {
+            name: model
+            for name, model in self._models.items()
+            if model.last_report_at >= cutoff
+        }
+
+    def _select_models_to_pause(self) -> set[str]:
+        if not any(model.waiting_for_replica for model in self._models.values()):
+            return set()
+        nodes = _read_node_occupancy(set(self._models))
+        models_with_placed_replicas = {
+            name for node in nodes for name in node.resources_held_by_model
+        }
+        models_waiting_for_room = sorted(
+            (
+                name
+                for name, model in self._models.items()
+                if model.waiting_for_replica
+                and name not in models_with_placed_replicas
+            ),
+            key=lambda name: self._models[name].waiting_since,
+        )
+        models_to_pause: set[str] = set()
+        for name in models_waiting_for_room:
+            models_to_pause |= self._fewest_idle_models_to_pause_for(
+                self._models[name].replica_resources, nodes, models_to_pause
+            )
+        return models_to_pause
+
+    def _fewest_idle_models_to_pause_for(
+        self,
+        required_resources: _RayResources,
+        nodes: list[_NodeOccupancy],
+        already_pausing: set[str],
+    ) -> 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
+            )
+            if models_to_pause is not None and (
+                fewest_models_to_pause is None
+                or len(models_to_pause) < len(fewest_models_to_pause)
+            ):
+                fewest_models_to_pause = models_to_pause
+        return set(fewest_models_to_pause or ())
+
+    def _least_recently_used_idle_models_freeing(
+        self,
+        required_resources: _RayResources,
+        node: _NodeOccupancy,
+        already_pausing: set[str],
+    ) -> 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
+            ),
+            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]
+            )
+            if _has_room_for(required_resources, resources_free_after_pause):
+                return models_to_pause
+        return None
+
+
+def _read_node_occupancy(
+    scheduled_replica_class_names: set[str],
+) -> list[_NodeOccupancy]:
+    occupancy_by_node_id = {
+        node["NodeID"]: _NodeOccupancy(free_resources=dict(node["Resources"]))
+        for node in ray.nodes()
+        if node["Alive"]
+    }
+    for actor in list_actors(
+        filters=[("state", "=", "ALIVE")],
+        detail=True,
+        limit=_STATE_API_RESULT_LIMIT,
+        raise_on_missing_output=False,
+    ):
+        actor_fields = vars(actor)
+        occupancy = occupancy_by_node_id.get(actor_fields["node_id"])
+        if occupancy is None:
+            continue
+        reserved_resources = actor_fields["required_resources"] or {}
+        _subtract_resources(occupancy.free_resources, reserved_resources)
+        replica_class_name = actor_fields["class_name"]
+        if replica_class_name in scheduled_replica_class_names:
+            _add_resources(
+                occupancy.resources_held_by_model.setdefault(replica_class_name, {}),
+                reserved_resources,
+            )
+    return list(occupancy_by_node_id.values())
+
+
+def _add_resources(target: _RayResources, amounts: _RayResources) -> None:
+    for name, amount in amounts.items():
+        target[name] = target.get(name, 0.0) + amount
+
+
+def _subtract_resources(target: _RayResources, amounts: _RayResources) -> None:
+    for name, amount in amounts.items():
+        target[name] = target.get(name, 0.0) - amount
+
+
+def _has_room_for(required: _RayResources, available: _RayResources) -> bool:
+    return all(
+        available.get(name, 0.0) + _RESOURCE_FLOAT_TOLERANCE >= amount
+        for name, amount in required.items()
+    )

{cortexgrid-0.3.8 → cortexgrid-0.3.10}/cortexgrid/_serve_entry.py

@@ -27,6 +27,7 @@ from typing import Any
 from ray import serve
 from ray.serve.deployment import Application
 
+from cortexgrid._model_scheduler import model_autoscaling_config
 from cortexgrid.serve import ingress_app
 
 
@@ -74,7 +75,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.
-        num_replicas=args.get("num_replicas", 1),
+        autoscaling_config=model_autoscaling_config(
+            args.get("num_replicas", 1), args.get("ray_actor_options", {})
+        ),
         # 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.8 → cortexgrid-0.3.10}/cortexgrid/checkpoint.py

@@ -1,7 +1,8 @@
 """Durable checkpointing for cortexgrid jobs.
 
-Save arbitrary state (primitives, torch tensors, state_dicts) via MLflow
-artifacts and resume from the latest checkpoint on retry.
+Save arbitrary state (primitives, torch tensors, state_dicts) to S3, with a
+manifest recorded by the jobs control plane, and resume from the latest
+checkpoint on retry.
 
 Usage (save)::
 
@@ -20,24 +21,22 @@ Usage (resume)::
 
 from __future__ import annotations
 
-import json
 import logging
 import tempfile
 from pathlib import Path
 from typing import Any
 
 import cloudpickle  # type: ignore
-from mlflow.tracking import MlflowClient
 
-from cortexgrid import s3_util
-from cortexgrid.experiment import Experiment, get_mlflow_tracking_uri
+from cortexgrid import s3_util, state
+from cortexgrid.experiment import Experiment
 
 log = logging.getLogger(__name__)
 _CORTEXGRID_JOB_ID: str | None = None
 
 
 class Checkpoint:
-    """Attribute-based checkpoint persisted via MLflow artifacts.
+    """Attribute-based checkpoint persisted to S3.
 
     Assign any cloudpickle-compatible value to an attribute and it will be
     saved when the context manager exits::
@@ -112,9 +111,8 @@ class Checkpoint:
             self._persist()
 
     def _persist(self) -> None:
-        """Upload attr blobs to MinIO; log manifest.json via MLflow."""
+        """Upload attr blobs to MinIO; record the manifest with the control plane."""
         exp = Experiment.get_instance()
-        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
         tmpdir = Path(tempfile.mkdtemp())
 
         manifest: dict[str, Any] = {"attrs": {}}
@@ -127,28 +125,17 @@ class Checkpoint:
             )
             manifest["attrs"][name] = {"uri": uri}
 
-        (tmpdir / "manifest.json").write_text(json.dumps(manifest))
-        client.log_artifact(
-            exp.run_id, str(tmpdir / "manifest.json"), artifact_path=self._prefix
-        )
+        state.put("runs", exp.run_id, "checkpoints", self._prefix, body=manifest)
         log.info("Checkpoint saved: %s (%d attrs)", self._prefix, len(self._data))
 
     @classmethod
     def _load(cls, prefix: str) -> Checkpoint | None:
-        """Download manifest via MLflow; download attr blobs from MinIO."""
+        """Read the manifest from the control plane; download attr blobs from MinIO."""
         exp = Experiment.get_instance()
-        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-
-        manifest_rel = f"{prefix}/manifest.json"
-        if not any(
-            a.path == manifest_rel for a in client.list_artifacts(exp.run_id, prefix)
-        ):
-            return None
-
-        try:
-            manifest_path = client.download_artifacts(exp.run_id, manifest_rel)
-            manifest = json.loads(Path(manifest_path).read_text())
-        except Exception:
+        # An unreachable control plane raises rather than reading as "no
+        # checkpoint", which would restart a retried job from scratch.
+        manifest = state.get("runs", exp.run_id, "checkpoints", prefix)
+        if manifest is None:
             return None
 
         data: dict[str, Any] = {}