cortexgrid 0.2.97__tar.gz → 0.3.0__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.2.97 → cortexgrid-0.3.0}/PKG-INFO +57 -5
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/__init__.py +37 -8
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/_ray_job_driver.py +19 -2
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/jobs.py +234 -4
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/docs/cortexgrid/README.md +56 -4
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/pyproject.toml +1 -1
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/.gitignore +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/LICENSE +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/_bundle.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/_serve_entry.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/checkpoint.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/experiment.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/infra.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/mlflow_util.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/model_serving.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/model_storage.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/py.typed +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/ray_util.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/s3_util.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/secrets.py +0 -0
- {cortexgrid-0.2.97 → cortexgrid-0.3.0}/cortexgrid/serve.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: cortexgrid
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
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
|
|
@@ -107,11 +107,61 @@ def train_step(batch):
|
|
|
107
107
|
# MLflow and S3 env vars are injected automatically
|
|
108
108
|
return loss
|
|
109
109
|
|
|
110
|
-
|
|
111
|
-
print(f"Submitted: {job_id}")
|
|
110
|
+
job = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
|
|
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
|
|
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.
|
|
115
|
+
|
|
116
|
+
##### Blocking on the result
|
|
117
|
+
|
|
118
|
+
`JobFuture.result(timeout=None)` blocks until Ray reports the job terminal, then behaves like a local call: it returns whatever the function returned, or raises whatever the function raised.
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
def score(batch) -> dict[str, float]:
|
|
122
|
+
# runs on the DGX
|
|
123
|
+
return {"loss": 0.12}
|
|
124
|
+
|
|
125
|
+
# Fire-and-forget — the form training uses. Nothing blocks, and a retry=True
|
|
126
|
+
# job cannot be waited on at all (see below).
|
|
127
|
+
training = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
|
|
128
|
+
print(training.job_id, training.status().value) # "blue-42 running"
|
|
129
|
+
|
|
130
|
+
# Blocking — the value comes back as if score() had run locally.
|
|
131
|
+
job = cortexgrid.remote(score, batch, num_gpus=1)
|
|
132
|
+
metrics = job.result() # {"loss": 0.12}
|
|
133
|
+
|
|
134
|
+
# ...or with a deadline. On TimeoutError the job keeps running; wait again later.
|
|
135
|
+
metrics = job.result(timeout=600)
|
|
136
|
+
|
|
137
|
+
# From another process — the handle is gone, the job id is enough.
|
|
138
|
+
metrics = cortexgrid.get_job_result(job_id, timeout=600)
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
Failures arrive as exceptions, not as a status to inspect:
|
|
142
|
+
|
|
143
|
+
```python
|
|
144
|
+
try:
|
|
145
|
+
metrics = job.result()
|
|
146
|
+
except BadBatch as exc: # exactly what score() raised on the DGX
|
|
147
|
+
print(exc.__cause__) # JobFailed, carrying the remote traceback
|
|
148
|
+
except cortexgrid.JobFailed: # the job never got as far as returning
|
|
149
|
+
... # submission failed, driver died, or stopped
|
|
150
|
+
except cortexgrid.JobResultUnavailable:
|
|
151
|
+
... # it returned, but the value is not transportable
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
The driver cloudpickles the outcome to `job/{job_id}/result.pkl`, beside the payload manifest and the lifecycle, and the waiting side reads it back:
|
|
155
|
+
|
|
156
|
+
| Outcome | What `result()` does |
|
|
157
|
+
|---|---|
|
|
158
|
+
| The function returned | Returns its value |
|
|
159
|
+
| The function raised | Raises that exception, with the remote traceback attached as a chained `JobFailed` cause |
|
|
160
|
+
| The job never got as far as returning (submission failed, driver died, job stopped) | Raises `JobFailed` |
|
|
161
|
+
| The value or the exception did not survive cloudpickle | Raises `JobResultUnavailable` (value) or `JobFailed` (exception). A value that cannot be pickled never fails the job itself |
|
|
162
|
+
| `timeout` elapsed | Raises `TimeoutError`; the job keeps running |
|
|
163
|
+
|
|
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.
|
|
115
165
|
|
|
116
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.
|
|
117
167
|
|
|
@@ -197,7 +247,9 @@ The requirements are part of the model, not of the serve-app class: GPUs, RAM an
|
|
|
197
247
|
| `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
|
|
198
248
|
| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
|
|
199
249
|
| `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
|
|
200
|
-
| `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a
|
|
250
|
+
| `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a `JobFuture` |
|
|
251
|
+
| `JobFuture.status()` / `.done()` / `.result(timeout=None)` | Live status of a submitted job, and its function's return value (blocking) |
|
|
252
|
+
| `cortexgrid.get_job_result(job_id, timeout=None)` | Block on a job of the current run by id; returns its value or raises its exception |
|
|
201
253
|
| `cortexgrid.list_experiment_run_jobs(run_id)` | List `JobLifecycle` records for every cortexgrid job in a run |
|
|
202
254
|
| `cortexgrid.stop_experiment_run_jobs(run_id)` | Request every job in a run to stop (flips the `stop_requested` latch) |
|
|
203
255
|
| `cortexgrid.get_ray_job_status(ray_job_id)` | Live Ray status for a submission id |
|
|
@@ -2,14 +2,18 @@
|
|
|
2
2
|
|
|
3
3
|
import cortexgrid
|
|
4
4
|
|
|
5
|
-
# Fire-and-forget training on the DGX; returns a
|
|
5
|
+
# Fire-and-forget training on the DGX; returns a handle immediately.
|
|
6
6
|
# The jobs control plane picks up the submission and dispatches it to Ray.
|
|
7
|
-
|
|
8
|
-
print(f"Submitted: {job_id}")
|
|
7
|
+
training = cortexgrid.remote(my_train, config, num_gpus=1, retry=True)
|
|
8
|
+
print(f"Submitted: {training.job_id} ({training.status().value})")
|
|
9
9
|
|
|
10
|
-
#
|
|
11
|
-
|
|
12
|
-
|
|
10
|
+
# Or block on the function's return value, as if it had run locally.
|
|
11
|
+
# The job's own exception is what a failed job raises here.
|
|
12
|
+
job = cortexgrid.remote(score, batch, num_gpus=1)
|
|
13
|
+
loss = job.result(timeout=600)
|
|
14
|
+
|
|
15
|
+
# In another process, the job id is enough
|
|
16
|
+
loss = cortexgrid.get_job_result(job_id)
|
|
13
17
|
|
|
14
18
|
# Inside the training function — checkpoint after each epoch
|
|
15
19
|
with cortexgrid.checkpoint() as ckpt:
|
|
@@ -39,7 +43,12 @@ from cortexgrid.jobs import (
|
|
|
39
43
|
schedule_remote_job,
|
|
40
44
|
list_experiment_run_jobs,
|
|
41
45
|
stop_experiment_run_jobs,
|
|
46
|
+
wait_for_job_result,
|
|
47
|
+
JobFailed,
|
|
48
|
+
JobFuture,
|
|
42
49
|
JobLifecycle,
|
|
50
|
+
JobResult,
|
|
51
|
+
JobResultUnavailable,
|
|
43
52
|
LifecycleEvent,
|
|
44
53
|
Payload,
|
|
45
54
|
)
|
|
@@ -104,8 +113,11 @@ def remote(
|
|
|
104
113
|
num_cpus: int = 1,
|
|
105
114
|
retry: bool = False,
|
|
106
115
|
**kwargs: Any,
|
|
107
|
-
) ->
|
|
108
|
-
"""Submit a function to the control plane. Returns a job
|
|
116
|
+
) -> JobFuture:
|
|
117
|
+
"""Submit a function to the control plane. Returns a handle on the job.
|
|
118
|
+
|
|
119
|
+
The call does not block: the returned `JobFuture` carries the job id and
|
|
120
|
+
offers `status()`, `done()` and `result()`. Only `result()` waits."""
|
|
109
121
|
experiment = Experiment.get_instance()
|
|
110
122
|
return schedule_remote_job(
|
|
111
123
|
experiment.experiment_name,
|
|
@@ -119,6 +131,17 @@ def remote(
|
|
|
119
131
|
)
|
|
120
132
|
|
|
121
133
|
|
|
134
|
+
def get_job_result(job_id: str, timeout: float | None = None) -> Any:
|
|
135
|
+
"""Block until the job finishes, then return what its function returned.
|
|
136
|
+
|
|
137
|
+
For a job of the current Experiment's run, known only by its id — the
|
|
138
|
+
handle `remote` returned may live in another process. Raises whatever the
|
|
139
|
+
job's function raised, so the call reads like a local one; see
|
|
140
|
+
`wait_for_job_result` for the other errors and for `timeout`."""
|
|
141
|
+
experiment = Experiment.get_instance()
|
|
142
|
+
return wait_for_job_result(experiment.run_id, job_id, timeout)
|
|
143
|
+
|
|
144
|
+
|
|
122
145
|
def save_model(
|
|
123
146
|
weights_dir: str | Path,
|
|
124
147
|
serve_app: type,
|
|
@@ -173,11 +196,17 @@ __all__ = [
|
|
|
173
196
|
"delete_run",
|
|
174
197
|
# Ray / jobs
|
|
175
198
|
"remote",
|
|
199
|
+
"get_job_result",
|
|
200
|
+
"wait_for_job_result",
|
|
176
201
|
"get_ray_job_status",
|
|
177
202
|
"list_experiment_run_jobs",
|
|
178
203
|
"stop_experiment_run_jobs",
|
|
179
204
|
"JobStatus",
|
|
205
|
+
"JobFailed",
|
|
206
|
+
"JobFuture",
|
|
180
207
|
"JobLifecycle",
|
|
208
|
+
"JobResult",
|
|
209
|
+
"JobResultUnavailable",
|
|
181
210
|
"LifecycleEvent",
|
|
182
211
|
"Payload",
|
|
183
212
|
"get_ray_status",
|
|
@@ -9,11 +9,23 @@ from pathlib import Path
|
|
|
9
9
|
|
|
10
10
|
from cortexgrid.checkpoint import set_cortexgrid_job_id
|
|
11
11
|
from cortexgrid.experiment import Experiment
|
|
12
|
-
from cortexgrid.jobs import Payload
|
|
12
|
+
from cortexgrid.jobs import JobResult, Payload
|
|
13
13
|
|
|
14
14
|
log = logging.getLogger("ray-job-driver")
|
|
15
15
|
|
|
16
16
|
|
|
17
|
+
def record(result: JobResult) -> None:
|
|
18
|
+
"""Persist the job's outcome beside its payload.
|
|
19
|
+
|
|
20
|
+
Best effort: a job that ran must not be reported as failed because its
|
|
21
|
+
result could not be uploaded. A caller waiting on the result gets
|
|
22
|
+
JobResultUnavailable instead, and the reason is in this log."""
|
|
23
|
+
try:
|
|
24
|
+
result.save_to_mlflow()
|
|
25
|
+
except Exception:
|
|
26
|
+
log.exception("Failed to record result for job %s", result.job_id)
|
|
27
|
+
|
|
28
|
+
|
|
17
29
|
def main(payload_path: str) -> None:
|
|
18
30
|
logging.basicConfig(
|
|
19
31
|
level=logging.INFO,
|
|
@@ -41,7 +53,12 @@ def main(payload_path: str) -> None:
|
|
|
41
53
|
payload.args,
|
|
42
54
|
payload.kwargs,
|
|
43
55
|
)
|
|
44
|
-
|
|
56
|
+
try:
|
|
57
|
+
value = payload.fn(*payload.args, **payload.kwargs)
|
|
58
|
+
except BaseException as exc:
|
|
59
|
+
record(JobResult.from_exception(payload, exc))
|
|
60
|
+
raise
|
|
61
|
+
record(JobResult.from_value(payload, value))
|
|
45
62
|
|
|
46
63
|
|
|
47
64
|
if __name__ == "__main__":
|
|
@@ -12,18 +12,39 @@ from pathlib import Path
|
|
|
12
12
|
import sys
|
|
13
13
|
import tarfile
|
|
14
14
|
import tempfile
|
|
15
|
+
import time
|
|
16
|
+
import traceback
|
|
15
17
|
from typing import Any, Callable
|
|
16
18
|
|
|
17
19
|
from cortexgrid import s3_util
|
|
18
20
|
from cortexgrid._bundle import bundle, stage, worker_provides
|
|
19
21
|
from cortexgrid.infra import get_mlflow_tracking_uri
|
|
20
|
-
from cortexgrid.ray_util import
|
|
22
|
+
from cortexgrid.ray_util import (
|
|
23
|
+
JobStatus,
|
|
24
|
+
get_ray_job_id_for_cortexgrid_job,
|
|
25
|
+
get_ray_job_status,
|
|
26
|
+
)
|
|
21
27
|
from haikunator import Haikunator # type: ignore
|
|
22
28
|
from mlflow.tracking import MlflowClient
|
|
23
29
|
from pydantic import BaseModel, ConfigDict
|
|
24
30
|
|
|
25
31
|
log = logging.getLogger(__name__)
|
|
26
32
|
|
|
33
|
+
_JOB_POLL_INTERVAL_S = 5.0
|
|
34
|
+
_TERMINAL_JOB_STATES = (JobStatus.FINISHED, JobStatus.FAILED, JobStatus.STOPPED)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class JobFailed(RuntimeError):
|
|
38
|
+
"""A job did not produce a return value: its function raised, or the job
|
|
39
|
+
ended in a terminal state before the function ever returned. Carries the
|
|
40
|
+
remote traceback when there is one."""
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class JobResultUnavailable(RuntimeError):
|
|
44
|
+
"""The job's function returned, but its value cannot be handed back here:
|
|
45
|
+
it did not survive cloudpickle on the cluster, cannot be unpickled by this
|
|
46
|
+
process, or was never recorded."""
|
|
47
|
+
|
|
27
48
|
|
|
28
49
|
@dataclass
|
|
29
50
|
class LifecycleEvent:
|
|
@@ -218,6 +239,215 @@ class Payload(BaseModel):
|
|
|
218
239
|
return payload
|
|
219
240
|
|
|
220
241
|
|
|
242
|
+
def _try_dumps(obj: Any) -> tuple[bytes | None, str | None]:
|
|
243
|
+
"""cloudpickle ``obj``, or report why it could not be pickled.
|
|
244
|
+
|
|
245
|
+
Returns ``(blob, None)`` on success and ``(None, reason)`` on failure.
|
|
246
|
+
Nothing a job produces is allowed to turn into a job failure, so every
|
|
247
|
+
pickling error is captured rather than raised."""
|
|
248
|
+
try:
|
|
249
|
+
return cloudpickle.dumps(obj), None
|
|
250
|
+
except Exception as exc:
|
|
251
|
+
return None, f"{type(obj).__name__} did not survive cloudpickle: {exc!r}"
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
@dataclass
|
|
255
|
+
class JobResult:
|
|
256
|
+
"""What a job's function returned or raised, recorded by the driver.
|
|
257
|
+
|
|
258
|
+
Written to ``job/{job_id}/result.pkl``, beside the payload manifest and
|
|
259
|
+
the lifecycle. The value and the exception are pickled separately from
|
|
260
|
+
the envelope so an object that cannot be pickled costs only its own
|
|
261
|
+
field: a job whose return value does not survive cloudpickle still
|
|
262
|
+
finishes, and the reason is still readable here.
|
|
263
|
+
"""
|
|
264
|
+
|
|
265
|
+
experiment_name: str
|
|
266
|
+
run_id: str
|
|
267
|
+
job_id: str
|
|
268
|
+
ok: bool
|
|
269
|
+
value_pickle: bytes | None = None # cloudpickled return value
|
|
270
|
+
value_error: str | None = None # why value_pickle is None
|
|
271
|
+
exception_pickle: bytes | None = None # cloudpickled exception, when picklable
|
|
272
|
+
traceback: str | None = None # formatted remote traceback
|
|
273
|
+
|
|
274
|
+
@classmethod
|
|
275
|
+
def from_value(cls, payload: "Payload", value: Any) -> "JobResult":
|
|
276
|
+
blob, error = _try_dumps(value)
|
|
277
|
+
return cls(
|
|
278
|
+
experiment_name=payload.experiment_name,
|
|
279
|
+
run_id=payload.run_id,
|
|
280
|
+
job_id=payload.job_id,
|
|
281
|
+
ok=True,
|
|
282
|
+
value_pickle=blob,
|
|
283
|
+
value_error=error,
|
|
284
|
+
)
|
|
285
|
+
|
|
286
|
+
@classmethod
|
|
287
|
+
def from_exception(cls, payload: "Payload", exc: BaseException) -> "JobResult":
|
|
288
|
+
blob, _ = _try_dumps(exc)
|
|
289
|
+
return cls(
|
|
290
|
+
experiment_name=payload.experiment_name,
|
|
291
|
+
run_id=payload.run_id,
|
|
292
|
+
job_id=payload.job_id,
|
|
293
|
+
ok=False,
|
|
294
|
+
exception_pickle=blob,
|
|
295
|
+
traceback="".join(traceback.format_exception(exc)),
|
|
296
|
+
)
|
|
297
|
+
|
|
298
|
+
def save_to_mlflow(self) -> None:
|
|
299
|
+
artifact_path = f"job/{self.job_id}"
|
|
300
|
+
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
301
|
+
log.info("Recording result for job %s (ok=%s)", self.job_id, self.ok)
|
|
302
|
+
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
303
|
+
local_path = Path(tmp_dir, "result.pkl")
|
|
304
|
+
local_path.write_bytes(cloudpickle.dumps(self))
|
|
305
|
+
client.log_artifact(
|
|
306
|
+
self.run_id, str(local_path), artifact_path=artifact_path
|
|
307
|
+
)
|
|
308
|
+
|
|
309
|
+
@classmethod
|
|
310
|
+
def load_from_mlflow(cls, run_id: str, job_id: str) -> "JobResult":
|
|
311
|
+
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
312
|
+
result_rel = f"job/{job_id}/result.pkl"
|
|
313
|
+
if not any(
|
|
314
|
+
a.path == result_rel for a in client.list_artifacts(run_id, f"job/{job_id}")
|
|
315
|
+
):
|
|
316
|
+
raise FileNotFoundError(f"artifact {result_rel} not found in run {run_id}")
|
|
317
|
+
local_path = client.download_artifacts(run_id, result_rel)
|
|
318
|
+
result = cloudpickle.loads(Path(local_path).read_bytes())
|
|
319
|
+
if not isinstance(result, cls):
|
|
320
|
+
raise TypeError(f"{result_rel} in run {run_id} is not a JobResult")
|
|
321
|
+
return result
|
|
322
|
+
|
|
323
|
+
def unwrap(self) -> Any:
|
|
324
|
+
"""Return the value the job's function returned, or raise what it raised.
|
|
325
|
+
|
|
326
|
+
A failed job raises the original exception, rebuilt from its pickle, so
|
|
327
|
+
a blocking caller sees what a local call would have raised; the remote
|
|
328
|
+
traceback rides along as the chained ``JobFailed`` cause. When the
|
|
329
|
+
exception itself did not survive pickling, ``JobFailed`` is raised
|
|
330
|
+
instead.
|
|
331
|
+
"""
|
|
332
|
+
if not self.ok:
|
|
333
|
+
failure = JobFailed(f"job {self.job_id} failed:\n{self.traceback}")
|
|
334
|
+
remote_exc = self._unpickled_exception()
|
|
335
|
+
if remote_exc is None:
|
|
336
|
+
raise failure
|
|
337
|
+
raise remote_exc from failure
|
|
338
|
+
if self.value_pickle is None:
|
|
339
|
+
raise JobResultUnavailable(
|
|
340
|
+
f"job {self.job_id} finished, but its return value was not "
|
|
341
|
+
f"recorded: {self.value_error}"
|
|
342
|
+
)
|
|
343
|
+
try:
|
|
344
|
+
return cloudpickle.loads(self.value_pickle)
|
|
345
|
+
except Exception as exc:
|
|
346
|
+
raise JobResultUnavailable(
|
|
347
|
+
f"job {self.job_id} finished, but its return value could not be "
|
|
348
|
+
f"unpickled here: {exc!r}"
|
|
349
|
+
) from exc
|
|
350
|
+
|
|
351
|
+
def _unpickled_exception(self) -> BaseException | None:
|
|
352
|
+
"""The remote exception object, or None when it cannot be rebuilt here."""
|
|
353
|
+
if self.exception_pickle is None:
|
|
354
|
+
return None
|
|
355
|
+
try:
|
|
356
|
+
exc = cloudpickle.loads(self.exception_pickle)
|
|
357
|
+
except Exception:
|
|
358
|
+
log.warning("Job %s: remote exception could not be unpickled", self.job_id)
|
|
359
|
+
return None
|
|
360
|
+
return exc if isinstance(exc, BaseException) else None
|
|
361
|
+
|
|
362
|
+
|
|
363
|
+
def _missing_result(
|
|
364
|
+
lifecycle: JobLifecycle, ray_job_id: str | None, status: JobStatus
|
|
365
|
+
) -> Exception:
|
|
366
|
+
"""The error for a job that reached ``status`` leaving no result behind.
|
|
367
|
+
|
|
368
|
+
Two ways in: the job's function never ran (the driver died, or the control
|
|
369
|
+
plane could not hand the job to Ray), or it ran and the upload of its
|
|
370
|
+
result failed. Either way the detail lives in the Ray logs, and a
|
|
371
|
+
submission error lives on the lifecycle."""
|
|
372
|
+
error = next((e.error for e in reversed(lifecycle.history) if e.error), None)
|
|
373
|
+
detail = f": {error}" if error else f" (ray job {ray_job_id}; check its logs)"
|
|
374
|
+
if status == JobStatus.FINISHED:
|
|
375
|
+
return JobResultUnavailable(
|
|
376
|
+
f"job {lifecycle.job_id} finished but recorded no result{detail}"
|
|
377
|
+
)
|
|
378
|
+
return JobFailed(
|
|
379
|
+
f"job {lifecycle.job_id} ended in {status.value} before its function "
|
|
380
|
+
f"returned{detail}"
|
|
381
|
+
)
|
|
382
|
+
|
|
383
|
+
|
|
384
|
+
def wait_for_job_result(run_id: str, job_id: str, timeout: float | None = None) -> Any:
|
|
385
|
+
"""Block until the job reaches a terminal Ray state, then return what its
|
|
386
|
+
function returned — or raise what it raised, as a local call would.
|
|
387
|
+
|
|
388
|
+
Raises ``ValueError`` for a job submitted with ``retry=True``: retries are
|
|
389
|
+
unbounded by design, so a blocking wait would sit through however many
|
|
390
|
+
resubmissions the control plane makes. Raises ``TimeoutError`` once a
|
|
391
|
+
finite ``timeout`` elapses (``timeout=0`` polls once and gives up);
|
|
392
|
+
``timeout=None`` has no deadline. Raises ``JobFailed`` if the job never
|
|
393
|
+
got as far as returning, and ``JobResultUnavailable`` if it returned a
|
|
394
|
+
value that cannot be handed back here.
|
|
395
|
+
"""
|
|
396
|
+
deadline = None if timeout is None else time.monotonic() + timeout
|
|
397
|
+
while True:
|
|
398
|
+
lifecycle = JobLifecycle.load_from_mlflow(run_id, job_id)
|
|
399
|
+
if lifecycle.retry:
|
|
400
|
+
raise ValueError(
|
|
401
|
+
f"job {job_id} was submitted with retry=True, which resubmits "
|
|
402
|
+
"without bound; waiting on its result is not supported. Poll "
|
|
403
|
+
"list_experiment_run_jobs instead."
|
|
404
|
+
)
|
|
405
|
+
ray_job_id = lifecycle.get_ray_job_id()
|
|
406
|
+
status = get_ray_job_status(ray_job_id)
|
|
407
|
+
if status in _TERMINAL_JOB_STATES:
|
|
408
|
+
try:
|
|
409
|
+
result = JobResult.load_from_mlflow(run_id, job_id)
|
|
410
|
+
except FileNotFoundError:
|
|
411
|
+
raise _missing_result(lifecycle, ray_job_id, status) from None
|
|
412
|
+
return result.unwrap()
|
|
413
|
+
if deadline is not None and time.monotonic() >= deadline:
|
|
414
|
+
raise TimeoutError(
|
|
415
|
+
f"job {job_id} did not finish within {timeout}s "
|
|
416
|
+
f"(last status={status.value})"
|
|
417
|
+
)
|
|
418
|
+
time.sleep(_JOB_POLL_INTERVAL_S)
|
|
419
|
+
|
|
420
|
+
|
|
421
|
+
@dataclass(frozen=True)
|
|
422
|
+
class JobFuture:
|
|
423
|
+
"""A handle on a submitted job: its identity, its live status, and its result.
|
|
424
|
+
|
|
425
|
+
``schedule_remote_job`` returns one as soon as the job request is recorded.
|
|
426
|
+
Nothing here blocks until :meth:`result` is called.
|
|
427
|
+
"""
|
|
428
|
+
|
|
429
|
+
experiment_name: str
|
|
430
|
+
run_id: str
|
|
431
|
+
job_id: str
|
|
432
|
+
|
|
433
|
+
def status(self) -> JobStatus:
|
|
434
|
+
"""The job's live status, derived from Ray. ``PENDING`` covers both
|
|
435
|
+
'the control plane has not picked it up yet' and 'queued in Ray'."""
|
|
436
|
+
lifecycle = JobLifecycle.load_from_mlflow(self.run_id, self.job_id)
|
|
437
|
+
return get_ray_job_status(lifecycle.get_ray_job_id())
|
|
438
|
+
|
|
439
|
+
def done(self) -> bool:
|
|
440
|
+
"""True once Ray reports the job finished, failed or stopped."""
|
|
441
|
+
return self.status() in _TERMINAL_JOB_STATES
|
|
442
|
+
|
|
443
|
+
def result(self, timeout: float | None = None) -> Any:
|
|
444
|
+
"""Block until the job finishes and return its function's value.
|
|
445
|
+
|
|
446
|
+
Raises exactly what :func:`wait_for_job_result` raises — including the
|
|
447
|
+
job's own exception, so a blocking call reads like a local one."""
|
|
448
|
+
return wait_for_job_result(self.run_id, self.job_id, timeout)
|
|
449
|
+
|
|
450
|
+
|
|
221
451
|
def schedule_remote_job(
|
|
222
452
|
experiment_name: str,
|
|
223
453
|
run_id: str,
|
|
@@ -227,8 +457,8 @@ def schedule_remote_job(
|
|
|
227
457
|
num_cpus: int = 1,
|
|
228
458
|
retry: bool = False,
|
|
229
459
|
**kwargs: Any,
|
|
230
|
-
) ->
|
|
231
|
-
"""Submit a function to the control plane. Returns a job
|
|
460
|
+
) -> JobFuture:
|
|
461
|
+
"""Submit a function to the control plane. Returns a handle on the job."""
|
|
232
462
|
job_id = Haikunator().haikunate(token_length=2, token_chars="0123456789")
|
|
233
463
|
entry_file = Path(inspect.getfile(fn)).resolve()
|
|
234
464
|
driver_file = Path(__file__).with_name("_ray_job_driver.py")
|
|
@@ -263,7 +493,7 @@ def schedule_remote_job(
|
|
|
263
493
|
num_cpus=num_cpus,
|
|
264
494
|
pip_requirements=pip_requirements,
|
|
265
495
|
).save_to_mlflow()
|
|
266
|
-
return job_id
|
|
496
|
+
return JobFuture(experiment_name=experiment_name, run_id=run_id, job_id=job_id)
|
|
267
497
|
|
|
268
498
|
|
|
269
499
|
def list_experiment_run_jobs(run_id: str) -> list[JobLifecycle]:
|
|
@@ -79,11 +79,61 @@ def train_step(batch):
|
|
|
79
79
|
# MLflow and S3 env vars are injected automatically
|
|
80
80
|
return loss
|
|
81
81
|
|
|
82
|
-
|
|
83
|
-
print(f"Submitted: {job_id}")
|
|
82
|
+
job = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
|
|
83
|
+
print(f"Submitted: {job.job_id}")
|
|
84
84
|
```
|
|
85
85
|
|
|
86
|
-
`cortexgrid.remote` submits a job *request* (a pickled payload plus a `JobLifecycle` record) to MLflow and returns a
|
|
86
|
+
`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.
|
|
87
|
+
|
|
88
|
+
##### Blocking on the result
|
|
89
|
+
|
|
90
|
+
`JobFuture.result(timeout=None)` blocks until Ray reports the job terminal, then behaves like a local call: it returns whatever the function returned, or raises whatever the function raised.
|
|
91
|
+
|
|
92
|
+
```python
|
|
93
|
+
def score(batch) -> dict[str, float]:
|
|
94
|
+
# runs on the DGX
|
|
95
|
+
return {"loss": 0.12}
|
|
96
|
+
|
|
97
|
+
# Fire-and-forget — the form training uses. Nothing blocks, and a retry=True
|
|
98
|
+
# job cannot be waited on at all (see below).
|
|
99
|
+
training = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
|
|
100
|
+
print(training.job_id, training.status().value) # "blue-42 running"
|
|
101
|
+
|
|
102
|
+
# Blocking — the value comes back as if score() had run locally.
|
|
103
|
+
job = cortexgrid.remote(score, batch, num_gpus=1)
|
|
104
|
+
metrics = job.result() # {"loss": 0.12}
|
|
105
|
+
|
|
106
|
+
# ...or with a deadline. On TimeoutError the job keeps running; wait again later.
|
|
107
|
+
metrics = job.result(timeout=600)
|
|
108
|
+
|
|
109
|
+
# From another process — the handle is gone, the job id is enough.
|
|
110
|
+
metrics = cortexgrid.get_job_result(job_id, timeout=600)
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
Failures arrive as exceptions, not as a status to inspect:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
try:
|
|
117
|
+
metrics = job.result()
|
|
118
|
+
except BadBatch as exc: # exactly what score() raised on the DGX
|
|
119
|
+
print(exc.__cause__) # JobFailed, carrying the remote traceback
|
|
120
|
+
except cortexgrid.JobFailed: # the job never got as far as returning
|
|
121
|
+
... # submission failed, driver died, or stopped
|
|
122
|
+
except cortexgrid.JobResultUnavailable:
|
|
123
|
+
... # it returned, but the value is not transportable
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
The driver cloudpickles the outcome to `job/{job_id}/result.pkl`, beside the payload manifest and the lifecycle, and the waiting side reads it back:
|
|
127
|
+
|
|
128
|
+
| Outcome | What `result()` does |
|
|
129
|
+
|---|---|
|
|
130
|
+
| The function returned | Returns its value |
|
|
131
|
+
| The function raised | Raises that exception, with the remote traceback attached as a chained `JobFailed` cause |
|
|
132
|
+
| The job never got as far as returning (submission failed, driver died, job stopped) | Raises `JobFailed` |
|
|
133
|
+
| The value or the exception did not survive cloudpickle | Raises `JobResultUnavailable` (value) or `JobFailed` (exception). A value that cannot be pickled never fails the job itself |
|
|
134
|
+
| `timeout` elapsed | Raises `TimeoutError`; the job keeps running |
|
|
135
|
+
|
|
136
|
+
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.
|
|
87
137
|
|
|
88
138
|
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.
|
|
89
139
|
|
|
@@ -169,7 +219,9 @@ The requirements are part of the model, not of the serve-app class: GPUs, RAM an
|
|
|
169
219
|
| `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
|
|
170
220
|
| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
|
|
171
221
|
| `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
|
|
172
|
-
| `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a
|
|
222
|
+
| `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a `JobFuture` |
|
|
223
|
+
| `JobFuture.status()` / `.done()` / `.result(timeout=None)` | Live status of a submitted job, and its function's return value (blocking) |
|
|
224
|
+
| `cortexgrid.get_job_result(job_id, timeout=None)` | Block on a job of the current run by id; returns its value or raises its exception |
|
|
173
225
|
| `cortexgrid.list_experiment_run_jobs(run_id)` | List `JobLifecycle` records for every cortexgrid job in a run |
|
|
174
226
|
| `cortexgrid.stop_experiment_run_jobs(run_id)` | Request every job in a run to stop (flips the `stop_requested` latch) |
|
|
175
227
|
| `cortexgrid.get_ray_job_status(ray_job_id)` | Live Ray status for a submission id |
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|