Stream output from subgraphs while in-progress for sync stream

- previously implemented only for async stream

libs/langgraph/Makefile

@@ -40,16 +40,16 @@ start-postgres:
 stop-postgres:
 	docker compose -f tests/compose-postgres.yml down -v
 
-TEST_PATH ?= .
+TEST ?= .
 
 test:
-	make start-postgres && poetry run pytest $(TEST_PATH); \
+	make start-postgres && poetry run pytest $(TEST); \
 	EXIT_CODE=$$?; \
 	make stop-postgres; \
 	exit $$EXIT_CODE
 
 test_watch:
-	make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \
+	make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \
 	EXIT_CODE=$$?; \
 	make stop-postgres; \
 	exit $$EXIT_CODE

libs/langgraph/langgraph/pregel/__init__.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent
+import concurrent.futures
+import queue
 from collections import deque
 from functools import partial
 from typing import (
@@ -83,7 +86,6 @@ from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.utils.aio import Queue
 from langgraph.utils.config import (
     ensure_config,
     merge_configs,
@@ -92,6 +94,7 @@ from langgraph.utils.config import (
     patch_configurable,
 )
 from langgraph.utils.pydantic import create_model
+from langgraph.utils.queue import AsyncQueue, SyncQueue
 from langgraph.utils.runnable import RunnableCallable
 
 WriteValue = Union[Callable[[Input], Output], Any]
@@ -1162,11 +1165,14 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             ```
         """
 
-        stream = deque()
+        stream = SyncQueue()
 
         def output() -> Iterator:
-            while stream:
-                ns, mode, payload = stream.popleft()
+            while True:
+                try:
+                    ns, mode, payload = stream.get(block=False)
+                except queue.Empty:
+                    break
                 if subgraphs and isinstance(stream_mode, list):
                     yield (ns, mode, payload)
                 elif isinstance(stream_mode, list):
@@ -1210,7 +1216,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
 
             with SyncPregelLoop(
                 input,
-                stream=StreamProtocol(stream.append, stream_modes),
+                stream=StreamProtocol(stream.put, stream_modes),
                 config=config,
                 store=self.store,
                 checkpointer=checkpointer,
@@ -1228,6 +1234,22 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                 # enable subgraph streaming
                 if subgraphs:
                     loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream
+                    # we are careful to have a single waiter live at any one time
+                    # because on exit we increment semaphore count by exactly 1
+                    waiter: Optional[concurrent.futures.Future] = None
+                    # because sync futures cannot be cancelled, we instead
+                    # release the stream semaphore on exit, which will cause
+                    # a pending waiter to return immediately
+                    loop.stack.callback(stream._count.release)
+
+                    def get_waiter() -> asyncio.Task[None]:
+                        nonlocal waiter
+                        if waiter is None or waiter.done():
+                            return (waiter := loop.submit(stream.wait))
+                        else:
+                            return waiter
+                else:
+                    get_waiter = None
                 # Similarly to Bulk Synchronous Parallel / Pregel model
                 # computation proceeds in steps, while there are channel updates
                 # channel updates from step N are only visible in step N+1
@@ -1243,10 +1265,10 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                         loop.tasks.values(),
                         timeout=self.step_timeout,
                         retry_policy=self.retry_policy,
+                        get_waiter=get_waiter,
                     ):
                         # emit output
-                        for o in output():
-                            yield o
+                        yield from output()
             # emit output
             yield from output()
             # handle exit
@@ -1342,7 +1364,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
             ```
         """
 
-        stream = Queue()
+        stream = AsyncQueue()
         aioloop = asyncio.get_running_loop()
 
         def output() -> Iterator:

libs/langgraph/langgraph/pregel/runner.py

@@ -38,6 +38,7 @@ class PregelRunner:
         reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
+        get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
     ) -> Iterator[None]:
         tasks = tuple(tasks)
         # give control back to the caller
@@ -53,23 +54,30 @@ class PregelRunner:
                 if reraise:
                     raise
             return
+        # add waiter task if requested
+        if get_waiter is not None:
+            futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {
+                get_waiter(): None
+            }
+        else:
+            futures = {}
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
-        futures = {
-            self.submit(
-                run_with_retry,
-                task,
-                retry_policy,
-                __reraise_on_exit__=reraise,
-            ): task
-            for task in tasks
-            if not task.writes
-        }
+        for task in tasks:
+            if not task.writes:
+                futures[
+                    self.submit(
+                        run_with_retry,
+                        task,
+                        retry_policy,
+                        __reraise_on_exit__=reraise,
+                    )
+                ] = task
         all_futures = futures.copy()
         end_time = timeout + time.monotonic() if timeout else None
-        while futures:
-            done, _ = concurrent.futures.wait(
+        while len(futures) > (1 if get_waiter is not None else 0):
+            done, inflight = concurrent.futures.wait(
                 futures,
                 return_when=concurrent.futures.FIRST_COMPLETED,
                 timeout=(max(0, end_time - time.monotonic()) if end_time else None),
@@ -78,8 +86,13 @@ class PregelRunner:
                 break  # timed out
             for fut in done:
                 task = futures.pop(fut)
-                # task finished, commit writes
-                self.commit(task, _exception(fut))
+                if task is None:
+                    # waiter task finished, schedule another
+                    if inflight:
+                        futures[get_waiter()] = None
+                else:
+                    # task finished, commit writes
+                    self.commit(task, _exception(fut))
             else:
                 # remove references to loop vars
                 del fut, task
@@ -141,7 +154,7 @@ class PregelRunner:
         all_futures = futures.copy()
         end_time = timeout + loop.time() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
-            done, _ = await asyncio.wait(
+            done, inflight = await asyncio.wait(
                 futures,
                 return_when=asyncio.FIRST_COMPLETED,
                 timeout=(max(0, end_time - loop.time()) if end_time else None),
@@ -152,7 +165,8 @@ class PregelRunner:
                 task = futures.pop(fut)
                 if task is None:
                     # waiter task finished, schedule another
-                    futures[get_waiter()] = None
+                    if inflight:
+                        futures[get_waiter()] = None
                 else:
                     # task finished, commit writes
                     self.commit(task, _exception(fut))

libs/langgraph/langgraph/utils/aio.py

@@ -1,35 +0,0 @@
-import asyncio
-import sys
-
-PY_310 = sys.version_info >= (3, 10)
-
-
-class Queue(asyncio.Queue):
-    async def wait(self):
-        """If queue is empty, wait until an item is available.
-
-        Copied from Queue.get(), removing the call to .get_nowait(),
-        ie. this doesn't consume the item, just waits for it.
-        """
-        while self.empty():
-            if PY_310:
-                getter = self._get_loop().create_future()
-            else:
-                getter = self._loop.create_future()
-            self._getters.append(getter)
-            try:
-                await getter
-            except:
-                getter.cancel()  # Just in case getter is not done yet.
-                try:
-                    # Clean self._getters from canceled getters.
-                    self._getters.remove(getter)
-                except ValueError:
-                    # The getter could be removed from self._getters by a
-                    # previous put_nowait call.
-                    pass
-                if not self.empty() and not getter.cancelled():
-                    # We were woken up by put_nowait(), but can't take
-                    # the call.  Wake up the next in line.
-                    self._wakeup_next(self._getters)
-                raise

libs/langgraph/langgraph/utils/queue.py

@@ -0,0 +1,127 @@
+import asyncio
+import queue
+import sys
+import threading
+import types
+from collections import deque
+from time import monotonic
+
+PY_310 = sys.version_info >= (3, 10)
+
+
+class AsyncQueue(asyncio.Queue):
+    """Async unbounded FIFO queue with a wait() method.
+
+    Subclassed from asyncio.Queue, adding a wait() method."""
+
+    async def wait(self):
+        """If queue is empty, wait until an item is available.
+
+        Copied from Queue.get(), removing the call to .get_nowait(),
+        ie. this doesn't consume the item, just waits for it.
+        """
+        while self.empty():
+            if PY_310:
+                getter = self._get_loop().create_future()
+            else:
+                getter = self._loop.create_future()
+            self._getters.append(getter)
+            try:
+                await getter
+            except:
+                getter.cancel()  # Just in case getter is not done yet.
+                try:
+                    # Clean self._getters from canceled getters.
+                    self._getters.remove(getter)
+                except ValueError:
+                    # The getter could be removed from self._getters by a
+                    # previous put_nowait call.
+                    pass
+                if not self.empty() and not getter.cancelled():
+                    # We were woken up by put_nowait(), but can't take
+                    # the call.  Wake up the next in line.
+                    self._wakeup_next(self._getters)
+                raise
+
+
+class Semaphore(threading.Semaphore):
+    """Semaphore subclass with a wait() method."""
+
+    def wait(self, blocking: bool = True, timeout: float = None):
+        """Block until the semaphore can be acquired, but don't acquire it."""
+        if not blocking and timeout is not None:
+            raise ValueError("can't specify timeout for non-blocking acquire")
+        rc = False
+        endtime = None
+        with self._cond:
+            while self._value == 0:
+                if not blocking:
+                    break
+                if timeout is not None:
+                    if endtime is None:
+                        endtime = monotonic() + timeout
+                    else:
+                        timeout = endtime - monotonic()
+                        if timeout <= 0:
+                            break
+                self._cond.wait(timeout)
+            else:
+                rc = True
+        return rc
+
+
+class SyncQueue:
+    """Unbounded FIFO queue with a wait() method.
+    Adapted from pure Python implementation of queue.SimpleQueue.
+    """
+
+    def __init__(self):
+        self._queue = deque()
+        self._count = Semaphore(0)
+
+    def put(self, item, block=True, timeout=None):
+        """Put the item on the queue.
+
+        The optional 'block' and 'timeout' arguments are ignored, as this method
+        never blocks.  They are provided for compatibility with the Queue class.
+        """
+        self._queue.append(item)
+        self._count.release()
+
+    def get(self, block=True, timeout=None):
+        """Remove and return an item from the queue.
+
+        If optional args 'block' is true and 'timeout' is None (the default),
+        block if necessary until an item is available. If 'timeout' is
+        a non-negative number, it blocks at most 'timeout' seconds and raises
+        the Empty exception if no item was available within that time.
+        Otherwise ('block' is false), return an item if one is immediately
+        available, else raise the Empty exception ('timeout' is ignored
+        in that case).
+        """
+        if timeout is not None and timeout < 0:
+            raise ValueError("'timeout' must be a non-negative number")
+        if not self._count.acquire(block, timeout):
+            raise queue.Empty
+        try:
+            return self._queue.popleft()
+        except IndexError:
+            raise queue.Empty
+
+    def wait(self, block=True, timeout=None):
+        """If queue is empty, wait until an item maybe is available,
+        but don't consume it.
+        """
+        if timeout is not None and timeout < 0:
+            raise ValueError("'timeout' must be a non-negative number")
+        self._count.wait(block, timeout)
+
+    def empty(self):
+        """Return True if the queue is empty, False otherwise (not reliable!)."""
+        return len(self._queue) == 0
+
+    def qsize(self):
+        """Return the approximate size of the queue (not reliable!)."""
+        return len(self._queue)
+
+    __class_getitem__ = classmethod(types.GenericAlias)

libs/langgraph/tests/test_pregel.py

@@ -8362,6 +8362,86 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
     ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_stream_subgraphs_during_execution(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
+
+    class InnerState(TypedDict):
+        my_key: Annotated[str, operator.add]
+        my_other_key: str
+
+    def inner_1(state: InnerState):
+        return {"my_key": "got here", "my_other_key": state["my_key"]}
+
+    def inner_2(state: InnerState):
+        time.sleep(0.5)
+        return {
+            "my_key": " and there",
+            "my_other_key": state["my_key"],
+        }
+
+    inner = StateGraph(InnerState)
+    inner.add_node("inner_1", inner_1)
+    inner.add_node("inner_2", inner_2)
+    inner.add_edge("inner_1", "inner_2")
+    inner.set_entry_point("inner_1")
+    inner.set_finish_point("inner_2")
+
+    class State(TypedDict):
+        my_key: Annotated[str, operator.add]
+
+    def outer_1(state: State):
+        time.sleep(0.2)
+        return {"my_key": " and parallel"}
+
+    def outer_2(state: State):
+        return {"my_key": " and back again"}
+
+    graph = StateGraph(State)
+    graph.add_node("inner", inner.compile())
+    graph.add_node("outer_1", outer_1)
+    graph.add_node("outer_2", outer_2)
+
+    graph.add_edge(START, "inner")
+    graph.add_edge(START, "outer_1")
+    graph.add_edge(["inner", "outer_1"], "outer_2")
+    graph.add_edge("outer_2", END)
+
+    app = graph.compile(checkpointer=checkpointer)
+
+    start = time.perf_counter()
+    chunks: list[tuple[float, Any]] = []
+    config = {"configurable": {"thread_id": "2"}}
+    for c in app.stream({"my_key": ""}, config, subgraphs=True):
+        chunks.append((round(time.perf_counter() - start, 1), c))
+    for idx in range(len(chunks)):
+        elapsed, c = chunks[idx]
+        chunks[idx] = (round(elapsed - chunks[0][0], 1), c)
+
+    assert chunks == [
+        # arrives before "inner" finishes
+        (
+            0.0,
+            (
+                (AnyStr("inner:"),),
+                {"inner_1": {"my_key": "got here", "my_other_key": ""}},
+            ),
+        ),
+        (0.2, ((), {"outer_1": {"my_key": " and parallel"}})),
+        (
+            0.5,
+            (
+                (AnyStr("inner:"),),
+                {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
+            ),
+        ),
+        (0.5, ((), {"inner": {"my_key": "got here and there"}})),
+        (0.5, ((), {"outer_2": {"my_key": " and back again"}})),
+    ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_nested_graph_interrupts_parallel(
     request: pytest.FixtureRequest, checkpointer_name: str