Rename the import package app to fluksio
A wheel whose top-level module is `app` collides with anything else in a user's venv, so the package that is about to be published takes the name it is published under. Only the Python package moves; the repo, the Docker WORKDIR and the compose project keep theirs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -0,0 +1,706 @@
|
||||
"""Runs: a flow taken from its inputs to its outputs, once.
|
||||
|
||||
A cascade is what an always-on flow does when a value arrives — it has no
|
||||
beginning and no end worth recording. A *run* is the other shape the same
|
||||
engine can take: parameters go in, the graph executes until it drains, and
|
||||
what it produced is kept. That is what an ML experiment is, and what a
|
||||
CI-style job is, so both are this one entity.
|
||||
|
||||
Three things make a run different from a cascade, and each is deliberate:
|
||||
|
||||
* **Its own state.** Message names are global keys, so two runs of one flow
|
||||
sharing the engine's state would overwrite each other's values. A run gets a
|
||||
state backend namespaced to itself, which is a constructor argument rather
|
||||
than a change to the pipeline — every key the engine keeps, versions and
|
||||
node memory included, already goes through that backend.
|
||||
* **Its own record.** The event bus drops what it cannot keep up with, which
|
||||
is right for a live canvas and wrong for a result. The driver writes the
|
||||
run's rows itself, from the thread that is running it.
|
||||
* **Its own durability.** The queue wakes an engine up; from the moment a run
|
||||
starts, its database row is the truth. Redelivering hours of training
|
||||
because an acknowledgement was late is not recovery, so a run is
|
||||
acknowledged as soon as it is claimed and a stale lease — not an unacked
|
||||
stream entry — is what marks a run its engine died in the middle of.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import update
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlmodel import Session, col, select
|
||||
|
||||
from fluksio.core.db import engine as db_engine
|
||||
from fluksio.flow.controller import FlowController, RunContext
|
||||
from fluksio.flow.messages import qualify
|
||||
from fluksio.flow.pipeline import NodeOutcome, Pipeline
|
||||
from fluksio.flow.queue import WorkItem, WorkQueue
|
||||
from fluksio.flow.schemas import FlowDef
|
||||
from fluksio.flow.state import MemoryState, StateBackend
|
||||
from fluksio.models import Run, RunArtifact, RunMetric, RunNode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: How long a finished run's state is kept before Redis drops it. Long enough
|
||||
#: to look at what a failed run left behind, short enough not to accumulate.
|
||||
RUN_STATE_TTL = 24 * 3600
|
||||
#: How often a running run says it is still alive.
|
||||
LEASE_INTERVAL_S = 20.0
|
||||
#: A lease older than this belongs to an engine that is not coming back.
|
||||
LEASE_STALE_S = 90.0
|
||||
#: How often stale leases are looked for.
|
||||
SWEEP_INTERVAL_S = 30.0
|
||||
#: Runs driven at once. Node bodies are bounded by the worker pool anyway;
|
||||
#: this only bounds how many graphs are in flight.
|
||||
MAX_PARALLEL = 4
|
||||
CLAIM_COUNT = 4
|
||||
CLAIM_BLOCK_MS = 1000
|
||||
ERROR_CAP = 2000
|
||||
LOG_CAP = 8000
|
||||
#: Reported numbers held before they are written. A training loop reporting
|
||||
#: every step must not be a round trip every step.
|
||||
METRIC_BATCH = 500
|
||||
METRIC_FLUSH_S = 2.0
|
||||
#: How often a run waiting for a worker looks again.
|
||||
WAIT_RETRY_S = 15.0
|
||||
|
||||
#: Where a run's state lives, so it can never collide with the engine's own.
|
||||
RUN_NAMESPACE = "run"
|
||||
|
||||
|
||||
class RunRejected(ValueError):
|
||||
"""The run cannot be made: bad parameters, or a flow that cannot batch."""
|
||||
|
||||
|
||||
def new_run_id() -> str:
|
||||
"""Time-ordered, so the newest runs sort last without reading a column."""
|
||||
return f"{int(time.time() * 1000):013d}-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
|
||||
def digest_of(params: dict[str, Any], seed: int | None) -> str:
|
||||
"""What identifies a run's inputs: same parameters, same digest."""
|
||||
canonical = json.dumps(
|
||||
{"params": params, "seed": seed}, sort_keys=True, separators=(",", ":")
|
||||
)
|
||||
return hashlib.sha256(canonical.encode()).hexdigest()
|
||||
|
||||
|
||||
def batch_issues(flow: FlowDef) -> list[str]:
|
||||
"""Why this flow cannot be run as a batch, if it cannot.
|
||||
|
||||
One thing genuinely breaks: a port with a discretization interval holds
|
||||
values back for a timer to release, and a run has no timer — the engine
|
||||
would drop them instead of delaying them. On a *streaming* port that is
|
||||
exactly right and is how you keep a chart from being flooded: the run's
|
||||
history keeps every value, and the interval only thins what is published.
|
||||
Anywhere else it is a message quietly going missing.
|
||||
|
||||
A delay node is fine; without a queue to defer into it simply sleeps,
|
||||
which in a run is what was asked for.
|
||||
"""
|
||||
issues: list[str] = []
|
||||
for node in flow.nodes:
|
||||
for spec in list(node.requires) + list(node.provides):
|
||||
if spec.interval > 0 and not spec.stream:
|
||||
issues.append(
|
||||
f"Node '{node.id}' rate-limits '{spec.port or spec.name}'. "
|
||||
"A run has no timer to release what that holds back, so "
|
||||
"the value would be dropped. Remove the interval, or mark "
|
||||
"the port as streaming if it is a curve being thinned out."
|
||||
)
|
||||
return issues
|
||||
|
||||
|
||||
def required_labels(flow: FlowDef) -> list[str]:
|
||||
"""Worker labels this flow cannot run without.
|
||||
|
||||
Only the nodes that *require* their device. One that merely prefers it runs
|
||||
on the engine when no such worker is attached, so holding the whole run
|
||||
back to wait for one would be waiting for something it does not need.
|
||||
"""
|
||||
return sorted(
|
||||
{
|
||||
node.device
|
||||
for node in flow.nodes
|
||||
if node.device and node.device.strip() and node.device_policy == "require"
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
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 = pg_insert(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
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
self.controller = controller
|
||||
self.queue = queue
|
||||
# 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,
|
||||
) -> Run:
|
||||
"""Journal a run and wake an engine up for it. Never blocks on it."""
|
||||
flow = self.controller.store.read_flow(flow_name, draft=draft)
|
||||
issues = batch_issues(flow)
|
||||
if issues:
|
||||
raise RunRejected(" ".join(issues))
|
||||
params = params or {}
|
||||
# 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(),
|
||||
params=params,
|
||||
params_digest=digest_of(params, seed),
|
||||
seed=seed,
|
||||
group_id=group_id,
|
||||
cause=cause,
|
||||
status="queued",
|
||||
labels=required_labels(flow),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
actor=actor,
|
||||
)
|
||||
with Session(db_engine) as session:
|
||||
session.add(run)
|
||||
session.commit()
|
||||
session.refresh(run)
|
||||
|
||||
self.queue.add(WorkItem(kind="run", node="", flow=flow.name, run_id=run.id))
|
||||
return run
|
||||
|
||||
def cancel(self, run_id: str) -> bool:
|
||||
"""Stop a run: kill what it is executing, schedule nothing further."""
|
||||
with self._lock:
|
||||
pipeline = self._active.get(run_id)
|
||||
if pipeline is None:
|
||||
# Not running here — if it is still queued, refusing to start
|
||||
# is all the cancelling it needs.
|
||||
cancelled = self._finish_queued(run_id)
|
||||
if cancelled:
|
||||
self._cancelled.add(run_id)
|
||||
return cancelled
|
||||
self._cancelled.add(run_id)
|
||||
# The gate first, so nothing new is submitted while the running nodes
|
||||
# are being killed; a gated node is never handed to the executor, so
|
||||
# the graph drains instead of going further.
|
||||
pipeline.pause(self._flow_of(run_id) or "")
|
||||
workers = self.controller.workers
|
||||
if workers is not None:
|
||||
# Keyed by run, so a sweep cancelling one config leaves the others
|
||||
# training.
|
||||
workers.cancel_run(run_id)
|
||||
if self.controller.remote is not None:
|
||||
self.controller.remote.cancel_run(run_id)
|
||||
return True
|
||||
|
||||
def _flow_of(self, run_id: str) -> str | None:
|
||||
with Session(db_engine) as session:
|
||||
run = session.get(Run, run_id)
|
||||
return run.flow if run else None
|
||||
|
||||
def _finish_queued(self, run_id: str) -> bool:
|
||||
with Session(db_engine) as session:
|
||||
result = session.exec(
|
||||
update(Run)
|
||||
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
||||
.values(
|
||||
status="cancelled",
|
||||
finished_at=datetime.now(timezone.utc),
|
||||
status_reason="Cancelled before it started",
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
return bool(result.rowcount)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Threads
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _consume(self) -> None:
|
||||
failures = 0
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
# Runs put back to wait for a worker come due here. The claim
|
||||
# below blocks for a second, so this is about once a second.
|
||||
self.queue.move_due(time.time())
|
||||
items = self.queue.claim(CLAIM_COUNT, CLAIM_BLOCK_MS)
|
||||
failures = 0
|
||||
except Exception as exc:
|
||||
failures += 1
|
||||
logger.error("Could not claim runs: %s", exc)
|
||||
self._stop.wait(min(30.0, 2.0**failures))
|
||||
continue
|
||||
for item in items:
|
||||
if not item.run_id:
|
||||
self.queue.ack(item)
|
||||
continue
|
||||
missing = self._missing_labels(item.run_id)
|
||||
if missing:
|
||||
# Left in the queue rather than failed: submitting a run
|
||||
# before turning the GPU box on is a normal way to work, and
|
||||
# the run says what it is waiting for while it waits.
|
||||
self._waiting(item.run_id, missing)
|
||||
self._defer(item)
|
||||
continue
|
||||
# Acknowledged before it runs: from here on the row is the
|
||||
# record, and a lease that stops moving is what says otherwise.
|
||||
self.queue.ack(item)
|
||||
try:
|
||||
self._pool.submit(self._drive, item.run_id)
|
||||
except RuntimeError:
|
||||
logger.warning("Run %s not started: shutting down", item.run_id)
|
||||
|
||||
def _missing_labels(self, run_id: str) -> list[str]:
|
||||
"""Worker labels this run needs that nothing attached carries."""
|
||||
with Session(db_engine) as session:
|
||||
run = session.get(Run, run_id)
|
||||
needed = list(run.labels) if run else []
|
||||
if not needed:
|
||||
return []
|
||||
hub = self.controller.remote
|
||||
available = hub.labels() | {w.name for w in hub.workers()} if hub else set()
|
||||
# A node that only prefers its label runs locally instead, so it is not
|
||||
# a reason to hold the run back; that is decided per node at call time.
|
||||
return sorted(set(needed) - available)
|
||||
|
||||
def _waiting(self, run_id: str, missing: list[str]) -> None:
|
||||
reason = f"Waiting for a worker labelled {', '.join(missing)}"
|
||||
try:
|
||||
with Session(db_engine) as session:
|
||||
session.exec(
|
||||
update(Run)
|
||||
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
||||
.values(status_reason=reason)
|
||||
)
|
||||
session.commit()
|
||||
except Exception:
|
||||
logger.exception("Could not record what run %s is waiting for", run_id)
|
||||
|
||||
def _defer(self, item: WorkItem) -> None:
|
||||
"""Put an item back for later, and let go of this delivery."""
|
||||
try:
|
||||
self.queue.add_delayed(item, time.time() + WAIT_RETRY_S)
|
||||
self.queue.ack(item)
|
||||
except Exception:
|
||||
logger.exception("Could not defer run %s", item.run_id)
|
||||
|
||||
def _keep_leases(self) -> None:
|
||||
"""Say the local runs are alive, and clean up after engines that died."""
|
||||
last_sweep = 0.0
|
||||
while not self._stop.is_set():
|
||||
self._stop.wait(LEASE_INTERVAL_S)
|
||||
if self._stop.is_set():
|
||||
break
|
||||
with self._lock:
|
||||
mine = list(self._active)
|
||||
now = datetime.now(timezone.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(timezone.utc)
|
||||
with Session(db_engine) as session:
|
||||
result = session.exec(
|
||||
update(Run)
|
||||
.where(col(Run.id) == run_id, col(Run.status) == "queued")
|
||||
.values(
|
||||
status="running",
|
||||
started_at=now,
|
||||
lease_at=now,
|
||||
engine=self.engine_name,
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
if not result.rowcount:
|
||||
return None
|
||||
return session.exec(select(Run).where(col(Run.id) == run_id)).first()
|
||||
|
||||
def _drive(self, run_id: str) -> None:
|
||||
run = self._claim(run_id)
|
||||
if run is None:
|
||||
return
|
||||
started = time.perf_counter()
|
||||
status = "ok"
|
||||
reason = ""
|
||||
result: dict[str, Any] = {}
|
||||
state: StateBackend | None = None
|
||||
errors = 0
|
||||
|
||||
def observe(outcome: NodeOutcome) -> None:
|
||||
nonlocal errors
|
||||
if not outcome.ok:
|
||||
errors += 1
|
||||
self._record_node(run_id, outcome)
|
||||
|
||||
sink = MetricSink(run_id)
|
||||
try:
|
||||
flow = self.controller.store.read_flow(run.flow)
|
||||
state = self._state_factory(f"{RUN_NAMESPACE}:{run_id}")
|
||||
pipeline = self.controller.build_run_pipeline(
|
||||
flow,
|
||||
state=state,
|
||||
observer=observe,
|
||||
emission_observer=sink.handle,
|
||||
run=RunContext(run_id=run_id),
|
||||
)
|
||||
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:
|
||||
row = RunNode(
|
||||
run_id=run_id,
|
||||
node=outcome.node[:255],
|
||||
status="ok" if outcome.ok else "error",
|
||||
started_at=datetime.now(timezone.utc),
|
||||
duration_ms=outcome.duration_ms,
|
||||
error=outcome.error[:ERROR_CAP],
|
||||
logs=outcome.logs[:LOG_CAP],
|
||||
)
|
||||
try:
|
||||
with Session(db_engine) as session:
|
||||
session.merge(row)
|
||||
for message, ref in outcome.artifacts.items():
|
||||
session.merge(
|
||||
RunArtifact(
|
||||
run_id=run_id,
|
||||
name=message[:255],
|
||||
node=outcome.node[:255],
|
||||
digest=str(ref.get("digest") or "")[:71],
|
||||
size=int(ref.get("size") or 0),
|
||||
media_type=str(
|
||||
ref.get("media_type") or "application/octet-stream"
|
||||
)[:128],
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Could not record node '%s' of run %s", outcome.node, run_id
|
||||
)
|
||||
|
||||
def _finish(
|
||||
self,
|
||||
run_id: str,
|
||||
status: str,
|
||||
reason: str,
|
||||
result: dict[str, Any],
|
||||
duration_ms: float,
|
||||
) -> None:
|
||||
try:
|
||||
with Session(db_engine) as session:
|
||||
session.exec(
|
||||
update(Run)
|
||||
.where(col(Run.id) == run_id)
|
||||
.values(
|
||||
status=status,
|
||||
status_reason=reason,
|
||||
result=result,
|
||||
duration_ms=duration_ms,
|
||||
finished_at=datetime.now(timezone.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)
|
||||
Reference in New Issue
Block a user