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 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 @dataclass
class Preview: class Preview:
"""A draft as it would run, without running it.""" """A draft as it would run, without running it."""
@@ -397,7 +411,9 @@ class FlowController:
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
def _build_flows( 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]]: ) -> tuple[list[Node], dict[str, LoadedNode], dict[str, Any], dict[str, bool]]:
"""Instantiate the nodes of several flows, each published or draft.""" """Instantiate the nodes of several flows, each published or draft."""
nodes: list[Node] = [] nodes: list[Node] = []
@@ -408,7 +424,7 @@ class FlowController:
for flow, draft in flows: for flow, draft in flows:
for node_def in flow.nodes: 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 loaded[entry.id] = entry
if entry.node is not None: if entry.node is not None:
nodes.append(entry.node) nodes.append(entry.node)
@@ -423,7 +439,11 @@ class FlowController:
return nodes, loaded, initial_values, flow_inputs return nodes, loaded, initial_values, flow_inputs
def _build_node( 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: ) -> LoadedNode:
node_id = f"{flow}.{node_def.id}" node_id = f"{flow}.{node_def.id}"
entry = LoadedNode(id=node_id, flow=flow) entry = LoadedNode(id=node_id, flow=flow)
@@ -460,6 +480,8 @@ class FlowController:
code, code,
node_id=node_id, node_id=node_id,
timeout=node_def.timeout or settings.FLOW_NODE_TIMEOUT, 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( node = Node(
f=function, f=function,
@@ -848,6 +870,7 @@ class FlowController:
state: StateBackend, state: StateBackend,
draft: bool = False, draft: bool = False,
observer: Callable[[NodeOutcome], None] | None = None, observer: Callable[[NodeOutcome], None] | None = None,
run: RunContext | None = None,
) -> Pipeline: ) -> Pipeline:
"""Build one flow as a pipeline of its own, for a single run. """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 already holds. The state is the run's, which is what keeps two runs of
one flow from overwriting each other's messages. 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( pipeline = Pipeline(
nodes=nodes, nodes=nodes,
state=state, state=state,
+138 -19
View File
@@ -36,21 +36,20 @@ import uuid
from collections.abc import Callable from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any from typing import Any
from sqlalchemy import update from sqlalchemy import update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlmodel import Session, col, select from sqlmodel import Session, col, select
from app.core.db import engine as db_engine 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.messages import qualify
from app.flow.pipeline import NodeOutcome, Pipeline from app.flow.pipeline import NodeOutcome, Pipeline
from app.flow.queue import WorkItem, WorkQueue from app.flow.queue import WorkItem, WorkQueue
from app.flow.schemas import FlowDef from app.flow.schemas import FlowDef
from app.flow.state import MemoryState, StateBackend from app.flow.state import MemoryState, StateBackend
from app.models import Run, RunNode from app.models import Run, RunMetric, RunNode
if TYPE_CHECKING:
from app.flow.controller import FlowController
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -70,6 +69,12 @@ CLAIM_COUNT = 4
CLAIM_BLOCK_MS = 1000 CLAIM_BLOCK_MS = 1000
ERROR_CAP = 2000 ERROR_CAP = 2000
LOG_CAP = 8000 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. #: Where a run's state lives, so it can never collide with the engine's own.
RUN_NAMESPACE = "run" RUN_NAMESPACE = "run"
@@ -157,6 +162,113 @@ def collect_result(flow: FlowDef, state: StateBackend) -> dict[str, Any]:
return result 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: class RunService:
"""Accepts runs, drives them, and writes down what they did.""" """Accepts runs, drives them, and writes down what they did."""
@@ -260,13 +372,7 @@ class RunService:
return run return run
def cancel(self, run_id: str) -> bool: def cancel(self, run_id: str) -> bool:
"""Stop a run: nothing further is scheduled once its nodes return. """Stop a run: kill what it is executing, schedule nothing further."""
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.
"""
with self._lock: with self._lock:
pipeline = self._active.get(run_id) pipeline = self._active.get(run_id)
if pipeline is None: if pipeline is None:
@@ -277,9 +383,15 @@ class RunService:
self._cancelled.add(run_id) self._cancelled.add(run_id)
return cancelled return cancelled
self._cancelled.add(run_id) self._cancelled.add(run_id)
# Reusing the pause gate: a gated node is never submitted, so the graph # The gate first, so nothing new is submitted while the running nodes
# drains instead of going further. # 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 "") 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 return True
def _flow_of(self, run_id: str) -> str | None: def _flow_of(self, run_id: str) -> str | None:
@@ -415,11 +527,15 @@ class RunService:
errors += 1 errors += 1
self._record_node(run_id, outcome) self._record_node(run_id, outcome)
sink = MetricSink(run_id, publish=self._publish_event)
try: try:
flow = self.controller.store.read_flow(run.flow) flow = self.controller.store.read_flow(run.flow)
state = self._state_factory(f"{RUN_NAMESPACE}:{run_id}") state = self._state_factory(f"{RUN_NAMESPACE}:{run_id}")
pipeline = self.controller.build_run_pipeline( 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: with self._lock:
self._active[run_id] = pipeline self._active[run_id] = pipeline
@@ -444,6 +560,8 @@ class RunService:
with self._lock: with self._lock:
self._active.pop(run_id, None) self._active.pop(run_id, None)
self._cancelled.discard(run_id) 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) duration = round((time.perf_counter() - started) * 1000, 2)
self._finish(run_id, status, reason, result, duration) self._finish(run_id, status, reason, result, duration)
run.status = status run.status = status
@@ -501,10 +619,7 @@ class RunService:
logger.exception("Could not close run %s", run_id) logger.exception("Could not close run %s", run_id)
def _publish(self, run: Run, kind: str) -> None: def _publish(self, run: Run, kind: str) -> None:
events = self.controller.events self._publish_event(
if events is None:
return
events.publish(
{ {
"type": kind, "type": kind,
"flow": run.flow, "flow": run.flow,
@@ -514,3 +629,7 @@ class RunService:
"ts": time.time(), "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 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 native library, or a node printing during an import — lands in the server log
instead of corrupting the reply stream. 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 from __future__ import annotations
@@ -29,6 +36,7 @@ import contextlib
import hashlib import hashlib
import io import io
import json import json
import time
import traceback import traceback
from collections.abc import Callable from collections.abc import Callable
from types import ModuleType from types import ModuleType
@@ -38,6 +46,63 @@ from typing import Any, cast
#: pipe or the reply. #: pipe or the reply.
MAX_LOG = 16 * 1024 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]: def load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
"""Compile a node's source and return the function to run. """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: 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. # Everything the node writes to the real stdout now goes to the server log.
os.dup2(2, 1) os.dup2(2, 1)
_install_reporter()
cache: dict[tuple[str, str], Any] = {} cache: dict[tuple[str, str], Any] = {}
for line in sys.stdin: for line in sys.stdin:
if not line.strip(): if not line.strip():
continue continue
request = json.loads(line) request = json.loads(line)
_CALL_ID = str(request.get("call_id") or "")
captured = _Capped() captured = _Capped()
response: dict[str, Any] = {"id": request.get("id")} response: dict[str, Any] = {"call_id": _CALL_ID}
try: try:
with ( with (
contextlib.redirect_stdout(captured), contextlib.redirect_stdout(captured),
@@ -167,8 +236,8 @@ def main() -> None:
"traceback": _node_traceback(exc), "traceback": _node_traceback(exc),
} }
response["logs"] = captured.getvalue() response["logs"] = captured.getvalue()
rpc.write(json.dumps(response) + "\n") _RPC.write(json.dumps(response) + "\n")
rpc.flush() _RPC.flush()
if __name__ == "__main__": if __name__ == "__main__":
+105 -35
View File
@@ -102,6 +102,11 @@ class _Worker:
def __init__(self, python: str, generation: int) -> None: def __init__(self, python: str, generation: int) -> None:
self.generation = generation self.generation = generation
self.cancelled = False 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( self.proc = subprocess.Popen(
[python, str(WORKER_MAIN)], [python, str(WORKER_MAIN)],
stdin=subprocess.PIPE, stdin=subprocess.PIPE,
@@ -119,11 +124,15 @@ class _Worker:
self.proc.stdin.flush() self.proc.stdin.flush()
def read_line(self, deadline: float) -> str | None: 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 assert self.proc.stdout is not None
fd = self.proc.stdout.fileno() fd = self.proc.stdout.fileno()
buffer = bytearray()
while True: 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() remaining = deadline - time.monotonic()
if remaining <= 0: if remaining <= 0:
return None return None
@@ -133,11 +142,10 @@ class _Worker:
chunk = os.read(fd, 65536) chunk = os.read(fd, 65536)
if not chunk: if not chunk:
# Partial output before the pipe closed is a half-written # 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 "" return ""
buffer += chunk self._buffer += chunk
if buffer.endswith(b"\n"):
return buffer.decode(errors="replace")
def kill(self) -> None: def kill(self) -> None:
"""SIGKILL: user code has no cleanup we can trust to run. """SIGKILL: user code has no cleanup we can trust to run.
@@ -174,9 +182,10 @@ class PythonWorkerPool:
self.size = size self.size = size
self.events = events self.events = events
self._idle: queue.Queue[_Worker | None] = queue.Queue() self._idle: queue.Queue[_Worker | None] = queue.Queue()
# ponytail: _running is last-wins; two concurrent runs of one node mean # Keyed by (run, node): a sweep has the same node executing in several
# cancel kills the newest. Key by run id if that ever matters. # runs at once, and cancelling one of them must not kill the others.
self._running: dict[str, _Worker] = {} # The live pipeline's own executions carry an empty run.
self._running: dict[tuple[str, str], _Worker] = {}
self._generation = 0 self._generation = 0
self._stopped = False self._stopped = False
self._lock = threading.Lock() self._lock = threading.Lock()
@@ -269,17 +278,24 @@ class PythonWorkerPool:
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
def _request( 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]: ) -> dict[str, Any]:
worker = self._acquire() worker = self._acquire()
key = (run_id, node_id)
if node_id: if node_id:
self._running[node_id] = worker self._running[key] = worker
self._publish( self._publish(
{ {
"type": "node_started", "type": "node_started",
# A flow name cannot contain a dot, so this is exact. # A flow name cannot contain a dot, so this is exact.
"flow": node_id.split(".", 1)[0], "flow": node_id.split(".", 1)[0],
"node": node_id, "node": node_id,
"run": run_id,
"ts": time.time(), "ts": time.time(),
} }
) )
@@ -289,30 +305,43 @@ class PythonWorkerPool:
except OSError as exc: except OSError as exc:
raise RemoteError(f"worker died: {exc}") from exc raise RemoteError(f"worker died: {exc}") from exc
line = worker.read_line(time.monotonic() + timeout) while True:
if line: # Reset per line, so the timeout measures how long the node has
try: # been silent rather than how long it has been working. A node
return dict(json.loads(line)) # that reports nothing is still held to it, which is what keeps
except (TypeError, ValueError) as exc: # the deadline meaningful for the ones that never report.
# A reply we cannot read leaves this worker out of step: line = worker.read_line(time.monotonic() + timeout)
# whatever is still in its pipe would be taken by the next if line:
# request on this slot as its own answer. Retire it. try:
worker.cancelled = True message = dict(json.loads(line))
raise RemoteError( except (TypeError, ValueError) as exc:
f"unreadable reply from the worker: {exc}" # A reply we cannot read leaves this worker out of step:
) from exc # whatever is still in its pipe would be taken by the
if worker.cancelled: # next request on this slot as its own answer. Retire it.
raise NodeCancelled("cancelled while it was running") worker.cancelled = True
if line is None: raise RemoteError(
worker.kill() f"unreadable reply from the worker: {exc}"
raise NodeTimeout(f"exceeded {timeout}s and was killed") ) from exc
raise RemoteError("worker died") 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: finally:
# Under the lock, so a cancel that has already read this worker out # 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. # of _running cannot kill it after another node has taken the slot.
with self._lock: with self._lock:
if node_id: if node_id:
self._running.pop(node_id, None) self._running.pop(key, None)
self._release(worker) self._release(worker)
def compile(self, flow: str, node: str, source: str) -> str | None: def compile(self, flow: str, node: str, source: str) -> str | None:
@@ -338,18 +367,24 @@ class PythonWorkerPool:
params: dict[str, Any] | None, params: dict[str, Any] | None,
node_id: str, node_id: str,
timeout: float, timeout: float,
run_id: str = "",
on_event: Callable[[dict[str, Any]], None] | None = None,
) -> Any: ) -> Any:
response = self._request( response = self._request(
{ {
"op": "run", "op": "run",
"call_id": f"{run_id}:{node_id}" if run_id else node_id,
"flow": flow, "flow": flow,
"node": node, "node": node,
"source": source, "source": source,
"kwargs": kwargs, "kwargs": kwargs,
"params": params or {}, "params": params or {},
"run": {"id": run_id} if run_id else None,
}, },
timeout=timeout, timeout=timeout,
node_id=node_id, 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 # 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 # what puts a node's prints in the log panel, so it has to happen
@@ -366,25 +401,60 @@ class PythonWorkerPool:
) )
def proxy( 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]: ) -> 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: 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 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.""" """Stop a node that is running now. False when there was nothing to stop."""
with self._lock: with self._lock:
worker = self._running.get(node_id) worker = self._running.get((run_id, node_id))
if worker is None: if worker is None:
return False return False
worker.cancelled = True worker.cancelled = True
worker.kill() worker.kill()
return True 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: def _publish(self, event: dict[str, Any]) -> None:
if self.events is not None: if self.events is not None:
self.events.publish(event) 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"): with pytest.raises(Exception, match="shutting down"):
run(pool, "def process(params):\n return {'out': 1}\n") 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)