"""The ``Stage`` contract: autonomous, contract-bound units of *what* to compute.
Interface + pure data only -- no SMACT, no PyXtal, no MACE, no subprocess. The
concrete stages plug in at Phase 3, each as a thin ``Stage`` subclass plus a
**standalone runner module** (``python -m matmelt.core.stages.<name>_runner``)
that can be executed and debugged by hand, without the orchestrator.
Two theses meet here:
* **What vs where.** A stage describes work as a portable ``JobSpec`` --
``build_jobspec`` is the *only* place a command is constructed -- and never
executes anything itself, names no host, and imports no concrete backend.
* **Stages are autonomous and contract-bound.** Every stage *declares* what must
be true of its inputs (``requires``) and what it guarantees to produce
(``produces``), and *validates both itself*. That is what makes starting the
pipeline mid-way ("just screen these structures I already have") safe: the
stage checks the contract and refuses clearly -- naming exactly which files are
missing -- rather than producing silent garbage.
Failure vocabulary (three planes, failing independently):
* ``StageInputError`` / ``StageOutputError`` -- a *contract* violation for one
item: its inputs aren't satisfied, or the job "succeeded" but the declared
outputs aren't there. Per-item; the driver records it and moves on.
* ``StageVerdict(ok=False)`` -- the stage judged one item's result and it failed.
``terminal=True`` marks a deterministic *science verdict* a re-run cannot
change (no retry budget spent); ``False`` marks a *fault* a retry may fix.
* ``StageError`` -- the stage as a whole cannot run (programming/configuration
fault, e.g. an unsatisfiable backend). Never used for one bad item.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from matmelt.core.jobspec import JobResult, JobSpec
#: Manifest a fan-out stage's job writes into the parent item's directory: a
#: JSON list of child ids. The per-child seed files live under
#: ``children/<id>/`` next to it. The orchestrator validates and adopts them;
#: the runner never learns the workspace layout (which is what keeps runners
#: remote-safe — they only ever write into their own cwd).
CHILDREN_MANIFEST = "children.json"
#: Directory (inside the parent item's dir) holding one subdirectory of seed
#: files per child named in ``CHILDREN_MANIFEST``.
CHILDREN_DIR = "children"
[docs]
class StageError(Exception):
"""The stage as a whole cannot run -- a programming or configuration fault.
One item failing is never a ``StageError``: that is a ``StageVerdict`` with
``ok=False`` (or an input/output contract error), so the rest of the batch
continues.
"""
[docs]
class StageOutputError(StageError):
"""A finished item is missing files the stage's ``produces`` contract promised.
Raised by ``validate_outputs`` after a job completes but *before* the item is
marked done -- a job that exits 0 without producing its declared outputs must
not advance.
"""
[docs]
@dataclass(frozen=True)
class Requirement:
"""One runtime dependency a stage's runner needs in its interpreter.
Declared by the stage (the science stack is user-owned, never bundled);
checked by ``core/preflight.py`` against the stage's ``runner_python``
*before* any work starts, so a missing package is reported with the exact
install command instead of surfacing as N failed jobs.
Attributes
----------
module : str
The import name probed in the runner's interpreter.
pip_name : str
The exact pip install token. This differs from ``module`` often enough
that guessing one from the other would hand users a broken command --
``mace`` is installed as ``mace-torch``.
purpose : str
One human sentence: what the stage uses the package for.
Examples
--------
>>> req = Requirement("mace", "mace-torch", "ML relaxation")
>>> req.module, req.pip_name
('mace', 'mace-torch')
"""
module: str
pip_name: str
purpose: str
[docs]
@dataclass(frozen=True)
class WorkItem:
"""One unit of pipeline work: a material (or seed) and its directory.
Frozen snapshot: handing one to a stage must never let the stage mutate
coordination state.
Attributes
----------
id : str
The stable identifier used everywhere -- state store keys, log lines,
result records. Unique per material.
dir : pathlib.Path
The item's working directory. Every per-item file the stage reads or
writes lives under it, and stages address files **only** relative to
this directory -- never via an absolute path of their own construction.
"""
id: str
dir: Path
[docs]
@dataclass(frozen=True)
class StageOutputSpec:
"""What this stage guarantees an item's directory contains on success.
Checked by ``validate_outputs`` after the job finishes. A downstream stage's
``requires`` will typically name a subset of these files -- that hand-off is
the pipeline's real interface.
Attributes
----------
files : tuple of str
Filenames, relative to ``WorkItem.dir``, the stage guarantees exist and
are non-empty once an item has succeeded.
description : str
A human sentence for error messages and docs.
"""
files: tuple[str, ...] = ()
description: str = ""
[docs]
@dataclass(frozen=True)
class StageVerdict:
"""The stage's judgment of one item's finished job.
Attributes
----------
ok : bool
Whether the item advanced. Judged from the ``JobResult`` (exit code,
streams) and the item's own output files -- never by the backend.
reason : str
Short human-readable explanation when ``ok=False``. Persisted with the
failure and shown by ``status`` / ``results``.
terminal : bool
Classifies an ``ok=False`` verdict, and is meaningless when ``ok=True``.
``True`` marks a deterministic **science verdict** a re-run cannot
change ("screening rejected: unstable by 0.4 eV/atom") -- fail the item
on the spot, spend no retry budget. ``False`` (the default) marks a
**fault** (subprocess died, missing output, unknown) that a retry may
genuinely fix. When in doubt leave it ``False``: a "missing file" can be
an unmounted disk, and terminal-failing that would mass-drop healthy
items.
Examples
--------
>>> StageVerdict(ok=True)
StageVerdict(ok=True, reason='', terminal=False)
>>> StageVerdict(ok=False, reason="unstable by 0.4 eV/atom", terminal=True).terminal
True
"""
ok: bool
reason: str = ""
terminal: bool = False
[docs]
class Stage(ABC):
"""One pipeline stage: declares its contract, builds jobs, judges results.
The lifecycle a driver (the orchestrator, a test, or a user running one
stage standalone) puts an item through::
stage.validate_inputs(item) # refuse clearly, or proceed
spec = stage.build_jobspec(item) # the ONLY place a command is built
handle = backend.submit(spec) # the backend runs it, wherever
... poll until terminal, fetch ...
verdict = stage.interpret_result(item, result)
if verdict.ok:
stage.validate_outputs(item) # promised files actually exist
A stage never executes anything itself, never calls ``backend.release``
(the driver releases in ``finally``), and never touches coordination state
(advancing/failing an item is the state store's job, driven by the
orchestrator from the verdict).
**Fan-out stages** (``fan_out = True``): the funnel's 1 -> N steps (one
element combination -> N compositions; one composition -> M polymorphs).
Instead of the item advancing, its successful job must leave a
``CHILDREN_MANIFEST`` (a JSON list of child ids) plus one
``CHILDREN_DIR/<id>/`` directory of seed files per child, all inside the
item's own directory -- declare the manifest in ``produces``. The driver
adopts the children into the pipeline (each recording this item as its
``parent_id``) and retires the item as terminal ``EXPANDED``: it fanned
out, it did not fail, and it never advances past this stage itself.
"""
#: Whether this stage expands its item into children (see class docstring)
#: instead of advancing it. Class attribute -- override per subclass.
fan_out: bool = False
#: The science packages this stage's runner imports, checked by the
#: preflight before any work starts. Class attribute -- a science stage
#: overrides it; the default (nothing beyond stdlib) is itself a
#: contract, e.g. extract's.
runtime_requirements: tuple[Requirement, ...] = ()
#: The interpreter that executes this stage's runner on its backend.
#: Stages with a ``python`` constructor option overwrite this per
#: instance; the preflight probes ``runtime_requirements`` against it.
runner_python: str = "python3"
# -- The declared contract -------------------------------------------------
@property
@abstractmethod
def name(self) -> str:
"""Stable stage identifier -- also the per-item idempotency key in state.
Returns
-------
str
The canonical name (``"compose"``, ``"screen"``, ...). It is what a
config's ``stages:`` block names, what the stage registry maps, and
what the state store's completion ledger records per item, so
completing a stage twice is a no-op.
"""
@property
@abstractmethod
def requires(self) -> StageInputSpec:
"""What must be true of an item's directory before this stage runs.
Returns
-------
StageInputSpec
The declared input contract, enforced by ``validate_inputs``.
"""
@property
@abstractmethod
def produces(self) -> StageOutputSpec:
"""What this stage guarantees an item's directory contains on success.
Returns
-------
StageOutputSpec
The declared output contract, enforced by ``validate_outputs``.
"""
# -- The two abstract halves every stage implements --------------------------
[docs]
@abstractmethod
def build_jobspec(self, item: WorkItem) -> JobSpec:
"""Build the portable job description for one item.
The **only** place in a stage where a command is constructed. The spec
names no host, queue, or transport; paths point into ``item.dir``; any
stage options (model path, device, thresholds) were injected at stage
construction from config -- never hardcoded here.
Parameters
----------
item
The item to build a job for. Its inputs have already passed
``validate_inputs``.
Returns
-------
JobSpec
A portable job description any backend can execute.
Raises
------
StageError
If the stage is misconfigured such that no job can be built at all
-- a whole-stage fault, never a per-item verdict.
"""
[docs]
@abstractmethod
def interpret_result(self, item: WorkItem, result: JobResult) -> StageVerdict:
"""Judge one item's finished job.
Reads the faithful ``JobResult`` (exit code -- remember ``None`` means
"ran, code unknown", never success) and the item's own files.
Parameters
----------
item
The item whose job finished.
result
What the backend reported: exit code, streams, and the declared
output files it actually found.
Returns
-------
StageVerdict
The stage's judgment. Set ``terminal=True`` only for a
deterministic science rejection a re-run provably cannot change;
everything else is a retryable fault.
"""
# -- Concrete contract enforcement (override only to *extend*) --------------
[docs]
def validate_outputs(self, item: WorkItem) -> None:
"""Confirm the ``produces`` promise before an item is marked done.
Default: every declared ``produces.files`` entry must exist under
``item.dir`` and be non-empty. Override to add deeper checks; call
``super().validate_outputs(item)`` first.
Parameters
----------
item
The item whose job the stage judged successful.
Raises
------
StageOutputError
If any promised output file is missing or empty -- a job that
exited 0 without producing what it promised must not advance.
"""
missing = self._missing_files(item, self.produces.files)
if missing:
raise StageOutputError(
f"stage {self.name!r} finished item {item.id!r} but its promised "
f"output file(s) {missing} are missing or empty in {item.dir}"
)
@staticmethod
def _missing_files(item: WorkItem, names: tuple[str, ...]) -> list[str]:
"""The declared files absent or empty under ``item.dir``.
Parameters
----------
item
The item whose directory is inspected.
names
Declared filenames, relative to ``item.dir``.
Returns
-------
list of str
The subset of ``names`` that is missing or zero-byte, kept in
declared order so error messages read the way the contract does.
"""
missing: list[str] = []
for name in names:
path = item.dir / name
if not path.is_file() or path.stat().st_size == 0:
missing.append(name)
return missing