"""Dependency preflight: prove each stage is runnable BEFORE any work starts (3.6).
The machinery behind three user-facing promises (policy 2026-07-08): the
science stack is user-owned and never bundled; ``init``/``run`` tell the user
exactly what a run needs and the exact install command for anything missing,
before the first job; ``doctor`` shows per-stage readiness. This module is
the thin hook those (Phase-5) commands call -- it computes, it never renders
(all human output lives in ``ui/console.py``) and it never raises for a
missing dependency: absence is a *result* here, not an error.
Checks probe each stage's **declared runner interpreter**
(``stage.runner_python``) with a subprocess import -- the orchestrator's own
venv is deliberately irrelevant, because the runner executes in the user's
backend python, which is where the science stack must live. Without a probe,
a missing package surfaces only as N failed jobs with exit-3 stderr; with
it, the run refuses up front with one actionable line per problem.
Also here: the seed-enumeration report (how many combinations the configured
``elements`` x ``combination_sizes`` fan out into -- the n-ary blast radius,
visible before any work) and the MACE model status/download helpers (the
consent itself belongs to the interactive command; ``download_foundation_model``
documents that it must never be called without it).
"""
from __future__ import annotations
import math
import subprocess
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING
from matmelt.core.jobspec import HealthStatus
from matmelt.core.stages.base import Requirement, Stage
if TYPE_CHECKING:
from matmelt.backends.base import ComputeBackend
#: An import probe should be near-instant; a hung interpreter is a finding.
_PROBE_TIMEOUT_S = 60.0
#: Where mace caches downloaded foundation models (its own convention).
MACE_CACHE_DIR = "~/.cache/mace"
[docs]
@dataclass(frozen=True)
class RequirementStatus:
"""One stage requirement, checked against the stage's runner interpreter.
Attributes
----------
stage : str
The stage that declared the requirement.
python : str
The interpreter probed -- the stage's ``runner_python``, never ours.
requirement : Requirement or None
What was checked. ``None`` marks the single entry a stage with no
requirements yields: "needs nothing beyond the standard library" is a
checked, reportable fact, not an omission.
satisfied : bool
Whether the probe succeeded.
detail : str
Human-readable evidence -- the import's own last error line when it
failed.
install_command : str or None
The exact fix when unsatisfied: ``<python> -m pip install <name>``,
ssh-prefixed when the probe ran on a remote. ``None`` when satisfied.
where : str
Where the probe actually ran -- ``"local"``, or the remote backend's
target. An answer is only as good as its location.
"""
stage: str
python: str
requirement: Requirement | None
satisfied: bool
detail: str
install_command: str | None = None
where: str = "local"
[docs]
def check_import(python: str, module: str) -> tuple[bool, str]:
"""Ask whether ``python`` can import ``module``. Never raises.
Probes via ``<python> -c "import <module>"``, so the answer is about the
*runner's* interpreter rather than ours.
Parameters
----------
python
The interpreter to probe.
module
The import name to try.
Returns
-------
tuple
``(satisfied, detail)``. A missing interpreter, a failing import, and a
probe that hung past ``_PROBE_TIMEOUT_S`` all come back as
``(False, why)`` -- absence is a result here, never an exception.
"""
try:
proc = subprocess.run(
[python, "-c", f"import {module}"],
capture_output=True,
text=True,
timeout=_PROBE_TIMEOUT_S,
)
except FileNotFoundError:
return False, f"interpreter {python!r} not found"
except subprocess.TimeoutExpired:
return False, f"import probe timed out after {_PROBE_TIMEOUT_S:.0f}s"
except OSError as exc:
return False, f"cannot run interpreter {python!r}: {exc}"
if proc.returncode == 0:
return True, f"{module} importable"
stderr_lines = (proc.stderr or "").strip().splitlines()
detail = stderr_lines[-1] if stderr_lines else f"import exited {proc.returncode}"
return False, detail
[docs]
def preflight_stages(
stages: Sequence[Stage],
stage_backends: Mapping[str, ComputeBackend] | None = None,
) -> list[RequirementStatus]:
"""Check every stage's declared requirements, one entry per requirement.
Pure bookkeeping over the probes -- no rendering, no raising. The caller
(``run``'s gate, ``doctor``'s table) decides what an unsatisfied entry
means.
Parameters
----------
stages
The stages to check, in any order.
stage_backends
The ``Runtime`` mapping of stage name to backend, which makes the probe
**location-aware**: a stage wired to a backend exposing the duck-typed
``check_remote_import(python, module)`` hook (``ssh_exec`` does) is
probed on that remote -- the science stack must exist in the
interpreter that will actually run the staged runner, not in ours --
and its install command is ssh-prefixed accordingly. Without the
mapping, or for a backend without the hook, the probe runs locally,
which is exactly right for the ``local`` backend.
Returns
-------
list of RequirementStatus
One entry per declared requirement, plus one ``requirement=None`` entry
for each stage that declares none.
"""
statuses: list[RequirementStatus] = []
for stage in stages:
backend = (stage_backends or {}).get(stage.name)
remote_probe = getattr(backend, "check_remote_import", None)
if not stage.runtime_requirements:
statuses.append(
RequirementStatus(
stage=stage.name,
python=stage.runner_python,
requirement=None,
satisfied=True,
detail="standard library only -- always ready",
)
)
continue
for requirement in stage.runtime_requirements:
install = f"{stage.runner_python} -m pip install {requirement.pip_name}"
if callable(remote_probe):
ok, detail = remote_probe(stage.runner_python, requirement.module)
where = str(getattr(backend, "probe_target", "remote"))
install = f"ssh {where} {install}"
else:
ok, detail = check_import(stage.runner_python, requirement.module)
where = "local"
statuses.append(
RequirementStatus(
stage=stage.name,
python=stage.runner_python,
requirement=requirement,
satisfied=ok,
detail=detail,
install_command=None if ok else install,
where=where,
)
)
return statuses
[docs]
def is_ready(statuses: Sequence[RequirementStatus]) -> bool:
"""Decide whether a run may start.
Parameters
----------
statuses
What ``preflight_stages`` returned.
Returns
-------
bool
``True`` when every checked requirement is satisfied. An empty list is
vacuously ready.
"""
return all(status.satisfied for status in statuses)
# ---------------------------------------------------------------------------
# Backend reachability (the doctor's other half, 4.2)
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class BackendHealth:
"""One configured backend's reachability, as reported by ``health_check``.
Attributes
----------
backend : str
The backend's configured name.
kind : str
The implementing class's name, which is what ``doctor`` shows as the
backend's type.
healthy : bool
Whether it can take work right now.
detail : str
The backend's own explanation, secret-free by contract.
"""
backend: str
kind: str
healthy: bool
detail: str
[docs]
def check_backends(backends: Mapping[str, ComputeBackend]) -> list[BackendHealth]:
"""Ask every configured backend whether it is usable. Never raises.
Parameters
----------
backends
The configured backends by name. Iteration order is the mapping's,
which is the config file's.
Returns
-------
list of BackendHealth
One entry per backend, in that same order. ``health_check`` never
raises by contract; a backend that breaks the contract anyway is
reported unhealthy with the violation named -- ``doctor`` must never
crash on the thing it exists to diagnose.
"""
healths: list[BackendHealth] = []
for name, backend in backends.items():
try:
status = backend.health_check()
except Exception as exc: # contract violation: report it, never crash
status = HealthStatus(
ok=False, detail=f"health_check raised {exc!r} (backend contract violation)"
)
healths.append(
BackendHealth(
backend=name,
kind=type(backend).__name__,
healthy=status.ok,
detail=status.detail,
)
)
return healths
[docs]
@dataclass(frozen=True)
class DoctorReport:
"""Everything ``doctor`` shows: backend reachability + per-stage readiness.
Attributes
----------
backends : list of BackendHealth
One entry per configured backend, in config order.
requirements : list of RequirementStatus
One entry per declared stage requirement.
"""
backends: list[BackendHealth] = field(default_factory=list)
requirements: list[RequirementStatus] = field(default_factory=list)
@property
def ok(self) -> bool:
"""Whether the environment is fit to run.
Returns
-------
bool
``True`` when every backend is healthy and every requirement is
satisfied. It is what the ``doctor`` command's exit code is built
from.
"""
return all(health.healthy for health in self.backends) and is_ready(self.requirements)
[docs]
def doctor_report(
stages: Sequence[Stage],
stage_backends: Mapping[str, ComputeBackend],
backends: Mapping[str, ComputeBackend],
) -> DoctorReport:
"""Diagnose the whole environment in one call -- what ``doctor`` renders.
Parameters
----------
stages
The configured stages.
stage_backends
Which backend runs each stage, so readiness is probed where the work
will happen.
backends
Every configured backend, checked for reachability.
Returns
-------
DoctorReport
Backend reachability first -- an unreachable remote also explains its
stages' failed remote probes -- then location-aware stage readiness.
"""
return DoctorReport(
backends=check_backends(backends),
requirements=preflight_stages(stages, stage_backends),
)
# ---------------------------------------------------------------------------
# Seed enumeration: the n-ary blast radius, visible before any work
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class EnumerationReport:
"""How many seeds ``elements`` x ``combination_sizes`` fan out into.
Mirrors ``core.seeds.element_combinations`` exactly, elements deduped and
oversized sizes yielding zero.
Attributes
----------
n_elements : int
How many distinct elements were configured.
per_size : dict
``{size: count}`` for each requested combination size.
total : int
How many seeds a run would enqueue -- the n-ary blast radius, visible
before any work starts.
oversized : tuple of int
Configured sizes larger than the element set. These are silent no-ops
during enumeration, and are named here so the user learns why, say,
``[2, 5]`` over three elements produced only pairs.
"""
n_elements: int
per_size: dict[int, int]
total: int
oversized: tuple[int, ...]
[docs]
def enumeration_report(elements: Sequence[str], sizes: Sequence[int]) -> EnumerationReport:
"""Count the seeds per combination size without materialising them.
Parameters
----------
elements
The configured element symbols; duplicates are ignored.
sizes
The configured combination sizes.
Returns
-------
EnumerationReport
Counts computed with ``math.comb``, so a large element set is reported
instantly rather than enumerated.
Examples
--------
>>> report = enumeration_report(["Na", "Cl", "O"], [2, 5])
>>> report.total, report.oversized
(3, (5,))
"""
n = len(set(elements))
per_size = {size: math.comb(n, size) for size in sorted(set(sizes))}
return EnumerationReport(
n_elements=n,
per_size=per_size,
total=sum(per_size.values()),
oversized=tuple(size for size in sorted(set(sizes)) if size > n),
)
# ---------------------------------------------------------------------------
# MACE model status + (consented) download -- the screen stage's special case
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class ModelStatus:
"""Where the screen stage's MACE model will come from.
Three cases: configured and present is ``ok=True`` with nothing to do;
configured but missing is ``ok=False``, a config error to fix; not
configured at all is ``ok=True`` with ``needs_consent=True``.
Attributes
----------
ok : bool
Whether screening can proceed as configured.
needs_consent : bool
``True`` when the MACE-MP foundation model would be downloaded on first
use. The interactive command must surface ``detail`` and obtain
consent -- a download is never silent.
detail : str
The full disclosure or the exact problem, shown to the user verbatim.
"""
ok: bool
needs_consent: bool
detail: str
[docs]
def mace_model_status(
model_path: str,
model_size: str = "medium",
backend: ComputeBackend | None = None,
) -> ModelStatus:
"""Classify the configured (or absent) MACE model path.
Parameters
----------
model_path
The screen stage's configured ``model_path``. Empty means "use the
foundation model".
model_size
Which foundation model size the disclosure should name.
backend
The backend the screen stage is wired to, which makes the check
**location-aware**, mirroring ``preflight_stages``. When it exposes the
duck-typed ``check_remote_file(path)`` hook (``ssh_exec`` does), the
file is probed on that remote -- the model lives in the remote's cache,
not on our disk. Without the hook, or without a backend, the check is a
local ``stat``, exactly right for the ``local`` backend.
Returns
-------
ModelStatus
The classification, whose ``detail`` is written to be shown verbatim.
"""
if model_path:
remote_probe = getattr(backend, "check_remote_file", None)
if callable(remote_probe):
present, detail = remote_probe(model_path)
where = str(getattr(backend, "probe_target", "remote"))
if present:
return ModelStatus(
ok=True, needs_consent=False, detail=f"MACE model on {where}: {model_path}"
)
return ModelStatus(
ok=False,
needs_consent=False,
detail=(
f"screen.model_path is set to {model_path!r} but no such file "
f"exists on {where} ({detail}) -- fix the path, or clear it to use "
f"the MACE-MP foundation model"
),
)
path = Path(model_path).expanduser()
if path.is_file():
return ModelStatus(ok=True, needs_consent=False, detail=f"local MACE model: {path}")
return ModelStatus(
ok=False,
needs_consent=False,
detail=(
f"screen.model_path is set to {model_path!r} but no such file exists -- "
f"fix the path, or clear it to use the MACE-MP foundation model"
),
)
return ModelStatus(
ok=True,
needs_consent=True,
detail=(
f"no model_path configured: the MACE-MP foundation model "
f"(size={model_size}) will be downloaded on first use into the mace "
f"cache ({MACE_CACHE_DIR}) -- a one-time download on the order of "
f"hundreds of MB. It will NOT be downloaded without your consent."
),
)
[docs]
def download_foundation_model(
python: str, model_size: str = "medium", timeout_s: float = 1800.0
) -> tuple[bool, str]:
"""Fetch the MACE-MP foundation model into the mace cache, via the user's python.
**The caller owns consent.** This must only run after the user explicitly
agreed to the download disclosed by ``mace_model_status`` -- it is the
mechanism behind that offer, never an automatic step.
Parameters
----------
python
The interpreter to download with -- the one the screen runner will use,
so the model lands in the cache that run will read.
model_size
Which foundation model to fetch.
timeout_s
How long to allow. The default is generous: this is a
hundreds-of-megabytes download.
Returns
-------
tuple
``(ok, detail)``. Never raises. Loading ``mace_mp`` downloads the model
when absent and is a no-op when already cached, so a success also
verifies the cache. A missing ``mace`` is detected and the detail names
``pip install mace-torch``.
"""
snippet = (
f"from mace.calculators import mace_mp; "
f"mace_mp(model={model_size!r}, device='cpu', default_dtype='float64')"
)
try:
proc = subprocess.run(
[python, "-c", snippet],
capture_output=True,
text=True,
timeout=timeout_s,
)
except FileNotFoundError:
return False, f"interpreter {python!r} not found"
except subprocess.TimeoutExpired:
return False, f"model download timed out after {timeout_s:.0f}s"
except OSError as exc:
return False, f"cannot run interpreter {python!r}: {exc}"
if proc.returncode == 0:
return True, f"foundation model (size={model_size}) ready in {MACE_CACHE_DIR}"
stderr_lines = (proc.stderr or "").strip().splitlines()
detail = stderr_lines[-1] if stderr_lines else f"download exited {proc.returncode}"
if "No module named" in detail:
detail = f"{detail} -- install it first: {python} -m pip install mace-torch"
return False, detail
# ---------------------------------------------------------------------------
# AiiDA readiness -- the calculate stage's special case (the mace precedent)
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class AiidaStatus:
"""Whether the calculate stage's AiiDA *configuration* is usable.
The requirement probes already answer "is aiida importable?". This is the
config-level check they cannot see, following the mace-model precedent:
does the named profile load, and does the configured DFTB+ code label
exist? Both are probed in the stage's own runner interpreter -- the AiiDA
install that will actually submit the WorkChains, not ours.
Attributes
----------
ok : bool
Whether the profile loads and the code is found.
detail : str
The problem to fix, or the confirmation. Never a secret: only profile
and code labels, plus the last line of the probe's own error.
"""
ok: bool
detail: str
[docs]
def aiida_status(
python: str, profile: str | None, code_label: str, timeout_s: float = _PROBE_TIMEOUT_S
) -> AiidaStatus:
"""Probe ``load_profile`` + ``load_code`` via the runner interpreter. Never raises.
Intended to be called only when a ``calculate`` stage is configured; the
caller folds ``ok`` into the doctor verdict.
Parameters
----------
python
The calculate stage's runner interpreter, whose AiiDA install is the
one that will submit the WorkChains.
profile
The AiiDA profile to load, or ``None`` for the user's default.
code_label
The configured DFTB+ code label to resolve.
timeout_s
How long to allow the probe.
Returns
-------
AiidaStatus
A missing aiida, an unknown profile, and an unknown code label all come
back as ``ok=False`` with the reason, never as an exception.
"""
load = f"load_profile({profile!r})" if profile else "load_profile()"
snippet = f"from aiida import load_profile, orm\n{load}\norm.load_code({code_label!r})\n"
try:
proc = subprocess.run(
[python, "-c", snippet], capture_output=True, text=True, timeout=timeout_s
)
except FileNotFoundError:
return AiidaStatus(False, f"interpreter {python!r} not found")
except subprocess.TimeoutExpired:
return AiidaStatus(False, f"AiiDA probe timed out after {timeout_s:.0f}s")
except OSError as exc:
return AiidaStatus(False, f"cannot run interpreter {python!r}: {exc}")
if proc.returncode == 0:
return AiidaStatus(
True, f"AiiDA profile {profile or 'default'!r} loads; code {code_label!r} found"
)
stderr_lines = (proc.stderr or "").strip().splitlines()
detail = stderr_lines[-1] if stderr_lines else f"probe exited {proc.returncode}"
if "No module named" in detail:
detail = f"aiida not importable in {python!r} (see stage readiness above) -- {detail}"
return AiidaStatus(False, f"AiiDA not ready: {detail}")