stroblmeandClaude Opus 5
93374a310e
Four things the python SDK turned up, each fixed where every client sees it. A key no port declares is now an error rather than a silent drop, on the return, the yield and the emit alike — the contract the docs already stated. The SDK reads literal yields at sync time, so a typo fails before anything runs, and an emission of one fails the call rather than being logged where nobody looks. NaN and infinity are refused at the port. JSON cannot spell either, so one that travelled came back as a 500, a socket frame that stopped the canvas, or a metric batch the database dropped whole. An artifact input takes `@run:<id>.<output>` or a bare digest, resolved on the engine — so the CLI, the run dialog and a python caller mean the same thing, and a sweep can pass one at all. Node timeouts are off by default. The clock measured silence, which a training node is full of, and remote workers had already stopped enforcing it — their heartbeat reset it. Now a heartbeat proves the agent rather than the node, ninety seconds of nothing fails the call either way, and the engine touches work it is still running so a long node is not redelivered at sixty seconds. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019V5bsYGNxcgPs4xXmTPx69
912 lines
36 KiB
Python
912 lines
36 KiB
Python
"""Runs: a flow taken from its inputs to its outputs, once.
|
|
|
|
A cascade is what an always-on flow does when a value arrives — it has no
|
|
beginning and no end worth recording. A *run* is the other shape the same
|
|
engine can take: parameters go in, the graph executes until it drains, and
|
|
what it produced is kept. That is what an ML experiment is, and what a
|
|
CI-style job is, so both are this one entity.
|
|
|
|
Three things make a run different from a cascade, and each is deliberate:
|
|
|
|
* **Its own state.** Message names are global keys, so two runs of one flow
|
|
sharing the engine's state would overwrite each other's values. A run gets a
|
|
state backend namespaced to itself, which is a constructor argument rather
|
|
than a change to the pipeline — every key the engine keeps, versions and
|
|
node memory included, already goes through that backend.
|
|
* **Its own record.** The event bus drops what it cannot keep up with, which
|
|
is right for a live canvas and wrong for a result. The driver writes the
|
|
run's rows itself, from the thread that is running it.
|
|
* **Its own durability.** The queue wakes an engine up; from the moment a run
|
|
starts, its database row is the truth. Redelivering hours of training
|
|
because an acknowledgement was late is not recovery, so a run is
|
|
acknowledged as soon as it is claimed and a stale lease — not an unacked
|
|
stream entry — is what marks a run its engine died in the middle of.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import socket
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
|
|
from sqlalchemy import update
|
|
from sqlalchemy.dialects.sqlite import insert as upsert
|
|
from sqlmodel import Session, col, select
|
|
|
|
from fluksio.core.db import engine as db_engine
|
|
from fluksio.flow.artifacts import ArtifactStore, is_reference, valid_digest
|
|
from fluksio.flow.controller import FlowController, RunContext
|
|
from fluksio.flow.messages import DType, qualify
|
|
from fluksio.flow.pipeline import NodeOutcome, Pipeline
|
|
from fluksio.flow.queue import WorkItem, WorkQueue
|
|
from fluksio.flow.schemas import FlowDef
|
|
from fluksio.flow.state import MemoryState, StateBackend
|
|
from fluksio.models import Run, RunArtifact, RunMetric, RunNode
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
#: How long a finished run's state is kept before Redis drops it. Long enough
|
|
#: to look at what a failed run left behind, short enough not to accumulate.
|
|
RUN_STATE_TTL = 24 * 3600
|
|
#: How often a running run says it is still alive.
|
|
LEASE_INTERVAL_S = 20.0
|
|
#: A lease older than this belongs to an engine that is not coming back.
|
|
LEASE_STALE_S = 90.0
|
|
#: How often stale leases are looked for.
|
|
SWEEP_INTERVAL_S = 30.0
|
|
#: Runs driven at once. Node bodies are bounded by the worker pool anyway;
|
|
#: this only bounds how many graphs are in flight.
|
|
MAX_PARALLEL = 4
|
|
CLAIM_COUNT = 4
|
|
CLAIM_BLOCK_MS = 1000
|
|
ERROR_CAP = 2000
|
|
LOG_CAP = 8000
|
|
#: How much of a node's returned outputs is kept so a later run can restore
|
|
#: them. Past it the node simply is not cacheable — a cache is not a store,
|
|
#: and an artifact is how a large result is meant to travel.
|
|
OUTPUT_CAP = 256_000
|
|
#: Reported numbers held before they are written. A training loop reporting
|
|
#: every step must not be a round trip every step.
|
|
METRIC_BATCH = 500
|
|
METRIC_FLUSH_S = 2.0
|
|
#: How often a run waiting for a worker looks again.
|
|
WAIT_RETRY_S = 15.0
|
|
|
|
#: Where a run's state lives, so it can never collide with the engine's own.
|
|
RUN_NAMESPACE = "run"
|
|
|
|
|
|
class RunRejected(ValueError):
|
|
"""The run cannot be made: bad parameters, or a flow that cannot batch."""
|
|
|
|
|
|
def new_run_id() -> str:
|
|
"""Time-ordered, so the newest runs sort last without reading a column."""
|
|
return f"{int(time.time() * 1000):013d}-{uuid.uuid4().hex[:8]}"
|
|
|
|
|
|
def digest_of(params: dict[str, Any], seed: int | None) -> str:
|
|
"""What identifies a run's inputs: same parameters, same digest."""
|
|
canonical = json.dumps(
|
|
{"params": params, "seed": seed}, sort_keys=True, separators=(",", ":")
|
|
)
|
|
return hashlib.sha256(canonical.encode()).hexdigest()
|
|
|
|
|
|
def batch_issues(flow: FlowDef) -> list[str]:
|
|
"""Why this flow cannot be run as a batch, if it cannot.
|
|
|
|
One thing genuinely breaks: a port with a discretization interval holds
|
|
values back for a timer to release, and a run has no timer — the engine
|
|
would drop them instead of delaying them. On a *streaming* port that is
|
|
exactly right and is how you keep a chart from being flooded: the run's
|
|
history keeps every value, and the interval only thins what is published.
|
|
Anywhere else it is a message quietly going missing.
|
|
|
|
A delay node is fine; without a queue to defer into it simply sleeps,
|
|
which in a run is what was asked for.
|
|
"""
|
|
issues: list[str] = []
|
|
for node in flow.nodes:
|
|
for spec in list(node.requires) + list(node.provides):
|
|
if spec.interval > 0 and not spec.stream:
|
|
issues.append(
|
|
f"Node '{node.id}' rate-limits '{spec.port or spec.name}'. "
|
|
"A run has no timer to release what that holds back, so "
|
|
"the value would be dropped. Remove the interval, or mark "
|
|
"the port as streaming if it is a curve being thinned out."
|
|
)
|
|
return issues
|
|
|
|
|
|
def required_labels(flow: FlowDef) -> list[str]:
|
|
"""Worker labels this flow cannot run without.
|
|
|
|
Only the nodes that *require* their device. One that merely prefers it runs
|
|
on the engine when no such worker is attached, so holding the whole run
|
|
back to wait for one would be waiting for something it does not need.
|
|
"""
|
|
return sorted(
|
|
{
|
|
node.device
|
|
for node in flow.nodes
|
|
if node.device and node.device.strip() and node.device_policy == "require"
|
|
}
|
|
)
|
|
|
|
|
|
#: What a run's output is called from outside it: ``@run:<id>.<output>``.
|
|
RUN_REF_PREFIX = "@run:"
|
|
|
|
|
|
def resolve_references(
|
|
flow: FlowDef, params: dict[str, Any], artifacts: ArtifactStore | None = None
|
|
) -> dict[str, Any]:
|
|
"""Turn the text spellings of an artifact input into the reference itself.
|
|
|
|
A python caller hands one run's output straight to the next, because it has
|
|
the reference in its hand. A shell does not, and pasting the whole object
|
|
is not a command anyone wants to type — so an artifact input also takes
|
|
``@run:<id>.<output>``, naming what a run produced, or a bare
|
|
``sha256:...`` digest naming the bytes. Resolved here rather than in each
|
|
client, so the CLI, the browser and a python caller all mean the same
|
|
thing by the same string.
|
|
"""
|
|
wanted = {
|
|
declared.spec.name
|
|
for declared in flow.inputs
|
|
if declared.spec.dtype is DType.ARTIFACT
|
|
}
|
|
pending = {
|
|
key: value
|
|
for key, value in params.items()
|
|
if key in wanted and isinstance(value, str)
|
|
}
|
|
if not pending:
|
|
return params
|
|
|
|
resolved = dict(params)
|
|
with Session(db_engine) as session:
|
|
for key, text in pending.items():
|
|
if text.startswith(RUN_REF_PREFIX):
|
|
resolved[key] = _from_run(session, key, text[len(RUN_REF_PREFIX) :])
|
|
elif valid_digest(text):
|
|
resolved[key] = _from_digest(session, key, text, artifacts)
|
|
# Anything else is left alone: the type check names it better than
|
|
# a guess about what was meant would.
|
|
return resolved
|
|
|
|
|
|
def _from_run(session: Session, key: str, spelling: str) -> dict[str, Any]:
|
|
"""``<run id>.<output>`` as the reference that run produced."""
|
|
run_id, _, output = spelling.partition(".")
|
|
if not run_id or not output:
|
|
raise RunRejected(
|
|
f"Parameter '{key}': '{RUN_REF_PREFIX}{spelling}' names no output — "
|
|
f"write {RUN_REF_PREFIX}<run id>.<output>"
|
|
)
|
|
run = session.get(Run, run_id)
|
|
if run is None:
|
|
raise RunRejected(f"Parameter '{key}': there is no run '{run_id}'")
|
|
|
|
# The run's own result first: that is the reference as its producer made
|
|
# it, file name and all. The rows are the fallback, and they carry the
|
|
# message name instead — which loads the same bytes either way.
|
|
candidate = (run.result or {}).get(output)
|
|
if is_reference(candidate):
|
|
return dict(candidate)
|
|
|
|
rows = session.exec(
|
|
select(RunArtifact).where(col(RunArtifact.run_id) == run_id)
|
|
).all()
|
|
for row in rows:
|
|
if output in (row.name, row.name.rsplit(".", 1)[-1]):
|
|
return {
|
|
"digest": row.digest,
|
|
"size": row.size,
|
|
"media_type": row.media_type or "application/octet-stream",
|
|
"name": row.name,
|
|
}
|
|
known = ", ".join(sorted(row.name for row in rows)) or "none"
|
|
raise RunRejected(
|
|
f"Parameter '{key}': run '{run_id}' has no artifact '{output}' "
|
|
f"(it made: {known})"
|
|
)
|
|
|
|
|
|
def _from_digest(
|
|
session: Session, key: str, digest: str, artifacts: ArtifactStore | None
|
|
) -> dict[str, Any]:
|
|
"""A bare digest as a reference, with the size the store needs."""
|
|
row = session.exec(
|
|
select(RunArtifact)
|
|
.where(col(RunArtifact.digest) == digest)
|
|
.order_by(col(RunArtifact.run_id).desc())
|
|
).first()
|
|
if row is None:
|
|
raise RunRejected(
|
|
f"Parameter '{key}': no run has produced '{digest}', so there is "
|
|
"nothing here under that digest"
|
|
)
|
|
if artifacts is not None and artifacts.path(digest) is None:
|
|
raise RunRejected(
|
|
f"Parameter '{key}': '{digest}' is known but its bytes are gone "
|
|
"from this installation's store"
|
|
)
|
|
return {
|
|
"digest": row.digest,
|
|
"size": row.size,
|
|
"media_type": row.media_type or "application/octet-stream",
|
|
"name": row.name,
|
|
}
|
|
|
|
|
|
def seed_values(
|
|
flow: FlowDef, params: dict[str, Any], seed: int | None = None
|
|
) -> dict[str, Any]:
|
|
"""Turn a run's parameters into the messages the flow starts from.
|
|
|
|
A run's ``seed`` fills an input of that name when the flow declares one, so
|
|
the field that distinguishes two runs of one configuration is the same
|
|
number the flow draws from — otherwise it would only look like it was, and
|
|
a sweep over seeds would run the same experiment N times. An explicit
|
|
parameter still wins, and a flow that declares no ``seed`` ignores it.
|
|
"""
|
|
specs = {declared.spec.name: declared.spec for declared in flow.inputs}
|
|
if seed is not None and "seed" in specs and "seed" not in params:
|
|
params = {**params, "seed": seed}
|
|
values: dict[str, Any] = {}
|
|
for key, value in params.items():
|
|
spec = specs.get(key)
|
|
if spec is None:
|
|
known = ", ".join(sorted(specs)) or "none"
|
|
raise RunRejected(
|
|
f"'{key}' is not an input of flow '{flow.name}' (it declares: {known})"
|
|
)
|
|
try:
|
|
spec.check(value)
|
|
except TypeError as exc:
|
|
raise RunRejected(f"Parameter '{key}': {exc}") from exc
|
|
values[qualify(flow.name, spec.name)] = value
|
|
return values
|
|
|
|
|
|
def collect_result(flow: FlowDef, state: StateBackend) -> dict[str, Any]:
|
|
"""What the run produced, keyed by message name without the flow prefix."""
|
|
prefix = f"{flow.name}."
|
|
if flow.outputs:
|
|
names = [qualify(flow.name, name) for name in flow.outputs]
|
|
else:
|
|
# Everything the flow ended up holding. The engine's own bookkeeping is
|
|
# keyed by `__thing__:message`, so it never starts with the flow name.
|
|
names = sorted(key for key in state.keys() if key.startswith(prefix))
|
|
result: dict[str, Any] = {}
|
|
for name in names:
|
|
if name in state:
|
|
result[name[len(prefix) :] if name.startswith(prefix) else name] = state[
|
|
name
|
|
]
|
|
return result
|
|
|
|
|
|
class MetricSink:
|
|
"""Keeps the series a run's streaming outputs traced out.
|
|
|
|
A run's metrics are not logged; they are the numbers its nodes published on
|
|
the way to finishing. This watches the emissions, keeps the numeric ones,
|
|
and writes them in batches — synchronously rather than over the event bus,
|
|
which drops what it cannot keep up with, and a training curve with holes in
|
|
it is not a result.
|
|
|
|
The step is the count of emissions on that message. A node that publishes
|
|
every tenth training step therefore has steps 0, 1, 2 rather than 0, 10,
|
|
20 — a faithful x-axis of its own emissions, not of the loop inside it.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
run_id: str,
|
|
batch: int = METRIC_BATCH,
|
|
interval: float = METRIC_FLUSH_S,
|
|
) -> None:
|
|
self.run_id = run_id
|
|
self._batch = batch
|
|
self._interval = interval
|
|
self._rows: dict[tuple[str, int], RunMetric] = {}
|
|
self._steps: dict[str, int] = {}
|
|
self._last_flush = time.monotonic()
|
|
self._lock = threading.Lock()
|
|
|
|
def handle(self, node_id: str, outputs: dict[str, Any]) -> None:
|
|
"""One emission: every number in it belongs to this run's history."""
|
|
now = time.time()
|
|
rows: list[RunMetric] = []
|
|
with self._lock:
|
|
for name, value in outputs.items():
|
|
if not isinstance(value, (int, float)) or isinstance(value, bool):
|
|
# A checkpoint or a record is on the run some other way —
|
|
# as an artifact, or as its result. Only numbers are series.
|
|
continue
|
|
if not math.isfinite(value):
|
|
# The ports refuse these, so one here came from somewhere
|
|
# that does not go through them. A column that cannot hold
|
|
# it would take the whole batch down with it.
|
|
continue
|
|
step = self._steps.get(name, -1) + 1
|
|
self._steps[name] = step
|
|
row = RunMetric(
|
|
run_id=self.run_id,
|
|
name=name[:128],
|
|
step=step,
|
|
node=node_id[:255],
|
|
ts=now,
|
|
value=float(value),
|
|
)
|
|
self._rows[(row.name, row.step)] = row
|
|
due = (
|
|
len(self._rows) >= self._batch
|
|
or time.monotonic() - self._last_flush >= self._interval
|
|
)
|
|
if due:
|
|
rows = list(self._rows.values())
|
|
self._rows.clear()
|
|
self._last_flush = time.monotonic()
|
|
if rows:
|
|
self._write(rows)
|
|
|
|
def flush(self) -> None:
|
|
with self._lock:
|
|
rows = list(self._rows.values())
|
|
self._rows.clear()
|
|
self._last_flush = time.monotonic()
|
|
if rows:
|
|
self._write(rows)
|
|
|
|
def _write(self, rows: list[RunMetric]) -> None:
|
|
try:
|
|
with Session(db_engine) as session:
|
|
statement = upsert(RunMetric).values([row.model_dump() for row in rows])
|
|
session.exec(
|
|
statement.on_conflict_do_update(
|
|
index_elements=["run_id", "name", "step"],
|
|
set_={
|
|
"value": statement.excluded.value,
|
|
"ts": statement.excluded.ts,
|
|
"node": statement.excluded.node,
|
|
},
|
|
)
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
logger.exception(
|
|
"Could not write %d metric(s) of %s", len(rows), self.run_id
|
|
)
|
|
|
|
|
|
def origin_commit(flow: FlowDef) -> str:
|
|
"""The user repository's commit, for a flow declared in code elsewhere.
|
|
|
|
Marked dirty when the tree had uncommitted changes at sync, because then
|
|
the hash names something other than what actually ran.
|
|
"""
|
|
origin = flow.origin
|
|
if origin is None or not origin.commit:
|
|
return ""
|
|
return f"{origin.commit}-dirty" if origin.dirty else origin.commit
|
|
|
|
|
|
def _cacheable(outcome: NodeOutcome) -> str | None:
|
|
"""A node's outputs as stored, or None when it may not be reused."""
|
|
if not outcome.ok or not outcome.cache_key:
|
|
return None
|
|
try:
|
|
text = json.dumps(outcome.output_values, sort_keys=True, separators=(",", ":"))
|
|
except (TypeError, ValueError):
|
|
return None
|
|
return text if len(text) <= OUTPUT_CAP else None
|
|
|
|
|
|
class RunCache:
|
|
"""What earlier runs already worked out, looked up by cache key.
|
|
|
|
The rows are the run history itself — a node that ran successfully and
|
|
whose outputs were small enough to keep is a cache entry, without a store
|
|
of its own. An entry whose artifacts have gone from the content store is
|
|
not a hit: the reference would name bytes nobody can open.
|
|
"""
|
|
|
|
def __init__(self, artifacts: ArtifactStore | None = None) -> None:
|
|
self.artifacts = artifacts
|
|
|
|
def lookup(self, key: str) -> tuple[bool, dict[str, Any] | None]:
|
|
if not key:
|
|
return False, None
|
|
try:
|
|
with Session(db_engine) as session:
|
|
rows = session.exec(
|
|
select(RunNode)
|
|
.where(
|
|
col(RunNode.cache_key) == key,
|
|
col(RunNode.status).in_(("ok", "cached")),
|
|
col(RunNode.outputs).is_not(None),
|
|
)
|
|
# Run ids are time-ordered, so this is the most recent
|
|
# first. A handful is enough to get past entries whose
|
|
# artifacts have been collected.
|
|
.order_by(col(RunNode.run_id).desc())
|
|
.limit(5)
|
|
).all()
|
|
for row in rows:
|
|
outputs = json.loads(row.outputs or "null")
|
|
if outputs and self._unrestorable(outputs):
|
|
continue
|
|
return True, outputs
|
|
except Exception:
|
|
# A cache that cannot answer is a cache miss, never a failed run.
|
|
logger.exception("Cache lookup failed; running the node instead")
|
|
return False, None
|
|
|
|
def _unrestorable(self, outputs: dict[str, Any]) -> bool:
|
|
return any(
|
|
is_reference(value)
|
|
and (
|
|
self.artifacts is None
|
|
or self.artifacts.path(str(value["digest"])) is None
|
|
)
|
|
for value in outputs.values()
|
|
)
|
|
|
|
|
|
class RunService:
|
|
"""Accepts runs, drives them, and writes down what they did."""
|
|
|
|
def __init__(
|
|
self,
|
|
controller: FlowController,
|
|
queue: WorkQueue,
|
|
state_factory: Callable[[str], StateBackend] | None = None,
|
|
parallel: int = MAX_PARALLEL,
|
|
artifacts: ArtifactStore | None = None,
|
|
) -> None:
|
|
self.controller = controller
|
|
self.queue = queue
|
|
self._artifacts = artifacts
|
|
self._cache = RunCache(artifacts)
|
|
# Without one, a run gets a private in-memory state — which is exactly
|
|
# the isolation it wants, minus surviving the process.
|
|
self._state_factory = state_factory or (lambda _ns: MemoryState())
|
|
self.engine_name = f"{socket.gethostname()}-{os.getpid()}"[:64]
|
|
self._pool = ThreadPoolExecutor(max_workers=parallel, thread_name_prefix="run")
|
|
self._stop = threading.Event()
|
|
self._consumer: threading.Thread | None = None
|
|
self._keeper: threading.Thread | None = None
|
|
# Runs this process is driving, and the pipeline each is running, so a
|
|
# cancel has something to hold on to.
|
|
self._active: dict[str, Pipeline] = {}
|
|
self._cancelled: set[str] = set()
|
|
self._lock = threading.Lock()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Lifecycle
|
|
# -------------------------------------------------------------------------
|
|
|
|
def start(self) -> None:
|
|
if self._consumer is not None:
|
|
return
|
|
self._consumer = threading.Thread(
|
|
target=self._consume, name="run-consumer", daemon=True
|
|
)
|
|
self._consumer.start()
|
|
self._keeper = threading.Thread(
|
|
target=self._keep_leases, name="run-leases", daemon=True
|
|
)
|
|
self._keeper.start()
|
|
|
|
def stop(self) -> None:
|
|
self._stop.set()
|
|
for thread in (self._consumer, self._keeper):
|
|
if thread is not None:
|
|
thread.join(timeout=5)
|
|
self._consumer = None
|
|
self._keeper = None
|
|
self._pool.shutdown(wait=False)
|
|
self.queue.close()
|
|
|
|
def alive(self) -> bool:
|
|
return self._consumer is not None and self._consumer.is_alive()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Accepting work
|
|
# -------------------------------------------------------------------------
|
|
|
|
def submit(
|
|
self,
|
|
flow_name: str,
|
|
params: dict[str, Any] | None = None,
|
|
seed: int | None = None,
|
|
group_id: str | None = None,
|
|
cause: str = "api",
|
|
actor: str = "",
|
|
draft: bool = False,
|
|
no_cache: bool = False,
|
|
) -> Run:
|
|
"""Journal a run and wake an engine up for it. Never blocks on it."""
|
|
flow = self.controller.store.read_flow(flow_name, draft=draft)
|
|
issues = batch_issues(flow)
|
|
if issues:
|
|
raise RunRejected(" ".join(issues))
|
|
params = params or {}
|
|
# Resolved before it is stored, so what the run records is the same
|
|
# reference a python caller would have passed and every later reader —
|
|
# the digest, the cache, the run detail — sees one spelling.
|
|
params = resolve_references(flow, params, self._artifacts)
|
|
# Checked here rather than in the driver: a caller who mistyped a
|
|
# parameter should be told now, not by a run that fails in a minute.
|
|
seed_values(flow, params, seed)
|
|
|
|
run = Run(
|
|
id=new_run_id(),
|
|
flow=flow.name,
|
|
flow_version=flow.version,
|
|
commit=self.controller.store.head(),
|
|
origin_commit=origin_commit(flow),
|
|
params=params,
|
|
params_digest=digest_of(params, seed),
|
|
seed=seed,
|
|
group_id=group_id,
|
|
cause=cause,
|
|
no_cache=no_cache,
|
|
status="queued",
|
|
labels=required_labels(flow),
|
|
created_at=datetime.now(UTC),
|
|
actor=actor,
|
|
)
|
|
with Session(db_engine) as session:
|
|
session.add(run)
|
|
session.commit()
|
|
session.refresh(run)
|
|
|
|
self.queue.add(WorkItem(kind="run", node="", flow=flow.name, run_id=run.id))
|
|
return run
|
|
|
|
def cancel(self, run_id: str) -> bool:
|
|
"""Stop a run: kill what it is executing, schedule nothing further."""
|
|
with self._lock:
|
|
pipeline = self._active.get(run_id)
|
|
if pipeline is None:
|
|
# Not running here — if it is still queued, refusing to start
|
|
# is all the cancelling it needs.
|
|
cancelled = self._finish_queued(run_id)
|
|
if cancelled:
|
|
self._cancelled.add(run_id)
|
|
return cancelled
|
|
self._cancelled.add(run_id)
|
|
# The gate first, so nothing new is submitted while the running nodes
|
|
# are being killed; a gated node is never handed to the executor, so
|
|
# the graph drains instead of going further.
|
|
pipeline.pause(self._flow_of(run_id) or "")
|
|
workers = self.controller.workers
|
|
if workers is not None:
|
|
# Keyed by run, so a sweep cancelling one config leaves the others
|
|
# training.
|
|
workers.cancel_run(run_id)
|
|
if self.controller.remote is not None:
|
|
self.controller.remote.cancel_run(run_id)
|
|
return True
|
|
|
|
def _flow_of(self, run_id: str) -> str | None:
|
|
with Session(db_engine) as session:
|
|
run = session.get(Run, run_id)
|
|
return run.flow if run else None
|
|
|
|
def _finish_queued(self, run_id: str) -> bool:
|
|
with Session(db_engine) as session:
|
|
result = session.exec(
|
|
update(Run)
|
|
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
|
.values(
|
|
status="cancelled",
|
|
finished_at=datetime.now(UTC),
|
|
status_reason="Cancelled before it started",
|
|
)
|
|
)
|
|
session.commit()
|
|
return bool(result.rowcount)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Threads
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _consume(self) -> None:
|
|
failures = 0
|
|
while not self._stop.is_set():
|
|
try:
|
|
# Runs put back to wait for a worker come due here. The claim
|
|
# below blocks for a second, so this is about once a second.
|
|
self.queue.move_due(time.time())
|
|
items = self.queue.claim(CLAIM_COUNT, CLAIM_BLOCK_MS)
|
|
failures = 0
|
|
except Exception as exc:
|
|
failures += 1
|
|
logger.error("Could not claim runs: %s", exc)
|
|
self._stop.wait(min(30.0, 2.0**failures))
|
|
continue
|
|
for item in items:
|
|
if not item.run_id:
|
|
self.queue.ack(item)
|
|
continue
|
|
missing = self._missing_labels(item.run_id)
|
|
if missing:
|
|
# Left in the queue rather than failed: submitting a run
|
|
# before turning the GPU box on is a normal way to work, and
|
|
# the run says what it is waiting for while it waits.
|
|
self._waiting(item.run_id, missing)
|
|
self._defer(item)
|
|
continue
|
|
# Acknowledged before it runs: from here on the row is the
|
|
# record, and a lease that stops moving is what says otherwise.
|
|
self.queue.ack(item)
|
|
try:
|
|
self._pool.submit(self._drive, item.run_id)
|
|
except RuntimeError:
|
|
logger.warning("Run %s not started: shutting down", item.run_id)
|
|
|
|
def _missing_labels(self, run_id: str) -> list[str]:
|
|
"""Worker labels this run needs that nothing attached carries."""
|
|
with Session(db_engine) as session:
|
|
run = session.get(Run, run_id)
|
|
needed = list(run.labels) if run else []
|
|
if not needed:
|
|
return []
|
|
hub = self.controller.remote
|
|
available = hub.labels() | {w.name for w in hub.workers()} if hub else set()
|
|
# A node that only prefers its label runs locally instead, so it is not
|
|
# a reason to hold the run back; that is decided per node at call time.
|
|
return sorted(set(needed) - available)
|
|
|
|
def _waiting(self, run_id: str, missing: list[str]) -> None:
|
|
reason = f"Waiting for a worker labelled {', '.join(missing)}"
|
|
try:
|
|
with Session(db_engine) as session:
|
|
session.exec(
|
|
update(Run)
|
|
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
|
.values(status_reason=reason)
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
logger.exception("Could not record what run %s is waiting for", run_id)
|
|
|
|
def _defer(self, item: WorkItem) -> None:
|
|
"""Put an item back for later, and let go of this delivery."""
|
|
try:
|
|
self.queue.add_delayed(item, time.time() + WAIT_RETRY_S)
|
|
self.queue.ack(item)
|
|
except Exception:
|
|
logger.exception("Could not defer run %s", item.run_id)
|
|
|
|
def _keep_leases(self) -> None:
|
|
"""Say the local runs are alive, and clean up after engines that died."""
|
|
last_sweep = 0.0
|
|
while not self._stop.is_set():
|
|
self._stop.wait(LEASE_INTERVAL_S)
|
|
if self._stop.is_set():
|
|
break
|
|
with self._lock:
|
|
mine = list(self._active)
|
|
now = datetime.now(UTC)
|
|
try:
|
|
if mine:
|
|
with Session(db_engine) as session:
|
|
session.exec(
|
|
update(Run)
|
|
.where(col(Run.id).in_(mine))
|
|
.values(lease_at=now)
|
|
)
|
|
session.commit()
|
|
if time.monotonic() - last_sweep >= SWEEP_INTERVAL_S:
|
|
last_sweep = time.monotonic()
|
|
self._sweep_abandoned(now)
|
|
except Exception:
|
|
logger.exception("Could not refresh run leases")
|
|
|
|
def _sweep_abandoned(self, now: datetime) -> None:
|
|
cutoff = now - timedelta(seconds=LEASE_STALE_S)
|
|
with Session(db_engine) as session:
|
|
result = session.exec(
|
|
update(Run)
|
|
.where(
|
|
col(Run.status) == "running",
|
|
col(Run.lease_at) < cutoff,
|
|
)
|
|
.values(
|
|
status="abandoned",
|
|
finished_at=now,
|
|
status_reason="The engine running it stopped reporting",
|
|
)
|
|
)
|
|
session.commit()
|
|
if result.rowcount:
|
|
logger.warning("Marked %d run(s) abandoned", result.rowcount)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Driving one run
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _claim(self, run_id: str) -> Run | None:
|
|
"""Take the run, or leave it: whoever moves it out of `queued` owns it.
|
|
|
|
The compare-and-swap is what makes a redelivered item harmless — the
|
|
second engine to arrive updates nothing and walks away.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
with Session(db_engine) as session:
|
|
result = session.exec(
|
|
update(Run)
|
|
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
|
.values(
|
|
status="running",
|
|
started_at=now,
|
|
lease_at=now,
|
|
engine=self.engine_name,
|
|
)
|
|
)
|
|
session.commit()
|
|
if not result.rowcount:
|
|
return None
|
|
return session.exec(select(Run).where(col(Run.id) == run_id)).first()
|
|
|
|
def _drive(self, run_id: str) -> None:
|
|
run = self._claim(run_id)
|
|
if run is None:
|
|
return
|
|
started = time.perf_counter()
|
|
status = "ok"
|
|
reason = ""
|
|
result: dict[str, Any] = {}
|
|
state: StateBackend | None = None
|
|
errors = 0
|
|
|
|
def observe(outcome: NodeOutcome) -> None:
|
|
nonlocal errors
|
|
if not outcome.ok:
|
|
errors += 1
|
|
self._record_node(run_id, outcome)
|
|
|
|
sink = MetricSink(run_id)
|
|
try:
|
|
flow = self.controller.store.read_flow(run.flow)
|
|
state = self._state_factory(f"{RUN_NAMESPACE}:{run_id}")
|
|
pipeline = self.controller.build_run_pipeline(
|
|
flow,
|
|
state=state,
|
|
observer=observe,
|
|
emission_observer=sink.handle,
|
|
run=RunContext(run_id=run_id),
|
|
run_cache=None if run.no_cache else self._cache,
|
|
)
|
|
with self._lock:
|
|
self._active[run_id] = pipeline
|
|
if run_id in self._cancelled:
|
|
pipeline.pause(flow.name)
|
|
self._publish(run, "run_started")
|
|
pipeline.run(seed_values(flow, run.params, run.seed))
|
|
result = collect_result(flow, state)
|
|
with self._lock:
|
|
cancelled = run_id in self._cancelled
|
|
if cancelled:
|
|
status = "cancelled"
|
|
reason = "Cancelled while running"
|
|
elif errors:
|
|
status = "error"
|
|
reason = f"{errors} node(s) failed"
|
|
except Exception as exc:
|
|
logger.exception("Run %s failed", run_id)
|
|
status = "error"
|
|
reason = f"{type(exc).__name__}: {exc}"[:1024]
|
|
finally:
|
|
with self._lock:
|
|
self._active.pop(run_id, None)
|
|
self._cancelled.discard(run_id)
|
|
# Whatever the last batch was holding belongs to this run's record.
|
|
sink.flush()
|
|
duration = round((time.perf_counter() - started) * 1000, 2)
|
|
self._finish(run_id, status, reason, result, duration)
|
|
run.status = status
|
|
self._publish(run, "run_finished")
|
|
# Its values were only ever this run's; nothing reads them once it
|
|
# has a result. On Redis the namespace would expire anyway.
|
|
if state is not None and status != "error":
|
|
try:
|
|
state.clear()
|
|
except Exception:
|
|
logger.warning("Could not clear state of run %s", run_id)
|
|
|
|
def _record_node(self, run_id: str, outcome: NodeOutcome) -> None:
|
|
outputs = _cacheable(outcome)
|
|
row = RunNode(
|
|
run_id=run_id,
|
|
node=outcome.node[:255],
|
|
status="cached" if outcome.cached else ("ok" if outcome.ok else "error"),
|
|
started_at=datetime.now(UTC),
|
|
duration_ms=outcome.duration_ms,
|
|
error=outcome.error[:ERROR_CAP],
|
|
logs=outcome.logs[:LOG_CAP],
|
|
# Together or not at all: a row carrying a key must be one a
|
|
# lookup can actually restore from.
|
|
cache_key=outcome.cache_key if outputs is not None else "",
|
|
outputs=outputs,
|
|
)
|
|
try:
|
|
with Session(db_engine) as session:
|
|
session.merge(row)
|
|
for message, ref in outcome.artifacts.items():
|
|
session.merge(
|
|
RunArtifact(
|
|
run_id=run_id,
|
|
name=message[:255],
|
|
node=outcome.node[:255],
|
|
digest=str(ref.get("digest") or "")[:71],
|
|
size=int(ref.get("size") or 0),
|
|
media_type=str(
|
|
ref.get("media_type") or "application/octet-stream"
|
|
)[:128],
|
|
)
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
logger.exception(
|
|
"Could not record node '%s' of run %s", outcome.node, run_id
|
|
)
|
|
|
|
def _finish(
|
|
self,
|
|
run_id: str,
|
|
status: str,
|
|
reason: str,
|
|
result: dict[str, Any],
|
|
duration_ms: float,
|
|
) -> None:
|
|
try:
|
|
with Session(db_engine) as session:
|
|
session.exec(
|
|
update(Run)
|
|
.where(col(Run.id) == run_id)
|
|
.values(
|
|
status=status,
|
|
status_reason=reason,
|
|
result=result,
|
|
duration_ms=duration_ms,
|
|
finished_at=datetime.now(UTC),
|
|
)
|
|
)
|
|
session.commit()
|
|
except Exception:
|
|
logger.exception("Could not close run %s", run_id)
|
|
|
|
def _publish(self, run: Run, kind: str) -> None:
|
|
self._publish_event(
|
|
{
|
|
"type": kind,
|
|
"flow": run.flow,
|
|
"run": run.id,
|
|
"status": run.status,
|
|
"group": run.group_id or "",
|
|
"ts": time.time(),
|
|
}
|
|
)
|
|
|
|
def _publish_event(self, event: dict[str, Any]) -> None:
|
|
if self.controller.events is not None:
|
|
self.controller.events.publish(event)
|