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:
@@ -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
@@ -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)
|
||||
|
||||
@@ -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__":
|
||||
|
||||
+91
-21
@@ -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
|
||||
|
||||
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:
|
||||
return dict(json.loads(line))
|
||||
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.
|
||||
# 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"exceeded {timeout}s and was killed")
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user