"""The orchestrator: a single-threaded pass loop over durable state (2.2).
This is the seam where the three planes meet -- ``StateStore`` (coordination),
``Stage`` (what), ``ComputeBackend`` (where) -- and it is deliberately the
*only* place they meet. It is constructed from plain objects (an ordered list
of stages, a ``{stage_name: backend}`` mapping): it never reads config (that
wiring arrives at 2.3 via ``build_runtime``), never imports the CLI layer, and
never reaches inside a stage or a backend beyond their contracts.
The concurrency model (locked 2026-07-08): **parallelism is jobs in flight,
never threads.** Jobs run as OS processes wherever the backend puts them; the
orchestrator itself is one thread running short passes over the durable state
in SQLite. Every fact needed to resume -- who is where, which handles are
live -- is in the store, so the process can be killed at any point and a new
one continues the same run.
Each ``run_pass()`` is one non-blocking sweep, four steps:
1. **Stale-lease release** -- leases older than ``lease_timeout_s`` are
returned to their queue (materials owned by a recorded in-flight job are
never touched; the store guarantees it).
2. **Check phase** -- every job in the in-flight ledger is polled through its
recorded backend. Still-running jobs are left recorded (the pass never
waits). A terminal job is fetched, judged by its stage
(``interpret_result`` then ``validate_outputs``), attributed to the store
(``complete`` / ``fail`` -- or, for a ``fan_out`` stage, children adopted
from the manifest and the parent retired as ``EXPANDED``), its handle
released (best-effort), and its ledger record cleared.
3. **Finalize** -- materials that completed the *configured* funnel's last
stage are retired as first-class ``DONE`` via ``StateStore.finalize``
(idempotent, so a crash between check and finalize just retries next
pass). A run stopped mid-funnel (``stop_after`` or ``--only`` naming an
earlier stage) never finalizes: its materials rest at that stage's output
state, live and claimable, so a later run continues them.
4. **Submit phase** -- per stage in funnel order up to ``stop_after``:
backlog gate (a stage whose input state is empty is never considered) ->
health gate (an unhealthy backend pauses the stage; it resumes by itself
on a later pass) -> claim up to the backend's free slots
(``max_parallel`` minus its recorded in-flight jobs) -> per item:
``validate_inputs`` -> ``build_jobspec`` -> ``submit`` -> ``record_job``.
The handle is recorded *before* the pass ends, so a crash after submission
re-polls the job instead of resubmitting it; a job whose record cannot be
written is released (abandoned) so nothing runs untracked.
Failure attribution happens at exactly two helpers (the single attribution
point of the design):
* ``_complete`` -- a store fault while completing leaves the item leased; the
lease timeout recovers it. Never raises.
* ``_fail`` -- the verdict-vs-fault doctrine is applied here and only here: a
``StageVerdict`` with ``terminal=True`` is a deterministic science verdict a
re-run cannot change, so it bypasses the retry budget (failed with budget
1); everything else is a fault that spends one unit of ``retry_budget`` and
requeues while budget remains. ``StageInputError`` is budgeted the same way
(a missing input can be a not-yet-mounted disk). A raised ``StageError``
(the stage as a whole cannot run) fails *nothing*: the stage is skipped
this pass and its claimed items recover by lease timeout.
**A pass never raises.** Every phase guards its store/backend/stage calls;
an unexpected error is logged (stdlib ``logging``, logger
``matmelt.core.orchestrator``) and the pass moves on. ``run()`` loops
passes with a sleep between them and returns cleanly on ``KeyboardInterrupt``
-- in-flight jobs stay recorded and the next run re-attaches to them.
"""
from __future__ import annotations
import contextlib
import json
import logging
import shutil
import time
from dataclasses import asdict
from pathlib import Path
from typing import TYPE_CHECKING
from matmelt.core.jobspec import BackendError, Handle
from matmelt.core.stages.base import (
CHILDREN_DIR,
CHILDREN_MANIFEST,
Stage,
StageError,
StageInputError,
StageOutputError,
StageVerdict,
WorkItem,
)
from matmelt.core.stages.registry import STAGE_OUTPUT_STATE
from matmelt.core.state import (
LIFECYCLE,
MaterialState,
StateStore,
is_terminal,
)
if TYPE_CHECKING:
from collections.abc import Iterable, Mapping, Sequence
from matmelt.backends.base import ComputeBackend
from matmelt.core.state import InFlightJob
_log = logging.getLogger(__name__)
[docs]
class Orchestrator:
"""Drives materials through the funnel: check in-flight work, submit new work.
Parameters
----------
store
The durable coordination plane. The orchestrator is stateless beyond
it: two consecutive processes given the same store continue one run.
stages
The pipeline stages **in funnel order**. Each stage consumes the
previous stage's output state (the first consumes ``NEW``), and its
completion lands the material in its own output state: a name in
``STAGE_OUTPUT_STATE`` owns that lifecycle state by name -- so
omitting an optional stage never shifts its neighbours' states --
while any other name takes the next chain state positionally.
Output states must advance strictly (which bounds unmapped stages
at ``len(LIFECYCLE) - 1``). Names must be unique; the reserved
name ``"finalize"`` is not a legal stage name.
backends
``{stage_name: backend}`` -- which backend runs each stage's jobs.
Every stage must be wired (missing wiring is a construction error,
not a runtime surprise). Two stages may share one backend instance;
its free slots are then shared too.
work_dir
Root of the per-item working directories: item ``m`` works in
``work_dir/m`` (created at first submission). This is the only place
the id -> directory convention lives.
stop_after
Name of the last stage this run executes (a first-class knob, not a
debug hack): stages after it are never considered. Materials that
complete it **rest at its output state** -- live, claimable by the
next stage in a later run; they are finalized to terminal ``DONE``
only when the stop stage is the configured last stage (where resting
and finishing coincide). ``None`` means the full configured funnel.
only
Restrict the **submit phase** to exactly this stage (single-stage
runs on demand). The check phase is untouched -- any recorded
in-flight job still settles -- and finalize runs only when the only
stage is the configured last stage: a mid-funnel only-run must
never retire its output as ``DONE`` (a later run's stages still
consume it). Mutually exclusive with ``stop_after``.
materials
Restrict the run to this **selected lineage**: only materials whose
id -- or any ancestor's id, walking ``parent_id`` -- is in the set
are claimed by the submit phase or retired by finalize. Fan-out
children of a selected parent are therefore eligible automatically,
including on a rerun (the ancestry walk needs no memory of the
adoption). The check phase is untouched: any recorded in-flight job
still settles, whoever it belongs to. ``None`` = unrestricted; an
empty set is a construction error (it would silently do nothing).
retry_budget
Per-stage fault budget handed to ``StateStore.fail`` -- a faulted item
is requeued until the budget is exhausted, then goes terminal
``FAILED``. Terminal verdicts bypass it entirely.
lease_timeout_s
Age beyond which a lease is considered orphaned and released at the
start of every pass (in-flight-ledger materials excepted).
"""
def __init__(
self,
store: StateStore,
stages: Sequence[Stage],
backends: Mapping[str, ComputeBackend],
work_dir: str | Path,
*,
stop_after: str | None = None,
only: str | None = None,
materials: frozenset[str] | None = None,
retry_budget: int = 3,
lease_timeout_s: float = 900.0,
) -> None:
if not stages:
raise ValueError("stages must not be empty")
if materials is not None and not materials:
raise ValueError("materials must not be empty (None = unrestricted)")
names = [s.name for s in stages]
if len(set(names)) != len(names):
raise ValueError(f"duplicate stage names: {names}")
unwired = [n for n in names if n not in backends]
if unwired:
raise ValueError(f"stages with no backend wired: {unwired}")
if stop_after is not None and stop_after not in names:
raise ValueError(f"stop_after {stop_after!r} names no stage; stages: {names}")
if only is not None:
if only not in names:
raise ValueError(f"only {only!r} names no stage; stages: {names}")
if stop_after is not None:
raise ValueError("only and stop_after are mutually exclusive")
# Derive each stage's output state once, at construction. A name in
# STAGE_OUTPUT_STATE owns that lifecycle state; any other name takes
# the next chain state positionally. Outputs must advance strictly --
# a mapped stage out of funnel order, or more unmapped stages than
# chain states, is a construction error, not a runtime surprise.
outputs: list[MaterialState] = []
cursor = 0 # LIFECYCLE index of the previous stage's output (NEW to start)
for s in stages:
mapped = STAGE_OUTPUT_STATE.get(s.name)
if mapped is not None:
out = LIFECYCLE.index(mapped)
if out <= cursor:
raise ValueError(
f"stage {s.name!r} outputs {mapped.value}, which is not "
f"forward of {LIFECYCLE[cursor].value}; stages must be "
"in funnel order"
)
else:
out = cursor + 1
if out >= len(LIFECYCLE):
raise ValueError(
f"at most {len(LIFECYCLE) - 1} stages fit the lifecycle; got {len(stages)}"
)
outputs.append(LIFECYCLE[out])
cursor = out
self._store = store
self._backends = dict(backends)
self._work_dir = Path(work_dir)
self._retry_budget = int(retry_budget)
self._lease_timeout_s = float(lease_timeout_s)
stop_index = names.index(stop_after) if stop_after is not None else len(stages) - 1
#: The stages this run executes, in funnel order.
self._active: list[Stage] = list(stages[: stop_index + 1])
self._stages_by_name: dict[str, Stage] = {s.name: s for s in self._active}
#: Each active stage's input state: the previous stage's output (first = NEW).
self._input_states: list[MaterialState] = [LIFECYCLE[0], *outputs[:stop_index]]
#: Stage name -> the state its completion lands (``complete``'s target).
self._stage_outputs: dict[str, MaterialState] = {
s.name: outputs[i] for i, s in enumerate(self._active)
}
#: Completing the stop stage lands a material here; when the stop
#: stage is the configured last stage, finalize retires it as DONE.
self._stop_output: MaterialState = outputs[stop_index]
#: Single-stage restriction for the submit phase (None = all active stages).
self._only: str | None = only
#: Finalize runs only when this run executes through the configured
#: last stage. A mid-funnel stop -- ``--only`` or ``stop_after``
#: naming an earlier stage -- must never retire its output as
#: ``DONE``: later stages still consume it, in a later run.
last_configured = names[-1]
self._finalize_enabled: bool = (only is None or only == last_configured) and (
stop_after is None or stop_after == last_configured
)
#: The selected lineage (None = unrestricted); see ``_is_selected``.
self._materials: frozenset[str] | None = (
frozenset(materials) if materials is not None else None
)
#: Memoized eligibility verdicts. Safe for the process lifetime:
#: ``parent_id`` is immutable (enqueue is first-write-wins) and the
#: selection is fixed at construction.
self._eligibility: dict[str, bool] = {}
# -- Feeding ---------------------------------------------------------------
[docs]
def seed(self, material_ids: Iterable[str]) -> int:
"""Enqueue ``material_ids`` as ``NEW`` roots.
Parameters
----------
material_ids
The seed ids to offer, typically from
``core.seeds.element_combinations``.
Returns
-------
int
How many ids were offered -- not how many were new. The call is
idempotent end to end, since ``enqueue`` never regresses an
existing material, so replaying a seed list, or seeding while a run
is in progress, is always safe.
"""
n = 0
for mid in material_ids:
self._store.enqueue(mid)
n += 1
return n
# -- The pass loop -----------------------------------------------------------
[docs]
def run(self, *, max_passes: int | None = None, pass_interval_s: float = 1.0) -> None:
"""Run passes until a limit is reached or the user interrupts.
Parameters
----------
max_passes
How many passes to run. ``None`` runs until interrupted.
pass_interval_s
How long to sleep between passes.
Notes
-----
``KeyboardInterrupt`` is the supported way to stop an open-ended run:
it is caught, logged, and the method returns cleanly. In-flight jobs
stay recorded in the ledger, and the next ``run()`` re-attaches to them
instead of resubmitting.
"""
passes = 0
try:
while max_passes is None or passes < max_passes:
passes += 1
self.run_pass()
if max_passes is not None and passes >= max_passes:
break
time.sleep(pass_interval_s)
except KeyboardInterrupt:
_log.info(
"interrupted during pass %d; in-flight jobs stay recorded and "
"re-attach on the next run",
passes,
)
[docs]
def run_pass(self) -> dict[str, int]:
"""Sweep once, without blocking: release stale leases, check, finalize, submit.
Never raises.
Returns
-------
dict
What the pass did, with every key always present and zero when
nothing of that kind happened -- callers may index the summary
without checking:
``released``
Stale leases freed.
``completed``, ``failed``, ``requeued``
Materials attributed by the check phase or by submit-time input
validation.
``expanded``, ``children``
Fan-out parents retired as ``EXPANDED``, and the materials
those parents adopted.
``finalized``
Materials retired as ``DONE``.
``submitted``
New jobs recorded in the ledger.
"""
summary = {
"released": 0,
"completed": 0,
"failed": 0,
"requeued": 0,
"expanded": 0,
"children": 0,
"finalized": 0,
"submitted": 0,
}
try:
summary["released"] = self._store.release_stale_leases(self._lease_timeout_s)
except Exception:
_log.exception("stale-lease release failed; continuing the pass")
self._check_phase(summary)
self._finalize_phase(summary)
self._submit_phase(summary)
return summary
# -- Check phase -------------------------------------------------------------
def _check_phase(self, summary: dict[str, int]) -> None:
"""Poll every recorded in-flight job and settle the terminal ones.
Parameters
----------
summary
The pass summary, updated in place.
"""
try:
jobs = self._store.list_jobs()
except Exception:
_log.exception("check phase: could not list in-flight jobs; skipping")
return
for job in jobs:
try:
self._check_one_job(job, summary)
except Exception:
_log.exception(
"check phase: unexpected error on job %s; left recorded", job.job_key
)
def _check_one_job(self, job: InFlightJob, summary: dict[str, int]) -> None:
"""Poll one ledger record; if terminal, judge, attribute, release, clear.
Parameters
----------
job
The ledger record to resolve.
summary
The pass summary, updated in place.
Notes
-----
A job is deliberately *left recorded* when it cannot be resolved right
now: its stage is not part of this run, its stage is wired to a
different backend than the one that issued the handle, or the poll
itself failed. That last case is transient by assumption -- the durable
handle is exactly what makes waiting safe.
"""
stage = self._stages_by_name.get(job.stage)
if stage is None:
_log.warning(
"in-flight job %s belongs to stage %r, which this run does not "
"execute; left recorded",
job.job_key,
job.stage,
)
return
backend = self._backends[job.stage]
if job.backend_name != backend.name:
_log.warning(
"in-flight job %s was submitted to backend %r but stage %r is now "
"wired to %r; left recorded",
job.job_key,
job.backend_name,
job.stage,
backend.name,
)
return
raw = job.handle
handle = Handle(
backend_name=str(raw.get("backend_name", job.backend_name)),
job_id=str(raw.get("job_id", "")),
extra=dict(raw.get("extra") or {}),
)
try:
status = backend.poll(handle)
except BackendError as exc:
_log.warning(
"poll failed for job %s (%s); treated as transient, re-polled next pass",
job.job_key,
exc,
)
return
if not status.is_terminal:
return # still running: the pass never waits on it
try:
result = backend.fetch(handle)
except BackendError as exc:
for mid in job.material_ids:
self._fail(mid, job.stage, f"fetch failed: {exc}", terminal=False, summary=summary)
else:
for mid in job.material_ids:
item = self._item(mid)
try:
verdict = stage.interpret_result(item, result)
except Exception as exc:
verdict = StageVerdict(ok=False, reason=f"interpret_result raised: {exc}")
if verdict.ok:
try:
stage.validate_outputs(item)
except StageOutputError as exc:
self._fail(mid, job.stage, str(exc), terminal=False, summary=summary)
else:
if stage.fan_out:
self._settle_fan_out(stage, item, summary)
else:
self._complete(mid, job.stage, summary)
else:
self._fail(
mid,
job.stage,
verdict.reason or "stage verdict: not ok",
terminal=verdict.terminal,
summary=summary,
)
# Cleanup is best-effort; it must never mask the attribution above.
with contextlib.suppress(BackendError):
backend.release(handle)
try:
self._store.clear_job(job.job_key)
except Exception:
_log.exception("could not clear settled job %s; retried next pass", job.job_key)
# -- Fan-out (1 -> N): adopt children, retire the parent as EXPANDED -----------
def _settle_fan_out(self, stage: Stage, item: WorkItem, summary: dict[str, int]) -> None:
"""Adopt a fan-out stage's children, then retire the parent as ``EXPANDED``.
Parameters
----------
stage
The fan-out stage whose job just succeeded.
item
The parent item, whose directory holds the manifest and the child
seed directories.
summary
The pass summary, updated in place.
Notes
-----
The parent deliberately never transits an intermediate claimable state
-- no ``complete`` before ``expand``, because a crash in that window
would let the next stage claim the *parent* as if it were a child.
Instead: validate the manifest (a bad one is a budgeted fault, since it
may be a partial write a retry fixes), adopt every child idempotently
so a replay or a child shared with another parent is safe, and expand
the parent only once every child is in. A store fault on any child
leaves the parent leased; the lease timeout re-runs the stage, and
every adoption step tolerates the replay.
An empty manifest is *valid*: a seed with zero viable children still
retires as ``EXPANDED``, having fanned out into nothing, never as a
failure.
"""
try:
child_ids = self._read_children_manifest(item)
except ValueError as exc:
self._fail(item.id, stage.name, str(exc), terminal=False, summary=summary)
return
# Children inherit the funnel prefix up to and including this stage:
# completing those stage names walks a NEW child to this stage's
# output state, with the real stage names in its completions timeline.
lineage = [s.name for s in self._active[: self._active.index(stage) + 1]]
adopted = 0
for child_id in child_ids:
try:
self._adopt_child(item, child_id, lineage)
adopted += 1
except Exception:
_log.exception(
"could not adopt child %s of %s; parent stays leased for recovery",
child_id,
item.id,
)
if adopted < len(child_ids):
return # parent not expanded; lease timeout re-runs the (idempotent) settle
try:
self._store.expand(item.id, f"stage {stage.name!r} fanned out into {adopted} children")
except Exception:
_log.exception("store fault expanding %s; item stays leased for recovery", item.id)
return
summary["expanded"] += 1
summary["children"] += adopted
def _read_children_manifest(self, item: WorkItem) -> list[str]:
"""Parse and validate the fan-out manifest a stage's job wrote.
Parameters
----------
item
The parent item whose directory holds ``children.json``.
Returns
-------
list of str
The child ids, in manifest order.
Raises
------
ValueError
On any defect. Each id must be non-empty, unique, distinct from the
parent, a plain path segment -- a child id becomes a directory
name, so ``../evil`` must never escape the workspace -- and backed
by its ``children/<id>/`` seed directory.
"""
manifest = item.dir / CHILDREN_MANIFEST
try:
data = json.loads(manifest.read_text(encoding="utf-8"))
except (OSError, ValueError) as exc:
raise ValueError(f"unreadable children manifest {manifest}: {exc}") from None
if not isinstance(data, list) or not all(isinstance(c, str) and c for c in data):
raise ValueError(f"{manifest}: must be a JSON list of non-empty child ids")
if len(set(data)) != len(data):
raise ValueError(f"{manifest}: duplicate child ids")
for child_id in data:
if child_id == item.id:
raise ValueError(f"{manifest}: child id equals parent id {item.id!r}")
if child_id in {".", ".."} or Path(child_id).name != child_id:
raise ValueError(f"{manifest}: unsafe child id {child_id!r}")
if not (item.dir / CHILDREN_DIR / child_id).is_dir():
raise ValueError(
f"{manifest}: names child {child_id!r} but "
f"{CHILDREN_DIR}/{child_id}/ is missing"
)
return data
def _adopt_child(self, parent: WorkItem, child_id: str, lineage: list[str]) -> None:
"""Bring one child into the pipeline. Every step tolerates a replay.
Parameters
----------
parent
The item that fanned out, whose directory holds the child's seed
files.
child_id
The child's id, which is also its directory name in the workspace.
lineage
The funnel prefix to replay through the completions ledger, so the
child lands in the fan-out stage's output state with real stage
names in its timeline.
Notes
-----
Three steps, each safe to repeat: move the seed directory into the
workspace (skipped when the destination exists -- a replay, or a child
two parents both produced, where the first adoption wins); enqueue it
with the parent as lineage (idempotent, first ``parent_id`` wins); and
complete it through the funnel prefix (ledger-idempotent).
"""
src = parent.dir / CHILDREN_DIR / child_id
dst = self._work_dir / child_id
if not dst.exists():
shutil.move(str(src), str(dst))
self._store.enqueue(child_id, parent_id=parent.id)
for stage_name in lineage:
self._store.complete(child_id, stage_name, target_state=self._stage_outputs[stage_name])
# -- Lineage selection -----------------------------------------------------
def _is_selected(self, material_id: str) -> bool:
"""Decide whether a material belongs to the selected lineage.
Parameters
----------
material_id
The id to test.
Returns
-------
bool
``True`` when the id -- or any ancestor, walking ``parent_id``
chains through the store -- is in the selection, and always
``True`` when no selection is active. An unknown id or a broken
chain is simply not selected.
Notes
-----
Verdicts are memoized for every id on the walked chain. That is safe
for the process lifetime because ``parent_id`` is immutable (enqueue is
first-write-wins) and the selection is fixed at construction.
"""
if self._materials is None:
return True
chain: list[str] = []
verdict = False
current: str | None = material_id
while current is not None and current not in chain: # cycle-safe
cached = self._eligibility.get(current)
if cached is not None:
verdict = cached
break
if current in self._materials:
chain.append(current)
verdict = True
break
chain.append(current)
rec = self._store.get(current)
current = rec.parent_id if rec is not None else None
for mid in chain:
self._eligibility[mid] = verdict
return verdict
# -- Finalize phase ------------------------------------------------------------
def _finalize_phase(self, summary: dict[str, int]) -> None:
"""Retire materials that completed the configured funnel as first-class ``DONE``.
Parameters
----------
summary
The pass summary, updated in place.
Notes
-----
Query-then-finalize, both idempotent, so a crash anywhere here simply
repeats next pass. Runs *after* the check phase, so an item settled
this pass is finalized this pass. A mid-funnel stop finalizes nothing:
later stages still consume that output, in a later run.
"""
if not self._finalize_enabled:
return # mid-funnel stop: later stages still consume this output
if is_terminal(self._stop_output):
return # the funnel already ends in a terminal state; nothing to do
try:
ready = self._store.query({"state": self._stop_output})
except Exception:
_log.exception("finalize: query failed; skipping this pass")
return
for rec in ready:
if not self._is_selected(rec.material_id):
continue # selective run: non-selected items are never swept to DONE
try:
self._store.finalize(rec.material_id)
summary["finalized"] += 1
except Exception:
_log.exception("finalize failed for %s; retried next pass", rec.material_id)
# -- Submit phase -------------------------------------------------------------
def _submit_phase(self, summary: dict[str, int]) -> None:
"""Walk the active stages in funnel order, gating, claiming, and starting work.
Parameters
----------
summary
The pass summary, updated in place.
Notes
-----
A stage that raises ``StageError`` -- the stage as a whole cannot run
-- is skipped for this pass and fails nothing; its claimed items
recover by lease timeout. In only-mode exactly one stage submits.
"""
try:
backlog = self._store.count_by_state()
except Exception:
_log.exception("submit phase: count_by_state failed; skipping this pass")
return
for stage, needs in zip(self._active, self._input_states, strict=True):
if self._only is not None and stage.name != self._only:
continue # only-mode: exactly one stage submits new work
try:
self._submit_stage(stage, needs, backlog, summary)
except StageError as exc:
_log.error(
"stage %r cannot run (%s); skipped this pass, claimed items "
"recover by lease timeout",
stage.name,
exc,
)
except Exception:
_log.exception("submit phase: stage %r skipped this pass", stage.name)
def _submit_stage(
self,
stage: Stage,
needs: MaterialState,
backlog: dict[MaterialState, int],
summary: dict[str, int],
) -> None:
"""Apply the gates for one stage, then claim and submit up to free slots.
Parameters
----------
stage
The stage to consider.
needs
Its input state -- the previous stage's output.
backlog
The pass's single ``count_by_state`` reading, so the backlog gate
costs no extra query per stage.
summary
The pass summary, updated in place.
Notes
-----
Gate order is load-bearing. The backlog gate runs first, so a stage
with nothing to do costs nothing -- not even a health check. A
selective run adds a selection gate right behind it: when nothing in
the input state belongs to the selected lineage, the stage is equally
free, with no health check and no claim churn. Only then does an
unhealthy backend pause the stage, which resumes by itself on a later
pass.
"""
if backlog.get(needs, 0) <= 0:
return # backlog gate: an empty input queue means never considered
eligible: list[str] | None = None
if self._materials is not None:
eligible = [
rec.material_id
for rec in self._store.query({"state": needs})
if self._is_selected(rec.material_id)
]
if not eligible:
return # selection gate: nothing of ours here costs nothing
backend = self._backends[stage.name]
health = backend.health_check()
if not health.ok:
_log.warning(
"stage %r paused: backend %r unhealthy (%s); resumes when it recovers",
stage.name,
backend.name,
health.detail,
)
return
free = backend.capabilities().max_parallel - len(
self._store.list_jobs(backend_name=backend.name)
)
if free <= 0:
return # backend saturated; materials stay queued for a later pass
for rec in self._store.claim(needs, free, among=eligible):
self._submit_item(stage, backend, rec.material_id, summary)
def _submit_item(
self,
stage: Stage,
backend: ComputeBackend,
material_id: str,
summary: dict[str, int],
) -> None:
"""Validate, build, submit, and record one item's job.
Parameters
----------
stage
The stage doing the work.
backend
Where its job runs.
material_id
The claimed material.
summary
The pass summary, updated in place.
Raises
------
StageError
Upward, unhandled: the stage as a whole is broken, so the caller
pauses it and the item stays leased. Every *per-item* failure is
settled here through ``_fail`` and never propagates.
Notes
-----
The handle is recorded before the method returns, so a crash after
submission re-polls the job instead of resubmitting it. A job whose
record cannot be written is released, so nothing ever runs untracked.
"""
item = self._item(material_id)
item.dir.mkdir(parents=True, exist_ok=True)
try:
stage.validate_inputs(item)
except StageInputError as exc:
# Budgeted, not terminal: a missing input can be an infrastructure
# problem (unmounted disk), so it spends budget rather than
# mass-dropping items on the spot.
self._fail(material_id, stage.name, str(exc), terminal=False, summary=summary)
return
try:
spec = stage.build_jobspec(item)
except StageError:
raise # the stage as a whole is broken; the caller pauses it
except Exception as exc:
self._fail(
material_id,
stage.name,
f"build_jobspec raised: {exc}",
terminal=False,
summary=summary,
)
return
try:
handle = backend.submit(spec)
except BackendError as exc:
self._fail(
material_id,
stage.name,
f"backend {backend.name!r} rejected submission: {exc}",
terminal=False,
summary=summary,
)
return
job_key = f"{backend.name}:{handle.job_id}"
try:
self._store.record_job(job_key, stage.name, backend.name, asdict(handle), [material_id])
except Exception:
# An unrecorded job would be invisible to every future pass:
# abandon it so nothing runs untracked. The item stays leased and
# recovers by lease timeout.
_log.exception(
"could not record job %s; releasing the handle so nothing runs untracked",
job_key,
)
with contextlib.suppress(BackendError):
backend.release(handle)
return
summary["submitted"] += 1
# -- Attribution (the single point where results meet the store) ---------------
def _complete(self, material_id: str, stage_name: str, summary: dict[str, int]) -> None:
"""Advance one item to its stage's output state.
Parameters
----------
material_id
The item to advance.
stage_name
The stage that completed it.
summary
The pass summary, updated in place.
Notes
-----
Never raises: a store fault leaves the item leased, and the lease
timeout recovers it.
"""
try:
self._store.complete(
material_id, stage_name, target_state=self._stage_outputs[stage_name]
)
summary["completed"] += 1
except Exception:
_log.exception(
"store fault completing %s at stage %r; item stays leased for recovery",
material_id,
stage_name,
)
def _fail(
self,
material_id: str,
stage_name: str,
reason: str,
*,
terminal: bool,
summary: dict[str, int],
) -> None:
"""Fail one item -- the single point where the verdict-vs-fault doctrine applies.
Parameters
----------
material_id
The item that failed.
stage_name
Where it failed.
reason
Why, recorded verbatim and shown to the user.
terminal
``True`` for a deterministic science verdict a re-run cannot
change, which bypasses the retry budget by failing with a budget of
1. ``False`` for a fault, which spends one unit of ``retry_budget``
and requeues while budget remains.
summary
The pass summary, updated in place.
Notes
-----
Never raises: a store fault leaves the item leased for lease-timeout
recovery.
"""
budget = 1 if terminal else self._retry_budget
try:
outcome = self._store.fail(material_id, stage_name, reason, max_attempts=budget)
except Exception:
_log.exception(
"store fault failing %s at stage %r; item stays leased for recovery",
material_id,
stage_name,
)
return
if outcome is MaterialState.FAILED:
summary["failed"] += 1
_log.warning(
"material %s failed terminally at stage %r (%s): %s",
material_id,
stage_name,
"verdict" if terminal else "fault",
reason,
)
else:
summary["requeued"] += 1
# -- Plumbing -------------------------------------------------------------
def _item(self, material_id: str) -> WorkItem:
"""Resolve a material id to its work item.
The id-to-directory convention, defined exactly once.
Parameters
----------
material_id
The material to locate.
Returns
-------
WorkItem
The item, whose directory is ``work_dir/<material_id>``. The
directory is not created here; ``_submit_item`` creates it at first
submission.
"""
return WorkItem(id=material_id, dir=self._work_dir / material_id)