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.
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/PKG-INFO +11 -10
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/__init__.py +21 -31
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/checkpoint.py +13 -26
- cortexgrid-0.3.9/cortexgrid/experiment.py +238 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/infra.py +10 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/jobs.py +55 -163
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/mlflow_util.py +0 -14
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/model_serving.py +155 -67
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/model_storage.py +98 -142
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/ray_util.py +8 -37
- cortexgrid-0.3.9/cortexgrid/state.py +80 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/docs/cortexgrid/README.md +10 -9
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/pyproject.toml +6 -1
- cortexgrid-0.3.7/cortexgrid/experiment.py +0 -375
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/.gitignore +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/LICENSE +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_bundle.py +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_ray_job_driver.py +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/_serve_entry.py +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/py.typed +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/s3_util.py +0 -0
- {cortexgrid-0.3.7 → cortexgrid-0.3.9}/cortexgrid/secrets.py +0 -0
- {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.
|
|
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
|
|
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
|
|
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
|
|
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** —
|
|
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
|
|
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
|
|
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
|
|
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
|
-
|
|
|
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
|
|
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
|
|
202
|
-
`
|
|
203
|
-
|
|
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
|
-
|
|
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
|
|
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
|
-
|
|
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)
|
|
4
|
-
|
|
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
|
|
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
|
|
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;
|
|
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
|
-
(
|
|
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
|
-
"""
|
|
133
|
+
"""Read the manifest from the control plane; download attr blobs from MinIO."""
|
|
139
134
|
exp = Experiment.get_instance()
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
if
|
|
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
|
|