cortexgrid 0.2.96__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.96 → cortexgrid-0.3.0}/PKG-INFO +64 -8
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/__init__.py +51 -11
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/_ray_job_driver.py +19 -2
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/_serve_entry.py +11 -12
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/jobs.py +234 -4
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/model_serving.py +105 -9
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/model_storage.py +65 -5
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/docs/cortexgrid/README.md +63 -7
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/pyproject.toml +1 -1
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/.gitignore +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/LICENSE +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/_bundle.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/checkpoint.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/experiment.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/infra.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/mlflow_util.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/py.typed +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/ray_util.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/s3_util.py +0 -0
- {cortexgrid-0.2.96 → cortexgrid-0.3.0}/cortexgrid/secrets.py +0 -0
- {cortexgrid-0.2.96 → 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
|
|
|
@@ -165,19 +215,23 @@ app = FastAPI()
|
|
|
165
215
|
|
|
166
216
|
@serve.ingress(app)
|
|
167
217
|
class MyServeApp:
|
|
168
|
-
num_gpus = 1
|
|
169
|
-
|
|
170
218
|
def __init__(self, family: str, suffix: str, run_name: str) -> None:
|
|
171
219
|
self._weights_dir = cortexgrid.load_model(family, suffix, run_name)
|
|
172
220
|
|
|
173
221
|
@app.post("/complete")
|
|
174
222
|
async def complete(self, body: dict): ...
|
|
175
223
|
|
|
176
|
-
saved = cortexgrid.save_model(
|
|
224
|
+
saved = cortexgrid.save_model(
|
|
225
|
+
weights_dir, MyServeApp, family="qwen", suffix="instruct",
|
|
226
|
+
# What one replica needs; the model is deployed only on a host that has it.
|
|
227
|
+
requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
|
|
228
|
+
)
|
|
177
229
|
deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
|
|
178
230
|
print(deployed.url)
|
|
179
231
|
```
|
|
180
232
|
|
|
233
|
+
The requirements are part of the model, not of the serve-app class: GPUs, RAM and VRAM (GiB, 0 meaning no requirement) are stored with it and matched against what the cluster's hosts have free. Correct them later with `cortexgrid.set_model_requirements(family, suffix, run_name, requirements)` or on the model card in the dashboard; `cortexgrid.deploy_model(..., num_replicas=2)` chooses how many copies to run.
|
|
234
|
+
|
|
181
235
|
`save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
|
|
182
236
|
|
|
183
237
|
`save_model` is synchronous (registry lifecycle: `uploading` -> `ready`); `deploy_model` schedules the serving lifecycle (`deploying` -> `running`). With `wait=True` a failed deploy raises `cortexgrid.ModelDeployFailed`; `cortexgrid.wait_for_model_serving(family, suffix, run_name, timeout=...)` waits on a deploy started elsewhere, and re-deploying a failed model retries it from scratch. See [model-serving.md](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md) for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (`model_registry_status`, `model_serving_status`), and error handling.
|
|
@@ -193,7 +247,9 @@ print(deployed.url)
|
|
|
193
247
|
| `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
|
|
194
248
|
| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
|
|
195
249
|
| `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
|
|
196
|
-
| `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 |
|
|
197
253
|
| `cortexgrid.list_experiment_run_jobs(run_id)` | List `JobLifecycle` records for every cortexgrid job in a run |
|
|
198
254
|
| `cortexgrid.stop_experiment_run_jobs(run_id)` | Request every job in a run to stop (flips the `stop_requested` latch) |
|
|
199
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
|
)
|
|
@@ -80,12 +89,14 @@ from cortexgrid.model_storage import (
|
|
|
80
89
|
list_models,
|
|
81
90
|
load_model,
|
|
82
91
|
model_registry_status,
|
|
92
|
+
set_model_requirements,
|
|
83
93
|
)
|
|
84
94
|
from cortexgrid.model_storage import import_model as _import_model_storage
|
|
85
95
|
from cortexgrid.model_storage import save_model as _save_model_storage
|
|
86
96
|
from cortexgrid.model_serving import (
|
|
87
97
|
Deployment,
|
|
88
98
|
ModelDeployFailed,
|
|
99
|
+
ModelRequirements,
|
|
89
100
|
ServingStatus,
|
|
90
101
|
deploy_model,
|
|
91
102
|
list_deployed_models,
|
|
@@ -102,8 +113,11 @@ def remote(
|
|
|
102
113
|
num_cpus: int = 1,
|
|
103
114
|
retry: bool = False,
|
|
104
115
|
**kwargs: Any,
|
|
105
|
-
) ->
|
|
106
|
-
"""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."""
|
|
107
121
|
experiment = Experiment.get_instance()
|
|
108
122
|
return schedule_remote_job(
|
|
109
123
|
experiment.experiment_name,
|
|
@@ -117,11 +131,27 @@ def remote(
|
|
|
117
131
|
)
|
|
118
132
|
|
|
119
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
|
+
|
|
120
145
|
def save_model(
|
|
121
|
-
weights_dir: str | Path,
|
|
146
|
+
weights_dir: str | Path,
|
|
147
|
+
serve_app: type,
|
|
148
|
+
family: str,
|
|
149
|
+
suffix: str,
|
|
150
|
+
requirements: ModelRequirements | None = None,
|
|
122
151
|
) -> SavedModel:
|
|
123
152
|
"""Persist a weights directory under the current Experiment's run, paired
|
|
124
|
-
with the serve-app class that will front it at deploy time
|
|
153
|
+
with the serve-app class that will front it at deploy time and the
|
|
154
|
+
hardware one replica of it needs.
|
|
125
155
|
|
|
126
156
|
Every run saves a new copy under its own run_name - meant for weights the
|
|
127
157
|
run produced (e.g. a fine-tune). For a model produced elsewhere that should
|
|
@@ -134,6 +164,7 @@ def save_model(
|
|
|
134
164
|
family,
|
|
135
165
|
run_id=experiment.run_id,
|
|
136
166
|
run_name=experiment.run_name(),
|
|
167
|
+
requirements=requirements,
|
|
137
168
|
)
|
|
138
169
|
|
|
139
170
|
|
|
@@ -142,6 +173,7 @@ def import_model(
|
|
|
142
173
|
serve_app: type,
|
|
143
174
|
family: str,
|
|
144
175
|
suffix: str,
|
|
176
|
+
requirements: ModelRequirements | None = None,
|
|
145
177
|
) -> SavedModel:
|
|
146
178
|
"""Register a model produced elsewhere once, reuse it on every later call,
|
|
147
179
|
and record on the current Experiment's run which imported model it used.
|
|
@@ -151,7 +183,7 @@ def import_model(
|
|
|
151
183
|
`imported_model/<family>/<suffix>` holds the version's `created_at`, set
|
|
152
184
|
whether this call uploaded the model or reused it."""
|
|
153
185
|
experiment = Experiment.get_instance()
|
|
154
|
-
model = _import_model_storage(source, serve_app, family, suffix)
|
|
186
|
+
model = _import_model_storage(source, serve_app, family, suffix, requirements)
|
|
155
187
|
get_mlflow_client().set_tag(
|
|
156
188
|
experiment.run_id, f"imported_model/{family}/{suffix}", model.created_at
|
|
157
189
|
)
|
|
@@ -164,11 +196,17 @@ __all__ = [
|
|
|
164
196
|
"delete_run",
|
|
165
197
|
# Ray / jobs
|
|
166
198
|
"remote",
|
|
199
|
+
"get_job_result",
|
|
200
|
+
"wait_for_job_result",
|
|
167
201
|
"get_ray_job_status",
|
|
168
202
|
"list_experiment_run_jobs",
|
|
169
203
|
"stop_experiment_run_jobs",
|
|
170
204
|
"JobStatus",
|
|
205
|
+
"JobFailed",
|
|
206
|
+
"JobFuture",
|
|
171
207
|
"JobLifecycle",
|
|
208
|
+
"JobResult",
|
|
209
|
+
"JobResultUnavailable",
|
|
172
210
|
"LifecycleEvent",
|
|
173
211
|
"Payload",
|
|
174
212
|
"get_ray_status",
|
|
@@ -213,6 +251,8 @@ __all__ = [
|
|
|
213
251
|
"load_model",
|
|
214
252
|
"list_models",
|
|
215
253
|
"model_registry_status",
|
|
254
|
+
"ModelRequirements",
|
|
255
|
+
"set_model_requirements",
|
|
216
256
|
"delete_model",
|
|
217
257
|
# Model serving
|
|
218
258
|
"Deployment",
|
|
@@ -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__":
|
|
@@ -4,9 +4,9 @@ Ray Serve's REST `import_path` resolves to `cortexgrid._serve_entry:build`.
|
|
|
4
4
|
On the cluster replica, `build` imports the serve-app class bundled at
|
|
5
5
|
`save_model` time (its import path was stored as an MLflow tag), applies Ray's
|
|
6
6
|
ingress with the app it was marked with by `cortexgrid.serve.ingress` (again on
|
|
7
|
-
each replica, see `_IngressOnReplica`),
|
|
8
|
-
|
|
9
|
-
|
|
7
|
+
each replica, see `_IngressOnReplica`), wraps it as a Ray Serve deployment
|
|
8
|
+
with the replica count and Ray resource requests `deploy_model` derived from
|
|
9
|
+
the model's requirements, and binds it with the (family, suffix, run_name)
|
|
10
10
|
identifiers.
|
|
11
11
|
|
|
12
12
|
The serve-app owns everything about traffic: its own routes, request schemas,
|
|
@@ -14,11 +14,9 @@ streaming, and timeouts. cortexgrid does not interpose a request/response
|
|
|
14
14
|
contract - it only schedules the app and hands it the identifiers it needs to
|
|
15
15
|
fetch its own weights via `cortexgrid.load_model`.
|
|
16
16
|
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
declare resources in practice; revisit if plain class attributes prove too
|
|
21
|
-
limited.
|
|
17
|
+
The serve-app declares no resources: the hardware a replica needs belongs to
|
|
18
|
+
the model and is stored in the registry (`cortexgrid.ModelRequirements`), and
|
|
19
|
+
the replica count is chosen per `deploy_model`.
|
|
22
20
|
"""
|
|
23
21
|
|
|
24
22
|
from __future__ import annotations
|
|
@@ -72,12 +70,13 @@ def build(args: dict[str, Any]) -> Application:
|
|
|
72
70
|
{"_serve_app": serve_app},
|
|
73
71
|
)
|
|
74
72
|
serve_app = serve.ingress(app)(on_replica)
|
|
75
|
-
num_gpus = getattr(serve_app, "num_gpus", 0)
|
|
76
|
-
num_replicas = getattr(serve_app, "num_replicas", 1)
|
|
77
73
|
return serve.deployment(serve_app).options(
|
|
78
|
-
|
|
74
|
+
# This builder ships in the bundle, frozen at save time, while `args`
|
|
75
|
+
# come from the cortexgrid that deploys it; one older than the bundle
|
|
76
|
+
# sends neither key.
|
|
77
|
+
num_replicas=args.get("num_replicas", 1),
|
|
79
78
|
# Ray 2.32 lowered the default from 100 to 5; keep what serve-apps
|
|
80
79
|
# had on Ray 2.9.
|
|
81
80
|
max_ongoing_requests=_MAX_ONGOING_REQUESTS,
|
|
82
|
-
ray_actor_options=
|
|
81
|
+
ray_actor_options=args.get("ray_actor_options", {}),
|
|
83
82
|
).bind(args["family"], args["suffix"], args["run_name"])
|
|
@@ -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]:
|
|
@@ -8,7 +8,8 @@ uploaded to MinIO under
|
|
|
8
8
|
referenced via `runtime_env.working_dir` so Ray workers fetch it from there.
|
|
9
9
|
The bundle URL, class import path, pip list, and fingerprint are persisted as
|
|
10
10
|
MLflow tags on the ModelVersion so `deploy_model` can find them later without
|
|
11
|
-
the caller holding the class object.
|
|
11
|
+
the caller holding the class object. So are the model's `ModelRequirements`,
|
|
12
|
+
which `deploy_model` turns into the replica's Ray resource requests.
|
|
12
13
|
|
|
13
14
|
Naming: the Ray Serve application is named "<family>__<suffix>__<run_name>".
|
|
14
15
|
This relies on family/suffix/run_name not containing the literal "__".
|
|
@@ -206,10 +207,41 @@ def upload_bundle(
|
|
|
206
207
|
)
|
|
207
208
|
|
|
208
209
|
|
|
210
|
+
# Custom Ray resource each GPU worker advertises: the MiB of memory its GPUs
|
|
211
|
+
# have (see the ray-worker DaemonSet). A replica requests its vram_gb of it in
|
|
212
|
+
# MiB, so Ray places it only on a node with that much VRAM left. MiB because
|
|
213
|
+
# nvidia-smi reports MiB and a GPU's memory is not a whole number of GiB.
|
|
214
|
+
_VRAM_RESOURCE = "vram_mib"
|
|
215
|
+
|
|
216
|
+
_MIB_PER_GIB = 1024
|
|
217
|
+
|
|
218
|
+
_GIB = 1024**3
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
def _ray_actor_options(requirements: ModelRequirements) -> dict[str, Any]:
|
|
222
|
+
"""Translate ModelRequirements into a replica's Ray actor resource requests.
|
|
223
|
+
Ray places the replica only on a node with that much free and reserves it
|
|
224
|
+
there; a zero requirement requests nothing."""
|
|
225
|
+
options: dict[str, Any] = {"num_gpus": requirements.num_gpus}
|
|
226
|
+
if requirements.ram_gb > 0:
|
|
227
|
+
options["memory"] = int(requirements.ram_gb * _GIB)
|
|
228
|
+
if requirements.vram_gb > 0:
|
|
229
|
+
options["resources"] = {
|
|
230
|
+
_VRAM_RESOURCE: round(requirements.vram_gb * _MIB_PER_GIB)
|
|
231
|
+
}
|
|
232
|
+
return options
|
|
233
|
+
|
|
234
|
+
|
|
209
235
|
def _build_application_spec(
|
|
210
|
-
family: str,
|
|
236
|
+
family: str,
|
|
237
|
+
suffix: str,
|
|
238
|
+
run_name: str,
|
|
239
|
+
meta: BundleMetadata,
|
|
240
|
+
requirements: ModelRequirements,
|
|
241
|
+
num_replicas: int,
|
|
211
242
|
) -> dict[str, Any]:
|
|
212
|
-
"""Assemble a Ray Serve application schema from pre-bundled metadata
|
|
243
|
+
"""Assemble a Ray Serve application schema from pre-bundled metadata and
|
|
244
|
+
the model's requirements."""
|
|
213
245
|
# working_dir carries the serve-app's own source; Ray pip-installs the
|
|
214
246
|
# third-party distributions the image lacks into a per-node cached
|
|
215
247
|
# virtualenv layered on the image. No pip key when there are none, so Ray
|
|
@@ -230,6 +262,8 @@ def _build_application_spec(
|
|
|
230
262
|
"family": family,
|
|
231
263
|
"suffix": suffix,
|
|
232
264
|
"run_name": run_name,
|
|
265
|
+
"num_replicas": num_replicas,
|
|
266
|
+
"ray_actor_options": _ray_actor_options(requirements),
|
|
233
267
|
},
|
|
234
268
|
"runtime_env": runtime_env,
|
|
235
269
|
}
|
|
@@ -268,10 +302,66 @@ def metadata_from_tags(tags: dict[str, str]) -> BundleMetadata:
|
|
|
268
302
|
)
|
|
269
303
|
|
|
270
304
|
|
|
271
|
-
|
|
305
|
+
@dataclass
|
|
306
|
+
class ModelRequirements:
|
|
307
|
+
"""Hardware one replica of a model needs to be served, in GiB. Persisted as
|
|
308
|
+
tags on the ModelVersion next to the bundle metadata, so it is read without
|
|
309
|
+
touching the weights or importing the serve-app class.
|
|
310
|
+
|
|
311
|
+
Zero means no requirement: a model with no requirements is served on any
|
|
312
|
+
node, CPU-only included. Models saved before requirements existed carry
|
|
313
|
+
no tags and read as the defaults."""
|
|
314
|
+
|
|
315
|
+
num_gpus: int = 0
|
|
316
|
+
ram_gb: float = 0.0
|
|
317
|
+
# GPU memory across the replica's num_gpus GPUs, so it needs num_gpus >= 1.
|
|
318
|
+
vram_gb: float = 0.0
|
|
319
|
+
|
|
320
|
+
def __post_init__(self) -> None:
|
|
321
|
+
if self.num_gpus < 0 or self.ram_gb < 0 or self.vram_gb < 0:
|
|
322
|
+
raise ValueError(f"Model requirements cannot be negative: {self}")
|
|
323
|
+
if self.vram_gb > 0 and self.num_gpus == 0:
|
|
324
|
+
raise ValueError(
|
|
325
|
+
f"vram_gb={self.vram_gb} needs a GPU; set num_gpus >= 1"
|
|
326
|
+
)
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
# MLflow tag keys for the ModelRequirements.
|
|
330
|
+
_NUM_GPUS_TAG = "num_gpus"
|
|
331
|
+
_RAM_GB_TAG = "ram_gb"
|
|
332
|
+
_VRAM_GB_TAG = "vram_gb"
|
|
333
|
+
|
|
334
|
+
|
|
335
|
+
def has_requirement_tags(tags: dict[str, str]) -> bool:
|
|
336
|
+
"""Whether requirements were ever stored on the ModelVersion."""
|
|
337
|
+
return any(key in tags for key in (_NUM_GPUS_TAG, _RAM_GB_TAG, _VRAM_GB_TAG))
|
|
338
|
+
|
|
339
|
+
|
|
340
|
+
def requirements_to_tags(requirements: ModelRequirements) -> dict[str, str]:
|
|
341
|
+
"""Serialise ModelRequirements to MLflow tags. The inverse of
|
|
342
|
+
`requirements_from_tags`."""
|
|
343
|
+
return {
|
|
344
|
+
_NUM_GPUS_TAG: str(requirements.num_gpus),
|
|
345
|
+
_RAM_GB_TAG: str(requirements.ram_gb),
|
|
346
|
+
_VRAM_GB_TAG: str(requirements.vram_gb),
|
|
347
|
+
}
|
|
348
|
+
|
|
349
|
+
|
|
350
|
+
def requirements_from_tags(tags: dict[str, str]) -> ModelRequirements:
|
|
351
|
+
"""Deserialise ModelRequirements from a ModelVersion's MLflow tags; a
|
|
352
|
+
missing tag reads as no requirement."""
|
|
353
|
+
return ModelRequirements(
|
|
354
|
+
num_gpus=int(tags.get(_NUM_GPUS_TAG, "0")),
|
|
355
|
+
ram_gb=float(tags.get(_RAM_GB_TAG, "0")),
|
|
356
|
+
vram_gb=float(tags.get(_VRAM_GB_TAG, "0")),
|
|
357
|
+
)
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
def _load_deploy_metadata(
|
|
272
361
|
family: str, suffix: str, run_name: str
|
|
273
|
-
) -> BundleMetadata:
|
|
274
|
-
"""Read the bundle metadata `save_model` persisted on the
|
|
362
|
+
) -> tuple[BundleMetadata, ModelRequirements]:
|
|
363
|
+
"""Read the bundle metadata and requirements `save_model` persisted on the
|
|
364
|
+
ModelVersion."""
|
|
275
365
|
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
276
366
|
name = f"{family}__{suffix}"
|
|
277
367
|
versions = client.search_model_versions(
|
|
@@ -281,8 +371,9 @@ def _load_bundle_metadata(
|
|
|
281
371
|
raise ValueError(
|
|
282
372
|
f"No saved model for {family}/{suffix}/{run_name}; cannot deploy."
|
|
283
373
|
)
|
|
374
|
+
tags = versions[0].tags or {}
|
|
284
375
|
try:
|
|
285
|
-
return metadata_from_tags(
|
|
376
|
+
return metadata_from_tags(tags), requirements_from_tags(tags)
|
|
286
377
|
except KeyError as exc:
|
|
287
378
|
raise ValueError(
|
|
288
379
|
f"Saved model {family}/{suffix}/{run_name} is missing the deployment "
|
|
@@ -398,6 +489,7 @@ def deploy_model(
|
|
|
398
489
|
family: str,
|
|
399
490
|
suffix: str,
|
|
400
491
|
run_name: str,
|
|
492
|
+
num_replicas: int = 1,
|
|
401
493
|
wait: bool = False,
|
|
402
494
|
timeout: float | None = 300.0,
|
|
403
495
|
) -> Deployment:
|
|
@@ -408,6 +500,8 @@ def deploy_model(
|
|
|
408
500
|
|
|
409
501
|
The serve-app class is pulled from the MLflow ModelVersion tags `save_model`
|
|
410
502
|
wrote at save time; the caller does not need to hold the class object.
|
|
503
|
+
Each of the `num_replicas` replicas requests the model's `ModelRequirements`
|
|
504
|
+
from Ray, so it is placed only on a node that has them free.
|
|
411
505
|
|
|
412
506
|
A DEPLOY_FAILED app left by an earlier attempt is undeployed first, and it,
|
|
413
507
|
or an app still DELETING, is waited out before the new spec is PUT, so the
|
|
@@ -421,8 +515,10 @@ def deploy_model(
|
|
|
421
515
|
stuck in DEPLOYING) will hang forever.
|
|
422
516
|
"""
|
|
423
517
|
deadline = _deadline(timeout)
|
|
424
|
-
meta =
|
|
425
|
-
spec = _build_application_spec(
|
|
518
|
+
meta, requirements = _load_deploy_metadata(family, suffix, run_name)
|
|
519
|
+
spec = _build_application_spec(
|
|
520
|
+
family, suffix, run_name, meta, requirements, num_replicas
|
|
521
|
+
)
|
|
426
522
|
_clear_failed_application(family, suffix, run_name, timeout, deadline)
|
|
427
523
|
existing = [a for a in _current_application_specs() if a["name"] != spec["name"]]
|
|
428
524
|
# The controller registers the app, sets it DEPLOYING and stamps
|
|
@@ -8,6 +8,7 @@ Mapping cortexgrid taxonomy <-> MLflow Registry:
|
|
|
8
8
|
"s3://<bucket>/models/<run_name>/<family>/<suffix>/weights/"
|
|
9
9
|
run linkage -> ModelVersion.run_id (built-in MLflow field; unset
|
|
10
10
|
for imported models)
|
|
11
|
+
requirements -> ModelVersion.tags["num_gpus"], ["ram_gb"], ["vram_gb"]
|
|
11
12
|
|
|
12
13
|
Two ways in: `save_model` registers a fresh copy under the calling run's
|
|
13
14
|
run_name every time it runs (fine-tuned output); `import_model` registers a
|
|
@@ -36,10 +37,14 @@ from mlflow.tracking import MlflowClient
|
|
|
36
37
|
from cortexgrid import s3_util
|
|
37
38
|
from cortexgrid.infra import get_mlflow_tracking_uri, get_s3_bucket
|
|
38
39
|
from cortexgrid.model_serving import (
|
|
40
|
+
ModelRequirements,
|
|
39
41
|
build_bundle,
|
|
40
42
|
bundle_class,
|
|
43
|
+
has_requirement_tags,
|
|
41
44
|
metadata_from_tags,
|
|
42
45
|
metadata_to_tags,
|
|
46
|
+
requirements_from_tags,
|
|
47
|
+
requirements_to_tags,
|
|
43
48
|
upload_bundle,
|
|
44
49
|
)
|
|
45
50
|
|
|
@@ -81,6 +86,8 @@ class SavedModel:
|
|
|
81
86
|
# _UPLOAD_DEADLINE (writer presumed dead). Versions written before this tag
|
|
82
87
|
# existed report "ready".
|
|
83
88
|
phase: str
|
|
89
|
+
# Hardware one replica needs; defaults for versions stored without it.
|
|
90
|
+
requirements: ModelRequirements
|
|
84
91
|
|
|
85
92
|
|
|
86
93
|
def _phase_for(version: Any) -> str:
|
|
@@ -111,6 +118,7 @@ def _to_saved_model(version: Any) -> SavedModel:
|
|
|
111
118
|
data_blob_path=version.source,
|
|
112
119
|
size_bytes=int(version.tags.get("size_bytes", "0")),
|
|
113
120
|
phase=_phase_for(version),
|
|
121
|
+
requirements=requirements_from_tags(version.tags),
|
|
114
122
|
)
|
|
115
123
|
|
|
116
124
|
|
|
@@ -150,6 +158,7 @@ def save_model(
|
|
|
150
158
|
family: str,
|
|
151
159
|
run_id: str,
|
|
152
160
|
run_name: str,
|
|
161
|
+
requirements: ModelRequirements | None = None,
|
|
153
162
|
) -> SavedModel:
|
|
154
163
|
"""Upload a weights directory to S3 and register a new MLflow ModelVersion
|
|
155
164
|
paired with the serve-app that fronts it.
|
|
@@ -168,8 +177,13 @@ def save_model(
|
|
|
168
177
|
weights.
|
|
169
178
|
Its code is bundled and its import path, bundle URL, and pip list are
|
|
170
179
|
stored as tags on the ModelVersion so `deploy_model` can bind it later
|
|
171
|
-
without the caller holding the class object.
|
|
172
|
-
|
|
180
|
+
without the caller holding the class object.
|
|
181
|
+
|
|
182
|
+
`requirements` is the hardware one replica needs; None stores none, which
|
|
183
|
+
reads as no requirement. Change it later with `set_model_requirements`."""
|
|
184
|
+
return _upload_model(
|
|
185
|
+
weights_dir, serve_app, suffix, family, run_id, run_name, requirements
|
|
186
|
+
)
|
|
173
187
|
|
|
174
188
|
|
|
175
189
|
def import_model(
|
|
@@ -177,6 +191,7 @@ def import_model(
|
|
|
177
191
|
serve_app: type,
|
|
178
192
|
family: str,
|
|
179
193
|
suffix: str,
|
|
194
|
+
requirements: ModelRequirements | None = None,
|
|
180
195
|
) -> SavedModel:
|
|
181
196
|
"""Register a model produced elsewhere (e.g. a pretrained base model) under
|
|
182
197
|
the fixed key (family, suffix, IMPORTED), once.
|
|
@@ -190,7 +205,9 @@ def import_model(
|
|
|
190
205
|
If a version is already registered under the key:
|
|
191
206
|
- "ready": returns it without calling `source`. If `serve_app`'s code no
|
|
192
207
|
longer matches the stored bundle, it is re-bundled first and the
|
|
193
|
-
weights are kept (see `_refresh_bundle`).
|
|
208
|
+
weights are kept (see `_refresh_bundle`). `requirements` are stored
|
|
209
|
+
only if the version has none yet, so values changed since with
|
|
210
|
+
`set_model_requirements` are kept. To replace the weights,
|
|
194
211
|
`delete_model` it first.
|
|
195
212
|
- "uploading": raises RuntimeError - another process is importing it.
|
|
196
213
|
- "upload_failed" / "broken": deleted and imported again.
|
|
@@ -202,6 +219,10 @@ def import_model(
|
|
|
202
219
|
if existing is not None:
|
|
203
220
|
if existing.phase == _PHASE_READY:
|
|
204
221
|
_refresh_bundle(serve_app, family, suffix)
|
|
222
|
+
if requirements is not None:
|
|
223
|
+
existing.requirements = _set_missing_requirements(
|
|
224
|
+
family, suffix, requirements
|
|
225
|
+
)
|
|
205
226
|
return existing
|
|
206
227
|
if existing.phase == _PHASE_UPLOADING:
|
|
207
228
|
raise RuntimeError(
|
|
@@ -209,7 +230,9 @@ def import_model(
|
|
|
209
230
|
"another process"
|
|
210
231
|
)
|
|
211
232
|
delete_model(family, suffix, IMPORTED)
|
|
212
|
-
return _upload_model(
|
|
233
|
+
return _upload_model(
|
|
234
|
+
source, serve_app, suffix, family, None, IMPORTED, requirements
|
|
235
|
+
)
|
|
213
236
|
|
|
214
237
|
|
|
215
238
|
def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
|
|
@@ -233,6 +256,23 @@ def _refresh_bundle(serve_app: type, family: str, suffix: str) -> None:
|
|
|
233
256
|
client.set_model_version_tag(name, version.version, key, value)
|
|
234
257
|
|
|
235
258
|
|
|
259
|
+
def _set_missing_requirements(
|
|
260
|
+
family: str, suffix: str, requirements: ModelRequirements
|
|
261
|
+
) -> ModelRequirements:
|
|
262
|
+
"""Store `requirements` on an imported model whose version has none yet.
|
|
263
|
+
Returns the requirements the version holds afterwards."""
|
|
264
|
+
name = f"{family}__{suffix}"
|
|
265
|
+
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
266
|
+
version = client.search_model_versions(
|
|
267
|
+
f"name='{name}' and tags.run_name='{IMPORTED}'"
|
|
268
|
+
)[0]
|
|
269
|
+
if has_requirement_tags(version.tags):
|
|
270
|
+
return requirements_from_tags(version.tags)
|
|
271
|
+
for key, value in requirements_to_tags(requirements).items():
|
|
272
|
+
client.set_model_version_tag(name, version.version, key, value)
|
|
273
|
+
return requirements
|
|
274
|
+
|
|
275
|
+
|
|
236
276
|
def _upload_model(
|
|
237
277
|
weights: str | Path | Callable[[], str | Path],
|
|
238
278
|
serve_app: type,
|
|
@@ -240,6 +280,7 @@ def _upload_model(
|
|
|
240
280
|
family: str,
|
|
241
281
|
run_id: str | None,
|
|
242
282
|
run_name: str,
|
|
283
|
+
requirements: ModelRequirements | None,
|
|
243
284
|
) -> SavedModel:
|
|
244
285
|
"""Register a ModelVersion in "uploading", resolve `weights` to a directory
|
|
245
286
|
(calling it when it is a callable), upload the weights and the serve-app
|
|
@@ -255,7 +296,8 @@ def _upload_model(
|
|
|
255
296
|
# fetched and uploaded, and a concurrent `import_model` sees the import in
|
|
256
297
|
# flight for the whole download instead of starting one of its own. The
|
|
257
298
|
# size is stamped once the directory exists; the bundle tags and the flip
|
|
258
|
-
# to "ready" happen only after the upload lands.
|
|
299
|
+
# to "ready" happen only after the upload lands. No requirements leaves
|
|
300
|
+
# their tags unset, so a later `import_model` can still store them.
|
|
259
301
|
version = client.create_model_version(
|
|
260
302
|
name=name,
|
|
261
303
|
source=source,
|
|
@@ -265,6 +307,7 @@ def _upload_model(
|
|
|
265
307
|
"suffix": suffix,
|
|
266
308
|
"run_name": run_name,
|
|
267
309
|
_LIFECYCLE_TAG: _PHASE_UPLOADING,
|
|
310
|
+
**(requirements_to_tags(requirements) if requirements is not None else {}),
|
|
268
311
|
},
|
|
269
312
|
)
|
|
270
313
|
try:
|
|
@@ -333,6 +376,23 @@ def model_registry_status(
|
|
|
333
376
|
return _to_saved_model(versions[0]) if versions else None
|
|
334
377
|
|
|
335
378
|
|
|
379
|
+
def set_model_requirements(
|
|
380
|
+
family: str, suffix: str, run_name: str, requirements: ModelRequirements
|
|
381
|
+
) -> None:
|
|
382
|
+
"""Replace the hardware requirements stored on a model. Takes effect on
|
|
383
|
+
its next `deploy_model`; a replica already running keeps its placement.
|
|
384
|
+
Raises ValueError if the model was never registered."""
|
|
385
|
+
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
386
|
+
name = f"{family}__{suffix}"
|
|
387
|
+
versions = client.search_model_versions(
|
|
388
|
+
f"name='{name}' and tags.run_name='{run_name}'"
|
|
389
|
+
)
|
|
390
|
+
if not versions:
|
|
391
|
+
raise ValueError(f"No model {family}/{suffix}/{run_name}")
|
|
392
|
+
for key, value in requirements_to_tags(requirements).items():
|
|
393
|
+
client.set_model_version_tag(name, versions[0].version, key, value)
|
|
394
|
+
|
|
395
|
+
|
|
336
396
|
def delete_model(family: str, suffix: str, run_name: str) -> None:
|
|
337
397
|
"""Delete the ModelVersion in MLflow, its weights blob, and its serve bundle."""
|
|
338
398
|
client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
|
|
@@ -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
|
|
|
@@ -137,19 +187,23 @@ app = FastAPI()
|
|
|
137
187
|
|
|
138
188
|
@serve.ingress(app)
|
|
139
189
|
class MyServeApp:
|
|
140
|
-
num_gpus = 1
|
|
141
|
-
|
|
142
190
|
def __init__(self, family: str, suffix: str, run_name: str) -> None:
|
|
143
191
|
self._weights_dir = cortexgrid.load_model(family, suffix, run_name)
|
|
144
192
|
|
|
145
193
|
@app.post("/complete")
|
|
146
194
|
async def complete(self, body: dict): ...
|
|
147
195
|
|
|
148
|
-
saved = cortexgrid.save_model(
|
|
196
|
+
saved = cortexgrid.save_model(
|
|
197
|
+
weights_dir, MyServeApp, family="qwen", suffix="instruct",
|
|
198
|
+
# What one replica needs; the model is deployed only on a host that has it.
|
|
199
|
+
requirements=cortexgrid.ModelRequirements(num_gpus=1, ram_gb=8, vram_gb=16),
|
|
200
|
+
)
|
|
149
201
|
deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
|
|
150
202
|
print(deployed.url)
|
|
151
203
|
```
|
|
152
204
|
|
|
205
|
+
The requirements are part of the model, not of the serve-app class: GPUs, RAM and VRAM (GiB, 0 meaning no requirement) are stored with it and matched against what the cluster's hosts have free. Correct them later with `cortexgrid.set_model_requirements(family, suffix, run_name, requirements)` or on the model card in the dashboard; `cortexgrid.deploy_model(..., num_replicas=2)` chooses how many copies to run.
|
|
206
|
+
|
|
153
207
|
`save_model` saves a new copy under every run - meant for weights the run produced (e.g. a fine-tune). For a model produced elsewhere (e.g. a pretrained base model), `cortexgrid.import_model(source, MyServeApp, family, suffix)` uploads it once under `run_name=cortexgrid.IMPORTED` and on later runs only re-bundles `MyServeApp` if its code changed; deploy it with `deploy_model(family, suffix, cortexgrid.IMPORTED)`.
|
|
154
208
|
|
|
155
209
|
`save_model` is synchronous (registry lifecycle: `uploading` -> `ready`); `deploy_model` schedules the serving lifecycle (`deploying` -> `running`). With `wait=True` a failed deploy raises `cortexgrid.ModelDeployFailed`; `cortexgrid.wait_for_model_serving(family, suffix, run_name, timeout=...)` waits on a deploy started elsewhere, and re-deploying a failed model retries it from scratch. See [model-serving.md](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md) for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (`model_registry_status`, `model_serving_status`), and error handling.
|
|
@@ -165,7 +219,9 @@ print(deployed.url)
|
|
|
165
219
|
| `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
|
|
166
220
|
| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
|
|
167
221
|
| `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
|
|
168
|
-
| `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 |
|
|
169
225
|
| `cortexgrid.list_experiment_run_jobs(run_id)` | List `JobLifecycle` records for every cortexgrid job in a run |
|
|
170
226
|
| `cortexgrid.stop_experiment_run_jobs(run_id)` | Request every job in a run to stop (flips the `stop_requested` latch) |
|
|
171
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
|