Add tests for checkpointing

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

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
         )

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})

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}
         )

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")

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))

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))