"""Compile a saved recipe into a vibe-queue campaign, and follow it (#42).
A recipe (:class:`vibesketch.document.model.Recipe`) becomes a
``vq.campaign/1`` file in a directory of its own, with every node's sources
pinned, which ``vq campaign run`` runs: vibe-queue's controller owns
submission, idempotency, attempts, retries and verified artifacts, and its
journal is the record of what ran (see vibe-queue's ``docs/campaigns.md``).
Per species and per numerical stage there are two nodes:
- ``<species>-<stage>-run``, a ``compute`` node: the stage's pending QVF job
container (:func:`vibesketch.document.qvf.stage_container`) and the stage
driver (:mod:`vibesketch.stages.driver`), which starts from
``start/geometry.json`` and runs vibe-qc's container runner. Its outputs
are ``result.qvf`` and ``runner.log``.
- ``<species>-<stage>-check``, a ``validate`` node: the validator
(:mod:`vibesketch.stages.validator`) judges ``result.qvf`` and publishes
``accepted/`` (``verdict.json``, the accepted ``result.qvf`` and, for an
optimisation, ``geometry.json``), or fails.
A stage that starts from an earlier stage's geometry takes the ``accepted``
artifact of that stage's check node as its ``start`` input, a success edge,
so it can only ever start from a validated geometry of a succeeded producer
attempt, bound by vibe-queue to that attempt and manifest. A stage whose
parent is not its geometry source also waits for the parent's check node to
succeed. A stage that starts from the sketch's prepared conformer finds it in
its sources, written here; the ``prepare`` stage itself is the sketch's own
preparation and compiles to no node.
The campaign is immutable: it is compiled into a new directory, never over
an existing one, and its record on the recipe names the directory, vibe-queue's
digest of it and which node is which. A MACE model's licence is acknowledged
here, by the operator compiling the campaign; the acknowledgement goes into
the compiled stage containers, never into the recipe.
"""
from __future__ import annotations
import hashlib
import json
import os
import shutil
import stat
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from vibesketch.document.model import CampaignRecord, Document, Recipe, Stage
__all__ = [
"CampaignRefused",
"CampaignUnavailable",
"Compiled",
"UnknownCampaign",
"campaign_report",
"compile_campaign",
"retry_node",
"source_digest",
]
#: The run command's interpreter on a node's host, and vibe-qc's container runner.
DEFAULT_PYTHON = "python3"
DEFAULT_RUNNER = ("{python}", "-m", "vibeqc._cli", "run")
#: vibe-queue's labels: node ids and the campaign id.
_LABEL_LENGTH = 64
_RUN_DIR = "run"
[docs]
class CampaignRefused(ValueError):
"""The recipe cannot be compiled as asked; the reasons say why."""
def __init__(self, reasons: list[str]) -> None:
self.reasons = reasons
super().__init__("; ".join(reasons))
[docs]
class CampaignUnavailable(RuntimeError):
"""vibe-queue is not installed here, so its journal cannot be read."""
[docs]
class UnknownCampaign(KeyError):
"""The recipe has no campaign with that id."""
def __str__(self) -> str:
return f"the recipe has no campaign {self.args[0]!r}"
[docs]
@dataclass
class Compiled:
"""A compiled campaign: its record, the document that keeps it, and notes."""
document: Document
record: CampaignRecord
warnings: list[str] = field(default_factory=list)
[docs]
@property
def run_command(self) -> str:
directory = Path(self.record.directory)
return f"vq campaign run {directory / 'campaign.json'} {directory / _RUN_DIR}"
# ── vibe-queue's pins ───────────────────────────────────────────────────
def _canonical(payload: object) -> str:
text = json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False)
return hashlib.sha256(text.encode("utf-8")).hexdigest()
[docs]
def source_digest(root: Path) -> str:
"""vibe-queue's pin of a node's sources directory (``vq.campaign-sources/1``).
Every entry's relative path and type, and each file's size and SHA-256,
in name order, digested as canonical JSON: the algorithm of vibe-queue's
``vq.campaign.source_digest``, which a contract test checks when
vibe-queue is installed. Links and special files are refused.
"""
entries: list[dict[str, Any]] = []
def walk(directory: Path, prefix: str) -> None:
for child in sorted(os.scandir(directory), key=lambda e: e.name):
relative = f"{prefix}{child.name}"
info = child.stat(follow_symlinks=False)
if stat.S_ISDIR(info.st_mode):
entries.append({"path": relative, "type": "directory"})
walk(Path(child.path), relative + "/")
elif stat.S_ISREG(info.st_mode):
data = Path(child.path).read_bytes()
entries.append(
{
"path": relative,
"type": "file",
"size": len(data),
"sha256": hashlib.sha256(data).hexdigest(),
}
)
else:
raise CampaignRefused([f"sources {root} contain {relative}, not a regular file"])
walk(root, "")
return _canonical({"schema_name": "vq.campaign-sources/1", "entries": entries})
# ── Compiling ───────────────────────────────────────────────────────────
def _node_id(species: str, stage: str, role: str) -> str:
return f"{species}-{stage}-{role}"
def _start(recipe: Recipe, stage: Stage) -> dict[str, Any]:
"""Where ``stage`` starts: a conformer index, or the stage whose accepted geometry it takes."""
source = stage.geometry.stage
if source is not None:
origin = recipe.stage_by_id(source)
if origin.kind != "prepare":
return {"source": "stage", "stage": source, "conformer": None}
source_conformer = origin.geometry.conformer
return {"source": "conformer", "stage": source, "conformer": source_conformer}
return {"source": "conformer", "stage": None, "conformer": stage.geometry.conformer}
def discard_output(target: Path, existed: bool) -> None:
"""Undo a failed write into ``target``, which was checked new or empty beforehand.
A directory this call created goes; one the user gave keeps existing, and
only what was written into it is removed (it was empty when checked).
"""
if not existed:
shutil.rmtree(target, ignore_errors=True)
return
for child in list(target.iterdir()) if target.is_dir() else []:
if child.is_dir() and not child.is_symlink():
shutil.rmtree(child, ignore_errors=True)
else:
child.unlink(missing_ok=True)
def inside(root: Path, relative: str, what: str) -> Path:
"""``root / relative``, refused unless it stays inside ``root``.
Paths in a campaign's files and records may come from a document someone
else wrote; none of them may lead a reader outside the campaign (#69).
"""
candidate = (root / relative).resolve()
if Path(relative).is_absolute() or not candidate.is_relative_to(root.resolve()):
raise CampaignRefused([f"{what} {relative!r} leads outside {root}"])
return candidate
def _write_json(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
[docs]
def compile_campaign(
document: Document,
recipe_id: str,
directory: str | os.PathLike[str],
*,
campaign_id: str | None = None,
host: str | None = None,
program: str | None = "vibeqc",
python: str = DEFAULT_PYTHON,
runner: list[str] | None = None,
cpus: int = 1,
mem_mb: int | None = None,
accept_model_license: bool = False,
) -> Compiled:
"""Compile ``recipe_id`` into a vibe-queue campaign in ``directory``.
``directory`` must not exist or be empty. ``host`` and ``program`` are
the vibe-queue host and program the compute nodes run on (the program
whose interpreter ``python`` is); ``runner`` is the container runner's
command, ``{python}`` standing for the node's interpreter. A recipe with
a MACE stage needs ``accept_model_license``: the operator compiling the
campaign acknowledges the model's licence. Raises
:class:`CampaignRefused` with every reason the recipe cannot run:
refused stages, species without a prepared conformer, an unacknowledged
licence, ids too long for vibe-queue, a directory in use.
"""
from vibesketch.document.qvf import QvfError, qvf_available, stage_container
from vibesketch.inputs.recipe import resolve_stage
from vibesketch.stages import DRIVER, VALIDATOR
recipe = next((r for r in document.recipes if r.id == recipe_id), None)
if recipe is None:
raise CampaignRefused([f"the document has no recipe {recipe_id!r}"])
target = Path(directory).expanduser().resolve()
identifier = campaign_id or f"{recipe.id}-c{len(recipe.campaigns) + 1}"
reasons: list[str] = []
if any(c.id == identifier for c in recipe.campaigns):
reasons.append(f"recipe {recipe.id} already has a campaign {identifier!r}")
if target.exists() and (not target.is_dir() or any(target.iterdir())):
reasons.append(
f"{target} is not an empty directory; a compiled campaign is never rewritten"
)
if len(identifier) > _LABEL_LENGTH:
reasons.append(f"the campaign id {identifier!r} is longer than {_LABEL_LENGTH} characters")
if cpus < 1 or (mem_mb is not None and mem_mb < 1):
reasons.append("cpus and mem_mb must be positive")
from vibesketch.limits import MAX_CAMPAIGN_NODES, MAX_CPUS, MAX_MEM_MB
if cpus > MAX_CPUS or (mem_mb is not None and mem_mb > MAX_MEM_MB):
reasons.append(f"a node asks for at most {MAX_CPUS:,} cpus and {MAX_MEM_MB:,} MB")
if not qvf_available():
reasons.append("qvf-writer is not installed, so no stage container can be written")
species = {s.id: s for s in document.species}
numerical = [s for s in recipe.stages if s.kind != "prepare"]
if not numerical:
reasons.append(f"recipe {recipe.id} has no numerical stage to run")
nodes_wanted = 2 * len(numerical) * len(recipe.species)
if nodes_wanted > MAX_CAMPAIGN_NODES:
# Each node is a directory of files written before the queue sees it (#69).
reasons.append(
f"recipe {recipe.id} would compile to {nodes_wanted:,} nodes, more than the "
f"{MAX_CAMPAIGN_NODES:,} a campaign holds; split it over several recipes"
)
if any(s.kind == "mace" for s in numerical) and not accept_model_license:
reasons.append(
"the recipe runs a MACE model: an eligible operator acknowledges its licence when "
"compiling (accept_model_license / --accept-model-license); a recipe never carries it"
)
for species_id in recipe.species:
entry = species[species_id]
for stage in numerical:
start = _start(recipe, stage)
if start["source"] == "conformer" and len(entry.conformers) <= int(start["conformer"]):
reasons.append(
f"{species_id}: stage {stage.id} starts from conformer {start['conformer']}, "
"which the species does not have; embed it first"
)
resolved = resolve_stage(stage, entry)
reasons.extend(f"{species_id} {stage.id}: {why}" for why in resolved.refusals)
for role in ("run", "check"):
node = _node_id(species_id, stage.id, role)
if len(node) > _LABEL_LENGTH:
reasons.append(f"node id {node!r} is longer than {_LABEL_LENGTH} characters")
if reasons:
raise CampaignRefused(reasons)
# ``{python}`` stays in the command: the driver puts in the interpreter it runs under.
command = list(runner or DEFAULT_RUNNER)
sources = target / "sources"
nodes: list[dict[str, Any]] = []
roles: dict[str, dict[str, str]] = {}
existed = target.exists()
try:
target.mkdir(parents=True, exist_ok=True)
for species_id in recipe.species:
entry = species[species_id]
symbols = entry.conformers[0].symbols
for stage in numerical:
start = _start(recipe, stage)
resolved = resolve_stage(stage, entry)
run_id = _node_id(species_id, stage.id, "run")
check_id = _node_id(species_id, stage.id, "check")
run_dir = sources / run_id
run_dir.mkdir(parents=True)
container = stage_container(document, recipe.id, stage.id, species_id)
if stage.kind == "mace":
container = _acknowledged(container, accept_model_license)
(run_dir / "stage.qvf").write_bytes(container)
shutil.copyfile(DRIVER, run_dir / "run_stage.py")
_write_json(
run_dir / "stage.json",
{
"schema": "vibe-sketch.stage/1",
"recipe": recipe.id,
"stage": stage.id,
"species": species_id,
"kind": stage.kind,
"symbols": symbols,
"charge": entry.charge,
"multiplicity": entry.multiplicity,
"start": start,
"container": "stage.qvf",
"runner": command,
},
)
inputs: list[dict[str, Any]] = []
after: list[dict[str, Any]] = []
if start["source"] == "conformer":
conformer = entry.conformers[int(start["conformer"])]
_write_json(
run_dir / "start" / "geometry.json",
{
"schema": "vibe-sketch.accepted-geometry/1",
"recipe": recipe.id,
"stage": start["stage"],
"species": species_id,
"symbols": conformer.symbols,
"charge": entry.charge,
"multiplicity": entry.multiplicity,
"units": "angstrom",
"positions_angstrom": conformer.xyz_angstrom,
"method": conformer.method,
"conformer": start["conformer"],
"result_sha256": None,
},
)
else:
source = _node_id(species_id, str(start["stage"]), "check")
inputs.append(
{"name": "start", "source": f"{source}.accepted", "path": "start"}
)
parent = stage.parent
if parent is not None and recipe.stage_by_id(parent).kind != "prepare":
parent_check = _node_id(species_id, parent, "check")
if not any(i["source"].startswith(parent_check + ".") for i in inputs):
after.append({"node": parent_check, "on": "success"})
run_node: dict[str, Any] = {
"id": run_id,
"kind": "compute",
"command": [python, "run_stage.py"],
"cpus": cpus,
"sources": {"path": f"sources/{run_id}", "sha256": source_digest(run_dir)},
"inputs": inputs,
"outputs": [
{"name": "result", "path": "result.qvf", "kind": "file"},
{"name": "log", "path": "runner.log", "kind": "file"},
],
"after": after,
}
if host is not None:
run_node["host"] = host
if program is not None:
run_node["program"] = program
if mem_mb is not None:
run_node["mem_mb"] = mem_mb
check_dir = sources / check_id
check_dir.mkdir(parents=True)
shutil.copyfile(VALIDATOR, check_dir / "validate_stage.py")
_write_json(
check_dir / "check.json",
{
"schema": "vibe-sketch.check/1",
"recipe": recipe.id,
"stage": stage.id,
"species": species_id,
"kind": stage.kind,
"symbols": symbols,
"charge": entry.charge,
"multiplicity": entry.multiplicity,
"checks": resolved.checks,
},
)
check_node: dict[str, Any] = {
"id": check_id,
"kind": "validate",
"command": [python, "validate_stage.py"],
"cpus": 1,
"sources": {"path": f"sources/{check_id}", "sha256": source_digest(check_dir)},
"inputs": [
{"name": "result", "source": f"{run_id}.result", "path": "result.qvf"}
],
"outputs": [{"name": "accepted", "path": "accepted", "kind": "directory"}],
}
if host is not None:
check_node["host"] = host
nodes += [run_node, check_node]
roles[run_id] = {"species": species_id, "stage": stage.id, "role": "run"}
roles[check_id] = {"species": species_id, "stage": stage.id, "role": "check"}
campaign = {"schema_name": "vq.campaign/1", "campaign_id": identifier, "nodes": nodes}
_write_json(target / "campaign.json", campaign)
digest, warnings = _vibe_queue_digest(target / "campaign.json")
except (QvfError, OSError) as error:
discard_output(target, existed)
raise CampaignRefused([str(error)]) from None
except CampaignRefused:
discard_output(target, existed)
raise
record = CampaignRecord(
id=identifier,
directory=str(target),
campaign_sha256=digest,
created=datetime.now(UTC).replace(microsecond=0),
model_license_acknowledged=accept_model_license
and any(s.kind == "mace" for s in numerical),
nodes=roles,
)
recipes = [
r.model_copy(update={"campaigns": [*r.campaigns, record]}) if r.id == recipe.id else r
for r in document.recipes
]
rebuilt = document.model_copy(update={"recipes": recipes})
return Compiled(Document.model_validate(rebuilt.to_json_data()), record, warnings)
def _acknowledged(container: bytes, accepted: bool) -> bytes:
"""The MACE stage container with the operator's licence acknowledgement in its options."""
import io
import zipfile
source = zipfile.ZipFile(io.BytesIO(container))
manifest = json.loads(source.read("manifest.json"))
job = next(s for s in manifest["sections"] if s["kind"] == "job.spec")
member = job["members"]["spec"]
spec = json.loads(source.read(member["path"]))
spec.setdefault("options", {}).setdefault("mlip_options", {})["accept_academic_license"] = (
accepted
)
data = json.dumps(spec, indent=2).encode("utf-8")
member["sha256"] = hashlib.sha256(data).hexdigest()
if "size" in member:
member["size"] = len(data)
buffer = io.BytesIO()
with source, zipfile.ZipFile(buffer, "w") as result:
for info in source.infolist():
if info.filename == "manifest.json":
payload = json.dumps(manifest, indent=2).encode("utf-8")
elif info.filename == member["path"]:
payload = data
else:
payload = source.read(info.filename)
result.writestr(info, payload, compress_type=info.compress_type)
return buffer.getvalue()
def _vibe_queue_digest(campaign_file: Path) -> tuple[str | None, list[str]]:
"""vibe-queue's own check and digest of the campaign file, when it is installed."""
try:
from vq.campaign import campaign_digest, check_sources, load_campaign
except ImportError:
return None, [
"vibe-queue is not installed here, so the campaign was not checked by it and its "
f"digest is not recorded; run vq campaign validate {campaign_file}"
]
try:
campaign = load_campaign(campaign_file)
except ValueError as error:
raise CampaignRefused([f"vibe-queue refuses the compiled campaign: {error}"]) from None
problems = [
f"node {node}: {item['problem']}"
for node, item in check_sources(campaign, campaign_file.parent).items()
if item["problem"]
]
if problems:
raise CampaignRefused(problems)
return str(campaign_digest(campaign)), []
# ── Following ───────────────────────────────────────────────────────────
def _record(document: Document, recipe_id: str, campaign_id: str) -> tuple[Recipe, CampaignRecord]:
recipe = next((r for r in document.recipes if r.id == recipe_id), None)
if recipe is None:
from vibesketch.document.recipe import UnknownRecipe
raise UnknownRecipe(recipe_id)
record = next((c for c in recipe.campaigns if c.id == campaign_id), None)
if record is None:
raise UnknownCampaign(campaign_id)
return recipe, record
def _vq() -> Any:
try:
import vq.campaign as campaign_module
except ImportError:
raise CampaignUnavailable(
"vibe-queue is not installed, so its campaign journal cannot be read: install it "
"from a vibe-queue checkout (pip install '<vibe-queue checkout>')"
) from None
return campaign_module
[docs]
def campaign_report(document: Document, recipe_id: str, campaign_id: str) -> dict[str, Any]:
"""Where a compiled campaign stands: process state and scientific acceptance, apart.
Per species and stage: the run node's and the check node's state in
vibe-queue's journal (``waiting``, ``ready``, ``running``, ``succeeded``,
``failed``, ``blocked`` ...), their attempts with what each retried, and
``acceptance``: ``accepted`` with the verdict (review notes, energy,
frequencies, the software that ran) once the check node's artifact is
verified, ``rejected`` when the check failed, ``not run`` when the run
itself failed, else ``pending``. Raises :class:`CampaignUnavailable`
without vibe-queue.
"""
vq = _vq()
recipe, record = _record(document, recipe_id, campaign_id)
directory = Path(record.directory)
campaign = vq.load_campaign(directory / "campaign.json")
digest = vq.campaign_digest(campaign)
journal_path = directory / _RUN_DIR / "journal.json"
if journal_path.exists():
journal = vq.load_journal(journal_path)
else:
journal = vq.Journal.model_validate(
{"schema_name": "vq.campaign-journal/1", "campaign_sha256": digest}
)
status = vq.campaign_status(campaign, journal)
by_node = {entry["node"]: entry for entry in status["nodes"]}
stages: list[dict[str, Any]] = []
for species_id in recipe.species:
for stage in recipe.stages:
if stage.kind == "prepare":
entry = next(s for s in document.species if s.id == species_id)
stages.append(
{
"species": species_id, "stage": stage.id, "kind": stage.kind,
"acceptance": "prepared in the sketch",
"conformer": stage.geometry.conformer,
"method": entry.conformers[stage.geometry.conformer or 0].method
if entry.conformers else None,
}
) # fmt: skip
continue
run = by_node.get(_node_id(species_id, stage.id, "run"), {})
check = by_node.get(_node_id(species_id, stage.id, "check"), {})
row: dict[str, Any] = {
"species": species_id,
"stage": stage.id,
"kind": stage.kind,
"run": _process(run),
"check": _process(check),
"acceptance": "pending",
"verdict": None,
}
if run.get("status") == "failed":
row["acceptance"] = "not run"
if check.get("status") == "failed":
row["acceptance"] = "rejected"
elif check.get("status") == "succeeded":
attempt = len(check.get("attempts") or [])
verdict_path = (
directory / _RUN_DIR / "artifacts" / check["node"] / str(attempt)
/ "accepted" / "verdict.json"
) # fmt: skip
if verdict_path.is_file():
row["verdict"] = json.loads(verdict_path.read_text(encoding="utf-8"))
row["acceptance"] = "accepted" if row["verdict"].get("accepted") else "rejected"
else:
row["acceptance"] = "accepted, not yet published"
stages.append(row)
return {
"campaign": record.id,
"directory": str(directory),
"campaign_sha256": digest,
"recorded_sha256": record.campaign_sha256,
"status": status["status"],
"stages": stages,
"run_command": f"vq campaign run {directory / 'campaign.json'} {directory / _RUN_DIR}",
}
def _process(entry: dict[str, Any]) -> dict[str, Any]:
if not entry:
return {"node": None, "status": "absent"}
attempts = [
{
"attempt": a.get("attempt"),
"state": a.get("state"),
"retry_of": a.get("retry_of"),
"job": a.get("job"),
}
for a in entry.get("attempts") or []
]
return {
"node": entry["node"],
"status": entry["status"],
"reasons": entry.get("reasons") or [],
"next": (entry.get("next") or {}).get("action"),
"attempts": attempts,
}
[docs]
def retry_node(document: Document, recipe_id: str, campaign_id: str, node: str) -> int:
"""Authorise the next attempt of a failed node; returns its attempt number.
vibe-queue records the attempt as retrying the failed one, staged from the
producers' verified artifacts and under its own idempotency key; the next
``vq campaign run`` sends it. Raises :class:`CampaignRefused` when the
node is not a failed node of the campaign, and
:class:`CampaignUnavailable` without vibe-queue.
"""
vq = _vq()
_, record = _record(document, recipe_id, campaign_id)
if node not in record.nodes:
raise CampaignRefused([f"campaign {campaign_id} has no node {node!r}"])
from vq.campaign_controller import CliQueue, Controller, locked
directory = Path(record.directory)
run_dir = directory / _RUN_DIR
run_dir.mkdir(exist_ok=True)
try:
with locked(run_dir):
controller = Controller(
campaign=vq.load_campaign(directory / "campaign.json"),
run_dir=run_dir,
queue=CliQueue(),
sources_root=directory,
)
attempt = controller.authorize_retry(node)
except (vq.CampaignError, OSError) as error:
raise CampaignRefused([str(error)]) from None
return int(attempt.attempt)