diff --git a/backend/fluksio/flow/controller.py b/backend/fluksio/flow/controller.py index d4daaf0..39975ac 100644 --- a/backend/fluksio/flow/controller.py +++ b/backend/fluksio/flow/controller.py @@ -18,6 +18,7 @@ import asyncio import hashlib import json import logging +import threading import time import traceback from collections.abc import Callable @@ -169,6 +170,10 @@ class Wiring: requires: list[str] +#: Nothing held, as distinct from a node that yielded ``None``. +_NOTHING = object() + + class EmitSink: """Turns a worker's mid-call frames back into the node's own outputs. @@ -177,17 +182,78 @@ class EmitSink: soon as there is one. What arrives is a dict keyed by output port, which is the same thing a return value is — so it goes through the node, gets checked against the ports it declared, and is published from there. + + A generator's yields arrive here one at a time, and the last of them is the + node's *result* rather than an emission when the generator returns nothing + of its own. Which one is last is only known once the call ends, so a yield + is checked against the ports on arrival and published one behind; + :meth:`wrap` is what decides afterwards which the held one was. """ - __slots__ = ("node",) + __slots__ = ("node", "_held") def __init__(self) -> None: self.node: Node | None = None + # One thread serves one call, and the same node may be running on + # another thread for another run. + self._held = threading.local() + + def _take(self) -> Any: + """What is held for this thread, and nothing held afterwards.""" + held = getattr(self._held, "value", _NOTHING) + self._held.value = _NOTHING + return held def handle(self, event: dict[str, Any]) -> None: - if self.node is None or event.get("event") != "emit": + if self.node is None: return - self.node.emit(event.get("outputs") or {}) + kind = event.get("event") + if kind == "emit": + # ``fluksio.emit()``, which is a publication and nothing else — and + # every yield of a worker old enough to have drained its own. + self.node.emit(event.get("outputs") or {}) + return + if kind != "yield": + return + outputs = event.get("outputs") + held = self._take() + if held is not _NOTHING: + # A yield with one after it was an emission, whatever the node + # goes on to return. + self.node.emit(held) + # Checked here rather than where it is published: an undeclared port + # fails the call at the yield that produced it, not at the next one. + self.node._to_messages(outputs) + self._held.value = outputs + + def wrap(self, call: Callable[..., Any]) -> Callable[..., Any]: + """A worker-backed function, with its generator's last yield accounted for. + + The worker sends every yield as it happens and returns whatever its + generator returned, which is ``None`` for a node that only yields. So + the last yield is still here when the call ends, and this is where it + becomes either the node's result or one more emission. + """ + + def run(*args: Any, **kwargs: Any) -> Any: + # A call that died mid-generator cannot leave its last yield on + # this thread for the next one to adopt. + self._take() + try: + result = call(*args, **kwargs) + except BaseException: + self._take() + raise + held = self._take() + if held is _NOTHING: + return result + if result is None: + return held + if self.node is not None: + self.node.emit(held) + return result + + return run @dataclass @@ -979,6 +1045,9 @@ class FlowController: function if node_def.device_policy == "prefer" else None ), ) + # Outermost, so it sees the result whichever of the three + # above answered the call. + function = emissions.wrap(function) node = Node( f=with_settings(function, params), requires=_bound(node_def.requires), diff --git a/backend/fluksio/flow/nodes/base.py b/backend/fluksio/flow/nodes/base.py index 7ecab41..e959070 100644 --- a/backend/fluksio/flow/nodes/base.py +++ b/backend/fluksio/flow/nodes/base.py @@ -333,12 +333,19 @@ class Node: return self._to_messages(result) def _drain(self, generator: Iterator[Any]) -> Any: - """Publish each yield as it happens; the end of it is the result.""" + """Publish each yield as it happens; the end of it is the result. + + The last yield is the result when the generator returns nothing of its + own, and which one is last is only known once it ends — so a value is + published one behind but *checked* the moment it arrives, which is what + fails a mistyped port at the yield that produced it. + """ pending: Any = None have_pending = False try: while True: value = next(generator) + self._to_messages(value) if have_pending: self.emit(pending) pending, have_pending = value, True diff --git a/backend/tests/flow/test_runs.py b/backend/tests/flow/test_runs.py index 973c196..b53bd36 100644 --- a/backend/tests/flow/test_runs.py +++ b/backend/tests/flow/test_runs.py @@ -262,6 +262,25 @@ def test_an_emission_that_nothing_declares_names_what_was_emitted(): assert "study.undeclared" not in state +def test_a_mistyped_port_fails_at_the_yield_that_produced_it(): + """Not at the one after it: a training loop's second pass can be an hour.""" + passes = [] + + def stray(params): + passes.append(1) + yield {"undeclared": 1.0} + passes.append(2) + yield {"loss": 0.5} + + node = make_node("train", "study", stray, provides=[spec("loss", stream=True)]) + seen: list[NodeOutcome] = [] + Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run() + + assert not seen[0].ok + assert "undeclared" in seen[0].error + assert passes == [1] + + def test_emissions_reach_the_run_as_a_series_with_a_step_each(): sink = MetricSink("run-1", batch=1) sink.handle("study.train", {"study.loss": 1.0, "study.tag": "ignored"}) diff --git a/backend/tests/flow/test_workers.py b/backend/tests/flow/test_workers.py index a838d2c..043528e 100644 --- a/backend/tests/flow/test_workers.py +++ b/backend/tests/flow/test_workers.py @@ -9,7 +9,10 @@ import pytest from fluksio_worker.worker_main import ARTIFACT_DIR_ENV from fluksio.flow.artifacts import ArtifactStore +from fluksio.flow.controller import EmitSink from fluksio.flow.messages import DType, MessageSpec +from fluksio.flow.nodes import Node +from fluksio.flow.nodes.base import NodeOutputError from fluksio.flow.workers import NodeCancelled, NodeTimeout, PythonWorkerPool @@ -22,6 +25,24 @@ def pool() -> Iterator[PythonWorkerPool]: worker_pool.stop() +def _sink_node(node_id: str, port: str, dtype: DType) -> tuple[Node, list[dict]]: + """A node declaring one streaming port, and the list its emissions land in.""" + node = Node( + f=lambda: None, + requires=[], + provides=[MessageSpec(name=port, dtype=dtype, stream=True)], + name=node_id, + ) + node.assign_flow("demo", node_id) + published: list[dict] = [] + node._pipeline = type( + "Recorder", + (), + {"publish_emission": staticmethod(lambda _node, outputs: published.append(outputs))}, + )() + return node, published + + def run(pool: PythonWorkerPool, code: str, node: str = "demo", **kwargs): return pool.run("demo", node, code, kwargs, f"demo.{node}", timeout=5) @@ -224,22 +245,30 @@ def test_a_generator_node_publishes_each_yield_and_returns_the_end(pool): def test_without_a_return_the_last_yield_is_the_result(pool): - seen = [] - result = pool.run( - "demo", - "count", - "def process():\n" - " yield {'out': 1}\n" - " yield {'out': 2}\n" - " yield {'out': 3}\n", - {}, - "demo.count", - timeout=5, - on_event=seen.append, - ) + """The worker sends every yield; the engine is what holds the last one back.""" + node, published = _sink_node("count", "out", DType.INT) + sink = EmitSink() + sink.node = node + result = sink.wrap( + lambda: pool.run( + "demo", + "count", + "def process():\n" + " yield {'out': 1}\n" + " yield {'out': 2}\n" + " yield {'out': 3}\n", + {}, + "demo.count", + timeout=5, + on_event=sink.handle, + ) + )() + + # The last yield is the node's output rather than an emission, so it is not + # counted twice. assert result == {"out": 3} - assert [event["outputs"] for event in seen] == [{"out": 1}, {"out": 2}] + assert published == [{"demo.out": 1}, {"demo.out": 2}] def test_emit_reaches_the_same_ports_from_inside_a_callback(pool): @@ -438,31 +467,30 @@ def test_a_node_with_no_timeout_can_still_be_cancelled(pool): def test_an_emission_on_an_undeclared_port_fails_the_call(pool): """The engine's sink raises, and that has to reach the node's author. - A yield is held one behind — the last one is the return value when there is - no explicit return — so the mistake surfaces on the loop's second pass - rather than its first. Which is what a training loop does in a moment, and - a long way short of the hours it used to cost. + At the *first* yield, which is where the mistake is: the worker sends each + one as it happens, so a mistyped port is checked before the node goes back + to work rather than when the yield after it arrives. """ - from fluksio.flow.nodes.base import NodeOutputError - - def refuse(event): - raise NodeOutputError("'demo.gen' produced 'lss', which no port declares") + node, _ = _sink_node("gen", "loss", DType.FLOAT) + sink = EmitSink() + sink.node = node started = time.monotonic() with pytest.raises(NodeOutputError, match="lss"): - pool.run( - "demo", - "gen", - "import time\n\n\ndef process():\n" - " for _ in range(3):\n" - " yield {'lss': 1.0}\n" - " time.sleep(30)\n" - " return {'out': 1}\n", - {}, - "demo.gen", - timeout=0, - on_event=refuse, - ) + sink.wrap( + lambda: pool.run( + "demo", + "gen", + "import time\n\n\ndef process():\n" + " yield {'lss': 1.0}\n" + " time.sleep(30)\n" + " yield {'loss': 0.5}\n", + {}, + "demo.gen", + timeout=0, + on_event=sink.handle, + ) + )() # It did not wait out the node: the failure stopped the call. assert time.monotonic() - started < 10 diff --git a/worker/fluksio_worker/worker_main.py b/worker/fluksio_worker/worker_main.py index 59ffb1c..305875a 100644 --- a/worker/fluksio_worker/worker_main.py +++ b/worker/fluksio_worker/worker_main.py @@ -362,31 +362,21 @@ def _handle(request: dict[str, Any], cache: dict[tuple[str, str], Any]) -> Any: def _drain(generator: Any) -> Any: - """Run a generator node, publishing each yield as it happens. + """Run a generator node, sending each yield the moment it happens. Two shapes work, and they mean the same thing. Yield throughout and ``return`` the result at the end, which is the explicit one; or just yield, - and the last one is the result. Either way what the node *produces over - time* leaves through its ports while it is still running, and what it - *ends up with* is its return value. + and the last one is the result. Which of the two a node is cannot be known + until it ends, so the last yield has to be held back somewhere — and that + somewhere is the engine, which is the only side that knows what ports the + node declared. Holding it here instead cost a pass: a yield naming a port + that does not exist was only checked once the *next* one arrived. """ - pending: Any = None - have_pending = False try: while True: - value = next(generator) - # Held one behind: until the next yield arrives this might be the - # last one, and the last one is the result rather than an emission. - if have_pending: - _emit({"event": "emit", "outputs": pending}) - pending, have_pending = value, True + _emit({"event": "yield", "outputs": next(generator)}) except StopIteration as stop: - if stop.value is not None: - # It returned something, so every yield was an emission. - if have_pending: - _emit({"event": "emit", "outputs": pending}) - return stop.value - return pending if have_pending else None + return stop.value def main() -> None: