diff --git a/README.md b/README.md index 60611839c..e74f0b270 100644 --- a/README.md +++ b/README.md @@ -81,11 +81,11 @@ Check `examples` for more examples. - [ ] Add tests for subscribe_to_each().join() - [x] Add optional debug logging - [ ] Add an optional Diff value for Channels that implements `__add__`, returned by update(), yielded by Pregel for output channels. Add replacing_keys set to AddableDict. use an addabledict for yielding values. channels that dont implement it get marked with replacing_keys -- [ ] Implement checkpointing - - [ ] Use langchain.load dumps/loads functions (or use pickle?) - - [ ] Save checkpoints at end of each step - - [ ] Load checkpoint at start of invocation - - [ ] API to specify storage backend and save key +- [x] Implement checkpointing + - [x] Save checkpoints at end of each step/run + - [x] Load checkpoint at start of invocation + - [x] API to specify storage backend and save key + - [x] Tests - [ ] Add more examples - [ ] multi agent simulation - [ ] human in the loop @@ -93,6 +93,7 @@ Check `examples` for more examples. - [ ] agent executor (add current v total iterations info to read/write steps to enable doing a final update at the end) - [ ] run over dataset - [ ] Fault tolerance + - [ ] Expose a unique id to each step, hash of (app, chain, checkpoint) (include input updates for first step) - [ ] Retry individual processes in a step - [ ] Retry entire step? - [ ] Pregel.stream_log to contain additional keys specific to Pregel diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py index ef6ae43e5..c735d32c8 100644 --- a/permchain/checkpoint/base.py +++ b/permchain/checkpoint/base.py @@ -36,9 +36,9 @@ class BaseCheckpointAdapter(Serializable, ABC): ... async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None: - return asyncio.get_running_loop().run_in_executor(None, self.get, config) + return await asyncio.get_running_loop().run_in_executor(None, self.get, config) async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None: - return asyncio.get_running_loop().run_in_executor( + return await asyncio.get_running_loop().run_in_executor( None, self.put, config, checkpoint ) diff --git a/permchain/checkpoint/memory.py b/permchain/checkpoint/memory.py new file mode 100644 index 000000000..12d3a2932 --- /dev/null +++ b/permchain/checkpoint/memory.py @@ -0,0 +1,29 @@ +from typing import Any, Dict, Mapping, Sequence + +from langchain.pydantic_v1 import Field +from langchain.schema.runnable import RunnableConfig +from langchain.schema.runnable.utils import ConfigurableFieldSpec + +from permchain.checkpoint.base import BaseCheckpointAdapter + + +class MemoryCheckpoint(BaseCheckpointAdapter): + storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict) + + @property + def config_specs(self) -> Sequence[ConfigurableFieldSpec]: + return [ + ConfigurableFieldSpec( + id="thread_id", + annotation=str, + name="Thread ID", + description=None, + default="", + ), + ] + + def get(self, config: RunnableConfig) -> Mapping[str, Any] | None: + return self.storage.get(config["configurable"]["thread_id"], None) + + def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None: + return self.storage.update({config["configurable"]["thread_id"]: checkpoint}) diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py index abb693e63..66a20367c 100644 --- a/permchain/pregel/read.py +++ b/permchain/pregel/read.py @@ -89,6 +89,9 @@ class ChannelInvoke(RunnableBindingBase): ) def join(self, channels: Sequence[str]) -> ChannelInvoke: + assert isinstance(channels, list) or isinstance( + channels, tuple + ), "channels must be a list or tuple" joiner = RunnablePassthrough.assign( **{chan: ChannelRead(chan) for chan in channels} ) diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index afa51834e..927a78021 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -2,8 +2,14 @@ from typing import Any, Mapping, Sequence from permchain.channels.base import BaseChannel from permchain.channels.last_value import LastValue +from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION from permchain.pregel.read import ChannelBatch, ChannelInvoke +FORBIDDEN_CHANNEL_NAMES = { + CHECKPOINT_KEY_TS, + CHECKPOINT_KEY_VERSION, +} + def validate_chains_channels( chains: Mapping[str, ChannelInvoke | ChannelBatch], @@ -47,3 +53,7 @@ def validate_chains_channels( for chan in output: if chan not in channels: channels[chan] = LastValue(Any) + + for name in FORBIDDEN_CHANNEL_NAMES: + if name in channels: + raise ValueError(f"Channel name {name} is reserved") diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 619221270..d129c016e 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1,3 +1,4 @@ +import operator import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager @@ -9,9 +10,11 @@ from pytest_mock import MockerFixture from permchain import Channel, Pregel from permchain.channels.base import InvalidUpdateError +from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.context import Context from permchain.channels.last_value import LastValue from permchain.channels.topic import Topic +from permchain.checkpoint.memory import MemoryCheckpoint def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -221,6 +224,44 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] +def test_invoke_checkpoint(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + + def raise_if_above_10(input: int) -> int: + if input > 10: + raise ValueError("Input is too large") + return input + + chain_one = ( + Channel.subscribe_to(["input"]).join(["total"]) + | add_one + | Channel.write_to("output", "total") + | raise_if_above_10 + ) + + app = Pregel( + chains={"chain_one": chain_one}, + channels={"total": BinaryOperatorAggregate(int, operator.add)}, + checkpoint=MemoryCheckpoint(), + ) + + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 + assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 2 + # total is now 2, so output is 2+3=5 + assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 + assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7 + assert app.checkpoint.get({"configurable": {"thread_id": "2"}}).get("total") == 5 + + def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index b191b7438..091a11234 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1,4 +1,5 @@ import asyncio +import operator from contextlib import asynccontextmanager, contextmanager from typing import Any, AsyncGenerator, AsyncIterator, Generator @@ -8,9 +9,11 @@ from pytest_mock import MockerFixture from permchain import Channel, Pregel from permchain.channels.base import InvalidUpdateError +from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.context import Context from permchain.channels.last_value import LastValue from permchain.channels.topic import Topic +from permchain.checkpoint.memory import MemoryCheckpoint async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -230,6 +233,54 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] +async def test_invoke_checkpoint(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + + def raise_if_above_10(input: int) -> int: + if input > 10: + raise ValueError("Input is too large") + return input + + chain_one = ( + Channel.subscribe_to(["input"]).join(["total"]) + | add_one + | Channel.write_to("output", "total") + | raise_if_above_10 + ) + + app = Pregel( + chains={"chain_one": chain_one}, + channels={"total": BinaryOperatorAggregate(int, operator.add)}, + checkpoint=MemoryCheckpoint(), + ) + + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 + assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get( + "total" + ) == 2 + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get( + "total" + ) == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get( + "total" + ) == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 + assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get( + "total" + ) == 7 + assert (await app.checkpoint.aget({"configurable": {"thread_id": "2"}})).get( + "total" + ) == 5 + + async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))