cortexgrid 0.3.7__tar.gz → 0.3.9__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (23) hide show
  1. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/PKG-INFO +11 -10
  2. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/__init__.py +21 -31
  3. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/checkpoint.py +13 -26
  4. cortexgrid-0.3.9/cortexgrid/experiment.py +238 -0
  5. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/infra.py +10 -0
  6. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/jobs.py +55 -163
  7. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/mlflow_util.py +0 -14
  8. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/model_serving.py +155 -67
  9. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/model_storage.py +98 -142
  10. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/ray_util.py +8 -37
  11. cortexgrid-0.3.9/cortexgrid/state.py +80 -0
  12. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/docs/cortexgrid/README.md +10 -9
  13. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/pyproject.toml +6 -1
  14. cortexgrid-0.3.7/cortexgrid/experiment.py +0 -375
  15. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/.gitignore +0 -0
  16. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/LICENSE +0 -0
  17. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_bundle.py +0 -0
  18. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_ray_job_driver.py +0 -0
  19. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_serve_entry.py +0 -0
  20. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/py.typed +0 -0
  21. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/s3_util.py +0 -0
  22. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/secrets.py +0 -0
  23. {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/serve.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: cortexgrid
3
- Version: 0.3.7
3
+ Version: 0.3.9
4
4
  Summary: Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3
5
5
  Project-URL: Homepage, https://github.com/robodatalab/cortexgrid
6
6
  Project-URL: Repository, https://github.com/robodatalab/cortexgrid
@@ -55,7 +55,7 @@ import cortexgrid
55
55
  cortexgrid.init(experiment="weather-forecast")
56
56
  ```
57
57
 
58
- 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`.
58
+ 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`.
59
59
 
60
60
  **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.
61
61
 
@@ -80,7 +80,7 @@ No run-scoping context manager — `init()` starts the run, and every subsequent
80
80
 
81
81
  #### Checkpointing and resuming
82
82
 
83
- 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.
83
+ 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.
84
84
 
85
85
  ```python
86
86
  ckpt = cortexgrid.resume()
@@ -111,7 +111,7 @@ job = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
111
111
  print(f"Submitted: {job.job_id}")
112
112
  ```
113
113
 
114
- `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.
114
+ `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.
115
115
 
116
116
  ##### Blocking on the result
117
117
 
@@ -163,7 +163,7 @@ The driver cloudpickles the outcome to `job/{job_id}/result.pkl`, beside the pay
163
163
 
164
164
  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.
165
165
 
166
- 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.
166
+ 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.
167
167
 
168
168
  Each submission captures the code and dependencies the entry function needs automatically ([_bundle.py](https://github.com/robodatalab/cortexgrid/blob/main/cortexgrid/_bundle.py)):
169
169
  - `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
182
182
  cortexgrid.stop_experiment_run_jobs(run_id) # stops every job in the run
183
183
  ```
184
184
 
185
- `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.
185
+ `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.
186
186
 
187
187
  #### Object storage (S3/MinIO)
188
188
 
@@ -242,12 +242,12 @@ Anything else the serve-app has to know about the model - which model a provider
242
242
 
243
243
  | Function | Description |
244
244
  |----------|-------------|
245
- | `cortexgrid.init(experiment=None)` | Configure connections + start a new MLflow run inside the named experiment. One call per binary. |
245
+ | `cortexgrid.init(experiment=None)` | Configure connections + start a new run inside the named experiment. One call per binary. |
246
246
  | `cortexgrid.log_metric(key, value, step)` | Log a metric |
247
247
  | `cortexgrid.log_metrics(metrics, step)` | Log multiple metrics |
248
248
  | `cortexgrid.log_params(params)` | Log parameters |
249
249
  | `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
250
- | `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
250
+ | `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved on exit |
251
251
  | `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
252
252
  | `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a `JobFuture` |
253
253
  | `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
269
269
  | Service | Port | Purpose |
270
270
  |---------|------|---------|
271
271
  | Ray | 8265 | Dashboard + job submission (NodePort 30265) |
272
- | MLflow | 5000 | Experiment tracking, model registry |
272
+ | Jobs control plane | 8000 | cortexgrid's records (experiments, runs, jobs, model registry, deployments); schedules jobs on Ray (NodePort 30700) |
273
+ | MLflow | 5000 | Metrics and params of each run |
273
274
  | MinIO | 9000/9001 | S3-compatible artifact storage |
274
- | PostgreSQL | 5432 | MLflow metadata backend |
275
+ | PostgreSQL | 5432 | `cortexgrid` database (the control plane's records), MLflow metadata, UI notes |
275
276
  | Prometheus | 9090 | Metrics collection |
276
277
  | Grafana | 3000 | Dashboards (GPU, jobs, system) |
277
278
 
@@ -31,24 +31,18 @@ from __future__ import annotations
31
31
  from pathlib import Path
32
32
  from typing import Any, Callable
33
33
 
34
+ from cortexgrid import state
34
35
  from cortexgrid.checkpoint import checkpoint, resume
35
36
  from cortexgrid.experiment import (
36
37
  Experiment,
37
38
  delete_experiment,
38
39
  delete_run,
39
- experiment_has_active_runs,
40
- experiments_pending_deletion,
41
- finish_experiment_deletion,
42
- finish_run_deletion,
43
40
  list_experiments,
44
- runs_pending_deletion,
45
41
  )
46
42
  from cortexgrid.infra import get_ray_job_server_uri
47
43
  from cortexgrid.jobs import (
48
44
  schedule_remote_job,
49
45
  list_experiment_run_jobs,
50
- request_job_deletion,
51
- request_run_jobs_deletion,
52
46
  stop_experiment_run_jobs,
53
47
  wait_for_job_result,
54
48
  JobFailed,
@@ -74,17 +68,14 @@ from cortexgrid.mlflow_util import (
74
68
  list_run_metrics,
75
69
  get_metric_history,
76
70
  list_run_params,
77
- delete_run_artifacts,
78
71
  list_run_artifacts,
79
72
  )
80
73
  from cortexgrid.ray_util import (
81
74
  get_ray_status,
82
75
  get_ray_logs,
83
76
  get_ray_job_url,
84
- delete_ray_job,
85
77
  stop_ray_job,
86
78
  submit_ray_job,
87
- list_ray_job_statuses,
88
79
  list_ray_jobs_with_submission_id,
89
80
  get_ray_job_status,
90
81
  get_ray_job_id_for_cortexgrid_job,
@@ -198,16 +189,14 @@ def import_model(
198
189
  and record on the current Experiment's run which imported model it used.
199
190
 
200
191
  The model belongs to no run (see `cortexgrid.model_storage.import_model`),
201
- so the run keeps the link instead: the tag
202
- `imported_model/<family>/<suffix>` holds the version's `created_at`, set
203
- whether this call uploaded the model or reused it."""
192
+ so the run keeps the link instead: the run's record notes the model's
193
+ `created_at` under (family, suffix), whether this call uploaded the model
194
+ or reused it."""
204
195
  experiment = Experiment.get_instance()
205
196
  model = _import_model_storage(
206
197
  source, serve_app, family, suffix, requirements, config
207
198
  )
208
- get_mlflow_client().set_tag(
209
- experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
210
- )
199
+ _record_imported_model(experiment.run_id, family, suffix, model)
211
200
  return model
212
201
 
213
202
 
@@ -225,27 +214,33 @@ def register_model(
225
214
  `import_model` without the import: everything it needs beyond its code goes
226
215
  in `config`, which the serve-app reads with `model_config` at construction
227
216
  (see `cortexgrid.model_storage.register_model`). The model belongs to no
228
- run, so the run keeps the link the same way, under the same tag
229
- `imported_model/<family>/<suffix>`."""
217
+ run, so the run keeps the link the same way."""
230
218
  experiment = Experiment.get_instance()
231
219
  model = _register_model_storage(
232
220
  serve_app, family, suffix, requirements, config
233
221
  )
234
- get_mlflow_client().set_tag(
235
- experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
236
- )
222
+ _record_imported_model(experiment.run_id, family, suffix, model)
237
223
  return model
238
224
 
239
225
 
226
+ def _record_imported_model(
227
+ run_id: str, family: str, suffix: str, model: SavedModel
228
+ ) -> None:
229
+ """Note on the run's record which imported model it used."""
230
+ state.put(
231
+ "runs",
232
+ run_id,
233
+ "imported-models",
234
+ family,
235
+ suffix,
236
+ body={"created_at": model.created_at},
237
+ )
238
+
239
+
240
240
  __all__ = [
241
241
  "Experiment",
242
242
  "delete_experiment",
243
243
  "delete_run",
244
- "experiment_has_active_runs",
245
- "experiments_pending_deletion",
246
- "finish_experiment_deletion",
247
- "finish_run_deletion",
248
- "runs_pending_deletion",
249
244
  # Ray / jobs
250
245
  "remote",
251
246
  "get_job_result",
@@ -253,8 +248,6 @@ __all__ = [
253
248
  "get_ray_job_status",
254
249
  "get_ray_job_id_for_cortexgrid_job",
255
250
  "list_experiment_run_jobs",
256
- "request_job_deletion",
257
- "request_run_jobs_deletion",
258
251
  "stop_experiment_run_jobs",
259
252
  "JobStatus",
260
253
  "JobFailed",
@@ -267,11 +260,9 @@ __all__ = [
267
260
  "get_ray_status",
268
261
  "get_ray_logs",
269
262
  "get_ray_job_url",
270
- "delete_ray_job",
271
263
  "stop_ray_job",
272
264
  "submit_ray_job",
273
265
  "get_ray_job_server_uri",
274
- "list_ray_job_statuses",
275
266
  "list_ray_jobs_with_submission_id",
276
267
  "ray_submission_id",
277
268
  "get_ray_job_attempt",
@@ -284,7 +275,6 @@ __all__ = [
284
275
  "list_run_metrics",
285
276
  "get_metric_history",
286
277
  "list_run_params",
287
- "delete_run_artifacts",
288
278
  "list_run_artifacts",
289
279
  "list_experiments",
290
280
  # Checkpointing
@@ -1,7 +1,8 @@
1
1
  """Durable checkpointing for cortexgrid jobs.
2
2
 
3
- Save arbitrary state (primitives, torch tensors, state_dicts) via MLflow
4
- artifacts and resume from the latest checkpoint on retry.
3
+ Save arbitrary state (primitives, torch tensors, state_dicts) to S3, with a
4
+ manifest recorded by the jobs control plane, and resume from the latest
5
+ checkpoint on retry.
5
6
 
6
7
  Usage (save)::
7
8
 
@@ -20,24 +21,22 @@ Usage (resume)::
20
21
 
21
22
  from __future__ import annotations
22
23
 
23
- import json
24
24
  import logging
25
25
  import tempfile
26
26
  from pathlib import Path
27
27
  from typing import Any
28
28
 
29
29
  import cloudpickle # type: ignore
30
- from mlflow.tracking import MlflowClient
31
30
 
32
- from cortexgrid import s3_util
33
- from cortexgrid.experiment import Experiment, get_mlflow_tracking_uri
31
+ from cortexgrid import s3_util, state
32
+ from cortexgrid.experiment import Experiment
34
33
 
35
34
  log = logging.getLogger(__name__)
36
35
  _CORTEXGRID_JOB_ID: str | None = None
37
36
 
38
37
 
39
38
  class Checkpoint:
40
- """Attribute-based checkpoint persisted via MLflow artifacts.
39
+ """Attribute-based checkpoint persisted to S3.
41
40
 
42
41
  Assign any cloudpickle-compatible value to an attribute and it will be
43
42
  saved when the context manager exits::
@@ -112,9 +111,8 @@ class Checkpoint:
112
111
  self._persist()
113
112
 
114
113
  def _persist(self) -> None:
115
- """Upload attr blobs to MinIO; log manifest.json via MLflow."""
114
+ """Upload attr blobs to MinIO; record the manifest with the control plane."""
116
115
  exp = Experiment.get_instance()
117
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
118
116
  tmpdir = Path(tempfile.mkdtemp())
119
117
 
120
118
  manifest: dict[str, Any] = {"attrs": {}}
@@ -127,28 +125,17 @@ class Checkpoint:
127
125
  )
128
126
  manifest["attrs"][name] = {"uri": uri}
129
127
 
130
- (tmpdir / "manifest.json").write_text(json.dumps(manifest))
131
- client.log_artifact(
132
- exp.run_id, str(tmpdir / "manifest.json"), artifact_path=self._prefix
133
- )
128
+ state.put("runs", exp.run_id, "checkpoints", self._prefix, body=manifest)
134
129
  log.info("Checkpoint saved: %s (%d attrs)", self._prefix, len(self._data))
135
130
 
136
131
  @classmethod
137
132
  def _load(cls, prefix: str) -> Checkpoint | None:
138
- """Download manifest via MLflow; download attr blobs from MinIO."""
133
+ """Read the manifest from the control plane; download attr blobs from MinIO."""
139
134
  exp = Experiment.get_instance()
140
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
141
-
142
- manifest_rel = f"{prefix}/manifest.json"
143
- if not any(
144
- a.path == manifest_rel for a in client.list_artifacts(exp.run_id, prefix)
145
- ):
146
- return None
147
-
148
- try:
149
- manifest_path = client.download_artifacts(exp.run_id, manifest_rel)
150
- manifest = json.loads(Path(manifest_path).read_text())
151
- except Exception:
135
+ # An unreachable control plane raises rather than reading as "no
136
+ # checkpoint", which would restart a retried job from scratch.
137
+ manifest = state.get("runs", exp.run_id, "checkpoints", prefix)
138
+ if manifest is None:
152
139
  return None
153
140
 
154
141
  data: dict[str, Any] = {}
@@ -0,0 +1,238 @@
1
+ """cortexgrid's experiments and runs.
2
+
3
+ cortexgrid keeps its own record of each experiment and run with the jobs
4
+ control plane. Each maps onto an MLflow experiment and run, which exist only
5
+ so MLflow can display the run's metrics and params; the run id is MLflow's."""
6
+
7
+ from __future__ import annotations
8
+
9
+ import logging
10
+ import uuid
11
+ from dataclasses import dataclass
12
+
13
+ from cortexgrid import s3_util, state
14
+ from cortexgrid.infra import get_mlflow_tracking_uri
15
+ from cortexgrid.jobs import list_experiment_run_jobs, stop_experiment_run_jobs
16
+ from cortexgrid.ray_util import list_ray_jobs_with_submission_id, stop_ray_job
17
+ from cortexgrid.model_storage import delete_models_for_run
18
+ from haikunator import Haikunator # type: ignore
19
+ from mlflow.exceptions import MlflowException
20
+ from mlflow.tracking import MlflowClient
21
+
22
+
23
+ log = logging.getLogger(__name__)
24
+ _SINGLETON_EXPERIMENT: Experiment | None = None
25
+
26
+
27
+ @dataclass
28
+ class Experiment:
29
+ experiment_name: str
30
+ run_id: str
31
+
32
+ def run_name(self) -> str:
33
+ return state.get("runs", self.run_id)["run_name"]
34
+
35
+ @classmethod
36
+ def init(cls, name: str | None = None) -> "Experiment":
37
+ """Create a new run, in a new or existing experiment. Once per process."""
38
+ logging.basicConfig(
39
+ level=logging.INFO,
40
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(name)s: %(message)s",
41
+ datefmt="%H:%M:%S",
42
+ )
43
+
44
+ if _SINGLETON_EXPERIMENT is not None:
45
+ if name is not None and _SINGLETON_EXPERIMENT.experiment_name != name:
46
+ raise ValueError(
47
+ f"Active experiment has a different name {name} != "
48
+ f"{_SINGLETON_EXPERIMENT.experiment_name}"
49
+ )
50
+ return _SINGLETON_EXPERIMENT
51
+
52
+ experiment_name, run_id = _try_create_experiment_and_run(
53
+ experiment=name,
54
+ mlflow_tracking_uri=get_mlflow_tracking_uri(),
55
+ )
56
+ instance = cls(
57
+ experiment_name=experiment_name,
58
+ run_id=run_id,
59
+ )
60
+ set_instance(instance)
61
+ return instance
62
+
63
+ @classmethod
64
+ def from_experiment(cls, experiment_name: str, run_id: str) -> "Experiment":
65
+ """Bind to an existing experiment+run. Once per process."""
66
+ logging.basicConfig(
67
+ level=logging.INFO,
68
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(name)s: %(message)s",
69
+ datefmt="%H:%M:%S",
70
+ )
71
+
72
+ if _SINGLETON_EXPERIMENT is not None:
73
+ if (
74
+ _SINGLETON_EXPERIMENT.experiment_name != experiment_name
75
+ or _SINGLETON_EXPERIMENT.run_id != run_id
76
+ ):
77
+ raise ValueError(
78
+ f"Active experiment is different to the requested one: "
79
+ f"({experiment_name}, {run_id}) != "
80
+ f"({_SINGLETON_EXPERIMENT.experiment_name}, {_SINGLETON_EXPERIMENT.run_id})"
81
+ )
82
+ return _SINGLETON_EXPERIMENT
83
+
84
+ instance = cls(
85
+ experiment_name=experiment_name,
86
+ run_id=run_id,
87
+ )
88
+ set_instance(instance)
89
+ return instance
90
+
91
+ def get_jobs(self) -> list[str]:
92
+ """Return cortexgrid job IDs submitted against this experiment+run."""
93
+ return [job.job_id for job in list_experiment_run_jobs(self.run_id)]
94
+
95
+ @classmethod
96
+ def get_instance(cls) -> "Experiment":
97
+ if _SINGLETON_EXPERIMENT is None:
98
+ raise ValueError("Call Experiment.init or Experiment.from_experiment first")
99
+ return _SINGLETON_EXPERIMENT
100
+
101
+ @classmethod
102
+ def close(cls) -> None:
103
+ """Detach the active experiment so a different one can be init'd in this process."""
104
+ set_instance(None)
105
+
106
+
107
+ def set_instance(instance: Experiment | None) -> None:
108
+ global _SINGLETON_EXPERIMENT
109
+ _SINGLETON_EXPERIMENT = instance
110
+
111
+
112
+ def clear_instance() -> None:
113
+ # Use only in tests to clean between tests
114
+ global _SINGLETON_EXPERIMENT
115
+ _SINGLETON_EXPERIMENT = None
116
+
117
+
118
+ def _try_create_experiment_and_run(
119
+ experiment: str | None, mlflow_tracking_uri: str
120
+ ) -> tuple[str, str]:
121
+ name_gen = Haikunator()
122
+ if experiment is None:
123
+ experiment = name_gen.haikunate(token_length=2, token_chars="0123456789")
124
+
125
+ client = MlflowClient(tracking_uri=mlflow_tracking_uri)
126
+ record = state.get("experiments", experiment)
127
+ if record is None:
128
+ # A concurrent init of the same new experiment may win the insert; the
129
+ # record returned is the winner's, and its MLflow experiment is used.
130
+ record = state.put(
131
+ "experiments",
132
+ experiment,
133
+ body={"mlflow_experiment_id": _create_mlflow_experiment(client, experiment)},
134
+ )
135
+
136
+ run_name = name_gen.haikunate(token_length=2, token_chars="0123456789")
137
+ run = client.create_run(
138
+ experiment_id=record["mlflow_experiment_id"], run_name=run_name
139
+ )
140
+ state.put(
141
+ "runs",
142
+ run.info.run_id,
143
+ body={"run_name": run_name, "experiment_name": experiment},
144
+ )
145
+ return (experiment, run.info.run_id)
146
+
147
+
148
+ def _create_mlflow_experiment(client: MlflowClient, name: str) -> str:
149
+ """Create the MLflow experiment an experiment's runs log their metrics to.
150
+
151
+ MLflow reserves a name forever, deleted experiments included, so a name it
152
+ still holds from an earlier experiment gets a unique suffix. Nothing looks
153
+ an MLflow experiment up by name: cortexgrid's record keeps its id."""
154
+ try:
155
+ return client.create_experiment(name=name)
156
+ except MlflowException as exc:
157
+ if exc.error_code != "RESOURCE_ALREADY_EXISTS":
158
+ raise
159
+ return client.create_experiment(name=f"{name}__{uuid.uuid4().hex[:8]}")
160
+
161
+
162
+ def delete_run(run_id: str) -> None:
163
+ """Delete a run: cancel its Ray attempts, wipe its S3 job packages and
164
+ models, soft-delete its MLflow run, and drop cortexgrid's record of it
165
+ (its jobs, results and checkpoints go with it).
166
+
167
+ stop_experiment_run_jobs runs first so the control plane stops spawning
168
+ fresh Ray attempts for retry=True jobs before we tear the run down.
169
+
170
+ Idempotent: a run cortexgrid has no record of is treated as success."""
171
+ log.info("delete_run(%s): start", run_id)
172
+ if state.get("runs", run_id) is None:
173
+ log.info("delete_run(%s): early-exit, run not found", run_id)
174
+ return
175
+ stop_experiment_run_jobs(run_id)
176
+ log.info("delete_run(%s): stop_experiment_run_jobs done", run_id)
177
+ job_ids = [job.job_id for job in list_experiment_run_jobs(run_id)]
178
+ log.info("delete_run(%s): %d job(s) to clean", run_id, len(job_ids))
179
+ all_submissions = list_ray_jobs_with_submission_id()
180
+ for job_id in job_ids:
181
+ prefix = f"{run_id}-{job_id}-"
182
+ for sid in all_submissions:
183
+ if sid.startswith(prefix):
184
+ stop_ray_job(sid)
185
+ s3_util.delete_prefix(f"job/{job_id}/")
186
+ delete_models_for_run(run_id)
187
+ MlflowClient(tracking_uri=get_mlflow_tracking_uri()).delete_run(run_id)
188
+ state.delete("runs", run_id)
189
+ log.info("delete_run(%s): done", run_id)
190
+
191
+
192
+ def list_run_ids_in_experiment(name: str) -> list[str]:
193
+ """Return the run IDs of every run in the named experiment."""
194
+ return [
195
+ run["run_id"] for run in state.get("runs", params={"experiment_name": name})
196
+ ]
197
+
198
+
199
+ def delete_experiment(name: str) -> None:
200
+ """Delete every run in the experiment, then the experiment itself.
201
+
202
+ Its MLflow experiment is soft-deleted and keeps its name: MLflow never
203
+ frees one, so `Experiment.init(name)` gives a new experiment of the same
204
+ name an MLflow experiment under a suffixed name instead.
205
+
206
+ Idempotent: an experiment cortexgrid has no record of is treated as
207
+ success."""
208
+ log.info("delete_experiment(%r): start", name)
209
+ record = state.get("experiments", name)
210
+ if record is None:
211
+ log.info("delete_experiment(%r): early-exit, experiment not found", name)
212
+ return
213
+ run_ids = list_run_ids_in_experiment(name)
214
+ log.info("delete_experiment(%r): %d run(s) to delete", name, len(run_ids))
215
+ for run_id in run_ids:
216
+ delete_run(run_id)
217
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
218
+ client.delete_experiment(record["mlflow_experiment_id"])
219
+ state.delete("experiments", name)
220
+ log.info("delete_experiment(%r): done", name)
221
+
222
+
223
+ def list_experiments() -> list[Experiment]:
224
+ """Every run cortexgrid knows of, as (experiment_name, run_id) pairs."""
225
+ return [
226
+ Experiment(run["experiment_name"], run_id=run["run_id"])
227
+ for run in state.get("runs")
228
+ ]
229
+
230
+
231
+ def get_experiment_by_run_name(run_name: str) -> Experiment:
232
+ """Resolve a run by its haikunator name back to its (experiment_name, run_id) pair."""
233
+ runs = state.get("runs", params={"run_name": run_name})
234
+ if not runs:
235
+ raise ValueError(f"No run named {run_name!r}")
236
+ return Experiment(
237
+ experiment_name=runs[0]["experiment_name"], run_id=runs[0]["run_id"]
238
+ )
@@ -1,3 +1,5 @@
1
+ import functools
2
+
1
3
  from cortexgrid.secrets import get_secret
2
4
  from mlflow.tracking import MlflowClient
3
5
 
@@ -6,6 +8,14 @@ def get_mlflow_tracking_uri() -> str:
6
8
  return get_secret("MLFLOW_TRACKING_URI")
7
9
 
8
10
 
11
+ @functools.cache
12
+ def get_jobs_control_plane_uri() -> str:
13
+ """Base URL of the jobs control plane, which keeps cortexgrid's own records.
14
+ Cached: every job and registry lookup goes through it, and asking the
15
+ secrets server each time would add a round trip to each of them."""
16
+ return get_secret("JOBS_CONTROL_PLANE_URI")
17
+
18
+
9
19
  def get_ray_job_server_uri() -> str:
10
20
  return get_secret("RAY_JOB_SERVER_URI")
11
21