Source code for matmelt.core.state

"""The coordination plane: the funnel state machine + the SQLite ``StateStore`` (2.1).

Everything the orchestrator knows about "where each material is" lives here, in
two parts:

1. **The state machine** -- the fixed lifecycle a candidate material moves
   through, encoded exactly once (``MaterialState``, ``next_state``,
   ``is_terminal``, ``can_transition``) so there is a single definition of
   "what comes next" and "what is finished". The full funnel::

       NEW -> COMPOSED -> STRUCTURED -> SCREENED -> CALCULATED -> EXTRACTED -> DONE

   plus two terminal states outside the linear chain, each reachable from any
   live state: ``FAILED`` (a hard failure) and ``EXPANDED`` (a seed or parent
   retired after fanning out into children -- terminal but *not* a failure, so
   ``FAILED`` keeps meaning genuine failure in every downstream query).
   ``CALCULATED`` (Phase 6) is the output of the *optional* ``calculate``
   stage: a run that skips it walks SCREENED straight to EXTRACTED via
   ``complete``'s explicit ``target_state`` -- which stages a run executes is
   run configuration, not lifecycle physics (the ``finalize`` doctrine).

2. **The store** -- one concrete SQLite-backed ``StateStore`` (no ABC: this is a
   single-user tool; the contract lives in docstrings and tests, and an
   interface gets extracted only if a second database ever actually arrives).
   SQLite runs in WAL mode with a ``busy_timeout``, and every read-then-write
   sequence executes inside a ``BEGIN IMMEDIATE`` transaction, which is what
   makes ``claim`` an atomic lease even across separate store instances on the
   same file.

Three durability properties the store guarantees (the restart-safety thesis):

* **Atomic leasing.** ``claim`` selects unclaimed rows and stamps them inside
  one immediate transaction -- two claimers can never be handed the same item.
* **Idempotent progress.** ``complete`` is keyed on a ``(material_id, stage)``
  completions ledger: replaying a completion after a crash advances nothing
  twice. ``enqueue``, ``expand``, ``clear_job`` and terminal ``fail`` are
  likewise safe to replay.
* **The in-flight job ledger.** ``record_job`` persists every submitted backend
  handle *before* any polling, so a restarted run re-attaches to its live jobs
  instead of resubmitting them. A material referenced by a live record is never
  swept by ``release_stale_leases`` -- the recorded handle owns it until
  ``clear_job``. One live job per material is enforced by the schema itself
  (the link table's primary key).

This module never imports the config layer, the CLI layer, or a backend: the
handle it stores is an opaque JSON-safe dict, converted from/to the compute
plane's ``Handle`` at the orchestrator boundary (2.2).

**On the ``# nosec B608`` markers.** Several statements below build their SQL
with an f-string. In every one of them the interpolated part is *only* a run of
``?`` placeholders (or a fixed ``WHERE`` fragment assembled here) -- never a
value, never a caller-supplied identifier; the values themselves are always
passed to ``execute`` as bound parameters. That is the standard way to express
a variable-length ``IN (...)`` in sqlite3, and it is why bandit's SQL-injection
heuristic is suppressed at those exact lines and nowhere else.
"""

from __future__ import annotations

import json
import sqlite3
import threading
import time
from collections.abc import Collection, Iterator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# The state machine (defined once; every store method defers to it)
# ---------------------------------------------------------------------------


[docs] class MaterialState(str, Enum): """The lifecycle states of one candidate material. A ``str`` enum: members compare equal to their plain-string values and persist as text (``state.value`` into a column, ``MaterialState(text)`` back out) with no custom conversion. The seven forward states are listed in funnel order; ``FAILED`` and ``EXPANDED`` sit outside the chain, reachable from any live state. ``CALCULATED`` belongs to the *optional* ``calculate`` stage -- runs that skip it never park a material there. """ NEW = "NEW" COMPOSED = "COMPOSED" STRUCTURED = "STRUCTURED" SCREENED = "SCREENED" CALCULATED = "CALCULATED" EXTRACTED = "EXTRACTED" DONE = "DONE" FAILED = "FAILED" EXPANDED = "EXPANDED"
#: The forward funnel, in order. The successor of ``LIFECYCLE[i]`` is #: ``LIFECYCLE[i + 1]``; ``DONE`` (last) has none. ``FAILED`` / ``EXPANDED`` #: are intentionally absent -- they are sinks, not chain members. LIFECYCLE: tuple[MaterialState, ...] = ( MaterialState.NEW, MaterialState.COMPOSED, MaterialState.STRUCTURED, MaterialState.SCREENED, MaterialState.CALCULATED, MaterialState.EXTRACTED, MaterialState.DONE, ) #: Each live chain state -> its successor. Built from ``LIFECYCLE`` so the #: order is written in exactly one place. _SUCCESSOR: dict[MaterialState, MaterialState] = { LIFECYCLE[i]: LIFECYCLE[i + 1] for i in range(len(LIFECYCLE) - 1) } #: Each chain state -> its position, for the strictly-forward check on #: ``complete``'s explicit ``target_state``. ``FAILED`` / ``EXPANDED`` are #: absent -- they can never be a completion target. _CHAIN_INDEX: dict[MaterialState, int] = {state: i for i, state in enumerate(LIFECYCLE)} #: The end states: reached once, never left, never claimed again. _TERMINAL: frozenset[MaterialState] = frozenset( {MaterialState.DONE, MaterialState.FAILED, MaterialState.EXPANDED} )
[docs] def is_terminal(state: MaterialState) -> bool: """Say whether a state is an end state. Parameters ---------- state The state to classify. Returns ------- bool ``True`` for ``DONE``, ``FAILED`` and ``EXPANDED``. A terminal material has left the funnel: it has no successor and must never be claimed or advanced again. Examples -------- >>> is_terminal(MaterialState.SCREENED), is_terminal(MaterialState.EXPANDED) (False, True) """ return state in _TERMINAL
[docs] def next_state(state: MaterialState) -> MaterialState: """Give the state that follows when a stage completes a material. Parameters ---------- state The material's current state. Returns ------- MaterialState Its successor in ``LIFECYCLE``. Raises ------ ValueError For a terminal state -- asking for the successor of ``DONE``, ``FAILED`` or ``EXPANDED`` is a programming error, not a recoverable condition. Examples -------- >>> next_state(MaterialState.SCREENED) <MaterialState.CALCULATED: 'CALCULATED'> """ if is_terminal(state): raise ValueError(f"{state.value} is terminal and has no successor") return _SUCCESSOR[state]
[docs] def can_transition(current: MaterialState, target: MaterialState) -> bool: """Judge whether a stage-driven move is legal. The store consults this before every write, so an out-of-order replay can never corrupt state. Parameters ---------- current Where the material is now. target Where the caller proposes to move it. Returns ------- bool ``True`` for the three legal moves: the single forward step to ``current``'s successor (success); any live state to ``FAILED`` (hard failure); and any live state to ``EXPANDED`` (retired after fan-out). Everything else -- skipping a stage, going backwards, leaving a terminal state -- is ``False``. Notes ----- Two moves deliberately sit *outside* this function: ``finalize`` and ``complete``'s explicit ``target_state``. Which state ends a run, and which stages a run executes, are run configuration rather than lifecycle physics. """ if is_terminal(current): return False if target is MaterialState.FAILED or target is MaterialState.EXPANDED: return True return _SUCCESSOR[current] is target
# --------------------------------------------------------------------------- # Value objects handed out by the store (immutable snapshots) # ---------------------------------------------------------------------------
[docs] class StateError(Exception): """A coordination-plane contract violation. Raised for an unknown ``material_id``, an illegal transition (rejected by ``can_transition``), a bad query filter, or a ledger-invariant violation (empty material set, second live job for a material, unserialisable handle). Never used for "no work available" -- that is an empty list. """
[docs] @dataclass(frozen=True) class MaterialRecord: """One material's snapshot as handed back by ``claim`` / ``query`` / ``get``. Frozen: advancing a material happens only through ``complete`` / ``fail`` / ``expand``, never by mutating what a reader was handed. Attributes ---------- material_id : str The material's stable id. state : MaterialState Where it is in the funnel. claimed_at : float or None Epoch seconds when the current lease was taken, or ``None`` when unleased. The basis for ``release_stale_leases``. parent_id : str or None Lineage: the material this one was fanned out from, ``None`` for a root seed. A plain back-pointer, recorded at ``enqueue``. attempts : int Consecutive failed attempts at the current stage. Reset to 0 by ``complete``. error : str or None The latest attempt's failure reason, set on every ``fail`` including requeues, so a requeued material still shows why its last try died. fail_stage : str or None Which stage the material died at, set when it went terminal via ``fail``. fail_reason : str or None Why it died. Also reused by ``expand`` to carry its note, since a second note kind did not warrant another column. completions : dict ``{stage: completed_at}`` from the ledger -- the material's forward timeline. created_at : float Epoch seconds the row was inserted. updated_at : float Epoch seconds it was last written. Terminal rows are never written again, so a ``FAILED`` material's ``updated_at`` *is* its failure time; and for any material, ``now - updated_at`` is its time in the current state. That is the monitor's clock, with no extra columns needed. """ material_id: str state: MaterialState claimed_at: float | None = None parent_id: str | None = None attempts: int = 0 error: str | None = None fail_stage: str | None = None fail_reason: str | None = None completions: dict[str, float] = field(default_factory=dict) created_at: float = 0.0 updated_at: float = 0.0
[docs] @dataclass(frozen=True) class InFlightJob: """One submitted-but-unfinished backend job, as persisted by ``record_job``. Attributes ---------- job_key : str Caller-chosen unique identity, by convention ``"<backend_name>:<job_id>"``. Re-recording a key replaces the record. stage : str The stage whose work this job carries. backend_name : str Which configured backend the handle must be resolved against. handle : dict The backend ``Handle`` as an opaque JSON-safe dict. The orchestrator converts at the boundary, so this module never imports the compute plane. material_ids : tuple of str Every material whose result this job will produce. submitted_at : float Epoch seconds when the record was written. """ job_key: str stage: str backend_name: str handle: dict[str, Any] material_ids: tuple[str, ...] submitted_at: float
# --------------------------------------------------------------------------- # The SQLite store # --------------------------------------------------------------------------- #: Reserved completions-ledger stage name for the orchestrator's stop-stage #: bookkeeping transition (``finalize``). Not a real stage; a stage must never #: use it as its own name. FINALIZE_STAGE = "finalize" #: ``query`` filter keys, whitelisted so a caller-supplied field name is never #: interpolated into SQL unchecked. _QUERYABLE: frozenset[str] = frozenset({"material_id", "state", "parent_id"}) #: Stored text of the terminal states, for SQL exclusion lists. _TERMINAL_VALUES: tuple[str, ...] = tuple(s.value for s in sorted(_TERMINAL, key=lambda s: s.value)) #: Schema, applied idempotently on every open (every statement IF NOT EXISTS). #: Four tables: materials (state + lease + diagnostics), stage_completions #: (the idempotency ledger), jobs + job_materials (the in-flight handle ledger; #: ``job_materials.material_id`` being the PRIMARY KEY is what enforces "one #: live job per material" in the schema itself). _SCHEMA = """ CREATE TABLE IF NOT EXISTS materials ( material_id TEXT PRIMARY KEY, state TEXT NOT NULL, claimed_at REAL, -- epoch seconds; NULL = unleased parent_id TEXT, -- lineage back-pointer attempts INTEGER NOT NULL DEFAULT 0, -- failed tries at the current stage error TEXT, -- latest attempt's failure reason fail_stage TEXT, -- terminal diagnosis (fail/expand) fail_reason TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL ); -- Claiming filters on (state, claimed_at IS NULL); index the hot path. CREATE INDEX IF NOT EXISTS idx_materials_state_claim ON materials (state, claimed_at); CREATE INDEX IF NOT EXISTS idx_materials_parent ON materials (parent_id); CREATE TABLE IF NOT EXISTS stage_completions ( material_id TEXT NOT NULL, stage TEXT NOT NULL, completed_at REAL NOT NULL, PRIMARY KEY (material_id, stage), FOREIGN KEY (material_id) REFERENCES materials (material_id) ON DELETE CASCADE ); -- The monitor's rolling-window throughput reads filter on completed_at alone. CREATE INDEX IF NOT EXISTS idx_completions_time ON stage_completions (completed_at); CREATE TABLE IF NOT EXISTS jobs ( job_key TEXT PRIMARY KEY, stage TEXT NOT NULL, backend_name TEXT NOT NULL, handle TEXT NOT NULL, -- backend Handle as a JSON blob submitted_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS job_materials ( material_id TEXT PRIMARY KEY, -- one live job per material, by schema job_key TEXT NOT NULL, FOREIGN KEY (job_key) REFERENCES jobs (job_key) ON DELETE CASCADE, FOREIGN KEY (material_id) REFERENCES materials (material_id) ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS idx_job_materials_job ON job_materials (job_key); """
[docs] class StateStore: """The durable coordination plane, backed by one SQLite database file. One connection per instance, shared across threads behind an internal lock (``check_same_thread=False``); ``isolation_level=None`` gives explicit transaction control so read-then-write sequences run under ``BEGIN IMMEDIATE``. Cross-*instance* safety (two stores on one file) comes from SQLite itself: the immediate transaction takes the write lock up front, and ``busy_timeout`` absorbs brief contention instead of raising. Parameters ---------- db_path Path to the database file, created if absent. Opening an existing file is a safe no-op re-apply of the schema, so a run resumes against the same file it was interrupted on. busy_timeout_ms How long SQLite waits for a contended write lock before raising. read_only Open an *observer* store -- the ``monitor`` seam. The schema script is skipped (it is a write) and ``PRAGMA query_only=ON`` makes SQLite itself reject every write on the connection, so the no-corruption guarantee is enforced by the engine rather than by caller discipline. Mutating methods additionally raise up front, for a friendlier failure than ``SQLITE_READONLY``. Raises ------ StateError When ``read_only`` is set and no database exists at ``db_path``. An observer must never silently create an empty database. Notes ----- ``query_only`` is chosen over the URI ``?mode=ro`` deliberately: a WAL reader may need to create the ``-shm`` file -- after the last writer died uncleanly, for instance -- which ``mode=ro`` forbids and ``query_only`` allows. """ def __init__( self, db_path: str | Path, *, busy_timeout_ms: int = 5000, read_only: bool = False ) -> None: self._db_path = str(db_path) self._read_only = read_only if read_only and not Path(self._db_path).is_file(): raise StateError( f"cannot open read-only: no state database at {self._db_path} " f"-- start a run first ('matmelt run')" ) self._conn = sqlite3.connect(self._db_path, check_same_thread=False, isolation_level=None) self._conn.row_factory = sqlite3.Row self._lock = threading.Lock() self._conn.execute(f"PRAGMA busy_timeout={int(busy_timeout_ms)}") if read_only: # Engine-enforced write ban; journal_mode/foreign_keys are writer # concerns (the file is already WAL from its first writer). self._conn.execute("PRAGMA query_only=ON") return # WAL: readers never block the single writer; survives reopen. self._conn.execute("PRAGMA journal_mode=WAL") # Enforce the ledger foreign keys (off by default in sqlite3). self._conn.execute("PRAGMA foreign_keys=ON") with self._lock: self._conn.executescript(_SCHEMA)
[docs] def close(self) -> None: """Close the underlying connection. Safe to call more than once. A closed store is not reusable; open a new one against the same file, which is exactly what a resumed run does. """ with self._lock: self._conn.close()
def _assert_writable(self) -> None: """Refuse the call when this store was opened ``read_only``. Every mutation funnels through ``enqueue`` or ``_immediate``, so those two call sites make the guard complete by construction. Raises ------ StateError When the store is an observer. """ if self._read_only: raise StateError(f"store is read-only: {self._db_path} was opened with read_only=True") @contextmanager def _immediate(self) -> Iterator[None]: """Run the body inside ``BEGIN IMMEDIATE``: commit on exit, rollback on error. Takes SQLite's write lock up front, so select-then-stamp sequences are atomic against any other writer -- including a second store instance on the same file. Callers must already hold ``self._lock``. Yields ------ None Control, with the transaction open. Raises ------ StateError When the store is read-only. """ self._assert_writable() self._conn.execute("BEGIN IMMEDIATE") try: yield except BaseException: self._conn.execute("ROLLBACK") raise else: self._conn.execute("COMMIT") # -- Row mapping ---------------------------------------------------------- def _completions_for(self, material_ids: Sequence[str]) -> dict[str, dict[str, float]]: """Fetch the completion timelines for several materials at once. Parameters ---------- material_ids The materials to look up. An empty sequence skips the query. Returns ------- dict ``{material_id: {stage: completed_at}}``, built from one batched query so building N records costs two round-trips, not N + 1. """ if not material_ids: return {} marks = ",".join("?" for _ in material_ids) rows = self._conn.execute( f"SELECT material_id, stage, completed_at FROM stage_completions " # nosec B608 f"WHERE material_id IN ({marks}) ORDER BY completed_at", tuple(material_ids), ).fetchall() out: dict[str, dict[str, float]] = {} for row in rows: out.setdefault(row["material_id"], {})[row["stage"]] = row["completed_at"] return out def _rows_to_records(self, rows: Sequence[sqlite3.Row]) -> list[MaterialRecord]: """Build immutable snapshots from ``materials`` rows and their timelines. Parameters ---------- rows Rows selected from the ``materials`` table. Returns ------- list of MaterialRecord One record per row, in the given order, each carrying its completions timeline. """ timelines = self._completions_for([row["material_id"] for row in rows]) return [ MaterialRecord( material_id=row["material_id"], state=MaterialState(row["state"]), claimed_at=row["claimed_at"], parent_id=row["parent_id"], attempts=row["attempts"], error=row["error"], fail_stage=row["fail_stage"], fail_reason=row["fail_reason"], completions=timelines.get(row["material_id"], {}), created_at=row["created_at"], updated_at=row["updated_at"], ) for row in rows ] def _state_of(self, material_id: str) -> MaterialState | None: """Read one material's current state. The caller must hold the lock. Parameters ---------- material_id The material to look up. Returns ------- MaterialState or None Its state, or ``None`` when the store does not know the id. """ row = self._conn.execute( "SELECT state FROM materials WHERE material_id = ?", (material_id,) ).fetchone() return MaterialState(row["state"]) if row is not None else None # -- Enqueue / read -------------------------------------------------------
[docs] def enqueue(self, material_id: str, parent_id: str | None = None) -> None: """Add ``material_id`` to the funnel in state ``NEW``. Idempotent. Parameters ---------- material_id The id to enqueue. parent_id The material this one fanned out from, recorded once as lineage. ``None`` for a root seed. Raises ------ StateError When the store is read-only. Notes ----- Re-enqueueing an existing id touches nothing: its current state, attempts, and recorded ``parent_id`` all stand, first write wins. That is what lets a replayed seed or a re-run feeder never knock a material backwards or rewrite its lineage. """ self._assert_writable() now = time.time() with self._lock: self._conn.execute( "INSERT OR IGNORE INTO materials " "(material_id, state, parent_id, created_at, updated_at) " "VALUES (?, ?, ?, ?, ?)", (material_id, MaterialState.NEW.value, parent_id, now, now), )
[docs] def get(self, material_id: str) -> MaterialRecord | None: """Fetch one material's full snapshot. Read-only. Parameters ---------- material_id The id to fetch. Returns ------- MaterialRecord or None The snapshot, or ``None`` when the store does not know the id. """ with self._lock: row = self._conn.execute( "SELECT * FROM materials WHERE material_id = ?", (material_id,) ).fetchone() return self._rows_to_records([row])[0] if row is not None else None
[docs] def query(self, filters: dict[str, Any] | None = None) -> list[MaterialRecord]: """List materials matching ``filters``, oldest first. Read-only. No leasing, no writes. Parameters ---------- filters Field/value pairs, AND-combined. Keys must be one of ``material_id``, ``state`` or ``parent_id`` -- whitelisted, so a caller-supplied field name is never interpolated into SQL unchecked. A ``MaterialState`` value is coerced to its stored text. ``None`` or empty returns every material. Returns ------- list of MaterialRecord Matching snapshots, ordered by creation time then id. Raises ------ StateError For a filter key outside the whitelist. """ clauses: list[str] = [] params: list[Any] = [] for key, value in (filters or {}).items(): if key not in _QUERYABLE: raise StateError(f"unknown query filter field: {key!r}") clauses.append(f"{key} = ?") params.append(value.value if isinstance(value, MaterialState) else value) sql = "SELECT * FROM materials" if clauses: sql += " WHERE " + " AND ".join(clauses) sql += " ORDER BY created_at, material_id" with self._lock: rows = self._conn.execute(sql, params).fetchall() return self._rows_to_records(rows)
[docs] def count_by_state(self) -> dict[MaterialState, int]: """Tally the materials per state. Read-only. Returns ------- dict ``{state: count}`` from one ``GROUP BY``. States holding no materials are simply absent; callers zero-fill if they need the full set. """ with self._lock: rows = self._conn.execute( "SELECT state, COUNT(*) AS n FROM materials GROUP BY state" ).fetchall() return {MaterialState(row["state"]): row["n"] for row in rows}
[docs] def completion_counts_since(self, cutoff: float) -> dict[str, int]: """Count completions ledgered at or after a cutoff. Read-only. The rolling-window primitive behind the throughput and speed displays: the caller picks the window and divides. Parameters ---------- cutoff Epoch seconds. Pass ``0.0`` for all-time counts, or ``time.time() - window_s`` for a rolling window. Returns ------- dict ``{stage: count}``. Reserved ``FINALIZE_STAGE`` rows appear under their own name, leaving the caller to decide whether finalizations count as stage work. Stages with no completions in the window are simply absent. """ with self._lock: rows = self._conn.execute( "SELECT stage, COUNT(*) AS n FROM stage_completions " "WHERE completed_at >= ? GROUP BY stage", (cutoff,), ).fetchall() return {row["stage"]: row["n"] for row in rows}
[docs] def recent_failures(self, limit: int = 10) -> list[MaterialRecord]: """List the most recently failed materials, newest first. Read-only. Parameters ---------- limit How many to return. Zero or less returns nothing. Returns ------- list of MaterialRecord Ordered by ``updated_at`` descending. A terminal row is never written again, so its ``updated_at`` is exactly when it failed, and the record carries it so callers can render "failed at". """ if limit <= 0: return [] with self._lock: rows = self._conn.execute( "SELECT * FROM materials WHERE state = ? " "ORDER BY updated_at DESC, material_id LIMIT ?", (MaterialState.FAILED.value, limit), ).fetchall() return self._rows_to_records(rows)
[docs] def recent_successes(self, limit: int = 10) -> list[MaterialRecord]: """List the most recently finished materials, newest first. Read-only. The success mirror of ``recent_failures``. Parameters ---------- limit How many to return. Zero or less returns nothing. Returns ------- list of MaterialRecord ``DONE`` materials ordered by ``updated_at`` descending. A terminal row is never written again, so ``updated_at`` is exactly when the material finished and ``updated_at - created_at`` is how long its whole ride took; the monitor renders both. """ if limit <= 0: return [] with self._lock: rows = self._conn.execute( "SELECT * FROM materials WHERE state = ? " "ORDER BY updated_at DESC, material_id LIMIT ?", (MaterialState.DONE.value, limit), ).fetchall() return self._rows_to_records(rows)
[docs] def earliest_created_at(self) -> float | None: """Find when the first material entered the store. Read-only. Returns ------- float or None Epoch seconds of the earliest ``created_at``, or ``None`` on an empty store. This is the run's wall-clock anchor: elapsed time is ``now - earliest``. """ with self._lock: row = self._conn.execute("SELECT MIN(created_at) AS t FROM materials").fetchone() earliest = row["t"] return float(earliest) if earliest is not None else None
# -- Lease / advance / fail -----------------------------------------------
[docs] def claim( self, state: MaterialState, limit: int, *, among: Collection[str] | None = None ) -> list[MaterialRecord]: """Atomically lease up to ``limit`` unclaimed materials in ``state``. Selects the oldest unleased rows and stamps ``claimed_at`` inside one immediate transaction, so two concurrent claimers can never be handed the same material. Parameters ---------- state Which input state to claim from. limit The most materials to lease. Zero or less claims nothing. among Restrict the claim to these ids -- the selective-run seam, where the orchestrator precomputes the eligible lineage and leases nothing outside it. ``None`` is unrestricted; an empty collection claims nothing. Returns ------- list of MaterialRecord The leased snapshots, oldest first. Empty when nothing is available. Raises ------ StateError When ``state`` is terminal -- claiming one is a programming error -- or when the store is read-only. Notes ----- The ``among`` ids ride an SQL ``IN`` clause, which very old SQLite builds cap at 999 entries. That is fine at v1 scale; chunking is a later concern. """ if is_terminal(state): raise StateError(f"cannot claim terminal state {state.value}") if limit <= 0 or (among is not None and not among): return [] where = "WHERE state = ? AND claimed_at IS NULL" params: list[Any] = [state.value] if among is not None: eligible = sorted(set(among)) where += f" AND material_id IN ({','.join('?' for _ in eligible)})" params.extend(eligible) now = time.time() with self._lock, self._immediate(): rows = self._conn.execute( f"SELECT material_id FROM materials {where} " # nosec B608 f"ORDER BY created_at, material_id LIMIT ?", (*params, limit), ).fetchall() ids = [row["material_id"] for row in rows] if not ids: return [] marks = ",".join("?" for _ in ids) self._conn.execute( f"UPDATE materials SET claimed_at = ?, updated_at = ? " # nosec B608 f"WHERE material_id IN ({marks})", (now, now, *ids), ) claimed = self._conn.execute( f"SELECT * FROM materials WHERE material_id IN ({marks}) " # nosec B608 f"ORDER BY created_at, material_id", ids, ).fetchall() return self._rows_to_records(claimed)
[docs] def complete( self, material_id: str, stage: str, *, target_state: MaterialState | None = None ) -> None: """Advance ``material_id`` forward in the funnel and release its lease. Idempotent on the ``(material_id, stage)`` completions-ledger key: if this stage already completed this material the call is a no-op, so a crash-replay never advances twice. Otherwise the material moves to the target, the lease clears, the retry budget resets (``attempts = 0``, ``error = NULL``) and the ledger row is written -- all in one transaction. Parameters ---------- material_id The material to advance. stage The stage that completed it, which is also the ledger's idempotency key. target_state Where to move it. ``None`` takes the single step to its lifecycle successor. An explicit value moves it to any *strictly forward* chain state, which is how a run that skips an optional stage (say ``calculate``) walks its consumer's completion over the skipped state -- SCREENED straight to EXTRACTED. Like ``finalize``, the jump sits deliberately outside ``can_transition``: which stages a run executes is run configuration, not lifecycle physics. Raises ------ StateError For an unknown material; for one already terminal; for a backward, same-state, or off-chain (``FAILED`` / ``EXPANDED``) target; or when the store is read-only. """ now = time.time() with self._lock, self._immediate(): current = self._state_of(material_id) if current is None: raise StateError(f"unknown material_id: {material_id!r}") already = self._conn.execute( "SELECT 1 FROM stage_completions WHERE material_id = ? AND stage = ?", (material_id, stage), ).fetchone() if already is not None: return if is_terminal(current): raise StateError(f"illegal complete from {current.value} for {material_id!r}") if target_state is None: target = next_state(current) else: target_index = _CHAIN_INDEX.get(target_state) if target_index is None or target_index <= _CHAIN_INDEX[current]: raise StateError( f"illegal complete target {target_state.value} from " f"{current.value} for {material_id!r}: the target must be " f"strictly forward in the lifecycle chain" ) target = target_state self._conn.execute( "UPDATE materials SET state = ?, claimed_at = NULL, " "attempts = 0, error = NULL, updated_at = ? WHERE material_id = ?", (target.value, now, material_id), ) self._conn.execute( "INSERT INTO stage_completions (material_id, stage, completed_at) VALUES (?, ?, ?)", (material_id, stage, now), )
[docs] def fail( self, material_id: str, stage: str, reason: str, *, max_attempts: int = 1 ) -> MaterialState: """Spend one retry attempt; requeue while budget remains, else fail terminally. Increments ``attempts`` and records ``reason`` verbatim in ``error``. Parameters ---------- material_id The material that failed. stage Where it failed, recorded in ``fail_stage`` when the failure turns terminal. reason Why, stored verbatim. Shown by ``status`` and ``monitor``. max_attempts The stage's retry budget. The default of 1 makes the first failure terminal; the orchestrator passes the configured per-stage value. Returns ------- MaterialState The material's state after the call: unchanged when it was requeued (lease cleared, so the next claim re-runs the same stage), or ``FAILED`` when the budget ran out. A failed row stays in the database -- failed candidates are still data. Raises ------ StateError For an unknown material, one terminal as ``DONE`` or ``EXPANDED``, or a read-only store. Calling on an already-``FAILED`` material is idempotent: a no-op that spends no attempt. """ now = time.time() with self._lock, self._immediate(): current = self._state_of(material_id) if current is None: raise StateError(f"unknown material_id: {material_id!r}") if current is MaterialState.FAILED: return MaterialState.FAILED if not can_transition(current, MaterialState.FAILED): raise StateError(f"illegal fail from {current.value} for {material_id!r}") row = self._conn.execute( "SELECT attempts FROM materials WHERE material_id = ?", (material_id,) ).fetchone() attempts = row["attempts"] + 1 if attempts < max_attempts: self._conn.execute( "UPDATE materials SET attempts = ?, error = ?, " "claimed_at = NULL, updated_at = ? WHERE material_id = ?", (attempts, reason, now, material_id), ) return current self._conn.execute( "UPDATE materials SET state = ?, attempts = ?, error = ?, " "fail_stage = ?, fail_reason = ?, claimed_at = NULL, updated_at = ? " "WHERE material_id = ?", (MaterialState.FAILED.value, attempts, reason, stage, reason, now, material_id), ) return MaterialState.FAILED
[docs] def reset(self, material_id: str, target_state: MaterialState) -> None: """Un-fail a terminal ``FAILED`` material back to a live, claimable state. The recovery primitive behind ``matmelt retry``. The lease clears and all failure bookkeeping is wiped (``attempts = 0``, and ``error``, ``fail_stage``, ``fail_reason`` all set to NULL), so the next ``run`` re-claims the material and re-runs that stage. The **completions ledger is left untouched**: those earlier stages really did complete. Preserving the timeline is what lets the caller compute the correct ``target_state`` in the first place, and it keeps lineage and provenance intact. Parameters ---------- material_id The failed material to revive. target_state Where to put it back -- the state it was claimable from before it failed, i.e. its ``fail_stage``'s input state. The caller derives this from the material's completions timeline, because the state layer holds no stage-to-state map. Raises ------ StateError For an unknown material; for one that is not ``FAILED``, since this is a recovery primitive and not a general "set state"; for a terminal ``target_state``; or for a read-only store. """ now = time.time() with self._lock, self._immediate(): current = self._state_of(material_id) if current is None: raise StateError(f"unknown material_id: {material_id!r}") if current is not MaterialState.FAILED: raise StateError( f"cannot reset {material_id!r}: only FAILED materials may be reset, " f"but it is {current.value}" ) if is_terminal(target_state): raise StateError( f"illegal reset target {target_state.value} for {material_id!r}: " f"the target must be a live chain state, not a terminal one" ) self._conn.execute( "UPDATE materials SET state = ?, claimed_at = NULL, attempts = 0, " "error = NULL, fail_stage = NULL, fail_reason = NULL, updated_at = ? " "WHERE material_id = ?", (target_state.value, now, material_id), )
[docs] def expand(self, material_id: str, reason: str) -> None: """Retire ``material_id`` as ``EXPANDED`` -- it fanned out into children. Terminal like ``FAILED`` but *not* a failure. Keeping them distinct is what lets ``FAILED`` go on meaning genuine failure in every downstream query. The lease clears. Parameters ---------- material_id The parent to retire. reason A diagnostic note, stored in ``fail_reason`` -- a second note kind did not warrant another column. Raises ------ StateError For an unknown material, one terminal as ``DONE`` or ``FAILED``, or a read-only store. Calling on an already-``EXPANDED`` material is idempotent. """ now = time.time() with self._lock, self._immediate(): current = self._state_of(material_id) if current is None: raise StateError(f"unknown material_id: {material_id!r}") if current is MaterialState.EXPANDED: return if not can_transition(current, MaterialState.EXPANDED): raise StateError(f"illegal expand from {current.value} for {material_id!r}") self._conn.execute( "UPDATE materials SET state = ?, fail_reason = ?, " "claimed_at = NULL, updated_at = ? WHERE material_id = ?", (MaterialState.EXPANDED.value, reason, now, material_id), )
[docs] def finalize(self, material_id: str) -> None: """Retire a live material as first-class terminal ``DONE``. The orchestrator's bookkeeping move, not a stage transition: a material that completes the run's last stage has no consuming stage left, so this retires it as "went as far as this run was asked to go" and records a completions-ledger row under the reserved stage name ``FINALIZE_STAGE``. The lease clears and the retry budget resets. It sits deliberately outside ``can_transition``, which describes stage-driven moves only: finalize is legal from any live state, because which state is "the end" is run configuration rather than lifecycle physics. Parameters ---------- material_id The material to retire. Raises ------ StateError For an unknown material, one terminal as ``FAILED`` or ``EXPANDED``, or a read-only store. Calling on an already-``DONE`` material is idempotent. """ now = time.time() with self._lock, self._immediate(): current = self._state_of(material_id) if current is None: raise StateError(f"unknown material_id: {material_id!r}") if current is MaterialState.DONE: return if is_terminal(current): raise StateError(f"illegal finalize from {current.value} for {material_id!r}") self._conn.execute( "UPDATE materials SET state = ?, claimed_at = NULL, " "attempts = 0, error = NULL, updated_at = ? WHERE material_id = ?", (MaterialState.DONE.value, now, material_id), ) self._conn.execute( "INSERT OR IGNORE INTO stage_completions (material_id, stage, completed_at) " "VALUES (?, ?, ?)", (material_id, FINALIZE_STAGE, now), )
[docs] def release_stale_leases(self, timeout_s: float) -> int: """Clear leases older than ``timeout_s`` seconds. Run at startup and between passes: a lease older than the timeout belonged to a run that died mid-stage, and clearing it returns the material to its queue with its state unchanged. Parameters ---------- timeout_s Lease age, in seconds, past which a lease is considered abandoned. Returns ------- int How many leases were freed. Raises ------ StateError When the store is read-only. Notes ----- Only live, currently leased rows are touched, and a material referenced by a live in-flight job is **never** swept whatever its lease age: the recorded handle owns it until ``clear_job``. Releasing underneath a running job is exactly the double-submission hole the ledger closes. """ now = time.time() cutoff = now - timeout_s terminal_marks = ",".join("?" for _ in _TERMINAL_VALUES) with self._lock, self._immediate(): cur = self._conn.execute( f"UPDATE materials SET claimed_at = NULL, updated_at = ? " # nosec B608 f"WHERE claimed_at IS NOT NULL AND claimed_at < ? " f"AND state NOT IN ({terminal_marks}) " f"AND material_id NOT IN (SELECT material_id FROM job_materials)", (now, cutoff, *_TERMINAL_VALUES), ) return cur.rowcount
# -- The in-flight job ledger ----------------------------------------------
[docs] def record_job( self, job_key: str, stage: str, backend_name: str, handle: dict[str, Any], material_ids: Sequence[str], ) -> None: """Persist one submitted backend job: job row plus material links, atomically. Called immediately after ``backend.submit`` returns and *before* any polling -- this record is what survives a crash. Re-recording an existing ``job_key`` replaces the record and its links, so a replayed submit leaves exactly one record. Parameters ---------- job_key The job's unique identity, by convention ``"<backend_name>:<job_id>"``. stage Whose work the job carries. backend_name Which configured backend the handle resolves against. handle The backend ``Handle`` as a JSON-safe dict. material_ids Every material the job will produce a result for. Duplicates are dropped, caller order preserved. Raises ------ StateError When ``material_ids`` is empty; when the handle is not JSON-serialisable; when a material is already linked to a *different* live job, since the schema allows one live job per material; when a material is unknown to the store; or when the store is read-only. """ ids = list(dict.fromkeys(material_ids)) # de-dupe, preserve caller order if not ids: raise StateError(f"record_job {job_key!r}: material_ids must not be empty") try: handle_json = json.dumps(handle) except (TypeError, ValueError) as exc: raise StateError( f"record_job {job_key!r}: handle is not JSON-serialisable: {exc}" ) from None now = time.time() marks = ",".join("?" for _ in ids) with self._lock, self._immediate(): conflict = self._conn.execute( f"SELECT material_id, job_key FROM job_materials " # nosec B608 f"WHERE material_id IN ({marks}) AND job_key != ?", (*ids, job_key), ).fetchone() if conflict is not None: raise StateError( f"record_job {job_key!r}: material {conflict['material_id']!r} " f"is already in flight under job {conflict['job_key']!r}" ) self._conn.execute( "INSERT OR REPLACE INTO jobs " "(job_key, stage, backend_name, handle, submitted_at) VALUES (?, ?, ?, ?, ?)", (job_key, stage, backend_name, handle_json, now), ) # Rewrite the links explicitly; correctness must not depend on a # REPLACE-triggered cascade. self._conn.execute("DELETE FROM job_materials WHERE job_key = ?", (job_key,)) try: self._conn.executemany( "INSERT INTO job_materials (material_id, job_key) VALUES (?, ?)", [(mid, job_key) for mid in ids], ) except sqlite3.IntegrityError as exc: raise StateError( f"record_job {job_key!r}: unknown material among {ids!r}: {exc}" ) from None
[docs] def list_jobs( self, *, stage: str | None = None, backend_name: str | None = None ) -> list[InFlightJob]: """List the live in-flight jobs, oldest first. Read-only. Parameters ---------- stage Restrict to one stage's jobs. ``None`` lists every stage. backend_name Restrict to one backend's jobs. ``None`` lists every backend. Returns ------- list of InFlightJob Ordered by submission time then key, each carrying its material links. Notes ----- The check half of the pass loop lists everything; the submit half filters by ``backend_name`` to count live jobs against that backend's ``max_parallel``. """ clauses: list[str] = [] params: list[Any] = [] if stage is not None: clauses.append("stage = ?") params.append(stage) if backend_name is not None: clauses.append("backend_name = ?") params.append(backend_name) sql = "SELECT * FROM jobs" if clauses: sql += " WHERE " + " AND ".join(clauses) sql += " ORDER BY submitted_at, job_key" with self._lock: rows = self._conn.execute(sql, params).fetchall() jobs: list[InFlightJob] = [] for row in rows: links = self._conn.execute( "SELECT material_id FROM job_materials WHERE job_key = ? ORDER BY material_id", (row["job_key"],), ).fetchall() jobs.append( InFlightJob( job_key=row["job_key"], stage=row["stage"], backend_name=row["backend_name"], handle=json.loads(row["handle"]), material_ids=tuple(link["material_id"] for link in links), submitted_at=row["submitted_at"], ) ) return jobs
[docs] def clear_job(self, job_key: str) -> None: """Remove one in-flight record, making its materials sweepable again. Called after the job's results were attributed -- ``complete`` or ``fail`` per material -- and the backend handle released. Parameters ---------- job_key The record to remove. Clearing an unknown or already-cleared key is a no-op, so restart-time cleanup is always safe. Raises ------ StateError When the store is read-only. """ with self._lock, self._immediate(): # Explicit two-step delete: never rely on the FK pragma for the cascade. self._conn.execute("DELETE FROM job_materials WHERE job_key = ?", (job_key,)) self._conn.execute("DELETE FROM jobs WHERE job_key = ?", (job_key,))