stroblmeandClaude Opus 5
40f8ad378d
Slurm is not a machine that attaches and stays; it is a queue somebody else owns. So nothing here submits a node to it. It submits a job whose payload is an ordinary worker dialling back in, and everything downstream — the protocol, the artifacts, cancellation, the books — already worked and did not have to learn what Slurm is. The alternative, which Covalent takes, is to stage a serialized call and a runner onto the login node, poll squeue and copy the result back: a second way of running a node beside the one that exists. The cost of not doing that is one assumption, that a compute node can open a connection outward. Where that is false, _payload is the single method a staged variant would replace. Clusters are configured in provisioners.json beside the alerts, since this is infrastructure an operator writes rather than anything a flow says. The script is generated with the system ssh and no new dependency, and prerun owns the environment — deliberately no pip install, because what is on a cluster is somebody's decision. One outstanding request per profile, cancelled if it never attaches and on the way out. Nothing autoscales. The run gate needed the same hook: a run held before it starts never reaches the placer's own wait, so it would have queued forever on a machine nothing had asked for. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A6HeySA27EkGANZN95QySW
1191 lines
47 KiB
Python
1191 lines
47 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
|
|
|
|
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, 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 installation'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 installation'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.
|
|
|
|
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
|
|
|
|
|
|
#: 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
|
|
|
|
|
|
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 _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]
|
|
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,
|
|
idempotency_key: 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)
|
|
# 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=code_digest(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),
|
|
needs=required_resources(flow),
|
|
created_at=datetime.now(UTC),
|
|
actor=actor,
|
|
idempotency_key=idempotency_key,
|
|
)
|
|
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 _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_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)
|
|
try:
|
|
flow = self.controller.store.read_flow(run.flow)
|
|
# 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.
|
|
run.code_digest = self._restamp(run, code_digest(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, code_digest=run.code_digest),
|
|
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],
|
|
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():
|
|
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)
|