stroblmeandClaude Opus 5
058f16ec1d
Each was a loose end recorded under `### SDK` in the notepad. `serve` takes its own pidfile down on SIGTERM. uvicorn restores the handler it found and re-raises the signal it stopped on, so the default handler ended the process without unwinding and the `finally` never ran — which is what a stop sends, and what left `serve.pid` behind. `serve.log` is cut back past 5 MB by the engine rather than by the screen that started it, so an adopted engine is bounded too. Gated on its own stdout being an appended regular file, which is what makes the cut safe: the kernel then puts the next write at the new end. Cards are counted from `/dev/nvidia[0-9]*`, so `FLOW_GPUS`/`--gpus` of 0 means "work it out" the way `FLOW_CPUS` always has. The engine counts, not the accountant — a remote worker builds one of those from its own inventory, and detecting there would hand it the engine host's cards. The worker counts last: what a batch job says it was granted still wins. `GET /runs/metrics/names` is the distinct over a selection that `--list` and the terminal's metric picker were approximating by reading the newest run that had measured anything, which missed a name only an older run ever wrote. `MetricSink` announces each batch it has written (`run_metric`, carrying the names). Not a per-point event: one covers up to 500 points or two seconds of them, and the rows stay the record. The terminal comparison fills in as the first readings land instead of staying blank until reopened, and the browser refetches the run and any comparison rather than the list behind them. `retry --group` pages the list route by `before` instead of stopping at 500. The terminal dashboard takes the terminal's colours (`ansi-dark`), and the web UI can re-pair from Settings without disconnecting first. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PRQ9bmTvCbqCwXo9mxZzzV
1417 lines
57 KiB
Python
1417 lines
57 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 copy
|
|
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 pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlalchemy import update
|
|
from sqlalchemy.dialects.sqlite import insert as upsert
|
|
from sqlalchemy.exc import IntegrityError
|
|
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 NODE_TYPES, FlowController, RunContext
|
|
from fluksio.flow.messages import qualify
|
|
from fluksio.flow.pipeline import CacheHit, NodeOutcome, Pipeline
|
|
from fluksio.flow.queue import WorkItem, WorkQueue
|
|
from fluksio.flow.schemas import FlowDef, NodeDef
|
|
from fluksio.flow.state import MemoryState, StateBackend
|
|
from fluksio.flow.store import FlowStore
|
|
from fluksio.models import Run, RunArtifact, RunMetric, RunNode
|
|
from fluksio.sdk.imports import digest_of as files_digest
|
|
|
|
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, unless `FLOW_MAX_RUNS` says otherwise. Node bodies are
|
|
#: bounded by the worker pool anyway; this only bounds how many graphs are in
|
|
#: flight.
|
|
MAX_PARALLEL = 4
|
|
#: Runs claimed per poll, at least. Raising the pool raises this with it, so a
|
|
#: queue of ready runs fills the drivers in one pass rather than four a second.
|
|
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, store: FlowStore | None = None, draft: bool = False
|
|
) -> 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.
|
|
|
|
The other is a node with no body. The store answers a new node's template
|
|
when nothing was ever written for one, so such a node runs — and returns
|
|
``{}`` every time, without a word. `fluksio sync` writes every body before
|
|
it publishes, so this is unreachable from there; it is the other end that
|
|
is open.
|
|
"""
|
|
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."
|
|
)
|
|
if store is not None and _has_no_body(store, flow.name, node, draft):
|
|
issues.append(
|
|
f"Node '{node.id}' has no stored code. It would run as an "
|
|
"empty node and publish nothing, so the run is refused."
|
|
)
|
|
return issues
|
|
|
|
|
|
def _has_no_body(
|
|
store: FlowStore, flow: str, node: NodeDef, draft: bool = False
|
|
) -> bool:
|
|
"""A node that should carry its own source, and does not.
|
|
|
|
A shared node runs the library's copy, so it is not this: a missing library
|
|
fails the node loudly on its own.
|
|
"""
|
|
node_type = NODE_TYPES.get(node.type)
|
|
if node_type is None or not node_type.has_source or node.source_ref:
|
|
return False
|
|
return not store.has_node_source(flow, node.id, draft=draft)
|
|
|
|
|
|
def _runs_on_a_worker(node: NodeDef) -> bool:
|
|
"""Whether this node's body could go to a worker at all.
|
|
|
|
Only a node whose source travels can: a connector is an entry point loaded
|
|
in this process, so a device on one is a field nothing reads. Holding a run
|
|
for a worker that could never take that node is a run that never starts.
|
|
"""
|
|
node_type = NODE_TYPES.get(node.type)
|
|
return node_type is not None and node_type.has_source
|
|
|
|
|
|
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"
|
|
and _runs_on_a_worker(node)
|
|
}
|
|
)
|
|
|
|
|
|
def required_resources(flow: FlowDef) -> dict[str, Any] | None:
|
|
"""The largest single thing this flow needs a machine to have.
|
|
|
|
Dimension by dimension rather than per node, which is deliberately blunt:
|
|
it answers "could this run start at all", not "in what order". A node that
|
|
only *prefers* its device is left out — it runs here when nothing carries
|
|
the label, so it is never what a run is waiting for.
|
|
"""
|
|
cpus, gpus, ram = 0, 0, 0
|
|
device = ""
|
|
for node in flow.nodes:
|
|
wanted = node.resources
|
|
if wanted is None or not _runs_on_a_worker(node):
|
|
continue
|
|
cpus = max(cpus, wanted.cpus)
|
|
gpus = max(gpus, wanted.gpus)
|
|
ram = max(ram, wanted.ram or 0)
|
|
if node.device and node.device_policy == "require":
|
|
device = node.device
|
|
if not (cpus or gpus or ram):
|
|
return None
|
|
return {"cpus": cpus, "gpus": gpus, "ram_mb": ram, "device": device}
|
|
|
|
|
|
#: 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 spelling of a run's output into the value itself.
|
|
|
|
A python caller hands one run's output straight to the next, because it has
|
|
the value in its hand. A shell does not, and pasting a config object or an
|
|
artifact reference is not a command anyone wants to type — so any declared
|
|
input also takes ``@run:<id>.<output>``, naming what a run produced, and a
|
|
bare ``sha256:...`` digest naming 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.
|
|
|
|
Both spellings are therefore reserved on every input, ``str`` included: an
|
|
input that has to carry one of them literally is asking for a value this
|
|
engine reads as a name.
|
|
"""
|
|
declared = {one.spec.name for one in flow.inputs}
|
|
pending = {
|
|
key: value
|
|
for key, value in params.items()
|
|
if key in declared
|
|
and isinstance(value, str)
|
|
and (value.startswith(RUN_REF_PREFIX) or valid_digest(value))
|
|
}
|
|
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) :], artifacts
|
|
)
|
|
else:
|
|
resolved[key] = _from_digest(session, key, text, artifacts)
|
|
return resolved
|
|
|
|
|
|
def _from_run(
|
|
session: Session, key: str, spelling: str, artifacts: ArtifactStore | None = None
|
|
) -> Any:
|
|
"""``<run id>.<output>`` as the value 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 value as its producer made it,
|
|
# an artifact's file name and all. Copied, because what comes back is
|
|
# handed on as this run's parameter and must not alias the other run's row.
|
|
result = run.result or {}
|
|
if output in result:
|
|
return copy.deepcopy(result[output])
|
|
|
|
# The artifact rows are the fallback: bytes a node made that the flow never
|
|
# declared as an output. A row records what the node called the file, so
|
|
# the reference rebuilt here is the one its producer made — apart from the
|
|
# rows written before there was a column to keep it in, which answer with
|
|
# the message name as they always did.
|
|
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]):
|
|
if artifacts is not None and artifacts.path(row.digest) is None:
|
|
raise RunRejected(
|
|
f"Parameter '{key}': run '{run_id}' made '{output}', but its "
|
|
"bytes are gone from this instance's store"
|
|
)
|
|
return {
|
|
"digest": row.digest,
|
|
"size": row.size,
|
|
"media_type": row.media_type or "application/octet-stream",
|
|
"name": row.filename or row.name,
|
|
}
|
|
known = ", ".join(sorted({*result, *(row.name for row in rows)})) or "none"
|
|
raise RunRejected(
|
|
f"Parameter '{key}': run '{run_id}' has no output '{output}' (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 instance's store"
|
|
)
|
|
return {
|
|
"digest": row.digest,
|
|
"size": row.size,
|
|
"media_type": row.media_type or "application/octet-stream",
|
|
"name": row.filename or row.name,
|
|
}
|
|
|
|
|
|
#: How deep a stored value is walked looking for artifact references. A
|
|
#: reference nested past this is not something any port declares.
|
|
_REF_DEPTH = 8
|
|
#: How many state keys are read back at a time by the sweep.
|
|
_SWEEP_BATCH = 500
|
|
|
|
|
|
def _references_in(value: Any, found: set[str], depth: int = 0) -> None:
|
|
"""Every artifact digest inside one stored value."""
|
|
if is_reference(value):
|
|
found.add(str(value["digest"]))
|
|
return
|
|
if depth >= _REF_DEPTH:
|
|
return
|
|
if isinstance(value, dict):
|
|
items: Any = value.values()
|
|
elif isinstance(value, (list, tuple)):
|
|
items = value
|
|
else:
|
|
return
|
|
for item in items:
|
|
_references_in(item, found, depth + 1)
|
|
|
|
|
|
def sweep_artifacts(
|
|
store: ArtifactStore, state: StateBackend, grace_s: float = 0.0
|
|
) -> int:
|
|
"""Remove artifact bytes nothing refers to any more. Returns how many went.
|
|
|
|
Two things refer to an artifact: a run that recorded it, and a message
|
|
currently holding it. The second is what makes a stream of media
|
|
affordable — a camera publishing a frame a second replaces the reference
|
|
each time, so yesterday's frames are unreferenced by definition and the
|
|
store does not grow without bound. What a node *returns* is recorded
|
|
against its run and kept; what it emits along the way is not.
|
|
|
|
Skipped entirely while anything is running: a node that stores a checkpoint
|
|
an hour before it returns has neither a row nor a message naming it yet,
|
|
and a sweep in that window would take the bytes out from under it.
|
|
"""
|
|
with Session(db_engine) as session:
|
|
active = session.exec(
|
|
select(Run.id).where(col(Run.status).in_(("running", "queued"))).limit(1)
|
|
).first()
|
|
if active is not None:
|
|
logger.debug("Artifact sweep skipped: run %s is in flight", active)
|
|
return 0
|
|
keep = set(session.exec(select(col(RunArtifact.digest)).distinct()).all())
|
|
|
|
keys = state.keys()
|
|
for start in range(0, len(keys), _SWEEP_BATCH):
|
|
for value in state.get_multi(keys[start : start + _SWEEP_BATCH]).values():
|
|
_references_in(value, keep)
|
|
|
|
removed = store.collect(keep, grace_s)
|
|
if removed:
|
|
logger.info("Artifact sweep removed %d unreferenced artifacts", removed)
|
|
return removed
|
|
|
|
|
|
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.
|
|
|
|
A batch that has been written announces itself (`on_flush`), which is how a
|
|
screen watching a run knows there is something new to read. The event
|
|
carries names rather than numbers: it is a nudge, and the rows are the
|
|
record. One per batch, so it is nowhere near the per-point path.
|
|
|
|
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,
|
|
on_flush: Callable[[list[str]], None] | None = None,
|
|
) -> None:
|
|
self.run_id = run_id
|
|
self._batch = batch
|
|
self._interval = interval
|
|
self._on_flush = on_flush
|
|
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._written(rows)
|
|
|
|
def flush(self) -> None:
|
|
with self._lock:
|
|
rows = list(self._rows.values())
|
|
self._rows.clear()
|
|
self._last_flush = time.monotonic()
|
|
if rows:
|
|
self._written(rows)
|
|
|
|
def _written(self, rows: list[RunMetric]) -> None:
|
|
"""Write a batch, then say that it is there."""
|
|
self._write(rows)
|
|
if self._on_flush is None:
|
|
return
|
|
try:
|
|
self._on_flush(sorted({row.name for row in rows}))
|
|
except Exception:
|
|
logger.exception("Announcing a batch of run %s failed", self.run_id)
|
|
|
|
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
|
|
|
|
|
|
#: What a repository's python files hashed to, and the stat state that answer
|
|
#: was read from. Keyed by path: the walk is cheap, opening every file is not.
|
|
_digests: dict[str, tuple[tuple[tuple[str, int, int], ...], str]] = {}
|
|
|
|
#: Directories a project's own code is never in, and which are large.
|
|
_SKIP_DIRS = frozenset({"__pycache__", "node_modules"})
|
|
|
|
|
|
def repo_digest(repo: str) -> str:
|
|
"""What the python files under ``repo`` currently hash to.
|
|
|
|
A code-defined node is stored as a shim that imports the real function, so
|
|
the engine runs whatever is on disk when the worker starts. The commit
|
|
cannot name that: an uncommitted tree stamps every run `-dirty`, however
|
|
many times the code changed underneath. This can — it is read at run time,
|
|
from the files themselves.
|
|
|
|
Empty when there is no repository to read here, which is the honest answer
|
|
on an engine that does not share a filesystem with the author: a stamp
|
|
without a digest says the engine could not see the code, and the cache key
|
|
falls back to what it keyed on before.
|
|
"""
|
|
if not repo or not os.path.isdir(repo):
|
|
return ""
|
|
state: list[tuple[str, int, int]] = []
|
|
for root, dirs, files in os.walk(repo):
|
|
dirs[:] = [
|
|
d
|
|
for d in dirs
|
|
if not d.startswith(".")
|
|
and d not in _SKIP_DIRS
|
|
# A virtualenv carries this, and its site-packages is not the
|
|
# project's code — nor is it small.
|
|
and not os.path.exists(os.path.join(root, d, "pyvenv.cfg"))
|
|
]
|
|
for name in files:
|
|
if not name.endswith(".py"):
|
|
continue
|
|
path = os.path.join(root, name)
|
|
try:
|
|
info = os.stat(path)
|
|
except OSError:
|
|
continue
|
|
state.append((os.path.relpath(path, repo), info.st_mtime_ns, info.st_size))
|
|
state.sort()
|
|
stamp = tuple(state)
|
|
cached = _digests.get(repo)
|
|
if cached is not None and cached[0] == stamp:
|
|
return cached[1]
|
|
|
|
digest = hashlib.sha256()
|
|
for relative, _mtime, _size in state:
|
|
try:
|
|
body = Path(repo, relative).read_bytes()
|
|
except OSError:
|
|
continue
|
|
digest.update(relative.encode())
|
|
digest.update(b"\0")
|
|
digest.update(body)
|
|
digest.update(b"\0")
|
|
answer = digest.hexdigest()
|
|
_digests[repo] = (stamp, answer)
|
|
return answer
|
|
|
|
|
|
#: What a node's own modules hashed to, and the stat state behind that answer.
|
|
#: Keyed by the node's file list, the same trade `_digests` makes.
|
|
_node_digests: dict[tuple[str, ...], tuple[tuple[tuple[int, int], ...], str]] = {}
|
|
|
|
|
|
def node_digest(node: NodeDef, flow: FlowDef) -> str:
|
|
"""What the code this node reaches hashes to now.
|
|
|
|
`sync` recorded which project modules the node's function imports its way
|
|
to, so this is the narrow answer the repository-wide digest could not
|
|
give: it moves when the node's own code moves and stays put when a
|
|
notebook two directories away changes.
|
|
|
|
Read live, because the shim imports from disk whenever the worker starts
|
|
and a sweep may sit queued for hours. An engine that cannot see the files
|
|
keeps what sync recorded instead of nothing — a remote worker's runs used
|
|
to share one empty digest, and one key.
|
|
"""
|
|
files = node.code_files
|
|
if not files:
|
|
# Drawn on the canvas, or synced before this was recorded. The
|
|
# repository is the older, blunter answer and still the only one.
|
|
return code_digest(flow)
|
|
paths = tuple(sorted(files.values()))
|
|
state: list[tuple[int, int]] = []
|
|
for path in paths:
|
|
try:
|
|
info = os.stat(path)
|
|
except OSError:
|
|
return node.code_digest
|
|
state.append((info.st_mtime_ns, info.st_size))
|
|
stamp = tuple(state)
|
|
cached = _node_digests.get(paths)
|
|
if cached is not None and cached[0] == stamp:
|
|
return cached[1]
|
|
answer = files_digest(files) or node.code_digest
|
|
_node_digests[paths] = (stamp, answer)
|
|
return answer
|
|
|
|
|
|
def node_digests(flow: FlowDef) -> dict[str, str]:
|
|
"""Every node's digest, which is what its fingerprint is keyed on."""
|
|
return {node.id: node_digest(node, flow) for node in flow.nodes}
|
|
|
|
|
|
def code_digest(flow: FlowDef) -> str:
|
|
"""The digest of the repository a code-defined flow was declared in."""
|
|
origin = flow.origin
|
|
return repo_digest(origin.repo) if origin is not None else ""
|
|
|
|
|
|
def run_digest(flow: FlowDef) -> str:
|
|
"""What the run records as the code it ran: its nodes' digests, hashed.
|
|
|
|
One column over a flow's several nodes, so it changes when any node's own
|
|
code does and not when anything else in the repository does — which is
|
|
what makes it worth joining an exported table on.
|
|
"""
|
|
parts = node_digests(flow)
|
|
if not any(parts.values()):
|
|
return ""
|
|
canonical = json.dumps(parts, sort_keys=True, separators=(",", ":"))
|
|
return hashlib.sha256(canonical.encode()).hexdigest()
|
|
|
|
|
|
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) -> CacheHit | None:
|
|
if not key:
|
|
return 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
|
|
source = session.get(Run, row.run_id)
|
|
return CacheHit(
|
|
flow=source.flow if source is not None else "",
|
|
outputs=outputs,
|
|
# A row that was itself restored holds no series, so
|
|
# follow its pointer rather than adding a hop to it.
|
|
metrics_run=row.cached_from or row.run_id,
|
|
)
|
|
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 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]
|
|
# Taken as written: clamping a 0 up to 1 would hide a limit somebody
|
|
# set, and the pool below rejects an unusable one loudly anyway.
|
|
self.parallel = parallel
|
|
self._pool = ThreadPoolExecutor(
|
|
max_workers=self.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()
|
|
self._requeue_queued()
|
|
|
|
def _requeue_queued(self) -> None:
|
|
"""Wake the engine for runs an earlier process was holding.
|
|
|
|
A row is journaled `queued` and then a work item is added, and the two
|
|
are not one transaction — an in-memory queue loses the item with the
|
|
process, and a stream item nobody claimed is nobody's. The row is the
|
|
record, so this reads it back. Nothing is re-executed by it: claiming
|
|
is a compare-and-set out of `queued`, so an item that did survive is
|
|
claimed once whatever this adds.
|
|
"""
|
|
with Session(db_engine) as session:
|
|
waiting = session.exec(
|
|
select(Run.id, Run.flow).where(col(Run.status) == "queued")
|
|
).all()
|
|
for run_id, flow in waiting:
|
|
self.queue.add(WorkItem(kind="run", node="", flow=flow, run_id=run_id))
|
|
if waiting:
|
|
logger.info(
|
|
"Requeued %d run(s) left queued by an earlier engine", len(waiting)
|
|
)
|
|
|
|
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,
|
|
idempotency_key: str | None = None,
|
|
parent_id: str | None = None,
|
|
) -> Run:
|
|
"""Journal a run and wake an engine up for it. Never blocks on it."""
|
|
# Before anything else, including reading the flow: a caller retrying a
|
|
# submit it never got an answer for is owed the run that answer was
|
|
# about, whatever the flow says now.
|
|
if idempotency_key:
|
|
existing = self._by_key(idempotency_key)
|
|
if existing is not None:
|
|
return existing
|
|
flow = self.controller.store.read_flow(flow_name, draft=draft)
|
|
issues = batch_issues(flow, self.controller.store, draft=draft)
|
|
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)
|
|
# What the run actually starts from, not only what was passed: an input
|
|
# left out takes its declared value, and a row that records `{}` cannot
|
|
# say which. Folded literally — an initial is a value from the
|
|
# definition, never a reference to resolve.
|
|
declared = {
|
|
one.spec.name: one.initial for one in flow.inputs if one.initial is not None
|
|
}
|
|
# The run's own seed fills an input of that name, outranking what the
|
|
# flow declares and outranked by one passed as a parameter — the order
|
|
# `seed_values` applies, moved to where the record is written.
|
|
if seed is not None and any(one.spec.name == "seed" for one in flow.inputs):
|
|
declared["seed"] = seed
|
|
params = {**declared, **params}
|
|
# And back the other way, so the run-level column holds the seed the
|
|
# run actually used however it arrived. Otherwise `--seed 1` fills one
|
|
# column and a declared seed the other, and that is the single field an
|
|
# export still has to coalesce.
|
|
resolved_seed = params.get("seed")
|
|
if isinstance(resolved_seed, int) and not isinstance(resolved_seed, bool):
|
|
seed = resolved_seed
|
|
# 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),
|
|
code_digest=run_digest(flow),
|
|
params=params,
|
|
params_digest=digest_of(params, seed),
|
|
seed=seed,
|
|
group_id=group_id,
|
|
cause=cause,
|
|
no_cache=no_cache,
|
|
draft=draft,
|
|
status="queued",
|
|
labels=required_labels(flow),
|
|
needs=required_resources(flow),
|
|
created_at=datetime.now(UTC),
|
|
actor=actor,
|
|
idempotency_key=idempotency_key,
|
|
parent_id=parent_id,
|
|
)
|
|
with Session(db_engine) as session:
|
|
session.add(run)
|
|
try:
|
|
session.commit()
|
|
except IntegrityError:
|
|
# Two retries of one submit raced here. The unique index is
|
|
# what decided which of them is the run; this one reads it
|
|
# back rather than queueing a second execution of it.
|
|
session.rollback()
|
|
existing = self._by_key(idempotency_key) if idempotency_key else None
|
|
if existing is None:
|
|
raise
|
|
return existing
|
|
session.refresh(run)
|
|
|
|
self.queue.add(WorkItem(kind="run", node="", flow=flow.name, run_id=run.id))
|
|
return run
|
|
|
|
def _by_key(self, idempotency_key: str) -> Run | None:
|
|
"""The run a key already made, if it made one."""
|
|
with Session(db_engine) as session:
|
|
return session.exec(
|
|
select(Run).where(col(Run.idempotency_key) == idempotency_key)
|
|
).first()
|
|
|
|
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 retry(self, run_id: str, actor: str = "") -> Run:
|
|
"""Submit the same run again, as a run of its own.
|
|
|
|
What an engine that stopped mid-run leaves behind is a run marked
|
|
`abandoned`, and re-issuing the whole sweep to recover one config is
|
|
the blunt way back. This is the sharp one: the same flow, the same
|
|
inputs, the same group, so a sweep is completed rather than repeated.
|
|
The stage cache is what makes it cheap — the nodes that finished are
|
|
restored rather than run again.
|
|
|
|
A run of the working copy retries as one, which means whatever the
|
|
draft is now rather than what it was — there is no older draft to go
|
|
back to, and the point of retrying a draft run is the code on disk.
|
|
"""
|
|
with Session(db_engine) as session:
|
|
run = session.get(Run, run_id)
|
|
if run is None:
|
|
raise RunRejected(f"No run {run_id}")
|
|
if run.status in ("queued", "running"):
|
|
raise RunRejected("That run has not finished; cancel it first")
|
|
source = run.model_copy()
|
|
return self.submit(
|
|
source.flow,
|
|
params=dict(source.params),
|
|
seed=source.seed,
|
|
group_id=source.group_id,
|
|
cause="retry",
|
|
actor=actor,
|
|
no_cache=source.no_cache,
|
|
draft=source.draft,
|
|
parent_id=source.id,
|
|
)
|
|
|
|
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 —
|
|
# and asks the queue for nothing at all while nothing is
|
|
# waiting, which is the common case.
|
|
due = self.queue.next_due()
|
|
if due is not None and due <= time.time():
|
|
self.queue.move_due(time.time())
|
|
items = self.queue.claim(
|
|
max(CLAIM_COUNT, self.parallel), 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_requirements(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_requirements(self, run_id: str) -> list[str]:
|
|
"""What this run needs that nothing attached can give it yet."""
|
|
with Session(db_engine) as session:
|
|
run = session.get(Run, run_id)
|
|
needed = list(run.labels) if run else []
|
|
needs = dict(run.needs) if run and run.needs else None
|
|
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.
|
|
missing = [f"labelled {label}" for label in sorted(set(needed) - available)]
|
|
placer = self.controller.placer
|
|
if placer is not None:
|
|
# Only a wait something is going to end — a machine that could take
|
|
# this is being started for it. An ask nothing can grant is cut down
|
|
# to what is here and runs, so holding the run would be holding it
|
|
# for something that is not coming.
|
|
short = placer.satisfiable(needs)
|
|
if short:
|
|
missing.append(f"with {short}")
|
|
# Held here rather than at the placer's own wait, so this is
|
|
# where the machine has to be asked for. Asked again on every
|
|
# pass; one outstanding request is the provisioner's business.
|
|
placer.provision_for(needs)
|
|
return missing
|
|
|
|
def _waiting(self, run_id: str, missing: list[str]) -> None:
|
|
reason = f"Waiting for a worker {' and '.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 _restamp(self, run: Run, digest: str) -> str:
|
|
"""Record the code this run is about to execute, if it moved."""
|
|
if digest == run.code_digest:
|
|
return digest
|
|
with Session(db_engine) as session:
|
|
session.exec(
|
|
update(Run).where(col(Run.id) == run.id).values(code_digest=digest)
|
|
)
|
|
session.commit()
|
|
return digest
|
|
|
|
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,
|
|
on_flush=lambda names: self._publish_event(
|
|
{
|
|
"type": "run_metric",
|
|
"flow": run.flow,
|
|
"run": run_id,
|
|
"names": names,
|
|
"ts": time.time(),
|
|
}
|
|
),
|
|
)
|
|
try:
|
|
flow = self.controller.store.read_flow(run.flow, draft=run.draft)
|
|
# Read again, now that it is this run's turn: a sweep queues every
|
|
# run at once, and the code on disk is free to move in the hours
|
|
# before the last of them starts. What the record must name is the
|
|
# state that ran, not the state that was submitted. A draft that
|
|
# was published meanwhile reads as the published copy, which is
|
|
# the same document under a different name.
|
|
digests = node_digests(flow)
|
|
run.code_digest = self._restamp(run, run_digest(flow))
|
|
state = self._state_factory(f"{RUN_NAMESPACE}:{run_id}")
|
|
pipeline = self.controller.build_run_pipeline(
|
|
flow,
|
|
state=state,
|
|
draft=run.draft,
|
|
observer=observe,
|
|
emission_observer=sink.handle,
|
|
run=RunContext(
|
|
run_id=run_id,
|
|
code_digest=run.code_digest,
|
|
node_digests=digests,
|
|
),
|
|
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")
|
|
self._release_cards(run)
|
|
# 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:
|
|
if status != "error":
|
|
try:
|
|
state.clear()
|
|
except Exception:
|
|
logger.warning("Could not clear state of run %s", run_id)
|
|
# Even when the values are kept: on Redis this backend is a
|
|
# client and a connection pool of its own, built per run.
|
|
try:
|
|
state.close()
|
|
except Exception:
|
|
logger.warning("Could not close state of run %s", run_id)
|
|
|
|
def _release_cards(self, run: Run) -> None:
|
|
"""Hand a GPU run's device memory back when the run is over.
|
|
|
|
The accountant frees the card the moment the node returns, but the
|
|
worker that ran on it is kept warm and a library that preallocated
|
|
most of the VRAM never gives it up — so the next process to want the
|
|
card found it taken by one sitting idle.
|
|
"""
|
|
# ponytail: retires every CUDA pool rather than the ones this run used,
|
|
# which needs no bookkeeping — a concurrent GPU run's busy worker only
|
|
# dies when it returns, which is when its own memory should go back
|
|
# anyway. What it costs is the warm worker of a *live* flow's GPU node.
|
|
# Track the pools per run if that ever matters.
|
|
if not (run.needs or {}).get("gpus"):
|
|
return
|
|
pool = getattr(self.controller, "workers", None)
|
|
if pool is None:
|
|
return
|
|
try:
|
|
pool.retire_gpu_children()
|
|
except Exception:
|
|
logger.warning("Could not retire the GPU workers 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],
|
|
# The tail, not the head: a failure appends its traceback, and a
|
|
# chatty node would otherwise push it past the cap.
|
|
logs=outcome.logs[-LOG_CAP:],
|
|
cached_from=outcome.cached_from,
|
|
# 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():
|
|
# A run recording a reference is what "returned media is
|
|
# kept" means: copied out of the volatile ring, or the row
|
|
# would outlive the bytes it names by a few seconds.
|
|
if self._artifacts is not None:
|
|
self._artifacts.adopt(str(ref.get("digest") or ""))
|
|
session.merge(
|
|
RunArtifact(
|
|
run_id=run_id,
|
|
name=message[:255],
|
|
filename=str(ref.get("name") or "")[:255] or None,
|
|
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)
|