aind-codeocean-dispatch 0.2.0__py3-none-any.whl
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.
- aind_codeocean_dispatch/__init__.py +34 -0
- aind_codeocean_dispatch/py.typed +1 -0
- aind_codeocean_dispatch/runner.py +603 -0
- aind_codeocean_dispatch-0.2.0.dist-info/METADATA +73 -0
- aind_codeocean_dispatch-0.2.0.dist-info/RECORD +7 -0
- aind_codeocean_dispatch-0.2.0.dist-info/WHEEL +4 -0
- aind_codeocean_dispatch-0.2.0.dist-info/licenses/LICENSE +22 -0
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
"""Launch, monitor, and capture batches of Code Ocean capsule runs."""
|
|
2
|
+
|
|
3
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
4
|
+
|
|
5
|
+
from .runner import (
|
|
6
|
+
DispatchFailed,
|
|
7
|
+
DispatchInterrupted,
|
|
8
|
+
Job,
|
|
9
|
+
Outcome,
|
|
10
|
+
failed_without_results,
|
|
11
|
+
job_key,
|
|
12
|
+
key_tag,
|
|
13
|
+
raise_for_failures,
|
|
14
|
+
run_jobs,
|
|
15
|
+
succeeded,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
__all__ = [
|
|
19
|
+
"DispatchFailed",
|
|
20
|
+
"DispatchInterrupted",
|
|
21
|
+
"Job",
|
|
22
|
+
"Outcome",
|
|
23
|
+
"failed_without_results",
|
|
24
|
+
"job_key",
|
|
25
|
+
"key_tag",
|
|
26
|
+
"raise_for_failures",
|
|
27
|
+
"run_jobs",
|
|
28
|
+
"succeeded",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
try:
|
|
32
|
+
__version__ = version("aind-codeocean-dispatch")
|
|
33
|
+
except PackageNotFoundError:
|
|
34
|
+
__version__ = "0.0.0.dev0"
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1,603 @@
|
|
|
1
|
+
"""Launch a batch of Code Ocean runs, wait for each, and capture its results.
|
|
2
|
+
|
|
3
|
+
:func:`run_jobs` drives every job from one polling loop. Each job moves
|
|
4
|
+
pending → running → capturing → done, so the concurrency cap and launch pacing are
|
|
5
|
+
counters rather than locks, and no thread is parked per job.
|
|
6
|
+
|
|
7
|
+
Error policy
|
|
8
|
+
------------
|
|
9
|
+
The SDK's client retries nothing by default, so every 429 or 5xx surfaces here.
|
|
10
|
+
|
|
11
|
+
- **429** is retried on every call, honouring ``Retry-After``: the server rejected
|
|
12
|
+
the request without acting on it.
|
|
13
|
+
- **5xx and connection errors** are retried on reads only. A POST that fails this way
|
|
14
|
+
may already have created its computation or data asset, so retrying it could
|
|
15
|
+
duplicate one; the job fails instead, saying so.
|
|
16
|
+
- **Any other HTTP error** fails the job.
|
|
17
|
+
|
|
18
|
+
Policy beyond that -- concurrency cap, launch pacing, relaunching failed starts,
|
|
19
|
+
retrying failed captures, reusing captured results -- is opt-in.
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
from __future__ import annotations
|
|
23
|
+
|
|
24
|
+
import dataclasses
|
|
25
|
+
import logging
|
|
26
|
+
import re
|
|
27
|
+
import time
|
|
28
|
+
from collections.abc import Callable, Sequence
|
|
29
|
+
from dataclasses import dataclass, field
|
|
30
|
+
from enum import Enum, auto
|
|
31
|
+
from typing import TYPE_CHECKING, Any, Literal, TypeVar
|
|
32
|
+
|
|
33
|
+
import requests
|
|
34
|
+
from aind_code_ocean_pipeline_utils.cache import input_fingerprint
|
|
35
|
+
from aind_code_ocean_pipeline_utils.process import GracefulExit, check_shutdown
|
|
36
|
+
from codeocean.error import Error as CodeOceanError
|
|
37
|
+
from codeocean.models.computation import Computation, ComputationEndStatus, ComputationState, RunParams
|
|
38
|
+
from codeocean.models.data_asset import DataAsset, DataAssetParams, DataAssetSearchParams, DataAssetState
|
|
39
|
+
|
|
40
|
+
if TYPE_CHECKING:
|
|
41
|
+
from codeocean import CodeOcean
|
|
42
|
+
|
|
43
|
+
__all__ = [
|
|
44
|
+
"DispatchFailed",
|
|
45
|
+
"DispatchInterrupted",
|
|
46
|
+
"Job",
|
|
47
|
+
"Outcome",
|
|
48
|
+
"failed_without_results",
|
|
49
|
+
"job_key",
|
|
50
|
+
"key_tag",
|
|
51
|
+
"raise_for_failures",
|
|
52
|
+
"run_jobs",
|
|
53
|
+
"succeeded",
|
|
54
|
+
]
|
|
55
|
+
|
|
56
|
+
logger = logging.getLogger(__name__)
|
|
57
|
+
|
|
58
|
+
# Indirected so tests can drive the loop with a fake clock.
|
|
59
|
+
_now = time.monotonic
|
|
60
|
+
_sleep = time.sleep
|
|
61
|
+
|
|
62
|
+
_ASSET_POLL_S = 10.0
|
|
63
|
+
# Sleep in short chunks: a flag-only signal handler does not interrupt time.sleep.
|
|
64
|
+
_MAX_SLEEP_S = 5.0
|
|
65
|
+
_BACKOFF_BASE_S = 5.0
|
|
66
|
+
_BACKOFF_CAP_S = 300.0
|
|
67
|
+
_RATE_LIMITED = 429
|
|
68
|
+
_TRANSIENT_SERVER = frozenset({500, 502, 503, 504})
|
|
69
|
+
_KEY_PATTERN = re.compile(r"[A-Za-z0-9_.-]{1,64}")
|
|
70
|
+
_KEY_TAG_PREFIX = "dispatch-"
|
|
71
|
+
|
|
72
|
+
Status = Literal["captured", "completed", "reused", "failed"]
|
|
73
|
+
_T = TypeVar("_T")
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
@dataclass(frozen=True)
|
|
77
|
+
class Job:
|
|
78
|
+
"""One capsule or pipeline run to launch, and how to capture its results.
|
|
79
|
+
|
|
80
|
+
Parameters
|
|
81
|
+
----------
|
|
82
|
+
label : str
|
|
83
|
+
Unique within a batch; names the job in logs and outcomes.
|
|
84
|
+
run_params : RunParams
|
|
85
|
+
Passed verbatim to ``client.computations.run_capsule``.
|
|
86
|
+
capture : callable, optional
|
|
87
|
+
Builds the capture request from the finished computation, so the asset name
|
|
88
|
+
can depend on it. ``None`` stops at a successful computation.
|
|
89
|
+
key : str, optional
|
|
90
|
+
Resume key (see :func:`job_key`), added to the captured asset as
|
|
91
|
+
:func:`key_tag`. With ``reuse=True``, a Ready asset carrying it is returned
|
|
92
|
+
instead of launching. ``None`` never reuses.
|
|
93
|
+
"""
|
|
94
|
+
|
|
95
|
+
label: str
|
|
96
|
+
run_params: RunParams
|
|
97
|
+
capture: Callable[[Computation], DataAssetParams] | None = None
|
|
98
|
+
key: str | None = None
|
|
99
|
+
|
|
100
|
+
def __post_init__(self) -> None:
|
|
101
|
+
"""Reject keys that cannot round-trip through a Code Ocean tag search."""
|
|
102
|
+
if self.key is not None and not _KEY_PATTERN.fullmatch(self.key):
|
|
103
|
+
raise ValueError(f"job {self.label!r}: key {self.key!r} must match {_KEY_PATTERN.pattern}")
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
@dataclass(frozen=True)
|
|
107
|
+
class Outcome:
|
|
108
|
+
"""How one job ended.
|
|
109
|
+
|
|
110
|
+
``computation_ids`` lists every launch attempt in order, so a relaunched job's
|
|
111
|
+
failed starts stay traceable.
|
|
112
|
+
"""
|
|
113
|
+
|
|
114
|
+
job: Job
|
|
115
|
+
status: Status
|
|
116
|
+
asset_id: str | None = None
|
|
117
|
+
computation_ids: tuple[str, ...] = ()
|
|
118
|
+
error: str | None = None
|
|
119
|
+
|
|
120
|
+
@property
|
|
121
|
+
def ok(self) -> bool:
|
|
122
|
+
"""Whether the job produced what was asked of it."""
|
|
123
|
+
return self.status != "failed"
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
class DispatchFailed(RuntimeError):
|
|
127
|
+
"""Raised by :func:`raise_for_failures`; ``failures`` holds the failed outcomes."""
|
|
128
|
+
|
|
129
|
+
def __init__(self, failures: Sequence[Outcome]) -> None:
|
|
130
|
+
self.failures = tuple(failures)
|
|
131
|
+
lines = [f"{len(self.failures)} job(s) failed:"]
|
|
132
|
+
lines += [f" {o.job.label}: {o.error}" for o in self.failures]
|
|
133
|
+
super().__init__("\n".join(lines))
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
class DispatchInterrupted(GracefulExit):
|
|
137
|
+
"""A shutdown signal stopped :func:`run_jobs` before every job finished.
|
|
138
|
+
|
|
139
|
+
Computations already launched keep running: deleting them would discard work
|
|
140
|
+
that may be about to succeed. ``in_flight`` names them; ``finished`` holds the
|
|
141
|
+
outcomes reached before the signal.
|
|
142
|
+
"""
|
|
143
|
+
|
|
144
|
+
def __init__(self, signum: int, *, in_flight: Sequence[str], finished: Sequence[Outcome]) -> None:
|
|
145
|
+
super().__init__(signum)
|
|
146
|
+
self.in_flight = tuple(in_flight)
|
|
147
|
+
self.finished = tuple(finished)
|
|
148
|
+
|
|
149
|
+
def __str__(self) -> str:
|
|
150
|
+
"""List the computations left running."""
|
|
151
|
+
running = ", ".join(self.in_flight) or "none"
|
|
152
|
+
return (
|
|
153
|
+
f"interrupted by {self.signame}; {len(self.finished)} job(s) finished, computations left running: {running}"
|
|
154
|
+
)
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def succeeded(computation: Computation) -> bool:
|
|
158
|
+
"""Whether a finished computation ran to a clean exit."""
|
|
159
|
+
return (
|
|
160
|
+
computation.state == ComputationState.Completed
|
|
161
|
+
and computation.end_status not in (ComputationEndStatus.Failed, ComputationEndStatus.Stopped)
|
|
162
|
+
and computation.exit_code in (None, 0)
|
|
163
|
+
)
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def failed_without_results(computation: Computation) -> bool:
|
|
167
|
+
"""Whether a computation looks like it never started: failed, with no results.
|
|
168
|
+
|
|
169
|
+
A capsule that crashes before writing anything matches too, so a relaunch can be
|
|
170
|
+
spent on a real bug; ``start_retries`` bounds that cost.
|
|
171
|
+
"""
|
|
172
|
+
return computation.state == ComputationState.Failed and not computation.has_results
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def key_tag(key: str) -> str:
|
|
176
|
+
"""Return the data-asset tag that carries resume *key*."""
|
|
177
|
+
return f"{_KEY_TAG_PREFIX}{key}"
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
def job_key(run_params: RunParams, *, salt: str = "") -> str:
|
|
181
|
+
"""Return a resume key: 16 hex characters fingerprinting *run_params*.
|
|
182
|
+
|
|
183
|
+
Covers the capsule or pipeline id, version, attached assets, and parameters.
|
|
184
|
+
Attached assets and named parameters are order-insensitive; positional
|
|
185
|
+
parameters are not. Unset fields are dropped, so an SDK release that adds an
|
|
186
|
+
optional field does not change existing keys.
|
|
187
|
+
|
|
188
|
+
Parameters
|
|
189
|
+
----------
|
|
190
|
+
run_params : RunParams
|
|
191
|
+
The launch request.
|
|
192
|
+
salt : str, optional
|
|
193
|
+
Mixed in for what the request cannot show, such as the git commit of a
|
|
194
|
+
non-release capsule whose code changes without its id or version changing.
|
|
195
|
+
"""
|
|
196
|
+
spec = _prune(run_params.to_dict())
|
|
197
|
+
for name, sort_key in (("data_assets", ("id", "mount")), ("named_parameters", ("param_name",))):
|
|
198
|
+
if name in spec:
|
|
199
|
+
spec[name] = sorted(spec[name], key=lambda item: tuple(str(item.get(k)) for k in sort_key))
|
|
200
|
+
return input_fingerprint({"run_params": spec, "salt": salt}).removeprefix("sha256:")[:16]
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
def raise_for_failures(outcomes: Sequence[Outcome]) -> None:
|
|
204
|
+
"""Raise :class:`DispatchFailed` if any outcome failed."""
|
|
205
|
+
failures = [o for o in outcomes if not o.ok]
|
|
206
|
+
if failures:
|
|
207
|
+
raise DispatchFailed(failures)
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
def run_jobs(
|
|
211
|
+
client: CodeOcean,
|
|
212
|
+
jobs: Sequence[Job],
|
|
213
|
+
*,
|
|
214
|
+
max_in_flight: int | None = None,
|
|
215
|
+
launch_interval_s: float = 0.0,
|
|
216
|
+
start_retries: int = 1,
|
|
217
|
+
is_failed_start: Callable[[Computation], bool] = failed_without_results,
|
|
218
|
+
capture_retries: int = 0,
|
|
219
|
+
reuse: bool = False,
|
|
220
|
+
poll_interval_s: float = 60.0,
|
|
221
|
+
timeout_s: float | None = None,
|
|
222
|
+
) -> list[Outcome]:
|
|
223
|
+
"""Launch *jobs*, wait for each, capture its results, and report every outcome.
|
|
224
|
+
|
|
225
|
+
A job's failure never stops the others or raises; check the outcomes, or pass them
|
|
226
|
+
to :func:`raise_for_failures`.
|
|
227
|
+
|
|
228
|
+
Parameters
|
|
229
|
+
----------
|
|
230
|
+
client : CodeOcean
|
|
231
|
+
Authenticated client.
|
|
232
|
+
jobs : Sequence[Job]
|
|
233
|
+
Launched in order, subject to the cap and pacing.
|
|
234
|
+
max_in_flight : int, optional
|
|
235
|
+
Most computations running at once. ``None`` launches everything.
|
|
236
|
+
launch_interval_s : float, optional
|
|
237
|
+
Minimum seconds between launches. Rate limiting is handled by retrying 429s
|
|
238
|
+
regardless; this only spreads the burst.
|
|
239
|
+
start_retries : int, optional
|
|
240
|
+
Relaunches allowed per job when *is_failed_start* matches.
|
|
241
|
+
is_failed_start : callable, optional
|
|
242
|
+
Classifies a failed computation as a platform start failure.
|
|
243
|
+
capture_retries : int, optional
|
|
244
|
+
New capture attempts per job after a capture ends in the Failed state.
|
|
245
|
+
reuse : bool, optional
|
|
246
|
+
Return a Ready asset tagged with a job's key instead of launching it.
|
|
247
|
+
poll_interval_s : float, optional
|
|
248
|
+
Seconds between status checks of each running computation.
|
|
249
|
+
timeout_s : float, optional
|
|
250
|
+
Give up waiting on an attempt after this long. The computation is left
|
|
251
|
+
running, not stopped. ``None`` waits indefinitely.
|
|
252
|
+
|
|
253
|
+
Returns
|
|
254
|
+
-------
|
|
255
|
+
list[Outcome]
|
|
256
|
+
One per job, in input order.
|
|
257
|
+
|
|
258
|
+
Raises
|
|
259
|
+
------
|
|
260
|
+
ValueError
|
|
261
|
+
On duplicate labels or keys, or nonsensical limits.
|
|
262
|
+
DispatchInterrupted
|
|
263
|
+
When a shutdown signal arrives (handlers installed by the caller, see
|
|
264
|
+
:func:`aind_code_ocean_pipeline_utils.process.install_shutdown_handlers`).
|
|
265
|
+
"""
|
|
266
|
+
_validate(jobs, max_in_flight=max_in_flight, start_retries=start_retries, capture_retries=capture_retries)
|
|
267
|
+
return _Dispatcher(
|
|
268
|
+
client,
|
|
269
|
+
jobs,
|
|
270
|
+
max_in_flight=max_in_flight,
|
|
271
|
+
launch_interval_s=launch_interval_s,
|
|
272
|
+
start_retries=start_retries,
|
|
273
|
+
is_failed_start=is_failed_start,
|
|
274
|
+
capture_retries=capture_retries,
|
|
275
|
+
reuse=reuse,
|
|
276
|
+
poll_interval_s=poll_interval_s,
|
|
277
|
+
timeout_s=timeout_s,
|
|
278
|
+
).run()
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
def _validate(jobs: Sequence[Job], *, max_in_flight: int | None, start_retries: int, capture_retries: int) -> None:
|
|
282
|
+
labels = [job.label for job in jobs]
|
|
283
|
+
if len(set(labels)) != len(labels):
|
|
284
|
+
raise ValueError(f"job labels must be unique: {sorted(labels)}")
|
|
285
|
+
keys = [job.key for job in jobs if job.key is not None]
|
|
286
|
+
if len(set(keys)) != len(keys):
|
|
287
|
+
raise ValueError("job keys must be unique within a batch")
|
|
288
|
+
if max_in_flight is not None and max_in_flight < 1:
|
|
289
|
+
raise ValueError(f"max_in_flight must be at least 1, got {max_in_flight}")
|
|
290
|
+
if start_retries < 0 or capture_retries < 0:
|
|
291
|
+
raise ValueError("retry counts cannot be negative")
|
|
292
|
+
|
|
293
|
+
|
|
294
|
+
def _prune(value: Any) -> Any:
|
|
295
|
+
"""Drop ``None`` and empty containers, recursively."""
|
|
296
|
+
if isinstance(value, dict):
|
|
297
|
+
pruned = {k: _prune(v) for k, v in value.items()}
|
|
298
|
+
return {k: v for k, v in pruned.items() if v not in (None, [], {})}
|
|
299
|
+
if isinstance(value, list):
|
|
300
|
+
return [_prune(v) for v in value]
|
|
301
|
+
return value
|
|
302
|
+
|
|
303
|
+
|
|
304
|
+
def _retry_after(exc: CodeOceanError) -> float | None:
|
|
305
|
+
"""Seconds from a ``Retry-After`` header; HTTP-date values are ignored."""
|
|
306
|
+
response = getattr(exc.http_err, "response", None)
|
|
307
|
+
header = response.headers.get("Retry-After") if response is not None else None
|
|
308
|
+
try:
|
|
309
|
+
return max(0.0, float(header)) if header is not None else None
|
|
310
|
+
except ValueError:
|
|
311
|
+
return None
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
def _retry_delay(exc: Exception, attempt: int, *, idempotent: bool) -> float | None:
|
|
315
|
+
"""Seconds to wait before retrying the call that raised *exc*, or ``None`` if it must not be."""
|
|
316
|
+
backoff = min(_BACKOFF_CAP_S, _BACKOFF_BASE_S * 2**attempt)
|
|
317
|
+
if isinstance(exc, CodeOceanError):
|
|
318
|
+
if exc.status_code == _RATE_LIMITED:
|
|
319
|
+
after = _retry_after(exc)
|
|
320
|
+
return backoff if after is None else after
|
|
321
|
+
return backoff if idempotent and exc.status_code in _TRANSIENT_SERVER else None
|
|
322
|
+
return backoff if idempotent and isinstance(exc, requests.RequestException) else None
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
def _is_ambiguous(exc: Exception) -> bool:
|
|
326
|
+
"""Whether a failed POST may still have taken effect."""
|
|
327
|
+
if isinstance(exc, CodeOceanError):
|
|
328
|
+
return exc.status_code in _TRANSIENT_SERVER
|
|
329
|
+
return isinstance(exc, requests.RequestException)
|
|
330
|
+
|
|
331
|
+
|
|
332
|
+
def _describe(computation: Computation) -> str:
|
|
333
|
+
return (
|
|
334
|
+
f"computation {computation.id} ended state={computation.state} end_status={computation.end_status} "
|
|
335
|
+
f"exit_code={computation.exit_code} has_results={computation.has_results}"
|
|
336
|
+
)
|
|
337
|
+
|
|
338
|
+
|
|
339
|
+
class _Phase(Enum):
|
|
340
|
+
PENDING = auto()
|
|
341
|
+
RUNNING = auto()
|
|
342
|
+
CAPTURING = auto()
|
|
343
|
+
DONE = auto()
|
|
344
|
+
|
|
345
|
+
|
|
346
|
+
@dataclass
|
|
347
|
+
class _Slot:
|
|
348
|
+
job: Job
|
|
349
|
+
phase: _Phase = _Phase.PENDING
|
|
350
|
+
next_at: float = 0.0
|
|
351
|
+
computation_ids: list[str] = field(default_factory=list)
|
|
352
|
+
launched_at: float = 0.0
|
|
353
|
+
asset_id: str | None = None
|
|
354
|
+
# Consecutive retryable errors on the current call; drives backoff.
|
|
355
|
+
api_errors: int = 0
|
|
356
|
+
start_retries_used: int = 0
|
|
357
|
+
capture_retries_used: int = 0
|
|
358
|
+
reuse_checked: bool = False
|
|
359
|
+
outcome: Outcome | None = None
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
class _Dispatcher:
|
|
363
|
+
def __init__(
|
|
364
|
+
self,
|
|
365
|
+
client: CodeOcean,
|
|
366
|
+
jobs: Sequence[Job],
|
|
367
|
+
*,
|
|
368
|
+
max_in_flight: int | None,
|
|
369
|
+
launch_interval_s: float,
|
|
370
|
+
start_retries: int,
|
|
371
|
+
is_failed_start: Callable[[Computation], bool],
|
|
372
|
+
capture_retries: int,
|
|
373
|
+
reuse: bool,
|
|
374
|
+
poll_interval_s: float,
|
|
375
|
+
timeout_s: float | None,
|
|
376
|
+
) -> None:
|
|
377
|
+
self.client = client
|
|
378
|
+
self.slots = [_Slot(job) for job in jobs]
|
|
379
|
+
self.max_in_flight = max_in_flight
|
|
380
|
+
self.launch_interval_s = launch_interval_s
|
|
381
|
+
self.start_retries = start_retries
|
|
382
|
+
self.is_failed_start = is_failed_start
|
|
383
|
+
self.capture_retries = capture_retries
|
|
384
|
+
self.reuse = reuse
|
|
385
|
+
self.poll_interval_s = poll_interval_s
|
|
386
|
+
self.timeout_s = timeout_s
|
|
387
|
+
self.next_launch_at = 0.0
|
|
388
|
+
|
|
389
|
+
def run(self) -> list[Outcome]:
|
|
390
|
+
try:
|
|
391
|
+
while any(slot.phase is not _Phase.DONE for slot in self.slots):
|
|
392
|
+
check_shutdown()
|
|
393
|
+
self._check_reuse()
|
|
394
|
+
self._launch_ready()
|
|
395
|
+
for slot in self.slots:
|
|
396
|
+
if slot.phase is _Phase.RUNNING and slot.next_at <= _now():
|
|
397
|
+
self._poll_computation(slot)
|
|
398
|
+
elif slot.phase is _Phase.CAPTURING and slot.next_at <= _now():
|
|
399
|
+
self._poll_asset(slot)
|
|
400
|
+
self._sleep_until_due()
|
|
401
|
+
except GracefulExit as exc:
|
|
402
|
+
in_flight = [slot.computation_ids[-1] for slot in self.slots if slot.phase is _Phase.RUNNING]
|
|
403
|
+
finished = [slot.outcome for slot in self.slots if slot.outcome is not None]
|
|
404
|
+
logger.error("interrupted; computations left running: %s", ", ".join(in_flight) or "none")
|
|
405
|
+
raise DispatchInterrupted(exc.signum, in_flight=in_flight, finished=finished) from exc
|
|
406
|
+
return [slot.outcome for slot in self.slots if slot.outcome is not None]
|
|
407
|
+
|
|
408
|
+
# -- scheduling ---------------------------------------------------------
|
|
409
|
+
|
|
410
|
+
def _in_flight(self) -> int:
|
|
411
|
+
return sum(slot.phase is _Phase.RUNNING for slot in self.slots)
|
|
412
|
+
|
|
413
|
+
def _launch_ready(self) -> None:
|
|
414
|
+
for slot in self.slots:
|
|
415
|
+
if slot.phase is not _Phase.PENDING or slot.next_at > _now():
|
|
416
|
+
continue
|
|
417
|
+
if self.reuse and slot.job.key is not None and not slot.reuse_checked:
|
|
418
|
+
continue # still waiting on its reuse search
|
|
419
|
+
if self.max_in_flight is not None and self._in_flight() >= self.max_in_flight:
|
|
420
|
+
return
|
|
421
|
+
if _now() < self.next_launch_at:
|
|
422
|
+
return
|
|
423
|
+
self._launch(slot)
|
|
424
|
+
|
|
425
|
+
def _sleep_until_due(self) -> None:
|
|
426
|
+
now = _now()
|
|
427
|
+
pending = any(slot.phase is _Phase.PENDING for slot in self.slots)
|
|
428
|
+
due = [slot.next_at for slot in self.slots if slot.phase is not _Phase.DONE and slot.next_at > now]
|
|
429
|
+
if pending and self.next_launch_at > now:
|
|
430
|
+
due.append(self.next_launch_at)
|
|
431
|
+
if not due:
|
|
432
|
+
return
|
|
433
|
+
_sleep(min(_MAX_SLEEP_S, min(due) - now))
|
|
434
|
+
|
|
435
|
+
def _finish(self, slot: _Slot, status: Status, *, error: str | None = None) -> None:
|
|
436
|
+
slot.phase = _Phase.DONE
|
|
437
|
+
slot.outcome = Outcome(
|
|
438
|
+
job=slot.job,
|
|
439
|
+
status=status,
|
|
440
|
+
asset_id=slot.asset_id,
|
|
441
|
+
computation_ids=tuple(slot.computation_ids),
|
|
442
|
+
error=error,
|
|
443
|
+
)
|
|
444
|
+
if error is None:
|
|
445
|
+
logger.info("%s: %s", slot.job.label, status)
|
|
446
|
+
else:
|
|
447
|
+
logger.error("%s: failed: %s", slot.job.label, error)
|
|
448
|
+
|
|
449
|
+
def _call(self, slot: _Slot, what: str, fn: Callable[[], _T], *, idempotent: bool) -> _T | None:
|
|
450
|
+
"""Run one API call for *slot*; on error, schedule a retry or fail the job.
|
|
451
|
+
|
|
452
|
+
Returns ``None`` when the call did not succeed; the slot is then either
|
|
453
|
+
rescheduled (``next_at`` set) or done.
|
|
454
|
+
"""
|
|
455
|
+
try:
|
|
456
|
+
result = fn()
|
|
457
|
+
except (CodeOceanError, requests.RequestException) as exc:
|
|
458
|
+
delay = _retry_delay(exc, slot.api_errors, idempotent=idempotent)
|
|
459
|
+
if delay is None:
|
|
460
|
+
note = " It may have taken effect anyway; check before rerunning." if _is_ambiguous(exc) else ""
|
|
461
|
+
self._finish(slot, "failed", error=f"{what} failed: {exc}.{note}")
|
|
462
|
+
return None
|
|
463
|
+
slot.api_errors += 1
|
|
464
|
+
slot.next_at = _now() + delay
|
|
465
|
+
logger.warning("%s: %s failed (%s); retrying in %.0f s", slot.job.label, what, exc, delay)
|
|
466
|
+
return None
|
|
467
|
+
slot.api_errors = 0
|
|
468
|
+
return result
|
|
469
|
+
|
|
470
|
+
# -- phases -------------------------------------------------------------
|
|
471
|
+
|
|
472
|
+
def _check_reuse(self) -> None:
|
|
473
|
+
"""Resolve pending keyed jobs to existing captures, without using launch capacity."""
|
|
474
|
+
if not self.reuse:
|
|
475
|
+
return
|
|
476
|
+
for slot in self.slots:
|
|
477
|
+
if slot.phase is not _Phase.PENDING or slot.reuse_checked or slot.job.key is None:
|
|
478
|
+
continue
|
|
479
|
+
if slot.next_at > _now():
|
|
480
|
+
continue
|
|
481
|
+
tag = key_tag(slot.job.key)
|
|
482
|
+
params = DataAssetSearchParams(query=f"tag:{tag}", limit=100)
|
|
483
|
+
try:
|
|
484
|
+
results = self.client.data_assets.search_data_assets(params)
|
|
485
|
+
except (CodeOceanError, requests.RequestException) as exc:
|
|
486
|
+
delay = _retry_delay(exc, slot.api_errors, idempotent=True)
|
|
487
|
+
if delay is not None:
|
|
488
|
+
slot.api_errors += 1
|
|
489
|
+
slot.next_at = _now() + delay
|
|
490
|
+
continue
|
|
491
|
+
# Reuse is an optimisation: launching is still correct.
|
|
492
|
+
logger.warning("%s: reuse search failed (%s); launching", slot.job.label, exc)
|
|
493
|
+
slot.reuse_checked = True
|
|
494
|
+
continue
|
|
495
|
+
slot.api_errors = 0
|
|
496
|
+
slot.reuse_checked = True
|
|
497
|
+
# Tag search may tokenize, so confirm the exact tag client-side.
|
|
498
|
+
matches: list[DataAsset] = [
|
|
499
|
+
asset for asset in results.results if asset.state == DataAssetState.Ready and tag in (asset.tags or [])
|
|
500
|
+
]
|
|
501
|
+
if matches:
|
|
502
|
+
slot.asset_id = max(matches, key=lambda asset: asset.created).id
|
|
503
|
+
logger.info("%s: reusing captured asset %s", slot.job.label, slot.asset_id)
|
|
504
|
+
self._finish(slot, "reused")
|
|
505
|
+
|
|
506
|
+
def _launch(self, slot: _Slot) -> None:
|
|
507
|
+
computation = self._call(
|
|
508
|
+
slot, "launch", lambda: self.client.computations.run_capsule(slot.job.run_params), idempotent=False
|
|
509
|
+
)
|
|
510
|
+
if computation is None:
|
|
511
|
+
if slot.phase is _Phase.PENDING:
|
|
512
|
+
# The rate limit is per user, so hold every launch, not just this one.
|
|
513
|
+
self.next_launch_at = max(self.next_launch_at, slot.next_at)
|
|
514
|
+
return
|
|
515
|
+
slot.computation_ids.append(computation.id)
|
|
516
|
+
slot.phase = _Phase.RUNNING
|
|
517
|
+
slot.launched_at = _now()
|
|
518
|
+
slot.next_at = slot.launched_at + self.poll_interval_s
|
|
519
|
+
self.next_launch_at = slot.launched_at + self.launch_interval_s
|
|
520
|
+
logger.info("%s: launched %s (attempt %d)", slot.job.label, computation.id, len(slot.computation_ids))
|
|
521
|
+
|
|
522
|
+
def _poll_computation(self, slot: _Slot) -> None:
|
|
523
|
+
computation_id = slot.computation_ids[-1]
|
|
524
|
+
computation = self._call(
|
|
525
|
+
slot, "status check", lambda: self.client.computations.get_computation(computation_id), idempotent=True
|
|
526
|
+
)
|
|
527
|
+
if computation is None:
|
|
528
|
+
return
|
|
529
|
+
if computation.state not in (ComputationState.Completed, ComputationState.Failed):
|
|
530
|
+
if self.timeout_s is not None and _now() - slot.launched_at > self.timeout_s:
|
|
531
|
+
self._finish(
|
|
532
|
+
slot,
|
|
533
|
+
"failed",
|
|
534
|
+
error=f"timed out after {self.timeout_s:.0f} s; computation {computation_id} left running",
|
|
535
|
+
)
|
|
536
|
+
else:
|
|
537
|
+
slot.next_at = _now() + self.poll_interval_s
|
|
538
|
+
return
|
|
539
|
+
if succeeded(computation):
|
|
540
|
+
if slot.job.capture is None:
|
|
541
|
+
self._finish(slot, "completed")
|
|
542
|
+
else:
|
|
543
|
+
self._start_capture(slot, computation)
|
|
544
|
+
elif self.is_failed_start(computation) and slot.start_retries_used < self.start_retries:
|
|
545
|
+
slot.start_retries_used += 1
|
|
546
|
+
slot.phase = _Phase.PENDING
|
|
547
|
+
slot.next_at = _now()
|
|
548
|
+
logger.warning("%s: %s; relaunching", slot.job.label, _describe(computation))
|
|
549
|
+
else:
|
|
550
|
+
self._finish(slot, "failed", error=_describe(computation))
|
|
551
|
+
|
|
552
|
+
def _start_capture(self, slot: _Slot, computation: Computation) -> None:
|
|
553
|
+
assert slot.job.capture is not None
|
|
554
|
+
try:
|
|
555
|
+
params = slot.job.capture(computation)
|
|
556
|
+
# Broad on purpose: one job's bad callback must not orphan the others' computations.
|
|
557
|
+
except Exception as exc:
|
|
558
|
+
logger.exception("%s: capture callback raised", slot.job.label)
|
|
559
|
+
self._finish(slot, "failed", error=f"capture callback raised for computation {computation.id}: {exc!r}")
|
|
560
|
+
return
|
|
561
|
+
if slot.job.key is not None and key_tag(slot.job.key) not in params.tags:
|
|
562
|
+
params = dataclasses.replace(params, tags=[*params.tags, key_tag(slot.job.key)])
|
|
563
|
+
asset = self._call(
|
|
564
|
+
slot,
|
|
565
|
+
f"capture of computation {computation.id}",
|
|
566
|
+
lambda: self.client.data_assets.create_data_asset(params),
|
|
567
|
+
idempotent=False,
|
|
568
|
+
)
|
|
569
|
+
# On a retryable error the slot stays RUNNING; the next poll sees the same
|
|
570
|
+
# finished computation and comes back here.
|
|
571
|
+
if asset is None:
|
|
572
|
+
return
|
|
573
|
+
slot.asset_id = asset.id
|
|
574
|
+
slot.phase = _Phase.CAPTURING
|
|
575
|
+
slot.next_at = _now() + _ASSET_POLL_S
|
|
576
|
+
logger.info("%s: capturing %s as %s", slot.job.label, computation.id, asset.id)
|
|
577
|
+
|
|
578
|
+
def _poll_asset(self, slot: _Slot) -> None:
|
|
579
|
+
asset_id = slot.asset_id
|
|
580
|
+
assert asset_id is not None
|
|
581
|
+
asset = self._call(
|
|
582
|
+
slot,
|
|
583
|
+
f"status check of asset {asset_id}",
|
|
584
|
+
lambda: self.client.data_assets.get_data_asset(asset_id),
|
|
585
|
+
idempotent=True,
|
|
586
|
+
)
|
|
587
|
+
if asset is None:
|
|
588
|
+
return
|
|
589
|
+
if asset.state == DataAssetState.Ready:
|
|
590
|
+
self._finish(slot, "captured")
|
|
591
|
+
elif asset.state == DataAssetState.Failed:
|
|
592
|
+
if slot.capture_retries_used < self.capture_retries:
|
|
593
|
+
slot.capture_retries_used += 1
|
|
594
|
+
slot.asset_id = None
|
|
595
|
+
slot.phase = _Phase.RUNNING
|
|
596
|
+
slot.next_at = _now()
|
|
597
|
+
logger.warning("%s: capture %s failed; capturing again", slot.job.label, asset_id)
|
|
598
|
+
else:
|
|
599
|
+
self._finish(
|
|
600
|
+
slot, "failed", error=f"capture {asset_id} of computation {slot.computation_ids[-1]} failed"
|
|
601
|
+
)
|
|
602
|
+
else:
|
|
603
|
+
slot.next_at = _now() + _ASSET_POLL_S
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: aind-codeocean-dispatch
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Launch, monitor, and capture batches of Code Ocean capsule runs
|
|
5
|
+
Project-URL: Repository, https://github.com/AllenNeuralDynamics/aind-codeocean-dispatch/
|
|
6
|
+
Project-URL: Issues, https://github.com/AllenNeuralDynamics/aind-codeocean-dispatch/issues
|
|
7
|
+
Project-URL: Changelog, https://github.com/AllenNeuralDynamics/aind-codeocean-dispatch/blob/main/CHANGELOG.md
|
|
8
|
+
Author: Allen Institute for Neural Dynamics
|
|
9
|
+
Author-email: Galen Lynch <galen@galenlynch.com>
|
|
10
|
+
License-Expression: MIT
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Requires-Python: >=3.11
|
|
14
|
+
Requires-Dist: aind-code-ocean-pipeline-utils>=0.7.1
|
|
15
|
+
Requires-Dist: codeocean>=0.17
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# AIND Code Ocean dispatch
|
|
19
|
+
|
|
20
|
+

|
|
21
|
+
[](https://pypi.org/project/aind-codeocean-dispatch/)
|
|
22
|
+
[](https://github.com/semantic-release/semantic-release)
|
|
23
|
+
[](LICENSE)
|
|
24
|
+
[](https://github.com/astral-sh/ruff)
|
|
25
|
+
[](https://github.com/astral-sh/uv)
|
|
26
|
+
[](https://github.com/copier-org/copier)
|
|
27
|
+
|
|
28
|
+
Launch, monitor, and capture batches of Code Ocean capsule runs
|
|
29
|
+
|
|
30
|
+
## Installation
|
|
31
|
+
|
|
32
|
+
If you choose to clone the repository, you can install the package by running the following command from the root directory of the repository:
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
pip install .
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
Otherwise, you can use pip:
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
pip install aind-codeocean-dispatch
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
To develop the code, run:
|
|
47
|
+
```bash
|
|
48
|
+
uv sync
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
## Development
|
|
52
|
+
|
|
53
|
+
Please test your changes using the full linting and testing suite:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
./scripts/run_linters_and_checks.sh -c
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Or run individual commands:
|
|
60
|
+
```bash
|
|
61
|
+
uv run --frozen ruff format # Code formatting
|
|
62
|
+
uv run --frozen ruff check # Linting
|
|
63
|
+
uv run --frozen mypy # Type checking
|
|
64
|
+
uv run --frozen interrogate -v src # Documentation coverage
|
|
65
|
+
uv run --frozen codespell --check-filenames # Spell checking
|
|
66
|
+
uv run --frozen pytest --cov aind_codeocean_dispatch # Tests with coverage
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
## License
|
|
72
|
+
|
|
73
|
+
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
aind_codeocean_dispatch/__init__.py,sha256=NcdbegNPvYKWVRArhlkdEL7HP9kof1OOiH8aRcrcszY,660
|
|
2
|
+
aind_codeocean_dispatch/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
|
|
3
|
+
aind_codeocean_dispatch/runner.py,sha256=hMeVmd5v5bWCF_2CQYkNV5HzD_Ywn5eyORzNqQGkHA4,24178
|
|
4
|
+
aind_codeocean_dispatch-0.2.0.dist-info/METADATA,sha256=vPbl6NmG66g2-gKqzmGsJkPIeP-wvb5K7ZaacsFL8Rc,2797
|
|
5
|
+
aind_codeocean_dispatch-0.2.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
|
|
6
|
+
aind_codeocean_dispatch-0.2.0.dist-info/licenses/LICENSE,sha256=zMZxYbR1i_XxgyJP_rPSWKcxY7QtyxUOzbIpowHNjCs,1093
|
|
7
|
+
aind_codeocean_dispatch-0.2.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Allen Institute for Neural Dynamics
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
22
|
+
|