cortexgrid 0.3.8__tar.gz → 0.3.9__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.3.8
+Version: 0.3.9
 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.9}/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.8 → cortexgrid-0.3.9}/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] = {}

{cortexgrid-0.3.8 → cortexgrid-0.3.9}/cortexgrid/experiment.py

@@ -1,16 +1,22 @@
+"""cortexgrid's experiments and runs.
+
+cortexgrid keeps its own record of each experiment and run with the jobs
+control plane. Each maps onto an MLflow experiment and run, which exist only
+so MLflow can display the run's metrics and params; the run id is MLflow's."""
+
 from __future__ import annotations
 
 import logging
+import uuid
 from dataclasses import dataclass
-from pathlib import Path
 
-from cortexgrid import s3_util
+from cortexgrid import s3_util, state
 from cortexgrid.infra import get_mlflow_tracking_uri
-from cortexgrid.jobs import stop_experiment_run_jobs
+from cortexgrid.jobs import list_experiment_run_jobs, stop_experiment_run_jobs
 from cortexgrid.ray_util import list_ray_jobs_with_submission_id, stop_ray_job
 from cortexgrid.model_storage import delete_models_for_run
 from haikunator import Haikunator  # type: ignore
-from mlflow.entities import Experiment as MlflowExperiment
+from mlflow.exceptions import MlflowException
 from mlflow.tracking import MlflowClient
 
 
@@ -24,12 +30,11 @@ class Experiment:
     run_id: str
 
     def run_name(self) -> str:
-        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-        return client.get_run(self.run_id).info.run_name or self.run_id
+        return state.get("runs", self.run_id)["run_name"]
 
     @classmethod
     def init(cls, name: str | None = None) -> "Experiment":
-        """Create a new MLflow experiment+run. Once per process."""
+        """Create a new run, in a new or existing experiment. Once per process."""
         logging.basicConfig(
             level=logging.INFO,
             format="%(asctime)s.%(msecs)03d %(levelname)s %(name)s: %(message)s",
@@ -57,7 +62,7 @@ class Experiment:
 
     @classmethod
     def from_experiment(cls, experiment_name: str, run_id: str) -> "Experiment":
-        """Bind to an existing MLflow experiment+run. Once per process."""
+        """Bind to an existing experiment+run. Once per process."""
         logging.basicConfig(
             level=logging.INFO,
             format="%(asctime)s.%(msecs)03d %(levelname)s %(name)s: %(message)s",
@@ -85,12 +90,7 @@ class Experiment:
 
     def get_jobs(self) -> list[str]:
         """Return cortexgrid job IDs submitted against this experiment+run."""
-        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-        return [
-            Path(f.path).name
-            for f in client.list_artifacts(self.run_id, path="job")
-            if f.is_dir
-        ]
+        return [job.job_id for job in list_experiment_run_jobs(self.run_id)]
 
     @classmethod
     def get_instance(cls) -> "Experiment":
@@ -123,64 +123,59 @@ def _try_create_experiment_and_run(
         experiment = name_gen.haikunate(token_length=2, token_chars="0123456789")
 
     client = MlflowClient(tracking_uri=mlflow_tracking_uri)
-    # get_experiment_by_name also returns deleted experiments. A deleted one
-    # cannot take new runs, so a fresh experiment is created under its name.
-    experiment_obj = client.get_experiment_by_name(name=experiment)
-    if experiment_obj is not None and experiment_obj.lifecycle_stage != "active":
-        _release_deleted_experiment_name(client, experiment_obj)
-        experiment_obj = None
-    if experiment_obj:
-        experiment_id = experiment_obj.experiment_id
-    else:
-        experiment_id = client.create_experiment(name=experiment)
+    record = state.get("experiments", experiment)
+    if record is None:
+        # A concurrent init of the same new experiment may win the insert; the
+        # record returned is the winner's, and its MLflow experiment is used.
+        record = state.put(
+            "experiments",
+            experiment,
+            body={"mlflow_experiment_id": _create_mlflow_experiment(client, experiment)},
+        )
 
     run_name = name_gen.haikunate(token_length=2, token_chars="0123456789")
-    run = client.create_run(experiment_id=experiment_id, run_name=run_name)
-
+    run = client.create_run(
+        experiment_id=record["mlflow_experiment_id"], run_name=run_name
+    )
+    state.put(
+        "runs",
+        run.info.run_id,
+        body={"run_name": run_name, "experiment_name": experiment},
+    )
     return (experiment, run.info.run_id)
 
 
-def _deleted_experiment_name(name: str, experiment_id: str) -> str:
-    """The name a deleted experiment is moved to, freeing `name` for reuse.
-    Unique because experiment ids are."""
-    return f"{name}__deleted__{experiment_id}"
-
+def _create_mlflow_experiment(client: MlflowClient, name: str) -> str:
+    """Create the MLflow experiment an experiment's runs log their metrics to.
 
-def _release_deleted_experiment_name(
-    client: MlflowClient, experiment: MlflowExperiment
-) -> None:
-    """Move a deleted experiment off its name, so a new experiment can take it.
-
-    MLflow keeps a deleted experiment's name reserved (experiment names are
-    unique across every lifecycle stage) and refuses to rename a deleted
-    experiment, so it is restored only for the rename and deleted again."""
-    log.info(
-        "Experiment %r (id %s) is deleted; renaming it to release the name",
-        experiment.name,
-        experiment.experiment_id,
-    )
-    client.restore_experiment(experiment.experiment_id)
-    client.rename_experiment(
-        experiment.experiment_id,
-        _deleted_experiment_name(experiment.name, experiment.experiment_id),
-    )
-    client.delete_experiment(experiment.experiment_id)
+    MLflow reserves a name forever, deleted experiments included, so a name it
+    still holds from an earlier experiment gets a unique suffix. Nothing looks
+    an MLflow experiment up by name: cortexgrid's record keeps its id."""
+    try:
+        return client.create_experiment(name=name)
+    except MlflowException as exc:
+        if exc.error_code != "RESOURCE_ALREADY_EXISTS":
+            raise
+        return client.create_experiment(name=f"{name}__{uuid.uuid4().hex[:8]}")
 
 
 def delete_run(run_id: str) -> None:
-    """Soft-delete a run in MLflow, cancel its Ray attempts, and wipe its
-    S3 job packages so it cannot be relaunched or re-read.
+    """Delete a run: cancel its Ray attempts, wipe its S3 job packages and
+    models, soft-delete its MLflow run, and drop cortexgrid's record of it
+    (its jobs, results and checkpoints go with it).
 
     stop_experiment_run_jobs runs first so the control plane stops spawning
-    fresh Ray attempts for retry=True jobs before we tear the run down."""
+    fresh Ray attempts for retry=True jobs before we tear the run down.
+
+    Idempotent: a run cortexgrid has no record of is treated as success."""
     log.info("delete_run(%s): start", run_id)
+    if state.get("runs", run_id) is None:
+        log.info("delete_run(%s): early-exit, run not found", run_id)
+        return
     stop_experiment_run_jobs(run_id)
     log.info("delete_run(%s): stop_experiment_run_jobs done", run_id)
-    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-    job_ids = [
-        Path(f.path).name for f in client.list_artifacts(run_id, path="job") if f.is_dir
-    ]
-    log.info("delete_run(%s): %d job artifact(s) to clean", run_id, len(job_ids))
+    job_ids = [job.job_id for job in list_experiment_run_jobs(run_id)]
+    log.info("delete_run(%s): %d job(s) to clean", run_id, len(job_ids))
     all_submissions = list_ray_jobs_with_submission_id()
     for job_id in job_ids:
         prefix = f"{run_id}-{job_id}-"
@@ -189,75 +184,55 @@ def delete_run(run_id: str) -> None:
                 stop_ray_job(sid)
         s3_util.delete_prefix(f"job/{job_id}/")
     delete_models_for_run(run_id)
-    client.delete_run(run_id)
+    MlflowClient(tracking_uri=get_mlflow_tracking_uri()).delete_run(run_id)
+    state.delete("runs", run_id)
     log.info("delete_run(%s): done", run_id)
 
 
 def list_run_ids_in_experiment(name: str) -> list[str]:
-    """Return the run IDs of every active run in the named experiment."""
-    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-    exp = client.get_experiment_by_name(name)
-    if exp is None:
-        return []
+    """Return the run IDs of every run in the named experiment."""
     return [
-        r.info.run_id for r in client.search_runs(experiment_ids=[exp.experiment_id])
+        run["run_id"] for run in state.get("runs", params={"experiment_name": name})
     ]
 
 
 def delete_experiment(name: str) -> None:
-    """Soft-delete every run in the experiment, then the experiment itself.
+    """Delete every run in the experiment, then the experiment itself.
 
-    The experiment is renamed (`<name>__deleted__<id>`) before it is deleted:
-    MLflow keeps a deleted experiment's name reserved, and the rename frees it
-    so `Experiment.init(name)` can create a new experiment under it.
+    Its MLflow experiment is soft-deleted and keeps its name: MLflow never
+    frees one, so `Experiment.init(name)` gives a new experiment of the same
+    name an MLflow experiment under a suffixed name instead.
 
-    Idempotent: already-deleted experiments are treated as success."""
+    Idempotent: an experiment cortexgrid has no record of is treated as
+    success."""
     log.info("delete_experiment(%r): start", name)
-    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-    exp = client.get_experiment_by_name(name)
-    if exp is None:
+    record = state.get("experiments", name)
+    if record is None:
         log.info("delete_experiment(%r): early-exit, experiment not found", name)
         return
-    if exp.lifecycle_stage != "active":
-        log.info(
-            "delete_experiment(%r): early-exit, lifecycle=%s", name, exp.lifecycle_stage
-        )
-        return
-    runs = list(client.search_runs(experiment_ids=[exp.experiment_id]))
-    log.info("delete_experiment(%r): %d active run(s) to delete", name, len(runs))
-    for run in runs:
-        delete_run(run.info.run_id)
-    client.rename_experiment(
-        exp.experiment_id, _deleted_experiment_name(name, exp.experiment_id)
-    )
-    client.delete_experiment(exp.experiment_id)
+    run_ids = list_run_ids_in_experiment(name)
+    log.info("delete_experiment(%r): %d run(s) to delete", name, len(run_ids))
+    for run_id in run_ids:
+        delete_run(run_id)
+    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+    client.delete_experiment(record["mlflow_experiment_id"])
+    state.delete("experiments", name)
     log.info("delete_experiment(%r): done", name)
 
 
 def list_experiments() -> list[Experiment]:
-    """Map MLflow experiment names to their run IDs."""
-    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-    result: list[Experiment] = []
-    for exp in client.search_experiments():
-        runs = client.search_runs(experiment_ids=[exp.experiment_id])
-        for run in runs:
-            result.append(Experiment(exp.name, run_id=run.info.run_id))
-    return result
+    """Every run cortexgrid knows of, as (experiment_name, run_id) pairs."""
+    return [
+        Experiment(run["experiment_name"], run_id=run["run_id"])
+        for run in state.get("runs")
+    ]
 
 
 def get_experiment_by_run_name(run_name: str) -> Experiment:
     """Resolve a run by its haikunator name back to its (experiment_name, run_id) pair."""
-    client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
-    experiment_ids = [e.experiment_id for e in client.search_experiments()]
-    if not experiment_ids:
-        raise ValueError(f"No run named {run_name!r}")
-    runs = client.search_runs(
-        experiment_ids=experiment_ids,
-        filter_string=f"attributes.run_name = '{run_name}'",
-        max_results=1,
-    )
+    runs = state.get("runs", params={"run_name": run_name})
     if not runs:
         raise ValueError(f"No run named {run_name!r}")
-    run = runs[0]
-    exp = client.get_experiment(run.info.experiment_id)
-    return Experiment(experiment_name=exp.name, run_id=run.info.run_id)
+    return Experiment(
+        experiment_name=runs[0]["experiment_name"], run_id=runs[0]["run_id"]
+    )

{cortexgrid-0.3.8 → cortexgrid-0.3.9}/cortexgrid/infra.py

@@ -1,3 +1,5 @@
+import functools
+
 from cortexgrid.secrets import get_secret
 from mlflow.tracking import MlflowClient
 
@@ -6,6 +8,14 @@ def get_mlflow_tracking_uri() -> str:
     return get_secret("MLFLOW_TRACKING_URI")
 
 
+@functools.cache
+def get_jobs_control_plane_uri() -> str:
+    """Base URL of the jobs control plane, which keeps cortexgrid's own records.
+    Cached: every job and registry lookup goes through it, and asking the
+    secrets server each time would add a round trip to each of them."""
+    return get_secret("JOBS_CONTROL_PLANE_URI")
+
+
 def get_ray_job_server_uri() -> str:
     return get_secret("RAY_JOB_SERVER_URI")