"""The calculate stage: one screened-in structure -> DFTB+ result via the user's AiiDA (6.5).
The thin ``Stage`` half of the pair; everything AiiDA lives entirely in
``calculate_runner.py`` (self-contained: stdlib + aiida/aiida_mydftbp/ase,
zero ``matmelt`` imports), which this class ships to the backend as an
input artifact of every job. This module itself imports **no science and no
AiiDA** -- it builds job descriptions and judges results, nothing more, so it
loads on any machine.
The runner submits one ``mydftbp.single`` WorkChain per material through the
**user's own AiiDA installation** (aiida-core + the aiida-mydftbp plugin,
whose WorkChain does the self-healing retries itself -- exit codes 300-370,
SCC/geometry/eigenvalue healing ladders, warm starts). Our package is the
launcher/extractor around it and never ships or imports AiiDA at package
level (the clean-room amendment: these are declared user-owned runtime
requirements, exactly the smact/pyxtal/mace model). 1:1, optional, and
skippable -- configs without a ``calculate:`` entry behave exactly as before.
The runner writes ``calc_result.json`` (converged flag, WorkChain exit
status, every healing attempt with SKF-stripped parameters). Verdicts follow
the G-2 doctrine: a WorkChain that *terminated* with a healing-exhausted or
configuration-fatal exit status (310, 320, 330, 340, 350, 360) is a
**terminal science verdict** -- the plugin already retried everything it
knows, so re-running reproduces it. Transient statuses (300, 370), unknown
statuses, and a broken/absent result file all stay retryable faults -- when
in doubt, fault.
``validate_inputs`` refuses clearly when ``screen_result.json`` does not say
``RELAXED`` (extract's precedent): calculate serves only screened-in
materials, and a mid-pipeline start on a screened-out item is told exactly
why before any job runs.
Every option comes from config via the constructor -- code label, AiiDA
profile, SKF locations, k-mesh, walltime, queue -- never hardcoded.
"""
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,
StageInputError,
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 calculate 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 = "calculate_runner.py"
#: Inputs from the screen stage -- declared as input artifacts so the staging
#: transport knows to push them (in-place locally; staged as a no-op).
RELAXED_NAME = "relaxed.cif"
SCREEN_NAME = "screen_result.json"
#: The result document this stage produces (attempt history included).
RESULT_NAME = "calc_result.json"
#: WorkChain exit statuses that are deterministic on retry: the plugin's
#: self-healing ladders ran to exhaustion (320/330/340/350), DFTB+ failed
#: fatally (310), or the SKF configuration is wrong (360). Re-running the
#: same WorkChain reproduces them -- terminal science verdicts, no budget.
_TERMINAL_EXIT_STATUSES = frozenset({310, 320, 330, 340, 350, 360})
#: How many characters of a failed runner's stderr to surface in the verdict.
_REASON_TAIL = 300
[docs]
class CalculateStage(Stage):
"""Run one screened-in material through DFTB+ via the user's own AiiDA.
Parameters
----------
code_label
**Required.** Label of the user's configured AiiDA code for DFTB+
(e.g. ``"dftbplus@cluster"``) -- ``verdi code list`` shows them.
profile
AiiDA profile to load; ``None`` (default) uses the user's default
profile.
skf_label
Label of the stored Slater-Koster ``FolderData`` node; ensured
automatically by the runner (find-or-create, loud failure).
skf_prefix
**Required.** Absolute directory of the ``.skf`` parameter files on
the machine where DFTB+ actually executes -- baked into the HSD
input the WorkChain runs.
skf_local_path
Local directory of ``.skf`` files used to *create* the labelled
FolderData node when no node with ``skf_label`` exists yet. ``None``
(default) means find-only: a missing node is a loud, actionable
failure, never a silent stall.
kpoints
SupercellFolding k-mesh diagonal, exactly three integers >= 1.
max_wallclock_seconds
Walltime request per WorkChain attempt, in seconds.
queue_name
Scheduler queue/partition for the WorkChain's jobs; ``None``
(default) uses the code's computer default.
poll_interval_s
Seconds between WorkChain status polls by the runner.
python
Interpreter to invoke on the backend this stage is wired to -- the
user's AiiDA environment (their venv holds aiida-core,
aiida-mydftbp and ase; never bundled).
Raises
------
StageError
For any invalid option. The two required ones, ``code_label`` and
``skf_prefix``, get messages naming ``verdi code list`` and the remote
SKF directory respectively, so the fix is actionable at wiring time
rather than after N failed WorkChains.
"""
runtime_requirements = (
Requirement(
module="aiida",
pip_name="aiida-core",
purpose="AiiDA engine the calculate runner submits WorkChains through",
),
Requirement(
module="aiida_mydftbp",
pip_name="aiida-mydftbp",
purpose="DFTB+ plugin providing the self-healing mydftbp.single WorkChain",
),
Requirement(
module="ase",
pip_name="ase",
purpose="CIF -> DFTB+ .gen geometry conversion in the calculate runner",
),
)
def __init__(
self,
code_label: str = "",
profile: str | None = None,
skf_label: str = "ptbp_complete_set",
skf_prefix: str = "",
skf_local_path: str | None = None,
kpoints: list[int] | None = None,
max_wallclock_seconds: int = 3600,
queue_name: str | None = None,
poll_interval_s: float = 30.0,
python: str = "python3",
) -> None:
if not isinstance(python, str) or not python.strip():
raise StageError(f"calculate: python must be a non-empty string, got {python!r}")
if not isinstance(code_label, str) or not code_label.strip():
raise StageError(
"calculate: code_label is required -- the label of your configured "
f"AiiDA code for DFTB+ (see 'verdi code list'), got {code_label!r}"
)
if profile is not None and (not isinstance(profile, str) or not profile.strip()):
raise StageError(
f"calculate: profile must be a non-empty string or omitted, got {profile!r}"
)
if not isinstance(skf_label, str) or not skf_label.strip():
raise StageError(f"calculate: skf_label must be a non-empty string, got {skf_label!r}")
if not isinstance(skf_prefix, str) or not skf_prefix.strip():
raise StageError(
"calculate: skf_prefix is required -- the absolute directory of the "
".skf parameter files on the machine where DFTB+ executes, got "
f"{skf_prefix!r}"
)
if skf_local_path is not None and (
not isinstance(skf_local_path, str) or not skf_local_path.strip()
):
raise StageError(
"calculate: skf_local_path must be a non-empty string or omitted, "
f"got {skf_local_path!r}"
)
if kpoints is None:
kpoints = [4, 4, 4]
if (
not isinstance(kpoints, (list, tuple))
or len(kpoints) != 3
or any(isinstance(k, bool) or not isinstance(k, int) or k < 1 for k in kpoints)
):
raise StageError(f"calculate: kpoints must be three integers >= 1, got {kpoints!r}")
if (
isinstance(max_wallclock_seconds, bool)
or not isinstance(max_wallclock_seconds, int)
or max_wallclock_seconds < 1
):
raise StageError(
"calculate: max_wallclock_seconds must be an integer >= 1, got "
f"{max_wallclock_seconds!r}"
)
if queue_name is not None and (not isinstance(queue_name, str) or not queue_name.strip()):
raise StageError(
f"calculate: queue_name must be a non-empty string or omitted, got {queue_name!r}"
)
if (
isinstance(poll_interval_s, bool)
or not isinstance(poll_interval_s, (int, float))
or poll_interval_s <= 0
):
raise StageError(
f"calculate: poll_interval_s must be a number > 0, got {poll_interval_s!r}"
)
self.runner_python = python
self._code_label = code_label
self._profile = profile
self._skf_label = skf_label
self._skf_prefix = skf_prefix
self._skf_local_path = skf_local_path
self._kpoints = tuple(kpoints)
self._max_wallclock_seconds = max_wallclock_seconds
self._queue_name = queue_name
self._poll_interval_s = float(poll_interval_s)
@property
def name(self) -> str:
return "calculate"
@property
def requires(self) -> StageInputSpec:
return StageInputSpec(
files=(RELAXED_NAME, SCREEN_NAME),
description="the relaxed structure + RELAXED screening verdict from the screen stage",
)
@property
def produces(self) -> StageOutputSpec:
return StageOutputSpec(
files=(RESULT_NAME,),
description="DFTB+ WorkChain result with the full healing-attempt history",
)
[docs]
def build_jobspec(self, item: WorkItem) -> JobSpec:
runner = Path(__file__).with_name(RUNNER_NAME)
command = [
self.runner_python,
RUNNER_NAME,
"--id",
item.id,
"--cif",
RELAXED_NAME,
"--code-label",
self._code_label,
"--skf-label",
self._skf_label,
"--skf-prefix",
self._skf_prefix,
"--kpoints",
*(str(k) for k in self._kpoints),
"--max-wallclock-seconds",
str(self._max_wallclock_seconds),
"--poll-interval",
str(self._poll_interval_s),
]
if self._profile is not None:
command += ["--profile", self._profile]
if self._skf_local_path is not None:
command += ["--skf-local-path", self._skf_local_path]
if self._queue_name is not None:
command += ["--queue-name", self._queue_name]
return JobSpec(
command=command,
input_artifacts={
RUNNER_NAME: str(runner),
RELAXED_NAME: str(item.dir / RELAXED_NAME),
SCREEN_NAME: str(item.dir / SCREEN_NAME),
},
output_artifacts=[RESULT_NAME],
cwd=str(item.dir),
scratch_inputs=frozenset({RUNNER_NAME}),
)
[docs]
def interpret_result(self, item: WorkItem, result: JobResult) -> StageVerdict:
"""Judge the finished calculation from its exit code and result file.
Parameters
----------
item
The material whose WorkChain finished; its directory holds the
result document.
result
What the backend reported for the runner process -- note this is
the *runner's* exit code, not the WorkChain's, which lives inside
the result document.
Returns
-------
StageVerdict
One of four outcomes:
- A non-zero or unknown runner exit is a retryable fault -- missing
aiida, a stopped daemon, an unloadable profile.
- Exit 0 with ``converged: true`` advances the item.
- Exit 0 with a terminal WorkChain exit status (310, 320, 330, 340,
350, 360) is a **terminal science verdict**: the plugin's own
self-healing already retried everything it knows, or the failure
is configuration-fatal, so re-running reproduces it.
- Exit 0 with a transient status (300, 370), an unknown status, or
a missing or corrupt result document is a fault -- 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"calculate runner exited {code}: {tail}")
result_path = item.dir / RESULT_NAME
try:
doc = json.loads(result_path.read_text(encoding="utf-8"))
except OSError:
return StageVerdict(
ok=False, reason=f"calculate runner exited 0 but wrote no {RESULT_NAME}"
)
except ValueError as exc:
return StageVerdict(ok=False, reason=f"{RESULT_NAME} unreadable: {exc}")
if not isinstance(doc, dict):
return StageVerdict(ok=False, reason=f"{RESULT_NAME} is not a JSON object")
if doc.get("converged") is True:
return StageVerdict(ok=True)
exit_status = doc.get("exit_status")
if exit_status in _TERMINAL_EXIT_STATUSES:
return StageVerdict(
ok=False,
reason=(
f"DFTB+ WorkChain failed with exit_status {exit_status}: the "
f"plugin's self-healing is exhausted (or the failure is "
f"configuration-fatal) -- re-running reproduces it"
),
terminal=True,
)
return StageVerdict(
ok=False,
reason=(
f"DFTB+ WorkChain did not converge (exit_status {exit_status!r}) -- "
f"transient or unknown ground, retrying"
),
)
register("calculate", CalculateStage)