diff --git a/backend/app/flow/controller.py b/backend/app/flow/controller.py index f0aa47b..6157db4 100644 --- a/backend/app/flow/controller.py +++ b/backend/app/flow/controller.py @@ -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, diff --git a/backend/app/flow/runs.py b/backend/app/flow/runs.py index d0ef1c7..e8b1b58 100644 --- a/backend/app/flow/runs.py +++ b/backend/app/flow/runs.py @@ -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) diff --git a/backend/app/flow/worker_main.py b/backend/app/flow/worker_main.py index 2dce883..ac47db7 100644 --- a/backend/app/flow/worker_main.py +++ b/backend/app/flow/worker_main.py @@ -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__": diff --git a/backend/app/flow/workers.py b/backend/app/flow/workers.py index 47e11f7..2d1887c 100644 --- a/backend/app/flow/workers.py +++ b/backend/app/flow/workers.py @@ -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) diff --git a/backend/tests/flow/test_workers.py b/backend/tests/flow/test_workers.py index f82c066..8c30ef1 100644 --- a/backend/tests/flow/test_workers.py +++ b/backend/tests/flow/test_workers.py @@ -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)