feat(langchain): deterministic durability for parallel StateGraph nodes (auto-enabled) - #220
Draft
jgrier wants to merge 5 commits into
Draft
feat(langchain): deterministic durability for parallel StateGraph nodes (auto-enabled)#220jgrier wants to merge 5 commits into
jgrier wants to merge 5 commits into
Conversation
RestateMiddleware linearizes create_agent's parallel tool-call batch. This adds
`durable_scope`, which generalizes the same guarantee to an arbitrary StateGraph:
parallel branches, Send/map-reduce, and nested parallel subgraphs.
LangGraph runs a superstep's nodes concurrently on asyncio; `ctx.run_typed`
creates a journal command synchronously at call time, so parallel nodes create
commands in I/O-completion order -- non-deterministic, and different on replay
(journaled results resolve instantly). Restate matches the journal by position,
so replay mismatches and the invocation cannot recover.
`durable_scope(ctx)` yields a context wrapper; nodes call `dctx.run_typed(...)`
unchanged (no node wrappers, no config threading). A per-invocation coordinator
batches durable ops and creates each round's commands in stable sorted order
(deterministic journal) while awaiting the I/O concurrently via `restate.gather`.
Task identity uses LangGraph's replay-stable name/path metadata, composed
hierarchically for nesting.
RFC/draft: membership + per-task keying come from lazily monkeypatching a private
LangGraph function (langgraph.pregel._runner.arun_with_retry). This is brittle
across LangGraph versions; a merge-ready version wants a supported LangGraph
runner/executor hook. Details and the crash-recovery verification are in the PR.
Test is guarded by importorskip('langgraph') since langgraph is an optional
`langchain` extra, not a default test dependency.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…turnstile
Two changes that together make Rushabh's parallel-node/subgraph repro
deterministic with zero user code (validated: baseline 570-poison -> fixed).
1) StateGraph determinism coordinator (_stategraph.py), rewritten from the
earlier node-wrapper prototype:
- Patches ServerInvocationContext.run_typed at the class level so
restate_context()/RestateMiddleware route parallel durable ops through a
per-invocation coordinator, which creates their journal commands in a
stable, hierarchical-key sorted order (langgraph task name|path, composed
for nested subgraphs) and awaits the I/O concurrently via restate.gather.
The ctx stays a real ServerInvocationContext (extension_data still works).
- Patches langgraph.pregel._runner.arun_with_retry to tag each task with its
replay-stable key.
- enable() auto-activates a coordinator per invocation by wrapping the SDK's
invoke_handler; __init__ calls enable() on import so merely importing the
integration turns it on. Opt out with RESTATE_STATEGRAPH_DETERMINISM=0.
2) Parallel-safe middleware turnstile (_state.py, _middleware.py):
- _State now holds turnstiles keyed by tool_call_id instead of one shared
slot. Parallel tool-using agents in concurrent LangGraph nodes were
clobbering the single slot (KeyError) once the coordinator synchronized
their model calls; keying by the unique tool_call_id fixes it and is
backward-compatible for the single-agent tool batch.
RFC/draft: still monkeypatches private langgraph internals (arun_with_retry);
a merge-ready version wants a supported runner hook. Determinism validated
end-to-end against Rushabh's exact repro with a stubbed model; see notes.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…dinator
Adds a coordinator test modeling a node's concurrent tool calls as separate
Pregel push tasks with distinct hierarchical keys (verified against LangGraph:
.../tools|('__pregel_push', N, False)), alongside sibling single-op nodes.
Asserts they flush in stable sorted-key order with no same-key collision.
This documents the actual mechanism for multi-tool-call parallel nodes: the
coordinator orders the distinct-key tool-call tasks directly; the middleware
turnstile is not what serializes them. (Empirically: create_agent v1 doesn't
invoke the turnstile for tool exec; deepagents does, trivially for a single
call. The per-tool_call_id turnstile keying still matters to prevent the
cross-node clobber KeyError when parallel nodes share the middleware state.)
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…lock) A node with multiple tool calls that invokes the turnstile (e.g. deepagents) DEADLOCKED under the StateGraph coordinator: the tool calls run as separate Pregel tasks (distinct keys), all registered as active leaves; the turnstile blocks tool calls 2..N in wait_for BEFORE they submit, so they never become pending, so the coordinator's "flush when every active leaf is pending" barrier never fires, so tool call 1 never completes, so the turnstile never releases the rest. Circular deadlock (confirmed by tracing). Fix: awrap_tool_call skips the turnstile when a coordinator is active (coordinator_active()). The coordinator already orders the tool-call tasks by key, so the turnstile is redundant there; the stock (no-coordinator) path keeps using it. Verified: deepagents x3 tool calls completes under the fix (turnstile skipped) and under stock (turnstile serializes); his full graph with a 3-tool-call deep node reproduces the 570 in baseline and is 0x570 with the fix. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
I have read the CLA Document and I hereby sign the CLA You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot. |
…ht import) - mypy runs on python/ (not just tests/): the class-level run_typed patch needs type: ignore[method-assign, assignment], not just [method-assign]. - invoke_handler is imported but not re-exported by server_context; use getattr/setattr so pyright doesn't flag reportPrivateImportUsage. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
|
||
| .. note:: | ||
| Limitations: (1) round batching imposes a cross-branch barrier per durable-op | ||
| round, so heterogeneous per-branch latencies run slower than native; (2) assumes |
There was a problem hiding this comment.
For point (2), I assume a case where multiple LLM calls and tool calls and API calls will be made in one node is rather common. Should we just handle this?
| round, so heterogeneous per-branch latencies run slower than native; (2) assumes | ||
| one durable op in flight per task at a time — a node that fires multiple durable | ||
| ops *concurrently* (intra-node ``gather``) needs op-level sub-keys, not handled; | ||
| (3) a routed ``run_typed`` returns a coroutine rather than a ``RestateDurableFuture``, |
There was a problem hiding this comment.
This one will be tricky to remember for all engineers who work with restate
Could silently break code that does f = ctx.run_typed(...); await restate.gather(f, other) inside a LangGraph node. Worth a runtime warning if someone tries to pass the result to restate.gather ideally we should check if this can be addressed.
| return await graph.ainvoke(state) # plain nodes, middleware, tools | ||
|
|
||
| .. warning:: | ||
| **RFC / draft.** Task keying is obtained by monkeypatching a *private* LangGraph |
There was a problem hiding this comment.
Add a version check or at minimum a runtime warning if _runner_mod.arun_with_retry doesn't exist at import time, rather than silently falling through.
| _, full = cur | ||
| # Defer command creation to the coordinator's sorted flush; the flush | ||
| # calls the ORIGINAL run_typed (no recursion) to create the command. | ||
| return coord.submit(full, lambda: _orig_run_typed(self, *args, **kwargs)) |
There was a problem hiding this comment.
Both enable() and durable_scope always pass settle_turns=0. The _flush_when_stable loop (lines ~142-156) implements a multi-turn settling protocol, but it's never exercised.
Comment on lines
+118
to
+129
| def checkout(self, full_key: str) -> None: | ||
| self.active.discard(full_key) | ||
| self.pending.pop(full_key, None) | ||
| self._try_flush() | ||
|
|
||
| def _leaves(self) -> set[str]: | ||
| act = self.active | ||
| return {k for k in act if not any(o != k and o.startswith(k + "/") for o in act)} | ||
|
|
||
| def _ready_set(self) -> Optional[set[str]]: | ||
| leaves = self._leaves() | ||
| if leaves and all(k in self.pending for k in leaves): |
There was a problem hiding this comment.
_try_flush creates an asyncio.create_task(self._flush_when_stable()) firewall. If a new task registers and submits between _ready_set() returning non-None in
_try_flush and the actual flush in _flush_when_stable, the flush could grab an incomplete batch. With settle_turns=0 this is less likely (the while loop
body runs 0 times), but with settle_turns > 0 and interleaved awaits this could race.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
RestateMiddlewaremakes acreate_agentagent durable, but aStateGraphwith parallel nodes (fan-out,Send/map-reduce, nested subgraphs) breaks. LangGraph runs a superstep's nodes concurrently on asyncio, andctx.run_typedcreates a journal command synchronously at call time, so parallel nodes create commands in I/O-completion order — nondeterministic, and different on replay (journaled results resolve instantly). Restate matches the journal by position, so replay mismatches and the invocation poisons withVMException (570) Journal mismatch.Reported with a minimal repro:
plan → [branch, branch, subgraph, subgraph, deepagents] → join, every node journaled byRestateMiddleware.Fix
Importing
restate.ext.langchainenables it automatically — no wrappers, no config, no changes to graph/agent/handler code (opt out withRESTATE_STATEGRAPH_DETERMINISM=0).__init__.pycallsenable()on import, which installs a per-invocation coordinator:ServerInvocationContext.run_typedis patched sorestate_context()/middleware route parallel durable ops through the coordinator. It batches ops and creates their journal commands in a stable, hierarchical-key sorted order (LangGraph taskname|path, composed for nested subgraphs), awaiting the I/O concurrently viarestate.gather. The context stays a realServerInvocationContext, soextension_dataetc. are unaffected.langgraph.pregel._runner.arun_with_retryis patched to tag each task with its replay-stable key.invoke_handleris wrapped so every invocation gets a fresh coordinator.Two middleware changes come with it:
ctx.extension_dataslot, which parallel agents clobbered (KeyError). It's now keyed bytool_call_id, backward compatible for the single-agent tool batch.awrap_tool_callnow skips the turnstile when the coordinator is active — the coordinator already orders the tool-call tasks by key; the stock path keeps the turnstile.Validation
Reproduced and resolved against the reported repro, end-to-end against
restate-serverwith a stubbed model.inactivity_timeout=0s(suspend/replay on every await) surfaces the 570 on a single invocation.Unit tests in
tests/langchain_stategraph.pycover the coordinator's sorted-key ordering (fan-out, concurrent tool-call tasks, multiple rounds) and thecoordinator_active()gate; guarded byimportorskip("langgraph").