"""The ``run`` command: preflight gate, seed, then drain the funnel (5.1d).
The productized ``mock_test.md`` driver. Order of operations:
1. **Config** -- ``load_config`` -> ``build_runtime`` + ``build_stages``
(exit 2 on any config problem, YAML paths included).
2. **Preflight gate** -- the enumeration report first (the n-ary blast radius
is visible before anything else), then ``preflight_stages`` over the
**active** stages only -- the same ``stop_after`` prefix the orchestrator
will execute, so a stage this run never reaches cannot block it. Not
ready -> exit 1 with one exact install command per problem, before any
work starts (the 3.6 promise).
3. **Model consent** -- when the screen stage is active and no ``model_path``
is configured, the foundation-model download is disclosed and consent
obtained (``--yes`` for non-interactive runs); declining aborts cleanly.
With consent, the cache is warmed up front via the user's own interpreter
-- except when screen runs on a remote backend, where the model must live
in the remote's cache and will be fetched there on first use (disclosed).
A configured-but-missing model path is a config problem: refuse.
4. **Seed + drain** -- ``element_combinations`` seeds the store
(idempotent -- rerunning a finished or interrupted run is always safe),
then a single-threaded drain loop: ``run_pass`` until every material is
terminal, one progress line per pass with activity. ``--max-passes`` is
the escape hatch for smoke tests; Ctrl-C exits cleanly (exit 0) -- state
and in-flight jobs are durable, and the same command resumes.
``--materials``/``--materials-file`` make the run **selective**: only the
named materials -- and their descendants, via the orchestrator's lineage
walk -- are seeded, claimed, and finalized. Unselected materials sit
untouched in whatever state they hold, so the drain condition switches to
only-mode's "a pass with zero activity and nothing in flight".
A **mid-funnel ``stop_after``** (naming any stage but the configured last
one) uses the same drain condition: its materials rest at the stop stage's
output state -- live, resumable by a later run -- and never finalize, so an
"every material terminal" drain would spin forever.
Exit codes: 0 = the drain completed or was cleanly interrupted/limited
(science ``FAILED`` verdicts are *results*, reported in the summary -- not a
process failure); 1 = refused to start; 2 = config error.
"""
from __future__ import annotations
import time
from pathlib import Path
from typing import TYPE_CHECKING
from matmelt.commands import STATE_DB_FILENAME, resolve_config_path
from matmelt.config import ConfigError, build_runtime, build_stages, load_config
from matmelt.core.orchestrator import Orchestrator
from matmelt.core.preflight import (
download_foundation_model,
enumeration_report,
is_ready,
mace_model_status,
preflight_stages,
)
from matmelt.core.seeds import element_combinations
from matmelt.core.state import StateStore, is_terminal
from matmelt.ui.console import (
confirm_model_download,
render_config_error,
render_enumeration,
render_error,
render_interrupted,
render_note,
render_pass_summary,
render_run_header,
render_run_result,
render_unready,
render_warning,
)
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from matmelt.backends.base import ComputeBackend
from matmelt.config import Config
from matmelt.core.stages.base import Stage
[docs]
def run_run(
config_path: str,
*,
assume_yes: bool = False,
max_passes: int | None = None,
only: str | None = None,
materials: str | None = None,
materials_file: str | None = None,
) -> int:
"""Run the funnel a config describes, to completion or until interrupted.
Parameters
----------
config_path
The config file to run.
assume_yes
Accept the foundation-model download consent without prompting, for
unattended runs.
max_passes
Stop after this many orchestrator passes. An escape hatch for
debugging; the default drains.
only
Run exactly one stage on demand. The preflight gates just that stage,
seeding is skipped unless it is ``compose`` -- fresh ``NEW`` items
would otherwise sit unconsumed forever -- and the drain stops when the
stage has no work left. Conflicts loudly with ``pipeline.stop_after``.
materials
Comma-separated material ids to restrict this run to.
materials_file
A file of ids, one per line, with blanks and ``#`` comments skipped.
Combines with ``materials``, and both combine freely with ``only``.
The selection covers each named material **and its descendants**, so
fan-out children are eligible automatically.
Returns
-------
int
``0`` when the run drained or was interrupted cleanly -- **science
failures are results, not process failures**, so a run whose materials
were all screened out still exits 0. ``1`` when the environment is not
ready or consent was declined; ``2`` for a config or usage error.
Notes
-----
Ctrl-C is a supported way to stop: in-flight jobs stay recorded, and the
next ``run`` re-attaches to them instead of resubmitting.
"""
config_path = resolve_config_path(config_path)
try:
config = load_config(config_path)
runtime = build_runtime(config)
stages = build_stages(config)
except ConfigError as exc:
render_config_error(str(exc))
return 2
selection: frozenset[str] | None = None
if materials is not None or materials_file is not None:
picked = {part.strip() for part in (materials or "").split(",") if part.strip()}
if materials_file is not None:
try:
lines = Path(materials_file).read_text(encoding="utf-8").splitlines()
except OSError as exc:
render_error(f"run: cannot read --materials-file {materials_file}: {exc}")
return 2
picked.update(
line.strip() for line in lines if line.strip() and not line.lstrip().startswith("#")
)
if not picked:
render_error(
"run: --materials/--materials-file given but the selection is empty -- "
"name at least one material id."
)
return 2
selection = frozenset(picked)
names = [stage.name for stage in stages]
stop_after = config.pipeline.stop_after
# A mid-funnel stop_after never finalizes (the orchestrator's rule):
# materials rest at the stop stage's output, so the drain condition and
# the closing message both need to know.
finalizing = stop_after is None or stop_after == names[-1]
if only is not None:
if stop_after is not None:
render_error(
"run: --only and pipeline.stop_after are mutually exclusive -- "
"remove stop_after from the config or drop --only."
)
return 2
if only not in names:
render_error(
f"run: --only {only!r} names no configured stage; configured: {', '.join(names)}"
)
return 2
# Exactly one stage may gate this run.
active = [stage for stage in stages if stage.name == only]
else:
# The same active prefix the orchestrator will compute from
# stop_after: only what this run actually executes may gate it.
stop_index = names.index(stop_after) if stop_after is not None else len(stages) - 1
active = list(stages[: stop_index + 1])
render_run_header(config.run.name, [stage.name for stage in active])
render_enumeration(enumeration_report(config.run.elements, config.run.combination_sizes))
statuses = preflight_stages(active, runtime.stage_backends)
if not is_ready(statuses):
render_unready([status for status in statuses if not status.satisfied])
return 1
exit_code = _model_gate(config, runtime.stage_backends, active, assume_yes=assume_yes)
if exit_code != 0:
return exit_code
work_dir = Path(config.run.work_dir)
work_dir.mkdir(parents=True, exist_ok=True)
store = StateStore(work_dir / STATE_DB_FILENAME)
try:
orchestrator = Orchestrator(
store,
stages,
runtime.stage_backends,
work_dir,
stop_after=stop_after,
only=only,
materials=selection,
retry_budget=config.pipeline.retry_budget,
lease_timeout_s=config.pipeline.lease_timeout_s,
)
if only is None or only == "compose":
seeds = element_combinations(config.run.elements, config.run.combination_sizes)
if selection is not None:
# A selective run seeds only the chosen combinations; anything
# else in the selection is an existing (or descendant) id.
seeds = [seed for seed in seeds if seed in selection]
if seeds:
orchestrator.seed(seeds)
render_note(
f"Seeded {len(seeds)} combination(s) matching --materials; "
f"starting the drain loop."
)
else:
render_note(
"--materials: no seed combination matched -- relying on "
"materials already in the store."
)
else:
orchestrator.seed(seeds)
render_note(f"Seeded {len(seeds)} combination(s); starting the drain loop.")
else:
# Seeding a mid-funnel only-run would enqueue NEW items nothing
# in this run consumes -- the drain loop would never exit.
render_note(
f"--only {only}: seeding skipped -- this run consumes materials "
f"already waiting for the {only} stage."
)
if selection is not None:
# Typo protection, never a refusal: a selected id may also be a
# descendant a future pass produces.
for unknown in sorted(mid for mid in selection if store.get(mid) is None):
render_warning(
f"--materials: {unknown!r} matches no material in the store "
f"(a typo, or a descendant not yet produced)."
)
drained = False
passes = 0
try:
while max_passes is None or passes < max_passes:
summary = orchestrator.run_pass()
passes += 1
render_pass_summary(passes, summary)
if only is None and selection is None and finalizing:
if all(is_terminal(state) for state in store.count_by_state()):
drained = True
break
# Only-mode, selective, and mid-funnel-stop runs: other
# states (unselected materials, or items resting at the stop
# stage's output) legitimately hold non-terminal items, so
# "drained" is a pass with no activity at all and nothing in
# flight -- no further progress is possible.
elif not any(summary.values()) and not store.list_jobs():
drained = True
break
time.sleep(config.pipeline.pass_interval_s)
except KeyboardInterrupt:
render_interrupted()
return 0
render_run_result(
store.count_by_state(),
drained=drained,
only=only,
selective=selection is not None,
stopped_early=None if finalizing else stop_after,
)
return 0
finally:
store.close()
def _model_gate(
config: Config,
stage_backends: Mapping[str, ComputeBackend],
active: Sequence[Stage],
*,
assume_yes: bool,
) -> int:
"""Settle where the screening model comes from, asking consent if needed.
Parameters
----------
config
The validated configuration.
stage_backends
Used to stat a configured ``model_path`` **where screen will run**, so
a model living in a remote's cache is not misjudged as missing from our
disk.
active
The stages this run executes. No screen stage means nothing to gate.
assume_yes
Skip the prompt and accept the download.
Returns
-------
int
``0`` to proceed, ``1`` to refuse -- a broken ``model_path``, a
declined download, or a download that failed.
Notes
-----
A foundation-model download is **never silent**: the size and cache
location are disclosed and consent is asked for, with EOF read as a
decline. When screen runs remotely the download is the remote's business,
and the user is told so rather than warming a local cache nothing will
read.
"""
screen_stage = next((stage for stage in active if stage.name == "screen"), None)
if screen_stage is None:
return 0
screen_cfg = config.stages["screen"]
model_size = screen_cfg.options.get("model_size", "medium")
backend = stage_backends["screen"]
# Stat a configured model_path where screen will run (remote-aware), so a
# model living in the remote's cache is not misjudged as missing on our disk.
model = mace_model_status(screen_cfg.options.get("model_path", ""), model_size, backend=backend)
if not model.ok:
render_error(f"run: {model.detail}")
return 1
if not model.needs_consent:
return 0
if not assume_yes and not confirm_model_download(model.detail):
render_error(
"run: model download declined -- set screen.model_path to a local "
"model file, or re-run and consent to the download."
)
return 1
if callable(getattr(backend, "check_remote_import", None)):
# Warming a remote cache from here is not possible; the runner fetches
# it into the remote's own cache on first use (consent covers it).
render_note(
"Consent recorded: the screen stage runs remotely -- the foundation "
"model will be fetched into the remote's cache on first use."
)
return 0
ok, detail = download_foundation_model(screen_stage.runner_python, model_size)
if not ok:
render_error(f"run: foundation-model download failed: {detail}")
return 1
render_note(detail)
return 0