Stage contract

The Stage ABC — the what half of the architecture — plus the funnel registry.

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.

matmelt.core.stages.base.CHILDREN_MANIFEST = 'children.json'

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).

Type:

Manifest a fan-out stage’s job writes into the parent item’s directory

matmelt.core.stages.base.CHILDREN_DIR = 'children'

Directory (inside the parent item’s dir) holding one subdirectory of seed files per child named in CHILDREN_MANIFEST.

exception matmelt.core.stages.base.StageError

Bases: 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.

exception matmelt.core.stages.base.StageInputError

Bases: StageError

An item’s inputs do not satisfy the stage’s declared requires contract.

Raised by validate_inputs before any job is built or submitted – the “refuse clearly, never silent garbage” half of the autonomous-stage thesis. The message names the stage, the item, and exactly which files are missing.

exception matmelt.core.stages.base.StageOutputError

Bases: 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.

class matmelt.core.stages.base.Requirement(module, pip_name, purpose)

Bases: object

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.

module

The import name probed in the runner’s interpreter.

Type:

str

pip_name

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.

Type:

str

purpose

One human sentence: what the stage uses the package for.

Type:

str

Examples

>>> req = Requirement("mace", "mace-torch", "ML relaxation")
>>> req.module, req.pip_name
('mace', 'mace-torch')

class matmelt.core.stages.base.WorkItem(id, dir)

Bases: object

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.

id

The stable identifier used everywhere – state store keys, log lines, result records. Unique per material.

Type:

str

dir

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.

Type:

pathlib.Path

class matmelt.core.stages.base.StageInputSpec(files=(), description='')

Bases: object

What must be true of an item before this stage may run on it.

Extend by overriding Stage.validate_inputs for checks richer than file presence – the declared files remain the machine-readable core of the contract.

files

Filenames, relative to WorkItem.dir, that must exist and be non-empty before the stage may run.

Type:

tuple of str

description

A human sentence for error messages and docs, e.g. “relaxed structures from the screen stage”.

Type:

str

class matmelt.core.stages.base.StageOutputSpec(files=(), description='')

Bases: object

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.

files

Filenames, relative to WorkItem.dir, the stage guarantees exist and are non-empty once an item has succeeded.

Type:

tuple of str

description

A human sentence for error messages and docs.

Type:

str

class matmelt.core.stages.base.StageVerdict(ok, reason='', terminal=False)

Bases: object

The stage’s judgment of one item’s finished job.

ok

Whether the item advanced. Judged from the JobResult (exit code, streams) and the item’s own output files – never by the backend.

Type:

bool

reason

Short human-readable explanation when ok=False. Persisted with the failure and shown by status / results.

Type:

str

terminal

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.

Type:

bool

Examples

>>> StageVerdict(ok=True)
StageVerdict(ok=True, reason='', terminal=False)
>>> StageVerdict(ok=False, reason="unstable by 0.4 eV/atom", terminal=True).terminal
True

class matmelt.core.stages.base.Stage

Bases: 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.

fan_out: bool = False

Whether this stage expands its item into children (see class docstring) instead of advancing it. Class attribute – override per subclass.

runtime_requirements: tuple[Requirement, ...] = ()

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.

runner_python: str = 'python3'

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.

abstract property name: str

Stable stage identifier – also the per-item idempotency key in state.

Returns:

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.

Return type:

str

abstract property requires: StageInputSpec

What must be true of an item’s directory before this stage runs.

Returns:

The declared input contract, enforced by validate_inputs.

Return type:

StageInputSpec

abstract property produces: StageOutputSpec

What this stage guarantees an item’s directory contains on success.

Returns:

The declared output contract, enforced by validate_outputs.

Return type:

StageOutputSpec

abstract build_jobspec(item)

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 (WorkItem) – The item to build a job for. Its inputs have already passed validate_inputs.

Returns:

A portable job description any backend can execute.

Return type:

JobSpec

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.

abstract interpret_result(item, result)

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 (WorkItem) – The item whose job finished.

• result (JobResult) – What the backend reported: exit code, streams, and the declared output files it actually found.

Returns:

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.

Return type:

StageVerdict

validate_inputs(item)

Refuse clearly when item does not satisfy requires.

Default: every declared requires.files entry must exist under item.dir and be non-empty (a zero-byte input is treated as missing – it is never a usable structure or manifest). Override to add deeper checks (parse a file, verify a schema); call super().validate_inputs(item) first.

Parameters:

item (WorkItem) – The item about to be worked on.

Raises:

StageInputError – If any declared input file is missing or empty. The message names the stage, the item, and exactly which files are missing, in declared order.

Return type:

None

validate_outputs(item)

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 (WorkItem) – 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.

Return type:

None

The stage registry: stage name -> Stage class (5.1a).

The single seam where config becomes a running stage – the stage-plane mirror of backends/registry.py. A user’s YAML names each stage entry by the stage’s own name (compose, structure, …); this module maps that string to the concrete class and instantiates it with the entry’s options.

Registration is an explicit call, not an import-time decorator, so there is no hidden import-order dependency: a concrete stage module calls register("compose", ComposeStage) at the bottom of its file, and whoever wires the tool imports the stage modules it intends to offer (the loader’s build_stages imports the built-ins).

Construction contract: a registered class is instantiated as cls(**options) – unlike backends (whose cls(name, options) keeps a free-form dict because every kind wants different keys), each stage already declares a typed keyword signature and validates every value itself, raising StageError for anything invalid. Unpacking the config options onto that signature means an unknown option fails loudly with the stage named, instead of riding along silently in a dict.

FUNNEL_ORDER is the funnel, written down exactly once: it is the reference against which build_stages checks a config’s stage set. OPTIONAL_STAGES names the funnel members a config may omit – the configured stages must follow the funnel order, skipping only these. STAGE_OUTPUT_STATE is its state-plane companion: each canonical stage name -> the lifecycle state its completion lands. The orchestrator derives every stage’s input state (the previous stage’s output; the first stage consumes NEW) and its completion target from this map – so the stage <-> state coupling lives here by name, never by list position, and omitting an optional stage from a run cannot shift its neighbours’ states.

matmelt.core.stages.registry.FUNNEL_ORDER: tuple[str, ...] = ('compose', 'structure', 'screen', 'calculate', 'extract')

The funnel, in orchestrator order. build_stages requires a config’s stages to follow this order, omitting only OPTIONAL_STAGES members.

matmelt.core.stages.registry.OPTIONAL_STAGES: frozenset[str] = frozenset({'calculate'})

Funnel members a config may leave out entirely. Every other funnel stage up to the furthest configured one is required – a hole at a required stage is refused by build_stages, never silently bridged.

matmelt.core.stages.registry.STAGE_OUTPUT_STATE: dict[str, MaterialState] = {'calculate': MaterialState.CALCULATED, 'compose': MaterialState.COMPOSED, 'extract': MaterialState.EXTRACTED, 'screen': MaterialState.SCREENED, 'structure': MaterialState.STRUCTURED}

Canonical stage name -> the lifecycle state its completion lands. Stage names not in this map (test doubles) take the next chain state positionally.

matmelt.core.stages.registry.register(name, cls)

Register cls as the implementation of stage name.

Called explicitly at the bottom of each concrete stage module.

Parameters:

• name (str) – The stage’s canonical name, which is also what a config’s stages: block keys on.

• cls (type[Stage]) – The implementing class, which must be a Stage subclass.

Raises:

StageError – If cls is not a Stage subclass, or if name is already registered to a different class. Re-registering the same class is idempotent, so a module re-import is a no-op, but silently shadowing an existing stage is refused.

Return type:

None

matmelt.core.stages.registry.get(name)

Return the class registered under name.

Parameters:

name (str) – The stage name to look up.

Returns:

The registered Stage subclass.

Return type:

type

Raises:

StageError – If nothing is registered under name. The message names the stages that are known, so a config typo produces an actionable message rather than a KeyError.

matmelt.core.stages.registry.create_stage(name, options)

Instantiate the stage name with its config options as kwargs.

The one seam where a config block becomes a live stage.

Parameters:

• name (str) – Which registered stage to construct.

• options (dict[str, Any]) – The entry’s options: block, unpacked onto the stage’s typed keyword signature.

Returns:

The constructed, option-validated stage.

Return type:

Stage

Raises:

StageError – If name is unknown; if the stage’s own __init__ validation rejects a value; or if an option the stage does not accept (or a non-string key, which cannot be a keyword) is passed – caught here and reported with the stage named, instead of riding along silently.

Examples

>>> from matmelt.core.stages import compose  # registers "compose"
>>> stage = create_stage("compose", {"python": "python3", "threshold": 2})
>>> stage.name, stage.runner_python
('compose', 'python3')

matmelt.core.stages.registry.known_stages()

List the stage names currently registered.

Returns:

Sorted names, used in diagnostics and error messages. Which stages are present depends on which stage modules have been imported.

Return type:

list of str