"""Running calculations through vibe-queue (roadmap WP 2.5).
vibe-queue is optional (the ``queue`` extra, installed from a vibe-queue
checkout). Its Python API is imported lazily, as vibe-view does:
``vq.submit.submit_local`` to submit, ``vq.status.show_status_json`` to
follow a run, ``vq.kill.kill_job`` to cancel, ``vq.config.load_config`` for
the programs registered with the queue, and
``vq.daemon_control.is_daemon_serving`` to tell whether anything will run.
What is submitted is what the input writers write:
- **ORCA**: the input in a directory of its own, run as
``[<orca binary>, <input>]`` with the binary vibe-queue knows as the
program ``orca`` (an absolute path; vibe-queue runs no shell), the input's
``%pal`` processes as CPUs, and memory for ``%maxcore`` per process with
ORCA's recommended quarter of headroom.
- **vibe-qc**: a QVF job container (``vibesketch.document.qvf.job_container``:
the structure, the ``job.spec``, the document and
``provenance.run_status: "pending"``), submitted through vibe-queue's
single-QVF route, which runs ``vibeqc run`` from the venv program
registered for vibe-qc and needs that program's name and a pin of its
runtime (``vq.submit._program_runtime_pin_for_submit``, the call vibe-view
#86 settled). The runner updates the container in place, so the result
keeps its sketch. When no container can be built or pinned (no qvf-writer,
no registered vibe-qc program with a git identity, an older vibe-queue),
the run script goes instead, run by the Python of the program vibe-queue
knows as ``vibeqc`` or by this interpreter when vibe-qc imports here, and
:attr:`Submission.fallbacks` says why.
Each submission is recorded on the document as a job with its runs (the
queue's working directory and, for a container, its file name), so a saved
document knows its calculations and where their results are. The runs one
call starts form a group (``g1``, ``g2``, ...): a reaction's or a whole
scheme's species under one specification, each distinct species once, so
the group is followed and cancelled as one (roadmap WP 4.1).
"""
from __future__ import annotations
import functools
import importlib.util
import json
import math
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from vibesketch.document.model import Document, Job, JobSpec, Run
__all__ = [
"QueueUnavailable",
"Submission",
"TERMINAL",
"cancel",
"cancel_group",
"group_states",
"queue_status",
"reaction_species",
"run_states",
"submit",
]
#: Programs whose inputs vibe-queue can run, by writer name.
PROGRAMS = ("orca", "vibeqc")
#: States after which a run changes no more (vq.spec.TERMINAL_STATES).
TERMINAL = frozenset(
{"completed", "failed", "killed", "interrupted", "oom_killed", "starved",
"time_exceeded", "aborted_by_queue"}
) # fmt: skip
#: The tag every vibe-sketch submission carries.
TAG = "vibe-sketch"
#: The tag a QVF job container carries besides, as vibe-view's do.
CONTAINER_TAG = "qvf-container"
HOST = "localhost"
[docs]
class QueueUnavailable(RuntimeError):
"""vibe-queue is not installed, or cannot run this program."""
def _installed() -> bool:
try:
return importlib.util.find_spec("vq") is not None
except (ImportError, ValueError):
return False
def _missing() -> QueueUnavailable:
return QueueUnavailable(
"vibe-queue is not installed: install it from its checkout (pip install "
"<path-to-vibe-queue>), then the queue extra is satisfied"
)
def _programs() -> dict[str, Any]:
from vq.config import load_config
try:
return dict(load_config().programs)
except Exception as error: # vibe-queue's own configuration errors
raise QueueUnavailable(f"vibe-queue's configuration cannot be read: {error}") from None
def _orca_binary() -> str | None:
program = _programs().get("orca")
binary = getattr(program, "binary", None)
return str(binary) if binary and Path(str(binary)).is_absolute() else None
@functools.cache
def _vibeqc_runs_here() -> bool:
"""Whether this interpreter can import vibe-qc, compiled core and all.
Finding the package is not enough: the naming engine's path makes a vibe-qc
checkout's Python sources findable without its compiled core. The import
runs in a child process, once, so a failing one cannot disturb this one.
"""
if importlib.util.find_spec("vibeqc") is None:
return False
try:
child = subprocess.run(
[sys.executable, "-c", "import vibeqc"], capture_output=True, timeout=120, check=False
)
except (OSError, subprocess.TimeoutExpired):
return False
return child.returncode == 0
def _vibeqc_python() -> str | None:
for name in ("vibeqc", "vibe-qc"):
python = getattr(_programs().get(name), "python", None)
if python:
return str(python)
return sys.executable if _vibeqc_runs_here() else None
def _imports_vibeqc(program: Any) -> bool:
"""Whether a vibe-queue venv program is a vibe-qc runtime, by the module its probe imports."""
resolve = getattr(program, "effective_import_check", None)
module = resolve() if callable(resolve) else getattr(program, "import_check", None)
return isinstance(module, str) and module.strip() == "vibeqc"
def _container_program() -> tuple[str, Any]:
"""The vibe-queue program that runs QVF job containers, and its runtime pin.
vibe-queue runs a single ``.qvf`` as ``<the program's python> -m
vibeqc._cli run`` and refuses it without the program's name and a pin of
its runtime's git identity, which ``vq submit --program`` computes
itself and an API caller supplies (vibe-view #86). The program is the
venv program registered as ``vibeqc`` or ``vibe-qc``, else the only one
whose import probe is ``vibeqc``. Raises :class:`QueueUnavailable` with
the reason, which :func:`submit` records when it falls back to a script.
"""
programs = _programs()
venvs = {n: p for n, p in programs.items() if getattr(p, "kind", None) == "venv"}
named = [n for n in ("vibeqc", "vibe-qc") if n in venvs]
name: str
if named:
name = named[0]
else:
candidates = [n for n, p in venvs.items() if _imports_vibeqc(p)]
if len(candidates) != 1:
found = ", ".join(sorted(candidates)) or "none"
raise QueueUnavailable(
"a QVF job container needs one registered vibe-qc venv program in vibe-queue "
f"(found: {found}); register vibe-qc's venv as the program 'vibeqc' in "
"vibe-queue's config.toml"
)
name = candidates[0]
try:
from vq.submit import _program_runtime_pin_for_submit # the call vq submit makes
except ImportError:
raise QueueUnavailable(
"this vibe-queue cannot pin a vibe-qc runtime for a QVF job container; install "
"vibe-queue 0.26.10 or newer"
) from None
from vq.config import load_config
try:
pin = _program_runtime_pin_for_submit(load_config(), name)
except Exception as error: # vibe-queue's own configuration or git errors
raise QueueUnavailable(
f"vibe-queue cannot pin the runtime of its program {name!r}: {error}"
) from None
if pin is None or not getattr(pin, "expected_git_sha", None):
raise QueueUnavailable(
f"vibe-queue's program {name!r} has no readable git identity, so a QVF job "
"container cannot be pinned to it; check vq programs"
)
return name, pin
def _workdir(run_id: str) -> str | None:
"""The working directory vibe-queue made for a run, or ``None`` when it does not say."""
from vq.status import show_status_json
try:
workdir = json.loads(show_status_json(HOST, run_id, tail=0)).get("workdir")
except Exception: # the record is not needed to submit
return None
return str(workdir) if workdir else None
[docs]
def queue_status() -> dict[str, Any]:
"""Whether vibe-queue is installed and serving, and which programs it can run."""
if not _installed():
return {"installed": False, "daemon": False, "programs": [], "message": str(_missing())}
from vq.daemon_control import is_daemon_serving
try:
programs = [
name
for name, ready in (("orca", _orca_binary()), ("vibeqc", _vibeqc_python()))
if ready is not None
]
serving = bool(is_daemon_serving())
except Exception as error: # a status must never break the caller
return {"installed": True, "daemon": False, "programs": [], "message": str(error)}
message = (
None if serving else "the vibe-queue daemon is not running: run vq daemon run in a terminal"
)
return {"installed": True, "daemon": serving, "programs": programs, "message": message}
[docs]
@dataclass
class Submission:
"""What :func:`submit` did."""
document: Document
#: ``{"job", "species", "run", "input"}`` per run started; ``input`` is
#: ``container`` (a QVF job container), ``script`` or ``directory`` (ORCA).
submitted: list[dict[str, str]] = field(default_factory=list)
#: Per species id, why it was not submitted.
refused: dict[str, list[str]] = field(default_factory=dict)
#: Per species id, why a run script went to vibe-queue instead of a QVF job container.
fallbacks: dict[str, str] = field(default_factory=dict)
#: The group the started runs form; ``None`` when nothing was submitted.
group: str | None = None
def _job_name(species: str, spec: JobSpec) -> str:
"""A vibe-queue job name: letters, digits, ``-_.``, at most 50 characters."""
name = re.sub(r"[^A-Za-z0-9_.-]+", "-", f"vsk-{species}-{spec.method}-{spec.task}")
return name.strip("-")[:50]
def _next_group_id(document: Document) -> str:
taken = {r.group for j in document.jobs for r in j.runs or [] if r.group}
taken |= {s.id for s in document.species} | {r.id for r in document.reactions}
taken |= {j.id for j in document.jobs}
number = 1
while f"g{number}" in taken:
number += 1
return f"g{number}"
[docs]
def reaction_species(document: Document, reactions: list[str] | None = None) -> list[str]:
"""The distinct species of ``reactions`` (all of them when ``None``), in document order.
Reactants, reagents and products all count: a scheme is computed when
every species it draws is. A species in several steps appears once.
Raises the writers' ``UnknownReference`` for a reaction the document lacks.
"""
from vibesketch.inputs import UnknownReference
known = {r.id: r for r in document.reactions}
wanted = list(reactions) if reactions is not None else list(known)
for identifier in wanted:
if identifier not in known:
raise UnknownReference("reaction", identifier)
chosen: set[str] = set()
for identifier in wanted:
reaction = known[identifier]
chosen.update(reaction.reactants, reaction.reagents, reaction.products)
return [s.id for s in document.species if s.id in chosen]
def _next_job_id(document: Document) -> str:
taken = {j.id for j in document.jobs} | {s.id for s in document.species}
taken |= {r.id for r in document.reactions}
number = 1
while f"j{number}" in taken:
number += 1
return f"j{number}"
def _resources(spec: JobSpec) -> tuple[int, int | None]:
"""CPUs, and memory in MB with ORCA's 25 % headroom over ``%maxcore``."""
if spec.resources is None:
return 1, None
cpus = spec.resources.nprocs
maxcore = spec.resources.maxcore_mb
return cpus, math.ceil(cpus * maxcore / 0.75) if maxcore else None
[docs]
def submit(
document: Document,
*,
program: str,
spec: JobSpec,
species: list[str] | None = None,
reactions: list[str] | None = None,
container: bool = True,
) -> Submission:
"""Write each species' input and submit it to vibe-queue, as one group.
``species`` names species; ``reactions`` names reactions whose species
(reactants, reagents and products, each once) join them; with neither,
every species goes. The runs started form one group, named in
:attr:`Submission.group` and on each run.
A vibe-qc job goes as a QVF job container, which keeps the sketch with
its result; when none can be built or pinned, or with ``container``
false, the run script goes instead, and :attr:`Submission.fallbacks`
says why. Species the writer refuses, or that have no 3D structure, are
listed in :attr:`Submission.refused`. Raises :class:`QueueUnavailable`
when vibe-queue is missing or cannot run ``program`` either way, and the
writers' ``UnknownWriter``/``UnknownReference`` (a reaction, a species)
as :func:`write_inputs` does.
"""
from vibesketch.document.qvf import job_container, qvf_available
from vibesketch.inputs import write_inputs
if program not in PROGRAMS:
raise QueueUnavailable(f"vibe-sketch submits ORCA and vibe-qc jobs, not {program!r}")
if reactions is not None:
from_reactions = set(reaction_species(document, reactions))
named = set(species or [])
species = [s.id for s in document.species if s.id in named or s.id in from_reactions]
if not _installed():
raise _missing()
from vq.submit import submit_local
runner: str | None = _orca_binary() if program == "orca" else _vibeqc_python()
containers: tuple[str, Any] | None = None
fallback: str | None = None
if program == "vibeqc" and container:
try:
if not qvf_available():
raise QueueUnavailable(
"qvf-writer is not installed, so no QVF job container can be written"
)
containers = _container_program()
except QueueUnavailable as error:
fallback = str(error)
if runner is None and containers is None:
reason = (
"vibe-queue knows no ORCA binary: register it as the program 'orca' "
"(an absolute path) in vibe-queue's config.toml"
if program == "orca"
else "no Python with vibe-qc: register vibe-qc's venv as the program 'vibeqc' "
"in vibe-queue's config.toml"
)
if fallback is not None:
reason = f"{fallback}; and {reason}"
raise QueueUnavailable(reason)
batch = write_inputs(document, program=program, spec=spec, species=species)
result = Submission(document, refused=dict(batch.refused))
jobs = list(document.jobs)
cpus, memory = _resources(spec)
group = _next_group_id(document) if batch.bundles else None
for bundle in batch.bundles:
job = next(
(j for j in jobs if j.species == bundle.species and j.program == program
and j.spec == spec),
None,
) # fmt: skip
if job is None:
job = Job(id=_next_job_id(document.model_copy(update={"jobs": jobs})),
species=bundle.species, program=program, spec=spec) # fmt: skip
jobs.append(job)
name = _job_name(bundle.species, spec)
archive: str | None = None
with tempfile.TemporaryDirectory(prefix="vibe-sketch-") as folder:
if containers is not None:
# The container carries the document as it is, this job included.
carried = Document.model_validate(
document.model_copy(update={"jobs": jobs}).to_json_data()
)
path = Path(folder) / f"{name}.qvf"
path.write_bytes(job_container(carried, job))
run_id = submit_local(
host=HOST, input_file=path, program=containers[0],
program_runtime_pin=containers[1], cpus=cpus, job_name=name,
tags=[TAG, CONTAINER_TAG],
) # fmt: skip
kind, archive = "container", path.name
else:
for file in bundle.files:
(Path(folder) / file.name).write_text(file.content, encoding="utf-8")
if program == "orca":
run_id = submit_local(
host=HOST, directory=Path(folder), command=[runner, bundle.files[0].name],
program="orca", cpus=cpus, mem_mb=memory, job_name=name, tags=[TAG],
) # fmt: skip
kind = "directory"
else:
run_id = submit_local(
host=HOST, input_file=Path(folder) / bundle.files[0].name, python=runner,
cpus=cpus, job_name=name, tags=[TAG],
) # fmt: skip
kind = "script"
if fallback is not None:
result.fallbacks[bundle.species] = fallback
run = Run(
id=str(run_id),
submitted=datetime.now(UTC).replace(microsecond=0),
directory=_workdir(str(run_id)),
container=archive,
group=group,
)
jobs[jobs.index(job)] = job.model_copy(update={"runs": [*(job.runs or []), run]})
result.submitted.append(
{"job": job.id, "species": bundle.species, "run": run.id, "input": kind}
)
rebuilt = document.model_copy(update={"jobs": jobs})
result.document = Document.model_validate(rebuilt.to_json_data())
result.group = group if result.submitted else None
return result
[docs]
def run_states(document: Document) -> list[dict[str, Any]]:
"""Every run of the document's jobs with its state in vibe-queue, newest first."""
if not _installed():
raise _missing()
from vq.status import show_status_json
rows: list[dict[str, Any]] = []
for job in document.jobs:
for run in job.runs or []:
row: dict[str, Any] = {
"job": job.id, "species": job.species, "program": job.program,
"task": job.spec.task, "method": job.spec.method, "run": run.id,
"submitted": run.submitted.isoformat(), "directory": run.directory,
"container": run.container, "group": run.group,
} # fmt: skip
try:
status = json.loads(show_status_json(run.host, run.id, tail=0))
except FileNotFoundError:
row["state"] = "unknown"
else:
row |= {
key: status.get(key)
for key in ("state", "exit_code", "started_at", "finished_at",
"failure_reason", "workdir")
} # fmt: skip
row["finished"] = row.get("state") in TERMINAL
rows.append(row)
rows.sort(key=lambda r: str(r["submitted"]), reverse=True)
return rows
[docs]
def group_states(document: Document) -> list[dict[str, Any]]:
"""The document's run groups with their progress, newest first.
Each row names the ``group``, its ``runs`` and ``species``, the
``program``, ``task`` and ``method`` they share, a tally of the runs'
``states``, and whether the group is ``finished``.
"""
groups: dict[str, dict[str, Any]] = {}
for row in run_states(document):
name = row.get("group")
if not name:
continue
entry = groups.setdefault(
name,
{
"group": name,
"runs": [],
"species": [],
"program": row["program"],
"task": row["task"],
"method": row["method"],
"states": {},
"submitted": row["submitted"],
"finished": True,
}, # fmt: skip
)
entry["runs"].append(row["run"])
entry["species"].append(row["species"])
state = str(row.get("state"))
entry["states"][state] = entry["states"].get(state, 0) + 1
entry["finished"] = entry["finished"] and bool(row["finished"])
entry["submitted"] = min(entry["submitted"], row["submitted"])
return sorted(groups.values(), key=lambda g: (str(g["submitted"]), _group_number(g["group"])),
reverse=True) # fmt: skip
def _group_number(name: str) -> int:
"""The number in ``g12``; the newest group wins a tie on the submission second."""
digits = name[1:] if name[:1] == "g" else ""
return int(digits) if digits.isdigit() else 0
[docs]
def cancel_group(document: Document, group: str) -> list[str]:
"""Cancel the runs of ``group`` that have not finished; the ids cancelled.
Finished runs are left alone. Raises ``KeyError`` when the document has
no run in that group.
"""
runs = [run for job in document.jobs for run in job.runs or [] if run.group == group]
if not runs:
raise KeyError(group)
states = {row["run"]: row for row in run_states(document)}
cancelled = []
for run in runs:
row = states.get(run.id, {})
if row.get("finished") or row.get("state") in (None, "unknown"):
continue
try:
cancel(run.id)
except (KeyError, ValueError): # gone from the queue, or finished meanwhile
continue
cancelled.append(run.id)
return cancelled
[docs]
def cancel(run_id: str) -> str:
"""Cancel a run; raises ``ValueError`` for a finished one, ``KeyError`` for none."""
if not _installed():
raise _missing()
from vq.kill import kill_job
try:
return str(kill_job(HOST, run_id, reason="cancelled in vibe-sketch", via="vibe-sketch"))
except FileNotFoundError:
raise KeyError(run_id) from None