WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-21 23:22:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1075 from langchain-ai/nc/19jul/nested-checkpoints

Browse Source
This commit is contained in:
Nuno Campos
2024-07-22 18:41:56 -07:00
committed by GitHub
parent ab80c113ac c268d695f0
commit bd4dccd9cb
14 changed files with 2885 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/channels/manager.py
+7 -3
View File
@@ -1,6 +1,6 @@
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from datetime import datetime, timezone
from typing import Any, AsyncGenerator, Generator, Mapping
from typing import Any, AsyncGenerator, Generator, Mapping, Optional
from langchain_core.runnables import RunnableConfig
@@ -43,7 +43,11 @@ async def AsyncChannelsManager(
def create_checkpoint(
checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], step: int
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
step: int,
*,
id: Optional[str] = None,
) -> Checkpoint:
"""Create a checkpoint for the given channels."""
ts = datetime.now(timezone.utc).isoformat()
@@ -56,7 +60,7 @@ def create_checkpoint(
return Checkpoint(
v=1,
ts=ts,
id=str(uuid6(clock_seq=step)),
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=checkpoint["channel_versions"],
versions_seen=checkpoint["versions_seen"],
libs/langgraph/langgraph/checkpoint/memory.py
+31 -4
View File
@@ -44,7 +44,7 @@ class MemorySaver(BaseCheckpointSaver):
asyncio.run(coro) # Output: 2
"""
storage: defaultdict[str, dict[str, tuple[bytes, bytes]]]
storage: defaultdict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]
def __init__(
self,
@@ -72,7 +72,7 @@ class MemorySaver(BaseCheckpointSaver):
thread_id = config["configurable"]["thread_id"]
if ts := config["configurable"].get("thread_ts"):
if saved := self.storage[thread_id].get(ts):
checkpoint, metadata = saved
checkpoint, metadata, parent_ts = saved
writes = self.writes[(thread_id, ts)]
return CheckpointTuple(
config=config,
@@ -81,11 +81,19 @@ class MemorySaver(BaseCheckpointSaver):
pending_writes=[
(id, c, self.serde.loads(v)) for id, c, v in writes
],
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
)
else:
if checkpoints := self.storage[thread_id]:
ts = max(checkpoints.keys())
checkpoint, metadata = checkpoints[ts]
checkpoint, metadata, parent_ts = checkpoints[ts]
writes = self.writes[(thread_id, ts)]
return CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
@@ -94,6 +102,14 @@ class MemorySaver(BaseCheckpointSaver):
pending_writes=[
(id, c, self.serde.loads(v)) for id, c, v in writes
],
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
)
def list(
@@ -120,7 +136,9 @@ class MemorySaver(BaseCheckpointSaver):
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
for thread_id in thread_ids:
for ts, (checkpoint, metadata_b) in self.storage[thread_id].items():
for ts, (checkpoint, metadata_b, parent_ts) in sorted(
self.storage[thread_id].items(), key=lambda x: x[0], reverse=True
):
# filter by thread_ts
if before and ts >= before["configurable"]["thread_ts"]:
continue
@@ -143,6 +161,14 @@ class MemorySaver(BaseCheckpointSaver):
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=metadata,
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
)
def put(
@@ -169,6 +195,7 @@ class MemorySaver(BaseCheckpointSaver):
checkpoint["id"]: (
self.serde.dumps(checkpoint),
self.serde.dumps(metadata),
config["configurable"].get("thread_ts"), # parent
)
}
)
libs/langgraph/langgraph/constants.py
+11 -1
View File
@@ -3,9 +3,19 @@ from typing import Any
INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
CONFIG_KEY_RESUMING = "__pregel_resuming"
INTERRUPT = "__interrupt__"
TASKS = "__pregel_tasks"
RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT}
RESERVED = {
INTERRUPT,
TASKS,
CONFIG_KEY_SEND,
CONFIG_KEY_READ,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_RESUMING,
INPUT,
}
TAG_HIDDEN = "langsmith:hidden"
START = "__start__"
libs/langgraph/langgraph/errors.py
+12
View File
@@ -28,3 +28,15 @@ class InvalidUpdateError(Exception):
"""Raised when attempting to update a channel with an invalid sequence of updates."""
pass
class GraphInterrupt(Exception):
"""Raised when a subgraph is interrupted."""
pass
class EmptyInputError(Exception):
"""Raised when graph receives an empty input."""
pass
libs/langgraph/langgraph/pregel/__init__.py
+25 -4
View File
@@ -63,7 +63,9 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.constants import (
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
INTERRUPT,
)
@@ -281,7 +283,13 @@ class Pregel(
)
)
# these are provided by the Pregel class
if spec.id not in [CONFIG_KEY_READ, CONFIG_KEY_SEND]
if spec.id
not in [
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_RESUMING,
]
]
@property
@@ -699,6 +707,7 @@ class Pregel(
Union[str, Sequence[str]],
Optional[Sequence[str]],
Optional[Sequence[str]],
Optional[BaseCheckpointSaver],
]:
debug = debug if debug is not None else self.debug
if output_keys is None:
@@ -710,15 +719,24 @@ class Pregel(
stream_mode = stream_mode if stream_mode is not None else self.stream_mode
if not isinstance(stream_mode, list):
stream_mode = [stream_mode]
if config is not None and config.get("configurable", {}).get(CONFIG_KEY_READ):
if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None:
# if being called as a node in another graph, always use values mode
stream_mode = ["values"]
if config is not None and config.get("configurable", {}).get(
CONFIG_KEY_CHECKPOINTER
):
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
CONFIG_KEY_CHECKPOINTER
]
else:
checkpointer = self.checkpointer
return (
debug,
stream_mode,
output_keys,
interrupt_before,
interrupt_after,
checkpointer,
)
def stream(
@@ -820,6 +838,7 @@ class Pregel(
output_keys,
interrupt_before,
interrupt_after,
checkpointer,
) = self._defaults(
config,
stream_mode=stream_mode,
@@ -830,7 +849,7 @@ class Pregel(
)
with SyncPregelLoop(
input, config=config, checkpointer=self.checkpointer, graph=self
input, config=config, checkpointer=checkpointer, graph=self
) as loop:
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
@@ -1063,6 +1082,7 @@ class Pregel(
output_keys,
interrupt_before,
interrupt_after,
checkpointer,
) = self._defaults(
config,
stream_mode=stream_mode,
@@ -1072,7 +1092,7 @@ class Pregel(
debug=debug,
)
async with AsyncPregelLoop(
input, config=config, checkpointer=self.checkpointer, graph=self
input, config=config, checkpointer=checkpointer, graph=self
) as loop:
aioloop = asyncio.get_event_loop()
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1201,6 +1221,7 @@ class Pregel(
read_channels(loop.channels, output_keys)
)
except BaseException as e:
# TODO use on_chain_end if exc is GraphInterrupt
await asyncio.shield(run_manager.on_chain_error(e))
raise
libs/langgraph/langgraph/pregel/algo.py
+21 -1
View File
@@ -26,9 +26,11 @@ from langchain_core.runnables.config import (
from langgraph.channels.base import BaseChannel
from langgraph.channels.context import Context
from langgraph.channels.manager import ChannelsManager, create_checkpoint
from langgraph.checkpoint.base import Checkpoint, copy_checkpoint
from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, copy_checkpoint
from langgraph.constants import (
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
INTERRUPT,
RESERVED,
@@ -213,6 +215,8 @@ def prepare_next_tasks(
config: RunnableConfig,
step: int,
for_execution: Literal[False],
is_resuming: bool = False,
checkpointer: Literal[None] = None,
manager: Literal[None] = None,
) -> list[PregelTaskDescription]:
...
@@ -227,6 +231,8 @@ def prepare_next_tasks(
config: RunnableConfig,
step: int,
for_execution: Literal[True],
is_resuming: bool,
checkpointer: Optional[BaseCheckpointSaver],
manager: Union[None, ParentRunManager, AsyncParentRunManager],
) -> list[PregelExecutableTask]:
...
@@ -241,6 +247,8 @@ def prepare_next_tasks(
step: int,
*,
for_execution: bool,
is_resuming: bool = False,
checkpointer: Optional[BaseCheckpointSaver] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
@@ -291,6 +299,8 @@ def prepare_next_tasks(
PregelTaskWrites(packet.node, writes, triggers),
config,
),
# in Send we can't checkpoint nested graphs
# as they could be running in parallel
},
),
triggers,
@@ -332,6 +342,12 @@ def prepare_next_tasks(
"langgraph_task_idx": len(tasks),
}
task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata)))
if parent_thread_id := config.get("configurable", {}).get(
"thread_id"
):
thread_id: Optional[str] = f"{parent_thread_id}-{name}"
else:
thread_id = None
writes = deque()
tasks.append(
PregelExecutableTask(
@@ -363,6 +379,10 @@ def prepare_next_tasks(
PregelTaskWrites(name, writes, triggers),
config,
),
CONFIG_KEY_CHECKPOINTER: checkpointer,
CONFIG_KEY_RESUMING: is_resuming,
"thread_id": thread_id,
"thread_ts": checkpoint["id"],
},
),
triggers,
libs/langgraph/langgraph/pregel/debug.py
+8 -6
View File
@@ -70,14 +70,15 @@ def map_debug_tasks(
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
metadata = config["metadata"].copy()
metadata.pop("thread_ts", None)
yield {
"type": "task",
"timestamp": ts,
"step": step,
"payload": {
"id": str(
uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"])))
),
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
"name": name,
"input": input,
"triggers": triggers,
@@ -95,14 +96,15 @@ def map_debug_task_results(
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
metadata = config["metadata"].copy()
metadata.pop("thread_ts", None)
yield {
"type": "task_result",
"timestamp": ts,
"step": step,
"payload": {
"id": str(
uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"])))
),
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
"name": name,
"result": [w for w in writes if w[0] in stream_channels_list],
},
libs/langgraph/langgraph/pregel/executor.py
+88 -49
View File
@@ -1,14 +1,14 @@
import asyncio
import concurrent.futures
import sys
from contextlib import contextmanager
from contextlib import ExitStack
from contextvars import copy_context
from types import TracebackType
from typing import (
AsyncContextManager,
Awaitable,
Callable,
Iterator,
ContextManager,
Optional,
Protocol,
TypeVar,
@@ -18,6 +18,8 @@ from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.config import get_executor_for_config
from typing_extensions import ParamSpec
from langgraph.errors import GraphInterrupt
P = ParamSpec("P")
T = TypeVar("T")
@@ -34,41 +36,63 @@ class Submit(Protocol[P, T]):
...
@contextmanager
def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]:
tasks: dict[concurrent.futures.Future, bool] = {}
with get_executor_for_config(config) as executor:
class BackgroundExecutor(ContextManager):
def __init__(self, config: RunnableConfig) -> None:
self.stack = ExitStack()
self.executor = self.stack.enter_context(get_executor_for_config(config))
self.tasks: dict[concurrent.futures.Future, bool] = {}
def done(task: concurrent.futures.Future) -> None:
try:
task.result()
except BaseException:
pass
else:
tasks.pop(task)
def submit(
fn: Callable[P, T],
*args: P.args,
__name__: Optional[str] = None, # currently not used in sync version
__cancel_on_exit__: bool = False,
**kwargs: P.kwargs,
) -> concurrent.futures.Future:
task = executor.submit(fn, *args, **kwargs)
tasks[task] = __cancel_on_exit__
task.add_done_callback(done)
return task
def submit(
self,
fn: Callable[P, T],
*args: P.args,
__name__: Optional[str] = None, # currently not used in sync version
__cancel_on_exit__: bool = False,
**kwargs: P.kwargs,
) -> concurrent.futures.Future[T]:
task = self.executor.submit(fn, *args, **kwargs)
self.tasks[task] = __cancel_on_exit__
task.add_done_callback(self.done)
return task
def done(self, task: concurrent.futures.Future) -> None:
try:
yield submit
finally:
for task, cancel in tasks.items():
if cancel:
task.cancel()
# executor waits for all tasks to finish on exit
for task in tasks:
# the first task to have raised an exception will be re-raised here
task.result()
task.result()
except GraphInterrupt:
# This exception is an interruption signal, not an error
# so we don't want to re-raise it on exit
self.tasks.pop(task)
except BaseException:
pass
else:
self.tasks.pop(task)
def __enter__(self) -> "submit":
return self.submit
def __exit__(
self,
exc_type: Optional[type[BaseException]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
# cancel all tasks that should be cancelled
for task, cancel in self.tasks.items():
if cancel:
task.cancel()
# wait for all tasks to finish
if tasks := {t for t in self.tasks if not t.done()}:
concurrent.futures.wait(tasks)
# shutdown the executor
self.stack.__exit__(exc_type, exc_value, traceback)
# re-raise the first exception that occurred in a task
if exc_type is None:
# if there's already an exception being raised, don't raise another one
for task in self.tasks:
try:
task.result()
except concurrent.futures.CancelledError:
pass
class AsyncBackgroundExecutor(AsyncContextManager):
@@ -97,24 +121,39 @@ class AsyncBackgroundExecutor(AsyncContextManager):
def done(self, task: asyncio.Task) -> None:
try:
task.result()
except GraphInterrupt:
# This exception is an interruption signal, not an error
# so we don't want to re-raise it on exit
self.tasks.pop(task)
except BaseException:
pass
else:
self.tasks.pop(task)
async def __aenter__(self) -> "submit":
async def __aenter__(self) -> Submit:
return self.submit
async def exit(self) -> None:
fut = asyncio.gather(*self.tasks, return_exceptions=True)
try:
rtns = await asyncio.shield(fut)
finally:
del self.tasks
for rtn in rtns:
# if this is ever changed to BaseException, need to ignore CancelledError
if isinstance(rtn, Exception):
raise rtn
async def exit(
self,
exc_type: Optional[type[BaseException]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> None:
# cancel all tasks that should be cancelled
for task, cancel in self.tasks.items():
if cancel:
task.cancel(self.sentinel)
# wait for all tasks to finish
if self.tasks:
await asyncio.wait(self.tasks)
# re-raise the first exception that occurred in a task
if exc_type is None:
# if there's already an exception being raised, don't raise another one
for task in self.tasks:
try:
task.result()
except asyncio.CancelledError:
pass
async def __aexit__(
self,
@@ -122,8 +161,8 @@ class AsyncBackgroundExecutor(AsyncContextManager):
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
for task, cancel in self.tasks.items():
if cancel:
task.cancel(self.sentinel)
# we cannot use `await` outside of asyncio.shield, as this code can run
# after owning task is cancelled, so pulling async logic to separate method
# wait for all background tasks to finish, shielded from cancellation
await asyncio.shield(self.exit())
await asyncio.shield(self.exit(exc_type, exc_value, traceback))
libs/langgraph/langgraph/pregel/loop.py
+141 -52
View File
@@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
from collections import deque
from contextlib import AsyncExitStack, ExitStack
from types import TracebackType
@@ -38,7 +39,8 @@ from langgraph.checkpoint.base import (
copy_checkpoint,
empty_checkpoint,
)
from langgraph.constants import INPUT, INTERRUPT
from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_RESUMING, INPUT, INTERRUPT
from langgraph.errors import EmptyInputError, GraphInterrupt
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValueMapping,
@@ -66,6 +68,8 @@ if TYPE_CHECKING:
V = TypeVar("V")
INPUT_DONE = object()
INPUT_RESUMING = object()
EMPTY_SEQ = ()
class PregelLoop:
@@ -76,8 +80,16 @@ class PregelLoop:
checkpointer_put_writes: Optional[
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
]
checkpointer_put: Optional[
Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any]
_checkpointer_put_after_previous: Optional[
Callable[
[
Optional[concurrent.futures.Future],
RunnableConfig,
Sequence[tuple[str, Any]],
str,
],
Any,
]
]
graph: "Pregel"
@@ -95,9 +107,27 @@ class PregelLoop:
]
tasks: Sequence[PregelExecutableTask]
stream: deque[Tuple[str, Any]]
is_nested: bool
# public
def __init__(
self,
input: Optional[Any],
*,
config: RunnableConfig,
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
self.input = input
self.config = config
self.checkpointer = checkpointer
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
"""Mark tasks as scheduled, to be used by queue-based executors."""
raise NotImplementedError
@@ -122,9 +152,9 @@ class PregelLoop:
def tick(
self,
*,
output_keys: Union[str, Sequence[str]] = None,
interrupt_after: Optional[Sequence[str]] = None,
interrupt_before: Optional[Sequence[str]] = None,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
) -> bool:
"""Execute a single iteration of the Pregel loop.
@@ -133,7 +163,7 @@ class PregelLoop:
if self.status != "pending":
raise RuntimeError("Cannot tick when status is no longer 'pending'")
if self.input is not INPUT_DONE:
if self.input not in (INPUT_DONE, INPUT_RESUMING):
self._first()
elif all(task.writes for task in self.tasks):
writes = [w for t in self.tasks for w in t.writes]
@@ -165,7 +195,10 @@ class PregelLoop:
# after execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
self.status = "interrupt_after"
return False
if self.is_nested:
raise GraphInterrupt(self)
else:
return False
else:
return False
@@ -184,6 +217,8 @@ class PregelLoop:
self.step,
for_execution=True,
manager=manager,
checkpointer=self.checkpointer,
is_resuming=self.input is INPUT_RESUMING,
)
# if no more tasks, we're done
@@ -199,12 +234,20 @@ class PregelLoop:
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks):
return self.tick()
return self.tick(
output_keys=output_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
)
# before execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
return False
if self.is_nested:
raise GraphInterrupt()
else:
return False
# produce debug output
self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks))
@@ -214,8 +257,23 @@ class PregelLoop:
# private
def _first(self) -> None:
# resuming from previous checkpoint requires
# - finding a previous checkpoint
# - receiving None input (outer graph) or RESUMING flag (subgraph)
is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
self.config.get("configurable", {}).get(CONFIG_KEY_RESUMING)
or self.input is None
)
# proceed past previous checkpoint
if is_resuming:
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
for k in self.channels:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
# map inputs to channel updates
if input_writes := deque(map_input(self.graph.input_channels, self.input)):
elif input_writes := deque(map_input(self.graph.input_channels, self.input)):
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
@@ -237,31 +295,33 @@ class PregelLoop:
# save input checkpoint
self._put_checkpoint({"source": "input", "writes": self.input})
else:
# no input is taken as signal to proceed past previous interrupt
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
for k in self.channels:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
raise EmptyInputError(f"Received no input for {self.graph.input_channels}")
# done with input
self.input = INPUT_DONE
self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
def _put_checkpoint(
self,
metadata: CheckpointMetadata,
) -> None:
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
# assign step
metadata["step"] = self.step
# bail if no checkpointer
if self.checkpointer_put is not None:
if self._checkpointer_put_after_previous is not None:
# create new checkpoint
self.checkpoint_metadata = metadata
self.checkpoint = create_checkpoint(
self.checkpoint, self.channels, self.step
self.checkpoint,
self.channels,
self.step,
# child graphs keep at most one checkpoint per parent checkpoint
# this is achieved by writing child checkpoints as progress is made
# (so that error recovery / resuming from interrupt don't lose work)
# but doing so always with an id equal to that of the parent checkpoint
id=self.config["configurable"]["thread_ts"] if self.is_nested else None,
)
# save it, without blocking
self.submit(
self.checkpointer_put,
# if there's a previous checkpoint save in progress, wait for it
# ensuring checkpointers receive checkpoints in order
self._put_checkpoint_fut = self.submit(
self._checkpointer_put_after_previous,
getattr(self, "_put_checkpoint_fut", None),
self.checkpoint_config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
@@ -287,6 +347,15 @@ class PregelLoop:
# increment step
self.step += 1
def _suppress_interrupt(
self,
exc_type: Optional[Type[BaseException]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
if exc_type is GraphInterrupt and not self.is_nested:
return True
class SyncPregelLoop(PregelLoop, ContextManager):
def __init__(
@@ -297,19 +366,29 @@ class SyncPregelLoop(PregelLoop, ContextManager):
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
self.stack = ExitStack()
self.input = input
self.config = config
self.checkpointer = checkpointer
self.checkpointer_get_next_version = (
checkpointer.get_next_version if checkpointer else increment
)
self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
self.checkpointer_put = checkpointer.put if checkpointer else None
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
self.stack.push(self._suppress_interrupt)
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.put_writes
else:
self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None
self.checkpointer_put_writes = None
def _checkpointer_put_after_previous(
self,
prev: Optional[concurrent.futures.Future],
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
try:
if prev is not None:
prev.result()
finally:
self.checkpointer.put(config, checkpoint, metadata)
# context manager
@@ -349,6 +428,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
# unwind stack
del self.graph
return self.stack.__exit__(exc_type, exc_value, traceback)
@@ -362,21 +442,29 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
) -> None:
self.stream = deque()
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
self.stack = AsyncExitStack()
self.input = input
self.config = config
self.checkpointer = checkpointer
self.checkpointer_get_next_version = (
checkpointer.get_next_version if checkpointer else increment
)
self.checkpointer_put_writes = (
checkpointer.aput_writes if checkpointer else None
)
self.checkpointer_put = checkpointer.aput if checkpointer else None
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
self.stack.push(self._suppress_interrupt)
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.aput_writes
else:
self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None
self.checkpointer_put_writes = None
async def _checkpointer_put_after_previous(
self,
prev: Optional[asyncio.Task],
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
) -> RunnableConfig:
try:
if prev is not None:
await prev
finally:
await self.checkpointer.aput(config, checkpoint, metadata)
# context manager
@@ -418,6 +506,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
# unwind stack
del self.graph
return await asyncio.shield(
self.stack.__aexit__(exc_type, exc_value, traceback)
libs/langgraph/poetry.lock
Generated
+15 -1
View File
@@ -2783,6 +2783,20 @@ pytest = ">=6.2.5"
[package.extras]
dev = ["pre-commit", "pytest-asyncio", "tox"]
[[package]]
name = "pytest-repeat"
version = "0.9.3"
description = "pytest plugin for repeating tests"
optional = false
python-versions = ">=3.7"
files = [
{file = "pytest_repeat-0.9.3-py3-none-any.whl", hash = "sha256:26ab2df18226af9d5ce441c858f273121e92ff55f5bb311d25755b8d7abdd8ed"},
{file = "pytest_repeat-0.9.3.tar.gz", hash = "sha256:ffd3836dfcd67bb270bec648b330e20be37d2966448c4148c4092d1e8aba8185"},
]
[package.dependencies]
pytest = "*"
[[package]]
name = "pytest-watcher"
version = "0.4.2"
@@ -4165,4 +4179,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "0d877d3879473de43aca1e1d36a8f420ff3f4b140807cb5cc24935d9114947be"
content-hash = "18b26895b05f2f7cdcd08d59ba164ac788e0e8005e3ec5d7089469b4f3a96aea"
libs/langgraph/pyproject.toml
+2 -1
View File
@@ -32,6 +32,7 @@ langchain-openai = ">=0.1.2"
langchain-anthropic = ">=0.1.8"
dataclasses-json = "^0.6.7"
pytest-xdist = {extras = ["psutil"], version = "^3.6.1"}
pytest-repeat = "^0.9.3"
[tool.poetry.group.dev]
optional = true
@@ -62,7 +63,7 @@ omit = ["tests/*"]
[tool.pytest-watcher]
now = true
delay = 0.1
runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
patterns = ["*.py"]
[build-system]
libs/langgraph/tests/memory_assert.py
+8 -1
View File
@@ -30,9 +30,11 @@ class MemorySaverAssertImmutable(MemorySaver):
self,
*,
serde: Optional[SerializerProtocol] = None,
put_sleep: Optional[float] = None,
) -> None:
super().__init__(serde=serde)
self.storage_for_copies = defaultdict(dict)
self.put_sleep = put_sleep
def put(
self,
@@ -40,6 +42,10 @@ class MemorySaverAssertImmutable(MemorySaver):
checkpoint: Checkpoint,
metadata: Optional[CheckpointMetadata] = None,
) -> None:
if self.put_sleep:
import time
time.sleep(self.put_sleep)
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
if saved := super().get(config):
@@ -85,7 +91,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
configurable = config["configurable"].copy()
# remove thread_ts to make testing simpler
configurable.pop("thread_ts", None)
thread_ts = configurable.pop("thread_ts", None)
self.storage[config["configurable"]["thread_id"]].update(
{
@@ -93,6 +99,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
self.serde.dumps(checkpoint),
# merge configurable fields and metadata
self.serde.dumps({**configurable, **metadata}),
thread_ts,
)
}
)
libs/langgraph/tests/test_pregel.py
+1248 -76
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+1268 -75
View File
File diff suppressed because it is too large Load Diff