"""The screen stage: one candidate structure -> keep (RELAXED) or drop (3.3).
The thin ``Stage`` half of the pair; the science lives entirely in
``screen_runner.py`` (self-contained: stdlib + mace/ase, zero
``matmelt`` imports), which this class ships to the backend as an input
artifact of every job. This module itself imports **no science** -- it builds
job descriptions and judges results, nothing more, so it loads on any machine.
The runner MACE vc-relaxes the item's ``structure.cif`` and writes a verdict
``screen_result.json`` (category ``RELAXED`` or ``FAILED`` + reasons +
metrics); on ``RELAXED`` it also writes ``relaxed.cif``. 1:1 -- this stage is
NOT fan-out: the item either advances or is screened out.
**This is the funnel's first terminal science verdict.** A parsed ``FAILED``
verdict is a deterministic rejection -- re-running the same relaxation with
the same model and settings reproduces it -- so ``interpret_result`` returns
``terminal=True``: the item fails on the spot and spends no retry budget
(G-2 verdict-vs-fault). Everything else (non-zero exit, missing/corrupt
verdict) stays a retryable fault: missing mace is an environment problem,
and a crashed relaxation could be an OOM -- when in doubt, fault.
Model path, device, dtype and the relaxation knobs come from config via the
constructor -- never hardcoded (the 3.6 preflight/auto-download offer hooks
into the same ``model_path`` option).
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import TYPE_CHECKING
from matmelt.core.jobspec import JobSpec
from matmelt.core.stages.base import (
Requirement,
Stage,
StageError,
StageInputSpec,
StageOutputSpec,
StageVerdict,
)
from matmelt.core.stages.registry import register
if TYPE_CHECKING:
from matmelt.core.jobspec import JobResult
from matmelt.core.stages.base import WorkItem
#: The self-contained script shipped with every screen job. Lives next to
#: this module inside the installed package; staged into the job's cwd by the
#: backend, so the command references it by bare name.
RUNNER_NAME = "screen_runner.py"
#: The input file a structure child carries -- declared as an input artifact
#: so the Phase-4 staging transport knows to push it (it already sits in the
#: item dir locally; backends stage in-place inputs as a no-op).
STRUCTURE_NAME = "structure.cif"
#: The verdict file the runner always writes on a completed screening.
VERDICT_NAME = "screen_result.json"
#: The relaxed structure the runner writes on a RELAXED verdict.
RELAXED_NAME = "relaxed.cif"
_MODEL_SIZES = ("small", "medium", "large")
#: float64 is the locked default -- required for a reliable stress tensor in
#: vc-relax (the MVP_V0 lesson); float32 stays available as an explicit,
#: eyes-open speed trade.
_DTYPES = ("float64", "float32")
#: How many characters of a failed runner's stderr to surface in the verdict.
_REASON_TAIL = 300
[docs]
class ScreenStage(Stage):
"""MACE vc-relax one candidate structure and keep or screen it out.
Parameters
----------
python
Interpreter to invoke on the backend this stage is wired to (the
user's own environment holds the science stack -- never bundled).
model_path
Local MACE model file. Empty (default) = the MACE-MP foundation
model, which mace downloads to its cache on first use -- the runner
discloses that before loading (3.6 adds the up-front consent offer).
model_size
Foundation model size when ``model_path`` is empty.
device
Compute device passed to MACE ("cpu" | "cuda" | "mps" | "auto").
dtype
Torch dtype; float64 (default) is required for a reliable stress
tensor in vc-relax.
fmax
Force convergence target for the relaxation (eV/A).
max_steps
Maximum optimizer steps for the relaxation.
Raises
------
StageError
For any invalid option, so a config mistake surfaces when the run is
wired rather than as N failed jobs later.
Notes
-----
The physical-sanity bounds that decide a verdict -- bond lengths, stress,
volume change, the energy window -- are deliberately **not** options. They
live as constants in the runner so verdicts stay comparable across a whole
database.
"""
runtime_requirements = (
Requirement(
module="mace",
pip_name="mace-torch",
purpose="MACE force-field vc-relaxation (screen runner; brings ase and torch)",
),
)
def __init__(
self,
python: str = "python3",
model_path: str = "",
model_size: str = "medium",
device: str = "cpu",
dtype: str = "float64",
fmax: float = 0.05,
max_steps: int = 200,
) -> None:
if not isinstance(python, str) or not python.strip():
raise StageError(f"screen: python must be a non-empty string, got {python!r}")
if not isinstance(model_path, str):
raise StageError(f"screen: model_path must be a string, got {model_path!r}")
if model_size not in _MODEL_SIZES:
raise StageError(
f"screen: model_size must be one of {_MODEL_SIZES}, got {model_size!r}"
)
if not isinstance(device, str) or not device.strip():
raise StageError(f"screen: device must be a non-empty string, got {device!r}")
if dtype not in _DTYPES:
raise StageError(f"screen: dtype must be one of {_DTYPES}, got {dtype!r}")
if isinstance(fmax, bool) or not isinstance(fmax, (int, float)) or fmax <= 0:
raise StageError(f"screen: fmax must be a number > 0, got {fmax!r}")
if isinstance(max_steps, bool) or not isinstance(max_steps, int) or max_steps < 1:
raise StageError(f"screen: max_steps must be an integer >= 1, got {max_steps!r}")
self.runner_python = python
self._model_path = model_path
self._model_size = model_size
self._device = device
self._dtype = dtype
self._fmax = float(fmax)
self._max_steps = max_steps
@property
def name(self) -> str:
return "screen"
@property
def requires(self) -> StageInputSpec:
return StageInputSpec(
files=(STRUCTURE_NAME,),
description="one candidate crystal from the structure stage",
)
@property
def produces(self) -> StageOutputSpec:
return StageOutputSpec(
files=(VERDICT_NAME, RELAXED_NAME),
description="screening verdict + the relaxed structure (on RELAXED)",
)
[docs]
def build_jobspec(self, item: WorkItem) -> JobSpec:
runner = Path(__file__).with_name(RUNNER_NAME)
command = [
self.runner_python,
RUNNER_NAME,
"--cif",
STRUCTURE_NAME,
"--model-size",
self._model_size,
"--device",
self._device,
"--dtype",
self._dtype,
"--fmax",
str(self._fmax),
"--max-steps",
str(self._max_steps),
]
if self._model_path:
command += ["--model-path", self._model_path]
return JobSpec(
command=command,
input_artifacts={
RUNNER_NAME: str(runner),
STRUCTURE_NAME: str(item.dir / STRUCTURE_NAME),
},
output_artifacts=[VERDICT_NAME, RELAXED_NAME],
cwd=str(item.dir),
scratch_inputs=frozenset({RUNNER_NAME}),
)
[docs]
def interpret_result(self, item: WorkItem, result: JobResult) -> StageVerdict:
"""Judge the finished screening from its exit code and verdict file.
The funnel's first terminal science verdict.
Parameters
----------
item
The structure whose job finished; its directory holds the verdict
file.
result
What the backend reported.
Returns
-------
StageVerdict
One of four outcomes:
- A non-zero or unknown exit is a retryable fault -- missing mace,
an out-of-memory kill, a broken interpreter.
- Exit 0 with a ``RELAXED`` verdict advances the item.
- Exit 0 with a ``FAILED`` verdict is a **terminal science
rejection**: the same relaxation reproduces it, so no retry
budget is spent and the recorded reason carries the failed
checks.
- Exit 0 with a missing, corrupt, or unknown verdict means the
runner broke its contract. That is a fault, not a rejection --
on unknown ground, let the budget apply.
"""
if result.exit_code != 0:
tail = (result.stderr or result.stdout).strip()[-_REASON_TAIL:]
code = "unknown" if result.exit_code is None else str(result.exit_code)
return StageVerdict(ok=False, reason=f"screen runner exited {code}: {tail}")
verdict_path = item.dir / VERDICT_NAME
try:
verdict = json.loads(verdict_path.read_text(encoding="utf-8"))
except OSError:
return StageVerdict(
ok=False, reason=f"screen runner exited 0 but wrote no {VERDICT_NAME}"
)
except ValueError as exc:
return StageVerdict(ok=False, reason=f"{VERDICT_NAME} unreadable: {exc}")
category = verdict.get("category") if isinstance(verdict, dict) else None
if category == "RELAXED":
return StageVerdict(ok=True)
if category == "FAILED":
reasons = verdict.get("failure_reasons") or []
joined = (
"; ".join(str(reason) for reason in reasons)
if isinstance(reasons, list)
else str(reasons)
)
return StageVerdict(
ok=False,
reason=f"screened out: {joined or 'no reason recorded'}",
terminal=True,
)
return StageVerdict(ok=False, reason=f"{VERDICT_NAME} has unknown category {category!r}")
register("screen", ScreenStage)