Runs report while they run, and the reader frames lines properly

A training loop has numbers worth keeping thousands of steps before it has a
result. Node code now imports fluksio and calls log_metric/progress, which
sends a line back without ending the call; the engine writes those to
run_metric in batches from the run's own driver rather than folding them off
the event bus, which drops what it cannot keep up with.

Two things fall out. Each event resets the worker deadline, so a node's
timeout measures silence rather than duration — which is what lets a two-hour
training keep a liveness contract instead of racing it. And the worker pool's
_running is now keyed by (run, node), so cancelling one config of a sweep
kills that training and leaves the rest alone.

Fixes a latent framing bug: read_line returned whatever a read had taken,
which was fine while a worker only ever sent one line per request and
unparseable as soon as it sent several. It now keeps the remainder.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AD8SfVhzXBG2nAfFcVh3iD
This commit is contained in:
2026-08-18 17:02:03 +02:00
co-authored by Claude Fable 5
parent 2dd9ce28f6
commit 3f5bb12ed0
5 changed files with 450 additions and 62 deletions
+29 -4
View File
@@ -91,6 +91,20 @@ class LoadedNode:
health_detail: str | None = None
@dataclass
class RunContext:
"""The run a set of nodes is being built for.
A run builds nodes of its own, so which run a node's worker call belongs to
is bound when the proxy is made rather than looked up while it runs. That
is what lets one run's node be cancelled without touching the same node in
another, and what gives a metric reported mid-training somewhere to go.
"""
run_id: str
on_event: Callable[[dict[str, Any]], None] | None = None
@dataclass
class Preview:
"""A draft as it would run, without running it."""
@@ -397,7 +411,9 @@ class FlowController:
# -------------------------------------------------------------------------
def _build_flows(
self, flows: list[tuple[FlowDef, bool]]
self,
flows: list[tuple[FlowDef, bool]],
run: RunContext | None = None,
) -> tuple[list[Node], dict[str, LoadedNode], dict[str, Any], dict[str, bool]]:
"""Instantiate the nodes of several flows, each published or draft."""
nodes: list[Node] = []
@@ -408,7 +424,7 @@ class FlowController:
for flow, draft in flows:
for node_def in flow.nodes:
entry = self._build_node(flow.name, node_def, draft=draft)
entry = self._build_node(flow.name, node_def, draft=draft, run=run)
loaded[entry.id] = entry
if entry.node is not None:
nodes.append(entry.node)
@@ -423,7 +439,11 @@ class FlowController:
return nodes, loaded, initial_values, flow_inputs
def _build_node(
self, flow: str, node_def: NodeDef, draft: bool = False
self,
flow: str,
node_def: NodeDef,
draft: bool = False,
run: RunContext | None = None,
) -> LoadedNode:
node_id = f"{flow}.{node_def.id}"
entry = LoadedNode(id=node_id, flow=flow)
@@ -460,6 +480,8 @@ class FlowController:
code,
node_id=node_id,
timeout=node_def.timeout or settings.FLOW_NODE_TIMEOUT,
run_id=run.run_id if run else "",
on_event=run.on_event if run else None,
)
node = Node(
f=function,
@@ -848,6 +870,7 @@ class FlowController:
state: StateBackend,
draft: bool = False,
observer: Callable[[NodeOutcome], None] | None = None,
run: RunContext | None = None,
) -> Pipeline:
"""Build one flow as a pipeline of its own, for a single run.
@@ -857,7 +880,9 @@ class FlowController:
already holds. The state is the run's, which is what keeps two runs of
one flow from overwriting each other's messages.
"""
nodes, _loaded, initial_values, _inputs = self._build_flows([(flow, draft)])
nodes, _loaded, initial_values, _inputs = self._build_flows(
[(flow, draft)], run=run
)
pipeline = Pipeline(
nodes=nodes,
state=state,
+138 -19
View File
@@ -36,21 +36,20 @@ import uuid
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any
from typing import Any
from sqlalchemy import update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlmodel import Session, col, select
from app.core.db import engine as db_engine
from app.flow.controller import FlowController, RunContext
from app.flow.messages import qualify
from app.flow.pipeline import NodeOutcome, Pipeline
from app.flow.queue import WorkItem, WorkQueue
from app.flow.schemas import FlowDef
from app.flow.state import MemoryState, StateBackend
from app.models import Run, RunNode
if TYPE_CHECKING:
from app.flow.controller import FlowController
from app.models import Run, RunMetric, RunNode
logger = logging.getLogger(__name__)
@@ -70,6 +69,12 @@ 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
#: Progress is for whoever is watching, so it is throttled hard.
PROGRESS_INTERVAL_S = 1.0
#: Where a run's state lives, so it can never collide with the engine's own.
RUN_NAMESPACE = "run"
@@ -157,6 +162,113 @@ def collect_result(flow: FlowDef, state: StateBackend) -> dict[str, Any]:
return result
class MetricSink:
"""Collects what a run's nodes report, and writes it down in batches.
Every call arrives on the thread of the node that made it, so this is
locked. It is written synchronously rather than published: three thousand
steps of a training curve is exactly the traffic the event bus is built to
drop, and a curve with holes in it is not a result.
"""
def __init__(
self,
run_id: str,
publish: Callable[[dict[str, Any]], None] | None = None,
batch: int = METRIC_BATCH,
interval: float = METRIC_FLUSH_S,
) -> None:
self.run_id = run_id
self._publish = publish
self._batch = batch
self._interval = interval
self._rows: dict[tuple[str, int], RunMetric] = {}
self._last_flush = time.monotonic()
self._last_progress = 0.0
self._lock = threading.Lock()
def handle(self, event: dict[str, Any]) -> None:
kind = event.get("event")
if kind == "metric":
self._metric(event)
elif kind == "progress":
self._progress(event)
def _metric(self, event: dict[str, Any]) -> None:
node = str(event.get("call_id") or "").split(":", 1)[-1]
row = RunMetric(
run_id=self.run_id,
name=str(event.get("name") or "")[:128],
step=int(event.get("step", -1)),
node=node[:255],
ts=float(event.get("ts") or time.time()),
value=float(event.get("value") or 0.0),
)
with self._lock:
# Same name and step twice is the later value; the primary key says
# so too, and colliding here is cheaper than colliding in Postgres.
self._rows[(row.name, row.step)] = row
due = (
len(self._rows) >= self._batch
or time.monotonic() - self._last_flush >= self._interval
)
rows = list(self._rows.values()) if due else []
if due:
self._rows.clear()
self._last_flush = time.monotonic()
if rows:
self._write(rows)
def _progress(self, event: dict[str, Any]) -> None:
"""Purely for whoever is watching: throttled, and never written down."""
if self._publish is None:
return
now = time.monotonic()
if now - self._last_progress < PROGRESS_INTERVAL_S:
return
self._last_progress = now
self._publish(
{
"type": "run_progress",
"run": self.run_id,
"node": str(event.get("call_id") or "").split(":", 1)[-1],
"fraction": event.get("fraction"),
"message": event.get("message") or "",
"ts": time.time(),
}
)
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."""
@@ -260,13 +372,7 @@ class RunService:
return run
def cancel(self, run_id: str) -> bool:
"""Stop a run: nothing further is scheduled once its nodes return.
A node already executing is left to finish. Killing one needs the
worker pool to know which run it belongs to, which is what the metric
stream adds next; until then a cancel that arrives mid-node is a stop
rather than an interruption.
"""
"""Stop a run: kill what it is executing, schedule nothing further."""
with self._lock:
pipeline = self._active.get(run_id)
if pipeline is None:
@@ -277,9 +383,15 @@ class RunService:
self._cancelled.add(run_id)
return cancelled
self._cancelled.add(run_id)
# Reusing the pause gate: a gated node is never submitted, so the graph
# drains instead of going further.
# 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)
return True
def _flow_of(self, run_id: str) -> str | None:
@@ -415,11 +527,15 @@ class RunService:
errors += 1
self._record_node(run_id, outcome)
sink = MetricSink(run_id, publish=self._publish_event)
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
flow,
state=state,
observer=observe,
run=RunContext(run_id=run_id, on_event=sink.handle),
)
with self._lock:
self._active[run_id] = pipeline
@@ -444,6 +560,8 @@ class RunService:
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
@@ -501,10 +619,7 @@ class RunService:
logger.exception("Could not close run %s", run_id)
def _publish(self, run: Run, kind: str) -> None:
events = self.controller.events
if events is None:
return
events.publish(
self._publish_event(
{
"type": kind,
"flow": run.flow,
@@ -514,3 +629,7 @@ class RunService:
"ts": time.time(),
}
)
def _publish_event(self, event: dict[str, Any]) -> None:
if self.controller.events is not None:
self.controller.events.publish(event)
+73 -4
View File
@@ -10,6 +10,13 @@ goes out on a private duplicate of fd 1 taken before anything else can write to
it. fd 1 itself is pointed at stderr, so a stray ``write(1, ...)`` — from a
native library, or a node printing during an import — lands in the server log
instead of corrupting the reply stream.
A node may also send lines back *while* it is still running: anything carrying
an ``event`` key is a report rather than the answer, and the engine keeps
reading. That is what makes a training curve visible during the hours it takes
to draw, and what tells the engine a long node is alive rather than hung —
each event resets its deadline, so the timeout measures silence rather than
duration. Node code reaches it by importing ``fluksio``.
"""
from __future__ import annotations
@@ -29,6 +36,7 @@ import contextlib
import hashlib
import io
import json
import time
import traceback
from collections.abc import Callable
from types import ModuleType
@@ -38,6 +46,63 @@ from typing import Any, cast
#: pipe or the reply.
MAX_LOG = 16 * 1024
#: The reply channel, opened by ``main``. Also what an event line goes down.
_RPC: Any = None
#: The call being served, so an event can say which one it belongs to.
_CALL_ID = ""
def _emit(event: dict[str, Any]) -> None:
"""Send one line back without ending the call."""
if _RPC is None:
return
event["call_id"] = _CALL_ID
event["ts"] = time.time()
_RPC.write(json.dumps(event) + "\n")
_RPC.flush()
class _Reporter(ModuleType):
"""``import fluksio`` — what node code says while it is still running.
Deliberately tiny and deliberately not a return value: a training loop has
numbers worth keeping thousands of steps before it has a result, and
holding them until it returns is how they get lost when it does not.
"""
def log_metric(self, name: str, value: float, step: int = -1) -> None:
"""Record one number, optionally at a step. Steps make a curve."""
_emit(
{
"event": "metric",
"name": str(name)[:128],
"value": float(value),
"step": int(step),
}
)
def log_metrics(self, values: dict[str, float], step: int = -1) -> None:
"""Several at once, which is what a training loop usually has."""
for name, value in values.items():
self.log_metric(name, value, step)
def progress(self, fraction: float | None = None, message: str = "") -> None:
"""How far along this node is, for whoever is watching it."""
_emit(
{
"event": "progress",
"fraction": None if fraction is None else float(fraction),
"message": str(message)[:200],
}
)
def _install_reporter() -> None:
"""Put ``fluksio`` on the import path of every node this worker runs."""
module = _Reporter("fluksio")
module.__doc__ = "Report metrics and progress from inside a node."
sys.modules["fluksio"] = module
def load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
"""Compile a node's source and return the function to run.
@@ -140,17 +205,21 @@ def _handle(request: dict[str, Any], cache: dict[tuple[str, str], Any]) -> Any:
def main() -> None:
rpc = os.fdopen(os.dup(1), "w")
global _RPC, _CALL_ID
_RPC = os.fdopen(os.dup(1), "w")
# Everything the node writes to the real stdout now goes to the server log.
os.dup2(2, 1)
_install_reporter()
cache: dict[tuple[str, str], Any] = {}
for line in sys.stdin:
if not line.strip():
continue
request = json.loads(line)
_CALL_ID = str(request.get("call_id") or "")
captured = _Capped()
response: dict[str, Any] = {"id": request.get("id")}
response: dict[str, Any] = {"call_id": _CALL_ID}
try:
with (
contextlib.redirect_stdout(captured),
@@ -167,8 +236,8 @@ def main() -> None:
"traceback": _node_traceback(exc),
}
response["logs"] = captured.getvalue()
rpc.write(json.dumps(response) + "\n")
rpc.flush()
_RPC.write(json.dumps(response) + "\n")
_RPC.flush()
if __name__ == "__main__":
+105 -35
View File
@@ -102,6 +102,11 @@ class _Worker:
def __init__(self, python: str, generation: int) -> None:
self.generation = generation
self.cancelled = False
# What a read took past the end of a line. A node reporting quickly
# puts several lines in one chunk, and the next one is the caller's to
# read — dropping it would lose a metric, returning it with the first
# would be unparseable.
self._buffer = bytearray()
self.proc = subprocess.Popen(
[python, str(WORKER_MAIN)],
stdin=subprocess.PIPE,
@@ -119,11 +124,15 @@ class _Worker:
self.proc.stdin.flush()
def read_line(self, deadline: float) -> str | None:
"""One reply line; ``None`` past the deadline, ``""`` if the pipe closed."""
"""One line; ``None`` past the deadline, ``""`` if the pipe closed."""
assert self.proc.stdout is not None
fd = self.proc.stdout.fileno()
buffer = bytearray()
while True:
end = self._buffer.find(b"\n")
if end >= 0:
line = bytes(self._buffer[: end + 1])
del self._buffer[: end + 1]
return line.decode(errors="replace")
remaining = deadline - time.monotonic()
if remaining <= 0:
return None
@@ -133,11 +142,10 @@ class _Worker:
chunk = os.read(fd, 65536)
if not chunk:
# Partial output before the pipe closed is a half-written
# reply, which is no more use than none at all.
# line, which is no more use than none at all.
self._buffer.clear()
return ""
buffer += chunk
if buffer.endswith(b"\n"):
return buffer.decode(errors="replace")
self._buffer += chunk
def kill(self) -> None:
"""SIGKILL: user code has no cleanup we can trust to run.
@@ -174,9 +182,10 @@ class PythonWorkerPool:
self.size = size
self.events = events
self._idle: queue.Queue[_Worker | None] = queue.Queue()
# ponytail: _running is last-wins; two concurrent runs of one node mean
# cancel kills the newest. Key by run id if that ever matters.
self._running: dict[str, _Worker] = {}
# Keyed by (run, node): a sweep has the same node executing in several
# runs at once, and cancelling one of them must not kill the others.
# The live pipeline's own executions carry an empty run.
self._running: dict[tuple[str, str], _Worker] = {}
self._generation = 0
self._stopped = False
self._lock = threading.Lock()
@@ -269,17 +278,24 @@ class PythonWorkerPool:
# -------------------------------------------------------------------------
def _request(
self, payload: dict[str, Any], timeout: float, node_id: str = ""
self,
payload: dict[str, Any],
timeout: float,
node_id: str = "",
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
) -> dict[str, Any]:
worker = self._acquire()
key = (run_id, node_id)
if node_id:
self._running[node_id] = worker
self._running[key] = worker
self._publish(
{
"type": "node_started",
# A flow name cannot contain a dot, so this is exact.
"flow": node_id.split(".", 1)[0],
"node": node_id,
"run": run_id,
"ts": time.time(),
}
)
@@ -289,30 +305,43 @@ class PythonWorkerPool:
except OSError as exc:
raise RemoteError(f"worker died: {exc}") from exc
line = worker.read_line(time.monotonic() + timeout)
if line:
try:
return dict(json.loads(line))
except (TypeError, ValueError) as exc:
# A reply we cannot read leaves this worker out of step:
# whatever is still in its pipe would be taken by the next
# request on this slot as its own answer. Retire it.
worker.cancelled = True
raise RemoteError(
f"unreadable reply from the worker: {exc}"
) from exc
if worker.cancelled:
raise NodeCancelled("cancelled while it was running")
if line is None:
worker.kill()
raise NodeTimeout(f"exceeded {timeout}s and was killed")
raise RemoteError("worker died")
while True:
# Reset per line, so the timeout measures how long the node has
# been silent rather than how long it has been working. A node
# that reports nothing is still held to it, which is what keeps
# the deadline meaningful for the ones that never report.
line = worker.read_line(time.monotonic() + timeout)
if line:
try:
message = dict(json.loads(line))
except (TypeError, ValueError) as exc:
# A reply we cannot read leaves this worker out of step:
# whatever is still in its pipe would be taken by the
# next request on this slot as its own answer. Retire it.
worker.cancelled = True
raise RemoteError(
f"unreadable reply from the worker: {exc}"
) from exc
if message.get("event"):
if on_event is not None:
try:
on_event(message)
except Exception:
logger.exception("Could not record a worker event")
continue
return message
if worker.cancelled:
raise NodeCancelled("cancelled while it was running")
if line is None:
worker.kill()
raise NodeTimeout(f"was silent for {timeout}s and was killed")
raise RemoteError("worker died")
finally:
# Under the lock, so a cancel that has already read this worker out
# of _running cannot kill it after another node has taken the slot.
with self._lock:
if node_id:
self._running.pop(node_id, None)
self._running.pop(key, None)
self._release(worker)
def compile(self, flow: str, node: str, source: str) -> str | None:
@@ -338,18 +367,24 @@ class PythonWorkerPool:
params: dict[str, Any] | None,
node_id: str,
timeout: float,
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
) -> Any:
response = self._request(
{
"op": "run",
"call_id": f"{run_id}:{node_id}" if run_id else node_id,
"flow": flow,
"node": node,
"source": source,
"kwargs": kwargs,
"params": params or {},
"run": {"id": run_id} if run_id else None,
},
timeout=timeout,
node_id=node_id,
run_id=run_id,
on_event=on_event,
)
# Into the tee, from the thread the engine is capturing on: this is
# what puts a node's prints in the log panel, so it has to happen
@@ -366,25 +401,60 @@ class PythonWorkerPool:
)
def proxy(
self, flow: str, node: str, source: str, node_id: str, timeout: float
self,
flow: str,
node: str,
source: str,
node_id: str,
timeout: float,
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
) -> Callable[..., Any]:
"""The callable a python node runs instead of its own compiled function."""
"""The callable a python node runs instead of its own compiled function.
A run builds its own nodes, so the run this proxy belongs to is bound
here rather than looked up — which is also what lets two runs of one
node be told apart when one of them is cancelled.
"""
def call(params: dict[str, Any] | None = None, **kwargs: Any) -> Any:
return self.run(flow, node, source, kwargs, params, node_id, timeout)
return self.run(
flow,
node,
source,
kwargs,
params,
node_id,
timeout,
run_id=run_id,
on_event=on_event,
)
return call
def cancel(self, node_id: str) -> bool:
def cancel(self, node_id: str, run_id: str = "") -> bool:
"""Stop a node that is running now. False when there was nothing to stop."""
with self._lock:
worker = self._running.get(node_id)
worker = self._running.get((run_id, node_id))
if worker is None:
return False
worker.cancelled = True
worker.kill()
return True
def cancel_run(self, run_id: str) -> int:
"""Stop every node this run has in a worker right now."""
with self._lock:
workers = [
worker
for (owner, _node), worker in self._running.items()
if owner == run_id
]
for worker in workers:
worker.cancelled = True
worker.kill()
return len(workers)
def _publish(self, event: dict[str, Any]) -> None:
if self.events is not None:
self.events.publish(event)
+105
View File
@@ -175,3 +175,108 @@ def test_a_pool_can_stop_while_a_node_is_running(pool):
with pytest.raises(Exception, match="shutting down"):
run(pool, "def process(params):\n return {'out': 1}\n")
# -----------------------------------------------------------------------------
# Reporting from inside a node that has not returned yet
# -----------------------------------------------------------------------------
def test_a_node_reports_metrics_while_it_is_still_running(pool):
seen = []
result = pool.run(
"demo",
"train",
"import fluksio\n"
"def process(params):\n"
" for step in range(3):\n"
" fluksio.log_metric('loss', 1.0 / (step + 1), step)\n"
" fluksio.progress(0.5, 'halfway')\n"
" return {'out': 1}\n",
{},
{},
"demo.train",
timeout=5,
run_id="r1",
on_event=seen.append,
)
assert result == {"out": 1}
metrics = [event for event in seen if event["event"] == "metric"]
assert [(m["name"], m["step"]) for m in metrics] == [
("loss", 0),
("loss", 1),
("loss", 2),
]
assert metrics[0]["value"] == 1.0
# Every event says which call it belongs to, so a sweep can tell them apart.
assert {m["call_id"] for m in metrics} == {"r1:demo.train"}
assert [event["event"] for event in seen if event["event"] == "progress"] == [
"progress"
]
def test_events_hold_off_the_timeout_but_silence_does_not(pool):
# The deadline measures silence: a node reporting every 0.05s stays alive
# well past a 0.3s timeout, which is what a two-hour training needs.
result = pool.run(
"demo",
"slow",
"import time, fluksio\n"
"def process(params):\n"
" for step in range(12):\n"
" time.sleep(0.05)\n"
" fluksio.log_metric('beat', step, step)\n"
" return {'done': True}\n",
{},
{},
"demo.slow",
timeout=0.3,
run_id="r2",
on_event=lambda _event: None,
)
assert result == {"done": True}
with pytest.raises(NodeTimeout):
pool.run(
"demo",
"quiet",
"import time\ndef process(params):\n time.sleep(2)\n return {}\n",
{},
{},
"demo.quiet",
timeout=0.3,
)
def test_cancelling_one_run_leaves_the_same_node_in_another_alone(pool):
# Keyed by (run, node): cancelling a config of a sweep must not kill the
# rest of it. With one slot the second run is not executing, so the check
# is that the pool refuses to find it rather than killing the wrong worker.
started = threading.Event()
def hold():
try:
pool.run(
"demo",
"hold",
"import time\ndef process(params):\n time.sleep(5)\n return {}\n",
{},
{},
"demo.hold",
timeout=10,
run_id="run-a",
)
except Exception:
pass
finally:
started.set()
thread = threading.Thread(target=hold, daemon=True)
thread.start()
time.sleep(0.5)
assert pool.cancel("demo.hold", run_id="run-b") is False
assert pool.cancel("demo.hold", run_id="run-a") is True
started.wait(timeout=5)
thread.join(timeout=5)