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.
@@ -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
+ ![CI](https://github.com/AllenNeuralDynamics/aind-codeocean-dispatch/actions/workflows/ci-call.yml/badge.svg)
21
+ [![PyPI - Version](https://img.shields.io/pypi/v/aind-codeocean-dispatch)](https://pypi.org/project/aind-codeocean-dispatch/)
22
+ [![semantic-release: angular](https://img.shields.io/badge/semantic--release-angular-e10079?logo=semantic-release)](https://github.com/semantic-release/semantic-release)
23
+ [![License](https://img.shields.io/badge/license-MIT-brightgreen)](LICENSE)
24
+ [![ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
25
+ [![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
26
+ [![Copier](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/copier-org/copier/master/img/badge/badge-grayscale-inverted-border.json)](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,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: hatchling 1.32.4
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
@@ -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
+