From 611ecdb1cdac1c6f50f6f19b67db18b0c48d384c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 6 May 2024 08:36:20 -0700 Subject: [PATCH] Checkpoint inputs before starting the first step for easier error recovery - this enables easier retrying, for any error just do .invoke(None, config) no matter which step the error happened on --- langgraph/checkpoint/sqlite.py | 12 ++- langgraph/pregel/__init__.py | 132 +++++++++++++++++++++++++-------- langgraph/pregel/types.py | 4 +- tests/test_pregel.py | 47 ++++++++---- tests/test_pregel_async.py | 36 +++++++-- 5 files changed, 175 insertions(+), 56 deletions(-) diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 54feff576..bde4a5d0a 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -1,5 +1,6 @@ import pickle import sqlite3 +import threading from contextlib import AbstractContextManager, contextmanager from types import TracebackType from typing import Any, Iterator, Optional @@ -94,6 +95,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): super().__init__(serde=serde) self.conn = conn self.is_setup = False + self.lock = threading.Lock() @classmethod def from_conn_string(cls, conn_string: str) -> "SqliteSaver": @@ -115,7 +117,13 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): memory = SqliteSaver.from_conn_string("checkpoints.sqlite") """ - return SqliteSaver(conn=sqlite3.connect(conn_string)) + return SqliteSaver( + conn=sqlite3.connect( + conn_string, + # https://ricardoanderegg.com/posts/python-sqlite-thread-safety/ + check_same_thread=False, + ) + ) def __enter__(self) -> Self: return self @@ -348,7 +356,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): saved_config ) # Output: {"configurable": {"thread_id": "1", "thread_ts": 2024-05-04T06:32:42.235444+00:00"}} """ - with self.cursor() as cur: + with self.lock, self.cursor() as cur: cur.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)", ( diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index de56f8925..4e5cf13a0 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -342,7 +342,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata, + saved.metadata if saved else None, ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: @@ -361,7 +361,7 @@ class Pregel( read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, - saved.metadata, + saved.metadata if saved else None, ) def get_state_history( @@ -623,6 +623,7 @@ class Pregel( run_id=config.get("run_id"), ) try: + bg: list[concurrent.futures.Future] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -656,6 +657,7 @@ class Pregel( else None ) checkpoint = saved.checkpoint if saved else empty_checkpoint() + start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint with ChannelsManager( self.channels, checkpoint @@ -668,6 +670,27 @@ class Pregel( ) # apply input writes _apply_writes(checkpoint, channels, input_writes) + # save input checkpoint + if self.checkpointer is not None: + checkpoint = create_checkpoint(checkpoint, channels) + bg.append( + executor.submit( + self.checkpointer.put, + checkpoint_config, + copy_checkpoint(checkpoint), + {"source": "input", "step": start}, + ) + ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"][ + "thread_id" + ], + "thread_ts": checkpoint["ts"], + } + } + # increment start to 0 + start += 1 else: # if received no input, take that as signal to proceed # past previous interrupt, if any @@ -681,7 +704,6 @@ class Pregel( # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - start = saved.metadata.get("step", -1) + 1 if saved else 0 stop = start + config["recursion_limit"] + 1 for step in range(start, stop): next_checkpoint, next_tasks = _prepare_next_tasks( @@ -786,21 +808,29 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = self.checkpointer.put( - checkpoint_config, - checkpoint, - {"source": "loop", "step": step}, - ) - if stream_mode == "debug": - yield map_debug_checkpoint( - step, + bg.append( + executor.submit( + self.checkpointer.put, checkpoint_config, - channels, - self.stream_channels_asis, + copy_checkpoint(checkpoint), + {"source": "loop", "step": step}, ) - elif stream_mode == "debug": + ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"][ + "thread_id" + ], + "thread_ts": checkpoint["ts"], + } + } + # yield debug checkpoint + if stream_mode == "debug": yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis + step, + checkpoint_config if self.checkpointer else None, + channels, + self.stream_channels_asis, ) # after execution, check if we should interrupt @@ -824,6 +854,12 @@ class Pregel( task.cancel() except NameError: pass + # wait for all background tasks to finish + done, _ = concurrent.futures.wait( + bg, return_when=concurrent.futures.ALL_COMPLETED + ) + for task in done: + task.result() async def astream( self, @@ -855,7 +891,7 @@ class Pregel( None, ) try: - tasks: list[asyncio.Task] = [] + bg: list[asyncio.Task] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -889,6 +925,7 @@ class Pregel( else None ) checkpoint = saved.checkpoint if saved else empty_checkpoint() + start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint async with AsyncChannelsManager(self.channels, checkpoint) as channels: # map inputs to channel updates @@ -899,6 +936,28 @@ class Pregel( ) # apply input writes _apply_writes(checkpoint, channels, input_writes) + # save input checkpoint + if self.checkpointer is not None: + checkpoint = create_checkpoint(checkpoint, channels) + bg.append( + asyncio.create_task( + self.checkpointer.aput( + checkpoint_config, + copy_checkpoint(checkpoint), + {"source": "input", "step": start}, + ) + ) + ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"][ + "thread_id" + ], + "thread_ts": checkpoint["ts"], + } + } + # increment start to 0 + start += 1 else: # if received no input, take that as signal to proceed # past previous interrupt, if any @@ -1027,21 +1086,30 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) - checkpoint_config = await self.checkpointer.aput( - checkpoint_config, - checkpoint, - {"source": "loop", "step": step}, - ) - if stream_mode == "debug": - yield map_debug_checkpoint( - step, - checkpoint_config, - channels, - self.stream_channels_asis, + bg.append( + asyncio.create_task( + self.checkpointer.aput( + checkpoint_config, + checkpoint, + {"source": "loop", "step": step}, + ) ) - elif stream_mode == "debug": + ) + checkpoint_config = { + "configurable": { + "thread_id": checkpoint_config["configurable"][ + "thread_id" + ], + "thread_ts": checkpoint["ts"], + } + } + # yield debug checkpoint + if stream_mode == "debug": yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis + step, + checkpoint_config if self.checkpointer else None, + channels, + self.stream_channels_asis, ) # after execution, check if we should interrupt @@ -1063,11 +1131,11 @@ class Pregel( try: for task in futures: task.cancel() - tasks.append(task) + bg.append(task) except NameError: pass - # wait for all tasks to finish - await asyncio.gather(*tasks, return_exceptions=True) + # wait for all background tasks to finish + await asyncio.gather(*bg) def invoke( self, diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index 4359b5994..d4dc92d40 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -3,6 +3,8 @@ from typing import Any, Literal, NamedTuple, Optional, Union from langchain_core.runnables import Runnable, RunnableConfig +from langgraph.checkpoint.base import CheckpointMetadata + class PregelTaskDescription(NamedTuple): name: str @@ -25,7 +27,7 @@ class StateSnapshot(NamedTuple): """Nodes to execute in the next step, if any""" config: RunnableConfig """Config used to fetch this snapshot""" - metadata: dict[str, Any] + metadata: CheckpointMetadata """Metadata associated with this snapshot""" parent_config: Optional[RunnableConfig] = None """Config used to fetch the parent snapshot, if any""" diff --git a/tests/test_pregel.py b/tests/test_pregel.py index edccfdd78..bf295801f 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2,6 +2,7 @@ import json import operator import time import warnings +from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import Annotated, Any, Generator, Literal, Optional, TypedDict, Union @@ -672,7 +673,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: if input > 10: @@ -681,7 +682,7 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: one = ( Channel.subscribe_to(["input"]).join(["total"]) - | add_one + | adder | Channel.write_to("output", "total") | raise_if_above_10 ) @@ -701,10 +702,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: thread_1 = {"configurable": {"thread_id": "1"}} # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, thread_1) == 2 + assert app.invoke(2, thread_1, debug=1) == 2 state = app.get_state(thread_1) assert state is not None assert state.values.get("total") == 2 + assert state.next == () assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"] # total is now 2, so output is 2+3=5 assert app.invoke(3, thread_1) == 5 @@ -719,14 +721,22 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: state = app.get_state(thread_1) assert state is not None assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert app.invoke(2, thread_1) == 9 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () thread_2 = {"configurable": {"thread_id": "2"}} # on a new thread, total starts out as 0, so output is 0+5=5 - assert app.invoke(5, thread_2) == 5 + assert app.invoke(5, thread_2, debug=True) == 5 state = app.get_state({"configurable": {"thread_id": "1"}}) assert state is not None - assert state.values.get("total") == 7 - assert state.next == () + assert state.values.get("total") == 16 + assert state.next == (), "checkpoint of other thread not touched" state = app.get_state(thread_2) assert state is not None assert state.values.get("total") == 5 @@ -735,8 +745,12 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: assert len(list(app.get_state_history(thread_1, limit=1))) == 1 # list all checkpoints for thread 1 thread_1_history = [c for c in app.get_state_history(thread_1)] - # there are 2: one for each successful ainvoke() - assert len(thread_1_history) == 2 + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } # sorted descending assert ( thread_1_history[0].config["configurable"]["thread_ts"] @@ -748,10 +762,10 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: ) assert len(cursored) == 1 assert cursored[0].config == thread_1_history[1].config - # the second checkpoint - assert thread_1_history[0].values["total"] == 7 - # the first checkpoint - assert thread_1_history[1].values["total"] == 2 + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 # can get each checkpoint using aget with config assert ( memory.get(thread_1_history[0].config)["ts"] @@ -769,7 +783,14 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: > thread_1_history[0].config["configurable"]["thread_ts"] ) # 1 more checkpoint in history - assert len(list(app.get_state_history(thread_1))) == 3 + assert len(list(app.get_state_history(thread_1))) == 8 + assert Counter( + c.metadata["source"] for c in app.get_state_history(thread_1) + ) == { + "update": 1, + "input": 4, + "loop": 3, + } # the latest checkpoint is the updated one assert app.get_state(thread_1) == app.get_state(thread_1_next_config) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 2a79de197..4a01b1631 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1,6 +1,7 @@ import asyncio import json import operator +from collections import Counter from contextlib import asynccontextmanager, contextmanager from typing import ( Annotated, @@ -712,13 +713,21 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: state = await app.aget_state(thread_1) assert state is not None assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1) == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () thread_2 = {"configurable": {"thread_id": "2"}} # on a new thread, total starts out as 0, so output is 0+5=5 assert await app.ainvoke(5, thread_2) == 5 state = await app.aget_state({"configurable": {"thread_id": "1"}}) assert state is not None - assert state.values.get("total") == 7 + assert state.values.get("total") == 16 assert state.next == () state = await app.aget_state(thread_2) assert state is not None @@ -728,8 +737,12 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 # list all checkpoints for thread 1 thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 2: one for each successful ainvoke() - assert len(thread_1_history) == 2 + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } # sorted descending assert ( thread_1_history[0].config["configurable"]["thread_ts"] @@ -744,10 +757,10 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: ] assert len(cursored) == 1 assert cursored[0].config == thread_1_history[1].config - # the second checkpoint - assert thread_1_history[0].values["total"] == 7 - # the first checkpoint - assert thread_1_history[1].values["total"] == 2 + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 # can get each checkpoint using aget with config assert (await memory.aget(thread_1_history[0].config))[ "ts" @@ -763,7 +776,14 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: > thread_1_history[0].config["configurable"]["thread_ts"] ) # 1 more checkpoint in history - assert len([h async for h in app.aget_state_history(thread_1)]) == 3 + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } # the latest checkpoint is the updated one assert await app.aget_state(thread_1) == await app.aget_state( thread_1_next_config