Compare commits

...
Author SHA1 Message Date
Will Fu-Hinthorn 685a755baf chore: use bounded queue for checkpointing 2026-04-09 17:33:50 -07:00
Will Fu-HinthornandClaude Opus 4.6 bde0c47cd9 fix: handle non-deterministic ordering in test_imp_nested
Nested tasks run concurrently so output order is non-deterministic.
Sort intermediate results before comparison, matching the pattern
already used by test_imp_task.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-08 18:04:55 -07:00
Will Fu-Hinthorn 7b034c2d25 chore: update conformance lint 2026-04-08 16:02:42 -07:00
7 changed files with 490 additions and 110 deletions
@@ -339,7 +339,6 @@ async def test_list_metadata_custom_keys(
assert results[0].metadata["run_id"] == "run-abc" assert results[0].metadata["run_id"] == "run-abc"
ALL_LIST_TESTS = [ ALL_LIST_TESTS = [
test_list_all, test_list_all,
test_list_by_thread, test_list_by_thread,
+1 -1
View File
@@ -263,7 +263,7 @@ wheels = [
[[package]] [[package]]
name = "langgraph-checkpoint-conformance" name = "langgraph-checkpoint-conformance"
version = "0.0.1" version = "0.0.2"
source = { editable = "." } source = { editable = "." }
dependencies = [ dependencies = [
{ name = "langgraph-checkpoint" }, { name = "langgraph-checkpoint" },
@@ -0,0 +1,215 @@
from __future__ import annotations
import asyncio
import os
import queue
import threading
from collections.abc import Callable
from contextlib import AbstractAsyncContextManager, AbstractContextManager
from dataclasses import dataclass
from types import TracebackType
from typing import Any
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
ChannelVersions,
Checkpoint,
CheckpointMetadata,
)
QUEUE_PUT_TIMEOUT = 0.05
CHECKPOINT_BACKLOG_ENV_VAR = "LANGGRAPH_CHECKPOINT_BACKLOG"
DEFAULT_CHECKPOINT_BACKLOG = 10
@dataclass(frozen=True)
class CheckpointRequest:
config: RunnableConfig
checkpoint: Checkpoint
metadata: CheckpointMetadata
new_versions: ChannelVersions
def _raise(error: BaseException) -> None:
raise error
def resolve_checkpoint_backlog() -> int:
if raw := os.getenv(CHECKPOINT_BACKLOG_ENV_VAR):
try:
backlog = int(raw)
except ValueError:
return DEFAULT_CHECKPOINT_BACKLOG
if backlog > 0:
return backlog
return DEFAULT_CHECKPOINT_BACKLOG
class SyncCheckpointWriter(AbstractContextManager):
def __init__(
self,
put: Callable[
[RunnableConfig, Checkpoint, CheckpointMetadata, ChannelVersions], Any
],
*,
max_pending: int | None = None,
) -> None:
self.put = put
max_pending = (
resolve_checkpoint_backlog() if max_pending is None else max_pending
)
self.queue: queue.Queue[CheckpointRequest | None] = queue.Queue(max_pending)
self.error: BaseException | None = None
self.closed = False
self.thread = threading.Thread(
target=self._run,
name="langgraph-checkpoint-writer",
daemon=True,
)
def __enter__(self) -> SyncCheckpointWriter:
self.thread.start()
return self
def submit(self, request: CheckpointRequest) -> None:
self._ensure_open()
while True:
self._raise_if_broken()
try:
self.queue.put(request, timeout=QUEUE_PUT_TIMEOUT)
except queue.Full:
continue
else:
self._raise_if_broken()
return
def _run(self) -> None:
while True:
item = self.queue.get()
if item is None:
return
try:
self.put(
item.config,
item.checkpoint,
item.metadata,
item.new_versions,
)
except BaseException as exc:
self.error = exc
return
def _ensure_open(self) -> None:
if self.closed:
raise RuntimeError("Checkpoint writer is closed")
def _raise_if_broken(self) -> None:
if self.error is not None:
_raise(self.error)
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
self.closed = True
while self.thread.is_alive():
if self.error is not None:
break
try:
self.queue.put(None, timeout=QUEUE_PUT_TIMEOUT)
except queue.Full:
continue
else:
break
self.thread.join()
if exc_type is None and self.error is not None:
_raise(self.error)
return None
class AsyncCheckpointWriter(AbstractAsyncContextManager):
def __init__(
self,
put: Callable[
[RunnableConfig, Checkpoint, CheckpointMetadata, ChannelVersions], Any
],
*,
max_pending: int | None = None,
) -> None:
self.put = put
max_pending = (
resolve_checkpoint_backlog() if max_pending is None else max_pending
)
self.queue: asyncio.Queue[CheckpointRequest | None] = asyncio.Queue(max_pending)
self.error: BaseException | None = None
self.closed = False
self.task: asyncio.Task[None] | None = None
async def __aenter__(self) -> AsyncCheckpointWriter:
self.task = asyncio.create_task(self._run(), name="langgraph-checkpoint-writer")
return self
async def submit(self, request: CheckpointRequest) -> None:
self._ensure_open()
while True:
self._raise_if_broken()
try:
await asyncio.wait_for(
self.queue.put(request),
timeout=QUEUE_PUT_TIMEOUT,
)
except asyncio.TimeoutError:
continue
else:
self._raise_if_broken()
return
async def _run(self) -> None:
while True:
item = await self.queue.get()
if item is None:
return
try:
await self.put(
item.config,
item.checkpoint,
item.metadata,
item.new_versions,
)
except BaseException as exc:
self.error = exc
return
def _ensure_open(self) -> None:
if self.closed:
raise RuntimeError("Checkpoint writer is closed")
def _raise_if_broken(self) -> None:
if self.error is not None:
_raise(self.error)
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
self.closed = True
while self.task is not None and not self.task.done():
if self.error is not None:
break
try:
await asyncio.wait_for(
self.queue.put(None),
timeout=QUEUE_PUT_TIMEOUT,
)
except asyncio.TimeoutError:
continue
else:
break
if self.task is not None:
await self.task
if exc_type is None and self.error is not None:
_raise(self.error)
+109 -90
View File
@@ -2,7 +2,6 @@ from __future__ import annotations
import asyncio import asyncio
import binascii import binascii
import concurrent.futures
from collections import defaultdict, deque from collections import defaultdict, deque
from collections.abc import Callable, Iterator, Mapping, Sequence from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import ( from contextlib import (
@@ -27,7 +26,6 @@ from langgraph.cache.base import BaseCache
from langgraph.checkpoint.base import ( from langgraph.checkpoint.base import (
WRITES_IDX_MAP, WRITES_IDX_MAP,
BaseCheckpointSaver, BaseCheckpointSaver,
ChannelVersions,
Checkpoint, Checkpoint,
CheckpointMetadata, CheckpointMetadata,
CheckpointTuple, CheckpointTuple,
@@ -92,6 +90,11 @@ from langgraph.pregel._checkpoint import (
create_checkpoint, create_checkpoint,
empty_checkpoint, empty_checkpoint,
) )
from langgraph.pregel._checkpoint_writer import (
AsyncCheckpointWriter,
CheckpointRequest,
SyncCheckpointWriter,
)
from langgraph.pregel._executor import ( from langgraph.pregel._executor import (
AsyncBackgroundExecutor, AsyncBackgroundExecutor,
BackgroundExecutor, BackgroundExecutor,
@@ -166,19 +169,6 @@ class PregelLoop:
checkpointer_get_next_version: GetNextVersion checkpointer_get_next_version: GetNextVersion
checkpointer_put_writes: Callable[[RunnableConfig, WritesT, str], Any] | None checkpointer_put_writes: Callable[[RunnableConfig, WritesT, str], Any] | None
checkpointer_put_writes_accepts_task_path: bool checkpointer_put_writes_accepts_task_path: bool
_checkpointer_put_after_previous: (
Callable[
[
concurrent.futures.Future | None,
RunnableConfig,
Checkpoint,
str,
ChannelVersions,
],
Any,
]
| None
)
_migrate_checkpoint: Callable[[Checkpoint], None] | None _migrate_checkpoint: Callable[[Checkpoint], None] | None
submit: Submit submit: Submit
channels: Mapping[str, BaseChannel] channels: Mapping[str, BaseChannel]
@@ -491,7 +481,7 @@ class PregelLoop:
) )
# produce debug output # produce debug output
if self._checkpointer_put_after_previous is not None: if self.checkpointer is not None:
self._emit( self._emit(
"checkpoints", "checkpoints",
map_debug_checkpoint, map_debug_checkpoint,
@@ -537,7 +527,7 @@ class PregelLoop:
return True return True
def after_tick(self) -> None: def _after_tick(self) -> CheckpointRequest | None:
# finish superstep # finish superstep
writes = [w for t in self.tasks.values() for w in t.writes] writes = [w for t in self.tasks.values() for w in t.writes]
# all tasks have finished # all tasks have finished
@@ -562,14 +552,14 @@ class PregelLoop:
# only replay (re-execute) done tasks on the first tick # only replay (re-execute) done tasks on the first tick
self.is_replaying = False self.is_replaying = False
# save checkpoint # save checkpoint
self._put_checkpoint({"source": "loop"}) return self._prepare_checkpoint({"source": "loop"})
# after execution, check if we should interrupt
def _finish_after_tick(self) -> None:
if self.interrupt_after and should_interrupt( if self.interrupt_after and should_interrupt(
self.checkpoint, self.interrupt_after, self.tasks.values() self.checkpoint, self.interrupt_after, self.tasks.values()
): ):
self.status = "interrupt_after" self.status = "interrupt_after"
raise GraphInterrupt() raise GraphInterrupt()
# unset resuming flag
self.config[CONF].pop(CONFIG_KEY_RESUMING, None) self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
def match_cached_writes(self) -> Sequence[PregelExecutableTask]: def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
@@ -619,7 +609,7 @@ class PregelLoop:
def _first( def _first(
self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None self, *, input_keys: str | Sequence[str], updated_channels: set[str] | None
) -> set[str] | None: ) -> tuple[set[str] | None, CheckpointRequest | None]:
# Resuming from a previous checkpoint requires two things: # Resuming from a previous checkpoint requires two things:
# 1. A prior checkpoint exists (channel_versions is non-empty) # 1. A prior checkpoint exists (channel_versions is non-empty)
# 2. The input signals continuation (not a fresh run with new input) # 2. The input signals continuation (not a fresh run with new input)
@@ -713,6 +703,7 @@ class PregelLoop:
) )
if updated_channels is not None: if updated_channels is not None:
updated_channels.update(null_updated_channels) updated_channels.update(null_updated_channels)
checkpoint_request = None
# proceed past previous checkpoint # proceed past previous checkpoint
if is_resuming: if is_resuming:
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
@@ -755,7 +746,7 @@ class PregelLoop:
) )
# save input checkpoint # save input checkpoint
self.updated_channels = updated_channels self.updated_channels = updated_channels
self._put_checkpoint({"source": "input"}) checkpoint_request = self._prepare_checkpoint({"source": "input"})
elif CONFIG_KEY_RESUMING not in configurable: elif CONFIG_KEY_RESUMING not in configurable:
raise EmptyInputError(f"Received no input for {input_keys}") raise EmptyInputError(f"Received no input for {input_keys}")
# Propagate resuming and replaying flags to subgraphs. # Propagate resuming and replaying flags to subgraphs.
@@ -785,9 +776,11 @@ class PregelLoop:
) )
# set flag # set flag
self.status = "pending" self.status = "pending"
return updated_channels return updated_channels, checkpoint_request
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: def _prepare_checkpoint(
self, metadata: CheckpointMetadata
) -> CheckpointRequest | None:
# assign step and parents # assign step and parents
exiting = metadata is self.checkpoint_metadata exiting = metadata is self.checkpoint_metadata
if exiting and self.checkpoint["id"] == self.checkpoint_id_saved: if exiting and self.checkpoint["id"] == self.checkpoint_id_saved:
@@ -797,8 +790,7 @@ class PregelLoop:
metadata["step"] = self.step metadata["step"] = self.step
metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})
self.checkpoint_metadata = metadata self.checkpoint_metadata = metadata
# do checkpoint? do_checkpoint = self.checkpointer is not None and (
do_checkpoint = self._checkpointer_put_after_previous is not None and (
exiting or self.durability != "exit" exiting or self.durability != "exit"
) )
# create new checkpoint # create new checkpoint
@@ -820,9 +812,8 @@ class PregelLoop:
for value in self.checkpoint["channel_values"][TASKS] for value in self.checkpoint["channel_values"][TASKS]
] ]
self.checkpoint["channel_values"][TASKS] = sanitized_tasks self.checkpoint["channel_values"][TASKS] = sanitized_tasks
# bail if no checkpointer request = None
if do_checkpoint:
if do_checkpoint and self._checkpointer_put_after_previous is not None:
self.prev_checkpoint_config = ( self.prev_checkpoint_config = (
self.checkpoint_config self.checkpoint_config
if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF] if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF]
@@ -844,17 +835,11 @@ class PregelLoop:
self.checkpoint_previous_versions, channel_versions self.checkpoint_previous_versions, channel_versions
) )
self.checkpoint_previous_versions = channel_versions self.checkpoint_previous_versions = channel_versions
request = CheckpointRequest(
# save it, without blocking config=self.checkpoint_config,
# if there's a previous checkpoint save in progress, wait for it checkpoint=copy_checkpoint(self.checkpoint),
# ensuring checkpointers receive checkpoints in order metadata=self.checkpoint_metadata,
self._put_checkpoint_fut = self.submit( new_versions=new_versions,
self._checkpointer_put_after_previous,
getattr(self, "_put_checkpoint_fut", None),
self.checkpoint_config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
new_versions,
) )
self.checkpoint_config = { self.checkpoint_config = {
**self.checkpoint_config, **self.checkpoint_config,
@@ -866,28 +851,18 @@ class PregelLoop:
if not exiting: if not exiting:
# increment step # increment step
self.step += 1 self.step += 1
return request
def _suppress_interrupt( def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
raise NotImplementedError
def _finalize_suppress(
self, self,
exc_type: type[BaseException] | None, exc_type: type[BaseException] | None,
exc_value: BaseException | None, exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None: ) -> bool | None:
# persist current checkpoint and writes
if self.durability == "exit" and (
# if it's a top graph
not self.is_nested
# or a nested graph with error or interrupt
or exc_value is not None
# or a nested graph with checkpointer=True
or all(NS_END not in part for part in self.checkpoint_ns)
):
self._put_checkpoint(self.checkpoint_metadata)
self._put_pending_writes()
# suppress interrupt
suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
if suppress: if suppress:
# emit one last "values" event, with pending writes applied
if ( if (
hasattr(self, "tasks") hasattr(self, "tasks")
and self.checkpoint_pending_writes and self.checkpoint_pending_writes
@@ -912,7 +887,6 @@ class PregelLoop:
[w for t in self.tasks.values() for w in t.writes], [w for t in self.tasks.values() for w in t.writes],
self.channels, self.channels,
) )
# emit INTERRUPT if exception is empty (otherwise emitted by put_writes)
if exc_value is not None and (not exc_value.args or not exc_value.args[0]): if exc_value is not None and (not exc_value.args or not exc_value.args[0]):
self._emit( self._emit(
"updates", "updates",
@@ -920,13 +894,26 @@ class PregelLoop:
[{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}] [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]
), ),
) )
# save final output
self.output = read_channels(self.channels, self.output_keys) self.output = read_channels(self.channels, self.output_keys)
# suppress interrupt
return True return True
elif exc_type is None: elif exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys) self.output = read_channels(self.channels, self.output_keys)
return None
def _suppress_interrupt(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
if self.durability == "exit" and (
not self.is_nested
or exc_value is not None
or all(NS_END not in part for part in self.checkpoint_ns)
):
self._put_checkpoint(self.checkpoint_metadata)
self._put_pending_writes()
return self._finalize_suppress(exc_type, exc_value)
def _emit( def _emit(
self, self,
@@ -1072,26 +1059,30 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
) )
else: else:
self.checkpointer_get_next_version = increment self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None # type: ignore[assignment]
self.checkpointer_put_writes = None self.checkpointer_put_writes = None
self.checkpointer_put_writes_accepts_task_path = False self.checkpointer_put_writes_accepts_task_path = False
self._checkpoint_writer: SyncCheckpointWriter | None = None
def _checkpointer_put_after_previous( def _dispatch_checkpoint_request(self, request: CheckpointRequest) -> None:
self, if self.durability == "async" and self._checkpoint_writer is not None:
prev: concurrent.futures.Future | None, self._checkpoint_writer.submit(request)
config: RunnableConfig, else:
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
try:
if prev is not None:
prev.result()
finally:
cast(BaseCheckpointSaver, self.checkpointer).put( cast(BaseCheckpointSaver, self.checkpointer).put(
config, checkpoint, metadata, new_versions request.config,
request.checkpoint,
request.metadata,
request.new_versions,
) )
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
if request := self._prepare_checkpoint(metadata):
self._dispatch_checkpoint_request(request)
def after_tick(self) -> None:
if request := self._after_tick():
self._dispatch_checkpoint_request(request)
self._finish_after_tick()
def match_cached_writes(self) -> Sequence[PregelExecutableTask]: def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
if self.cache is None: if self.cache is None:
return () return ()
@@ -1186,6 +1177,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
else [] else []
) )
self.submit = self.stack.enter_context(BackgroundExecutor(self.config)) self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
if self.checkpointer is not None and self.durability == "async":
self._checkpoint_writer = self.stack.enter_context(
SyncCheckpointWriter(self.checkpointer.put)
)
self.channels, self.managed = channels_from_checkpoint( self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint self.specs, self.checkpoint
) )
@@ -1194,12 +1189,14 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
self.step = self.checkpoint_metadata["step"] + 1 self.step = self.checkpoint_metadata["step"] + 1
self.stop = self.step + self.config["recursion_limit"] + 1 self.stop = self.step + self.config["recursion_limit"] + 1
self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy() self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
self.updated_channels = self._first( self.updated_channels, checkpoint_request = self._first(
input_keys=self.input_keys, input_keys=self.input_keys,
updated_channels=set(self.checkpoint.get("updated_channels")) # type: ignore[arg-type] updated_channels=set(self.checkpoint.get("updated_channels")) # type: ignore[arg-type]
if self.checkpoint.get("updated_channels") if self.checkpoint.get("updated_channels")
else None, else None,
) )
if checkpoint_request is not None:
self._dispatch_checkpoint_request(checkpoint_request)
return self return self
@@ -1268,26 +1265,26 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
) )
else: else:
self.checkpointer_get_next_version = increment self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None # type: ignore[assignment]
self.checkpointer_put_writes = None self.checkpointer_put_writes = None
self.checkpointer_put_writes_accepts_task_path = False self.checkpointer_put_writes_accepts_task_path = False
self._checkpoint_writer: AsyncCheckpointWriter | None = None
async def _checkpointer_put_after_previous( async def _dispatch_checkpoint_request(self, request: CheckpointRequest) -> None:
self, if self.durability == "async" and self._checkpoint_writer is not None:
prev: asyncio.Task | None, await self._checkpoint_writer.submit(request)
config: RunnableConfig, else:
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
try:
if prev is not None:
await prev
finally:
await cast(BaseCheckpointSaver, self.checkpointer).aput( await cast(BaseCheckpointSaver, self.checkpointer).aput(
config, checkpoint, metadata, new_versions request.config,
request.checkpoint,
request.metadata,
request.new_versions,
) )
async def aafter_tick(self) -> None:
if request := self._after_tick():
await self._dispatch_checkpoint_request(request)
self._finish_after_tick()
async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
if self.cache is None: if self.cache is None:
return [] return []
@@ -1332,6 +1329,22 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
}, },
) )
async def _asuppress_interrupt(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
if self.durability == "exit" and (
not self.is_nested
or exc_value is not None
or all(NS_END not in part for part in self.checkpoint_ns)
):
if request := self._prepare_checkpoint(self.checkpoint_metadata):
await self._dispatch_checkpoint_request(request)
self._put_pending_writes()
return self._finalize_suppress(exc_type, exc_value)
# context manager # context manager
async def __aenter__(self) -> Self: async def __aenter__(self) -> Self:
@@ -1387,20 +1400,26 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.submit = await self.stack.enter_async_context( self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config) AsyncBackgroundExecutor(self.config)
) )
if self.checkpointer is not None and self.durability == "async":
self._checkpoint_writer = await self.stack.enter_async_context(
AsyncCheckpointWriter(self.checkpointer.aput)
)
self.channels, self.managed = channels_from_checkpoint( self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint self.specs, self.checkpoint
) )
self.stack.push(self._suppress_interrupt) self.stack.push_async_exit(self._asuppress_interrupt)
self.status = "input" self.status = "input"
self.step = self.checkpoint_metadata["step"] + 1 self.step = self.checkpoint_metadata["step"] + 1
self.stop = self.step + self.config["recursion_limit"] + 1 self.stop = self.step + self.config["recursion_limit"] + 1
self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy() self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
self.updated_channels = self._first( self.updated_channels, checkpoint_request = self._first(
input_keys=self.input_keys, input_keys=self.input_keys,
updated_channels=set(self.checkpoint.get("updated_channels")) # type: ignore[arg-type] updated_channels=set(self.checkpoint.get("updated_channels")) # type: ignore[arg-type]
if self.checkpoint.get("updated_channels") if self.checkpoint.get("updated_channels")
else None, else None,
) )
if checkpoint_request is not None:
await self._dispatch_checkpoint_request(checkpoint_request)
return self return self
+1 -7
View File
@@ -2751,9 +2751,6 @@ class Pregel(
_state_mapper, _state_mapper,
) )
loop.after_tick() loop.after_tick()
# wait for checkpoint
if durability_ == "sync":
loop._put_checkpoint_fut.result()
# emit output # emit output
yield from _output( yield from _output(
stream_mode, stream_mode,
@@ -3143,10 +3140,7 @@ class Pregel(
_state_mapper, _state_mapper,
): ):
yield o yield o
loop.after_tick() await loop.aafter_tick()
# wait for checkpoint
if durability_ == "sync":
await cast(asyncio.Future, loop._put_checkpoint_fut)
finally: finally:
# ensure waiter doesn't remain pending on cancel/shutdown # ensure waiter doesn't remain pending on cancel/shutdown
if _cleanup_waiter is not None: if _cleanup_waiter is not None:
+108 -11
View File
@@ -54,6 +54,13 @@ from langgraph.pregel import (
NodeBuilder, NodeBuilder,
Pregel, Pregel,
) )
from langgraph.pregel._checkpoint_writer import (
CHECKPOINT_BACKLOG_ENV_VAR,
DEFAULT_CHECKPOINT_BACKLOG,
AsyncCheckpointWriter,
SyncCheckpointWriter,
resolve_checkpoint_backlog,
)
from langgraph.pregel._loop import SyncPregelLoop from langgraph.pregel._loop import SyncPregelLoop
from langgraph.pregel._runner import PregelRunner from langgraph.pregel._runner import PregelRunner
from langgraph.types import ( from langgraph.types import (
@@ -1329,20 +1336,22 @@ def test_imp_nested(
} }
thread1 = {"configurable": {"thread_id": "1"}} thread1 = {"configurable": {"thread_id": "1"}}
assert [*graph.stream([0, 1], thread1, durability=durability)] == [ result = [*graph.stream([0, 1], thread1, durability=durability)]
{"submapper": "0"}, # nested tasks run concurrently so output order is non-deterministic
assert sorted(result[:-1], key=lambda d: str(d)) == [
{"mapper": "00"}, {"mapper": "00"},
{"submapper": "1"},
{"mapper": "11"}, {"mapper": "11"},
{ {"submapper": "0"},
"__interrupt__": ( {"submapper": "1"},
Interrupt(
value="question",
id=AnyStr(),
),
)
},
] ]
assert result[-1] == {
"__interrupt__": (
Interrupt(
value="question",
id=AnyStr(),
),
)
}
assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [ assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [
"00answera", "00answera",
@@ -3724,6 +3733,7 @@ def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
"end": END, "end": END,
}, },
) )
workflow.add_conditional_edges( workflow.add_conditional_edges(
"Chart Generator", "Chart Generator",
router, router,
@@ -3747,6 +3757,93 @@ def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
def test_sync_durability_applies_checkpoint_backpressure() -> None:
first_put_started = threading.Event()
release_first_put = threading.Event()
put_calls = 0
visited: list[int] = []
result: dict[str, Any] = {}
error: dict[str, BaseException] = {}
class SlowFirstPutCheckpointer(InMemorySaver):
def put(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: Any,
) -> RunnableConfig:
nonlocal put_calls
put_calls += 1
if put_calls == 1:
first_put_started.set()
release_first_put.wait()
return super().put(config, checkpoint, metadata, new_versions)
class State(TypedDict):
counter: int
def increment(state: State) -> State:
visited.append(state["counter"])
return {"counter": state["counter"] + 1}
def should_continue(state: State) -> str:
return "loop" if state["counter"] < 4 else "done"
builder = StateGraph(State)
builder.add_node("increment", increment)
builder.add_edge(START, "increment")
builder.add_conditional_edges(
"increment", should_continue, {"loop": "increment", "done": END}
)
graph = builder.compile(checkpointer=SlowFirstPutCheckpointer())
def invoke() -> None:
try:
result["value"] = graph.invoke(
{"counter": 0},
{"configurable": {"thread_id": "1"}},
durability="async",
)
except BaseException as exc:
error["value"] = exc
thread = threading.Thread(target=invoke)
thread.start()
assert first_put_started.wait(timeout=1)
time.sleep(0.05)
assert thread.is_alive()
assert len(visited) <= DEFAULT_CHECKPOINT_BACKLOG + 1
release_first_put.set()
thread.join(timeout=1)
assert not thread.is_alive()
assert "value" not in error
assert result["value"] == {"counter": 4}
def test_checkpoint_backlog_uses_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv(CHECKPOINT_BACKLOG_ENV_VAR, "7")
assert resolve_checkpoint_backlog() == 7
assert SyncCheckpointWriter(lambda *_args: None).queue.maxsize == 7
assert AsyncCheckpointWriter(lambda *_args: None).queue.maxsize == 7
def test_checkpoint_backlog_invalid_env_uses_default(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv(CHECKPOINT_BACKLOG_ENV_VAR, "not-an-int")
assert resolve_checkpoint_backlog() == DEFAULT_CHECKPOINT_BACKLOG
monkeypatch.setenv(CHECKPOINT_BACKLOG_ENV_VAR, "0")
assert resolve_checkpoint_backlog() == DEFAULT_CHECKPOINT_BACKLOG
def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None: def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None:
"""This test verifies that a run's configurable fields are merged with the """This test verifies that a run's configurable fields are merged with the
previous checkpoint config for each step in the run. previous checkpoint config for each step in the run.
+56
View File
@@ -53,6 +53,7 @@ from langgraph.func import entrypoint, task
from langgraph.graph import END, START, StateGraph from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import MessagesState, add_messages from langgraph.graph.message import MessagesState, add_messages
from langgraph.pregel import NodeBuilder, Pregel from langgraph.pregel import NodeBuilder, Pregel
from langgraph.pregel._checkpoint_writer import DEFAULT_CHECKPOINT_BACKLOG
from langgraph.pregel._loop import AsyncPregelLoop from langgraph.pregel._loop import AsyncPregelLoop
from langgraph.pregel._runner import PregelRunner from langgraph.pregel._runner import PregelRunner
from langgraph.types import ( from langgraph.types import (
@@ -484,6 +485,61 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
assert False, "Task should be cancelled" assert False, "Task should be cancelled"
async def test_async_durability_applies_checkpoint_backpressure() -> None:
first_put_started = asyncio.Event()
release_first_put = asyncio.Event()
put_calls = 0
visited: list[int] = []
class SlowFirstPutCheckpointer(InMemorySaver):
async def aput(
self,
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
nonlocal put_calls
put_calls += 1
if put_calls == 1:
first_put_started.set()
await release_first_put.wait()
return await super().aput(config, checkpoint, metadata, new_versions)
class State(TypedDict):
counter: int
def increment(state: State) -> State:
visited.append(state["counter"])
return {"counter": state["counter"] + 1}
def should_continue(state: State) -> str:
return "loop" if state["counter"] < 4 else "done"
builder = StateGraph(State)
builder.add_node("increment", increment)
builder.add_edge(START, "increment")
builder.add_conditional_edges(
"increment", should_continue, {"loop": "increment", "done": END}
)
graph = builder.compile(checkpointer=SlowFirstPutCheckpointer())
task = asyncio.create_task(
graph.ainvoke(
{"counter": 0}, {"configurable": {"thread_id": "1"}}, durability="async"
)
)
await first_put_started.wait()
await asyncio.sleep(0.05)
assert not task.done()
assert len(visited) <= DEFAULT_CHECKPOINT_BACKLOG + 1
release_first_put.set()
assert await task == {"counter": 4}
async def test_node_cancellation_on_external_cancel() -> None: async def test_node_cancellation_on_external_cancel() -> None:
inner_task_cancelled = False inner_task_cancelled = False