Check a generator's yield at the yield that produced it

The worker held each yield one behind, because the last one is the node's
result when the generator returns nothing of its own. Only the engine knows
what ports a node declared, so the check happened when the *next* yield
arrived — a pass late, which for a training loop is however long one epoch
takes.

The worker now sends every yield as it happens and returns whatever its
generator returned; EmitSink holds the last one back and decides at the end
of the call what it was. Old "emit" frames are still handled, so a remote
agent that has not been restarted keeps working.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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
+        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),

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

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"})

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,8 +245,13 @@ 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 = []
+    """The worker sends every yield; the engine is what holds the last one back."""
-    result = pool.run(
+    node, published = _sink_node("count", "out", DType.INT)
+    sink = EmitSink()
+    sink.node = node
+
+    result = sink.wrap(
+        lambda: pool.run(
             "demo",
             "count",
             "def process():\n"
@@ -235,11 +261,14 @@ def test_without_a_return_the_last_yield_is_the_result(pool):
             {},
             "demo.count",
             timeout=5,
-        on_event=seen.append,
+            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
+    At the *first* yield, which is where the mistake is: the worker sends each
-    no explicit return — so the mistake surfaces on the loop's second pass
+    one as it happens, so a mistyped port is checked before the node goes back
-    rather than its first. Which is what a training loop does in a moment, and
+    to work rather than when the yield after it arrives.
-    a long way short of the hours it used to cost.
     """
-    from fluksio.flow.nodes.base import NodeOutputError
+    node, _ = _sink_node("gen", "loss", DType.FLOAT)
-
+    sink = EmitSink()
-    def refuse(event):
+    sink.node = node
-        raise NodeOutputError("'demo.gen' produced 'lss', which no port declares")
 
     started = time.monotonic()
     with pytest.raises(NodeOutputError, match="lss"):
-        pool.run(
+        sink.wrap(
+            lambda: 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",
+                "    yield {'loss': 0.5}\n",
                 {},
                 "demo.gen",
                 timeout=0,
-            on_event=refuse,
+                on_event=sink.handle,
             )
+        )()
     # It did not wait out the node: the failure stopped the call.
     assert time.monotonic() - started < 10
 

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
+    and the last one is the result. Which of the two a node is cannot be known
-    time* leaves through its ports while it is still running, and what it
+    until it ends, so the last yield has to be held back somewhere — and that
-    *ends up with* is its return value.
+    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)
+            _emit({"event": "yield", "outputs": 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
     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
 
 
 def main() -> None: