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>
This commit is contained in:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user