cortexgrid 0.3.6__tar.gz → 0.3.7__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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: cortexgrid
3
- Version: 0.3.6
3
+ Version: 0.3.7
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
@@ -36,12 +36,19 @@ from cortexgrid.experiment import (
36
36
  Experiment,
37
37
  delete_experiment,
38
38
  delete_run,
39
+ experiment_has_active_runs,
40
+ experiments_pending_deletion,
41
+ finish_experiment_deletion,
42
+ finish_run_deletion,
39
43
  list_experiments,
44
+ runs_pending_deletion,
40
45
  )
41
46
  from cortexgrid.infra import get_ray_job_server_uri
42
47
  from cortexgrid.jobs import (
43
48
  schedule_remote_job,
44
49
  list_experiment_run_jobs,
50
+ request_job_deletion,
51
+ request_run_jobs_deletion,
45
52
  stop_experiment_run_jobs,
46
53
  wait_for_job_result,
47
54
  JobFailed,
@@ -67,14 +74,17 @@ from cortexgrid.mlflow_util import (
67
74
  list_run_metrics,
68
75
  get_metric_history,
69
76
  list_run_params,
77
+ delete_run_artifacts,
70
78
  list_run_artifacts,
71
79
  )
72
80
  from cortexgrid.ray_util import (
73
81
  get_ray_status,
74
82
  get_ray_logs,
75
83
  get_ray_job_url,
84
+ delete_ray_job,
76
85
  stop_ray_job,
77
86
  submit_ray_job,
87
+ list_ray_job_statuses,
78
88
  list_ray_jobs_with_submission_id,
79
89
  get_ray_job_status,
80
90
  get_ray_job_id_for_cortexgrid_job,
@@ -231,6 +241,11 @@ __all__ = [
231
241
  "Experiment",
232
242
  "delete_experiment",
233
243
  "delete_run",
244
+ "experiment_has_active_runs",
245
+ "experiments_pending_deletion",
246
+ "finish_experiment_deletion",
247
+ "finish_run_deletion",
248
+ "runs_pending_deletion",
234
249
  # Ray / jobs
235
250
  "remote",
236
251
  "get_job_result",
@@ -238,6 +253,8 @@ __all__ = [
238
253
  "get_ray_job_status",
239
254
  "get_ray_job_id_for_cortexgrid_job",
240
255
  "list_experiment_run_jobs",
256
+ "request_job_deletion",
257
+ "request_run_jobs_deletion",
241
258
  "stop_experiment_run_jobs",
242
259
  "JobStatus",
243
260
  "JobFailed",
@@ -250,9 +267,11 @@ __all__ = [
250
267
  "get_ray_status",
251
268
  "get_ray_logs",
252
269
  "get_ray_job_url",
270
+ "delete_ray_job",
253
271
  "stop_ray_job",
254
272
  "submit_ray_job",
255
273
  "get_ray_job_server_uri",
274
+ "list_ray_job_statuses",
256
275
  "list_ray_jobs_with_submission_id",
257
276
  "ray_submission_id",
258
277
  "get_ray_job_attempt",
@@ -265,6 +284,7 @@ __all__ = [
265
284
  "list_run_metrics",
266
285
  "get_metric_history",
267
286
  "list_run_params",
287
+ "delete_run_artifacts",
268
288
  "list_run_artifacts",
269
289
  "list_experiments",
270
290
  # Checkpointing
@@ -4,19 +4,23 @@ import logging
4
4
  from dataclasses import dataclass
5
5
  from pathlib import Path
6
6
 
7
- from cortexgrid import s3_util
8
7
  from cortexgrid.infra import get_mlflow_tracking_uri
9
- from cortexgrid.jobs import stop_experiment_run_jobs
10
- from cortexgrid.ray_util import list_ray_jobs_with_submission_id, stop_ray_job
8
+ from cortexgrid.jobs import request_run_jobs_deletion
11
9
  from cortexgrid.model_storage import delete_models_for_run
12
10
  from haikunator import Haikunator # type: ignore
13
11
  from mlflow.entities import Experiment as MlflowExperiment
12
+ from mlflow.entities import Run as MlflowRun
14
13
  from mlflow.tracking import MlflowClient
15
14
 
16
15
 
17
16
  log = logging.getLogger(__name__)
18
17
  _SINGLETON_EXPERIMENT: Experiment | None = None
19
18
 
19
+ # Written on a run or an experiment that has been asked to go away. It is
20
+ # what keeps such a record out of every by-name lookup while the control
21
+ # plane, the one thing that reads it on purpose, tears the record down.
22
+ DELETE_REQUESTED_TAG = "cortexgrid.delete_requested"
23
+
20
24
 
21
25
  @dataclass
22
26
  class Experiment:
@@ -123,9 +127,19 @@ def _try_create_experiment_and_run(
123
127
  experiment = name_gen.haikunate(token_length=2, token_chars="0123456789")
124
128
 
125
129
  client = MlflowClient(tracking_uri=mlflow_tracking_uri)
130
+ # An experiment that asked to be deleted is invisible to the lookup
131
+ # below, so a new one would quietly take its name while the control
132
+ # plane is still tearing the old one down — and any run created here
133
+ # would be torn down with it. Refuse instead.
134
+ pending = get_experiment_by_name(experiment, include_deleting=True)
135
+ if pending is not None and _pending_deletion(pending.tags):
136
+ raise RuntimeError(
137
+ f"Experiment {experiment!r} is being deleted; wait for the "
138
+ "control plane to finish before creating it again"
139
+ )
126
140
  # get_experiment_by_name also returns deleted experiments. A deleted one
127
141
  # cannot take new runs, so a fresh experiment is created under its name.
128
- experiment_obj = client.get_experiment_by_name(name=experiment)
142
+ experiment_obj = get_experiment_by_name(experiment)
129
143
  if experiment_obj is not None and experiment_obj.lifecycle_stage != "active":
130
144
  _release_deleted_experiment_name(client, experiment_obj)
131
145
  experiment_obj = None
@@ -168,53 +182,48 @@ def _release_deleted_experiment_name(
168
182
 
169
183
 
170
184
  def delete_run(run_id: str) -> None:
171
- """Soft-delete a run in MLflow, cancel its Ray attempts, and wipe its
172
- S3 job packages so it cannot be relaunched or re-read.
173
-
174
- stop_experiment_run_jobs runs first so the control plane stops spawning
175
- fresh Ray attempts for retry=True jobs before we tear the run down."""
176
- log.info("delete_run(%s): start", run_id)
177
- stop_experiment_run_jobs(run_id)
178
- log.info("delete_run(%s): stop_experiment_run_jobs done", run_id)
179
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
180
- job_ids = [
181
- Path(f.path).name for f in client.list_artifacts(run_id, path="job") if f.is_dir
182
- ]
183
- log.info("delete_run(%s): %d job artifact(s) to clean", run_id, len(job_ids))
184
- all_submissions = list_ray_jobs_with_submission_id()
185
- for job_id in job_ids:
186
- prefix = f"{run_id}-{job_id}-"
187
- for sid in all_submissions:
188
- if sid.startswith(prefix):
189
- stop_ray_job(sid)
190
- s3_util.delete_prefix(f"job/{job_id}/")
185
+ """Ask for a run, and everything under it, to be deleted.
186
+
187
+ Intent only: every job in the run is latched and the run itself is
188
+ tagged. The control plane is the only thing that acts on either — it
189
+ stops the Ray attempts, wipes the job packages and artifacts, and
190
+ removes the run once nothing is left under it.
191
+
192
+ The tag takes the run out of every by-name lookup here, so a caller
193
+ that deletes a run and immediately looks one up never gets the one on
194
+ its way out. Models are not jobs and nothing else touches them, so
195
+ they are deleted outright.
196
+ """
197
+ log.info("delete_run(%s): requesting deletion", run_id)
198
+ request_run_jobs_deletion(run_id)
191
199
  delete_models_for_run(run_id)
192
- client.delete_run(run_id)
193
- log.info("delete_run(%s): done", run_id)
200
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
201
+ client.set_tag(run_id, DELETE_REQUESTED_TAG, "true")
194
202
 
195
203
 
196
204
  def list_run_ids_in_experiment(name: str) -> list[str]:
197
205
  """Return the run IDs of every active run in the named experiment."""
198
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
199
- exp = client.get_experiment_by_name(name)
206
+ exp = get_experiment_by_name(name)
200
207
  if exp is None:
201
208
  return []
202
- return [
203
- r.info.run_id for r in client.search_runs(experiment_ids=[exp.experiment_id])
204
- ]
209
+ return [r.info.run_id for r in search_runs([exp.experiment_id])]
205
210
 
206
211
 
207
212
  def delete_experiment(name: str) -> None:
208
- """Soft-delete every run in the experiment, then the experiment itself.
209
-
210
- The experiment is renamed (`<name>__deleted__<id>`) before it is deleted:
211
- MLflow keeps a deleted experiment's name reserved, and the rename frees it
212
- so `Experiment.init(name)` can create a new experiment under it.
213
-
214
- Idempotent: already-deleted experiments are treated as success."""
215
- log.info("delete_experiment(%r): start", name)
213
+ """Ask for an experiment, its runs and their jobs to be deleted.
214
+
215
+ The experiment is renamed off `name` (`<name>__deleted__<id>`) as the
216
+ last thing this does, so by the time the call returns the name is free
217
+ and `Experiment.init(name)` creates a fresh experiment instead of
218
+ attaching to the one being torn down. Everything else is intent: the
219
+ control plane deletes each run once its jobs are gone, and the
220
+ experiment once its runs are.
221
+
222
+ Idempotent: an experiment that is absent, already deleted or already
223
+ on its way out is treated as success."""
224
+ log.info("delete_experiment(%r): requesting deletion", name)
216
225
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
217
- exp = client.get_experiment_by_name(name)
226
+ exp = get_experiment_by_name(name)
218
227
  if exp is None:
219
228
  log.info("delete_experiment(%r): early-exit, experiment not found", name)
220
229
  return
@@ -223,41 +232,144 @@ def delete_experiment(name: str) -> None:
223
232
  "delete_experiment(%r): early-exit, lifecycle=%s", name, exp.lifecycle_stage
224
233
  )
225
234
  return
226
- runs = list(client.search_runs(experiment_ids=[exp.experiment_id]))
227
- log.info("delete_experiment(%r): %d active run(s) to delete", name, len(runs))
235
+ runs = search_runs([exp.experiment_id])
236
+ log.info("delete_experiment(%r): %d run(s) to request", name, len(runs))
228
237
  for run in runs:
229
238
  delete_run(run.info.run_id)
239
+ client.set_experiment_tag(exp.experiment_id, DELETE_REQUESTED_TAG, "true")
230
240
  client.rename_experiment(
231
241
  exp.experiment_id, _deleted_experiment_name(name, exp.experiment_id)
232
242
  )
233
- client.delete_experiment(exp.experiment_id)
234
- log.info("delete_experiment(%r): done", name)
243
+ log.info("delete_experiment(%r): requested", name)
235
244
 
236
245
 
237
- def list_experiments() -> list[Experiment]:
238
- """Map MLflow experiment names to their run IDs."""
246
+ def _pending_deletion(tags: dict[str, str] | None) -> bool:
247
+ return (tags or {}).get(DELETE_REQUESTED_TAG) == "true"
248
+
249
+
250
+ def search_experiments(include_deleting: bool = False) -> list[MlflowExperiment]:
251
+ """Every experiment, minus the ones on their way out.
252
+
253
+ This is the gate: reading experiments straight off MlflowClient
254
+ bypasses the deletion rule and brings records back from the dead, so
255
+ every reader — this module, the UI streams, the control plane — comes
256
+ through here. Only the control plane passes `include_deleting=True`,
257
+ because removing them is its job.
258
+ """
239
259
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
260
+ return [
261
+ e
262
+ for e in client.search_experiments()
263
+ if include_deleting or not _pending_deletion(e.tags)
264
+ ]
265
+
266
+
267
+ def get_experiment_by_name(
268
+ name: str, include_deleting: bool = False
269
+ ) -> MlflowExperiment | None:
270
+ """The experiment under `name`, unless it is on its way out.
271
+
272
+ The gate for by-name resolution. Like MlflowClient's own call it still
273
+ returns an experiment MLflow has deleted — `Experiment.init` has to see
274
+ one to release its name — but a record that asked to be deleted is,
275
+ to every caller but the control plane, already gone.
276
+ """
277
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
278
+ exp = client.get_experiment_by_name(name)
279
+ if exp is None:
280
+ return None
281
+ if not include_deleting and _pending_deletion(exp.tags):
282
+ return None
283
+ return exp
284
+
285
+
286
+ def search_runs(
287
+ experiment_ids: list[str],
288
+ include_deleting: bool = False,
289
+ filter_string: str = "",
290
+ ) -> list[MlflowRun]:
291
+ """Every run of those experiments, minus the ones on their way out.
292
+
293
+ The same gate, one level down, shaped like the client call it stands
294
+ in front of so callers keep MLflow's server-side filtering. See
295
+ `search_experiments`.
296
+ """
297
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
298
+ return [
299
+ r
300
+ for r in client.search_runs(
301
+ experiment_ids=experiment_ids, filter_string=filter_string
302
+ )
303
+ if include_deleting or not _pending_deletion(r.data.tags)
304
+ ]
305
+
306
+
307
+ def list_experiments(include_deleting: bool = False) -> list[Experiment]:
308
+ """Map MLflow experiment names to their run IDs.
309
+
310
+ Records that asked to be deleted are left out: they are on their way
311
+ to disappearing and must not be reachable any more. The control plane
312
+ passes `include_deleting=True`, because making them disappear is its
313
+ job and it cannot do it without seeing them."""
240
314
  result: list[Experiment] = []
241
- for exp in client.search_experiments():
242
- runs = client.search_runs(experiment_ids=[exp.experiment_id])
243
- for run in runs:
315
+ for exp in search_experiments(include_deleting):
316
+ for run in search_runs([exp.experiment_id], include_deleting):
244
317
  result.append(Experiment(exp.name, run_id=run.info.run_id))
245
318
  return result
246
319
 
247
320
 
321
+ def runs_pending_deletion() -> list[str]:
322
+ """Run ids tagged for deletion. The control plane is the only caller."""
323
+ return [
324
+ run.info.run_id
325
+ for exp in search_experiments(include_deleting=True)
326
+ for run in search_runs([exp.experiment_id], include_deleting=True)
327
+ if _pending_deletion(run.data.tags)
328
+ ]
329
+
330
+
331
+ def experiments_pending_deletion() -> list[str]:
332
+ """Experiment ids tagged for deletion. The control plane is the only caller."""
333
+ return [
334
+ exp.experiment_id
335
+ for exp in search_experiments(include_deleting=True)
336
+ if _pending_deletion(exp.tags)
337
+ ]
338
+
339
+
340
+ def experiment_has_active_runs(experiment_id: str) -> bool:
341
+ """Whether anything is still under the experiment. Control plane only.
342
+
343
+ Runs on their way out count: the experiment cannot go until they
344
+ have actually gone."""
345
+ return bool(search_runs([experiment_id], include_deleting=True))
346
+
347
+
348
+ def finish_run_deletion(run_id: str) -> None:
349
+ """Remove a run whose jobs are gone. The control plane is the only caller."""
350
+ log.info("finish_run_deletion(%s)", run_id)
351
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
352
+ client.delete_run(run_id)
353
+
354
+
355
+ def finish_experiment_deletion(experiment_id: str) -> None:
356
+ """Remove an experiment whose runs are gone. Control plane only."""
357
+ log.info("finish_experiment_deletion(%s)", experiment_id)
358
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
359
+ client.delete_experiment(experiment_id)
360
+
361
+
248
362
  def get_experiment_by_run_name(run_name: str) -> Experiment:
249
363
  """Resolve a run by its haikunator name back to its (experiment_name, run_id) pair."""
250
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
251
- experiment_ids = [e.experiment_id for e in client.search_experiments()]
364
+ experiment_ids = [e.experiment_id for e in search_experiments()]
252
365
  if not experiment_ids:
253
366
  raise ValueError(f"No run named {run_name!r}")
254
- runs = client.search_runs(
255
- experiment_ids=experiment_ids,
256
- filter_string=f"attributes.run_name = '{run_name}'",
257
- max_results=1,
367
+ runs = search_runs(
368
+ experiment_ids, filter_string=f"attributes.run_name = '{run_name}'"
258
369
  )
259
370
  if not runs:
260
371
  raise ValueError(f"No run named {run_name!r}")
261
372
  run = runs[0]
373
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
262
374
  exp = client.get_experiment(run.info.experiment_id)
263
375
  return Experiment(experiment_name=exp.name, run_id=run.info.run_id)
@@ -3,7 +3,7 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  import cloudpickle # type: ignore
6
- from dataclasses import asdict, dataclass, field
6
+ from dataclasses import asdict, dataclass, field, fields
7
7
  import inspect
8
8
  import io
9
9
  import json
@@ -67,6 +67,16 @@ class LifecycleEvent:
67
67
  error: str | None = None
68
68
 
69
69
 
70
+ def _known_fields_only(cls: type, data: dict[str, Any]) -> dict[str, Any]:
71
+ """`data` minus the keys `cls` has no field for.
72
+
73
+ Records outlive the code that wrote them and the code that reads them:
74
+ a control plane one release ahead writes fields an older client has
75
+ never seen, and that client has to keep reading its own jobs."""
76
+ known = {f.name for f in fields(cls)}
77
+ return {k: v for k, v in data.items() if k in known}
78
+
79
+
70
80
  @dataclass
71
81
  class JobLifecycle:
72
82
  """Static identity and latches for a job.
@@ -81,6 +91,7 @@ class JobLifecycle:
81
91
  run_id: str
82
92
  job_id: str
83
93
  stop_requested: bool = False # latch: False -> True, never cleared
94
+ delete_requested: bool = False # latch: False -> True, never cleared
84
95
  retry: bool = False # static flag set at job creation
85
96
  num_gpus: int = 0
86
97
  num_cpus: int = 1
@@ -90,13 +101,30 @@ class JobLifecycle:
90
101
  history: list[LifecycleEvent] = field(default_factory=list)
91
102
 
92
103
  def to_json(self) -> str:
93
- return json.dumps(asdict(self))
104
+ """The record as it goes to MLflow.
105
+
106
+ A latch that is still False is left out. Every reader of this file
107
+ pins the fields it knows — a released client parses it with an
108
+ exact constructor call — so a latch a reader has never heard of
109
+ must not appear until something actually set it.
110
+ """
111
+ data = asdict(self)
112
+ if not self.delete_requested:
113
+ data.pop("delete_requested")
114
+ return json.dumps(data)
94
115
 
95
116
  @classmethod
96
117
  def from_json(cls, text: str) -> "JobLifecycle":
97
- data = json.loads(text)
98
- data.pop("error", None)
99
- data["history"] = [LifecycleEvent(**e) for e in data.get("history", [])]
118
+ """Parse a record, including one a newer writer produced.
119
+
120
+ Fields this version does not know are dropped rather than raising,
121
+ so a client keeps working while the cluster runs ahead of it.
122
+ """
123
+ data = _known_fields_only(cls, json.loads(text))
124
+ data["history"] = [
125
+ LifecycleEvent(**_known_fields_only(LifecycleEvent, e))
126
+ for e in data.get("history", [])
127
+ ]
100
128
  return cls(**data)
101
129
 
102
130
  def get_ray_job_id(
@@ -147,17 +175,23 @@ class JobLifecycle:
147
175
 
148
176
  @classmethod
149
177
  def load_from_mlflow(cls, run_id: str, job_id: str) -> "JobLifecycle":
178
+ """Read a job's record.
179
+
180
+ One round trip: the fetch is the existence check. Listing the
181
+ directory first doubled the cost of every read, and a sweep over a
182
+ cluster's jobs pays it per job against a tracking server that
183
+ answers these one at a time.
184
+ """
150
185
  client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
151
186
  lifecycle_rel = f"job/{job_id}/lifecycle.json"
152
- if not any(
153
- a.path == lifecycle_rel
154
- for a in client.list_artifacts(run_id, f"job/{job_id}")
155
- ):
187
+ try:
188
+ local_path = client.download_artifacts(run_id, lifecycle_rel)
189
+ text = Path(local_path).read_text()
190
+ except Exception as err:
156
191
  raise FileNotFoundError(
157
192
  f"artifact {lifecycle_rel} not found in run {run_id}"
158
- )
159
- local_path = client.download_artifacts(run_id, lifecycle_rel)
160
- return cls.from_json(Path(local_path).read_text())
193
+ ) from err
194
+ return cls.from_json(text)
161
195
 
162
196
 
163
197
  class Payload(BaseModel):
@@ -496,22 +530,37 @@ def schedule_remote_job(
496
530
  return JobFuture(experiment_name=experiment_name, run_id=run_id, job_id=job_id)
497
531
 
498
532
 
533
+ def list_experiment_run_job_ids(run_id: str) -> list[str]:
534
+ """The ids of every job of this run — one listing, no records read.
535
+
536
+ Reading a job's record costs a download each; a caller that wants many
537
+ of them can take this list and fetch them as it pleases."""
538
+ client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
539
+ return [
540
+ Path(entry.path).name
541
+ for entry in client.list_artifacts(run_id, path="job")
542
+ if entry.is_dir
543
+ ]
544
+
545
+
546
+ def load_job(run_id: str, job_id: str) -> JobLifecycle | None:
547
+ """The job's record, or None when it has not landed yet.
548
+
549
+ A job whose lifecycle is still being written is a normal state, not an
550
+ error: the caller lists it again on the next poll."""
551
+ try:
552
+ return JobLifecycle.load_from_mlflow(run_id, job_id)
553
+ except Exception:
554
+ logging.getLogger(__name__).warning(
555
+ "Skipping job %s: missing lifecycle", job_id
556
+ )
557
+ return None
558
+
559
+
499
560
  def list_experiment_run_jobs(run_id: str) -> list[JobLifecycle]:
500
561
  """Return all jobs and their lifecycle states for this experiment+run."""
501
- client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
502
- entries = client.list_artifacts(run_id, path="job")
503
- result: list[JobLifecycle] = []
504
- for entry in entries:
505
- if not entry.is_dir:
506
- continue
507
- job_id = Path(entry.path).name
508
- try:
509
- result.append(JobLifecycle.load_from_mlflow(run_id, job_id))
510
- except Exception:
511
- logging.getLogger(__name__).warning(
512
- "Skipping job %s: missing lifecycle", job_id
513
- )
514
- return result
562
+ jobs = (load_job(run_id, job_id) for job_id in list_experiment_run_job_ids(run_id))
563
+ return [job for job in jobs if job is not None]
515
564
 
516
565
 
517
566
  def stop_experiment_run_jobs(run_id: str) -> None:
@@ -530,3 +579,30 @@ def stop_experiment_run_jobs(run_id: str) -> None:
530
579
  continue
531
580
  job.stop_requested = True
532
581
  job.save_to_mlflow()
582
+
583
+
584
+ def request_job_deletion(run_id: str, job_id: str) -> None:
585
+ """Ask for one job to be deleted by flipping the delete_requested latch.
586
+
587
+ This function never touches Ray, S3 or the job's artifacts: it
588
+ records the intent on the job's own record and returns. The control
589
+ plane observes the latch on its next poll and is the only thing that
590
+ tears the job down, so a job created and deleted between two polls
591
+ is simply never submitted.
592
+
593
+ Idempotent: an already-requested job is left alone.
594
+ """
595
+ lifecycle = JobLifecycle.load_from_mlflow(run_id, job_id)
596
+ if lifecycle.delete_requested:
597
+ return
598
+ lifecycle.delete_requested = True
599
+ lifecycle.save_to_mlflow()
600
+
601
+
602
+ def request_run_jobs_deletion(run_id: str) -> None:
603
+ """Ask for every job in the run to be deleted. See `request_job_deletion`."""
604
+ for job in list_experiment_run_jobs(run_id):
605
+ if job.delete_requested:
606
+ continue
607
+ job.delete_requested = True
608
+ job.save_to_mlflow()
@@ -14,6 +14,7 @@ import requests
14
14
  from cortexgrid.experiment import Experiment
15
15
  from cortexgrid.infra import get_mlflow_tracking_uri
16
16
  from mlflow.entities import Metric
17
+ from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
17
18
  from mlflow.tracking import MlflowClient
18
19
 
19
20
 
@@ -108,3 +109,16 @@ def list_run_artifacts(run_id: str, path: str = "") -> list[str]:
108
109
  """Return artifact paths for a run."""
109
110
  client = get_mlflow_client()
110
111
  return [a.path for a in client.list_artifacts(run_id, path=path)]
112
+
113
+
114
+ def delete_run_artifacts(run_id: str, path: str) -> None:
115
+ """Delete everything a run stores under `path`, directory and all.
116
+
117
+ MlflowClient has no delete: it goes through the run's artifact
118
+ repository, which the tracking server proxies to object storage."""
119
+ client = get_mlflow_client()
120
+ repo = get_artifact_repository(
121
+ client.get_run(run_id).info.artifact_uri,
122
+ tracking_uri=get_mlflow_tracking_uri(),
123
+ )
124
+ repo.delete_artifacts(path)
@@ -58,23 +58,43 @@ def get_ray_status(ray_job_id: str | None) -> str | None:
58
58
  return client.get_job_status(ray_job_id).value
59
59
 
60
60
 
61
+ def to_job_status(ray_status: str | None) -> JobStatus:
62
+ """Map one of Ray's words onto ours. `None` means never submitted."""
63
+ if ray_status is None:
64
+ return JobStatus.PENDING
65
+ if ray_status == "SUCCEEDED":
66
+ return JobStatus.FINISHED
67
+ if ray_status == "FAILED":
68
+ return JobStatus.FAILED
69
+ if ray_status == "STOPPED":
70
+ return JobStatus.STOPPED
71
+ return JobStatus.RUNNING
72
+
73
+
61
74
  def get_ray_job_status(ray_job_id: str | None) -> JobStatus:
62
75
  """Derive a job's observable status from a live Ray query.
63
76
 
64
77
  Returns ``PENDING`` both when ``ray_job_id is None`` (never submitted)
65
78
  and when Ray itself reports ``PENDING`` (queued). Callers that need
66
79
  to distinguish those two must check ``ray_job_id is None`` first.
80
+
81
+ One query per job. A caller with many jobs wants
82
+ ``list_ray_job_statuses`` instead, which answers for all of them in
83
+ a single call.
67
84
  """
68
- ray_status = get_ray_status(ray_job_id)
69
85
  if ray_job_id is None:
70
86
  return JobStatus.PENDING
71
- if ray_status == "SUCCEEDED":
72
- return JobStatus.FINISHED
73
- if ray_status == "FAILED":
74
- return JobStatus.FAILED
75
- if ray_status == "STOPPED":
76
- return JobStatus.STOPPED
77
- return JobStatus.RUNNING
87
+ return to_job_status(get_ray_status(ray_job_id))
88
+
89
+
90
+ def list_ray_job_statuses() -> dict[str, JobStatus]:
91
+ """Every ray job's status, by submission id, in one call."""
92
+ client = get_ray_job_submission_client()
93
+ return {
94
+ job.submission_id: to_job_status(job.status.value)
95
+ for job in client.list_jobs()
96
+ if job.submission_id is not None
97
+ }
78
98
 
79
99
 
80
100
  def get_ray_logs(ray_job_id: str | None) -> str | None:
@@ -101,6 +121,15 @@ def stop_ray_job(ray_job_id: str) -> None:
101
121
  client.stop_job(ray_job_id)
102
122
 
103
123
 
124
+ def delete_ray_job(ray_job_id: str) -> None:
125
+ """Drop a ray job and its data from Ray's job store.
126
+
127
+ Ray rejects the call for a job that has not settled, so callers stop
128
+ the job and let it reach a terminal state first."""
129
+ client = get_ray_job_submission_client()
130
+ client.delete_job(ray_job_id)
131
+
132
+
104
133
  def list_ray_jobs_with_submission_id() -> list[str]:
105
134
  """List all ray jobs, the ones that received submission id."""
106
135
  client = get_ray_job_submission_client()
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "cortexgrid"
3
- version = "0.3.6"
3
+ version = "0.3.7"
4
4
  description = "Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3"
5
5
  readme = "docs/cortexgrid/README.md"
6
6
  license = "Apache-2.0"
File without changes
File without changes