WIP

permchain/channels/base.py

@@ -14,11 +14,12 @@ from typing import (
 
 from typing_extensions import Self
 
+from permchain.checkpoint.base import Checkpoint
 from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
 
 Value = TypeVar("Value")
 Update = TypeVar("Update")
-Checkpoint = TypeVar("Checkpoint")
+C = TypeVar("C")
 
 
 class EmptyChannelError(Exception):
@@ -34,7 +35,7 @@ class InvalidUpdateError(Exception):
     pass
 
 
-class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
+class BaseChannel(Generic[Value, Update, C], ABC):
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -47,14 +48,12 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
     @contextmanager
     @abstractmethod
-    def empty(
-        self, checkpoint: Optional[Checkpoint] = None
-    ) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: Optional[C] = None) -> Generator[Self, None, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
 
     @asynccontextmanager
     async def aempty(
-        self, checkpoint: Optional[Checkpoint] = None
+        self, checkpoint: Optional[C] = None
     ) -> AsyncGenerator[Self, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
         with self.empty(checkpoint) as value:
@@ -74,7 +73,7 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
         Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
     @abstractmethod
-    def checkpoint(self) -> Checkpoint | None:
+    def checkpoint(self) -> C | None:
         """Return a string representation of the channel's current state.
 
         Raises EmptyChannelError if the channel is empty (never updated yet),
@@ -84,12 +83,11 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 @contextmanager
 def ChannelsManager(
     channels: Mapping[str, BaseChannel],
-    checkpoint: Optional[Mapping[str, Any]],
+    checkpoint: Checkpoint,
 ) -> Generator[Mapping[str, BaseChannel], None, None]:
     """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
     # TODO use https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
-    checkpoint = checkpoint or {}
-    empty = {k: v.empty(checkpoint.get(k)) for k, v in channels.items()}
+    empty = {k: v.empty(checkpoint["values"].get(k)) for k, v in channels.items()}
     try:
         yield {k: v.__enter__() for k, v in empty.items()}
     finally:
@@ -112,15 +110,12 @@ async def AsyncChannelsManager(
             await v.__aexit__(None, None, None)
 
 
-def create_checkpoint(channels: Mapping[str, BaseChannel]) -> Mapping[str, Any]:
+def create_checkpoint(checkpoint, channels: Mapping[str, BaseChannel]) -> Checkpoint:
     """Create a checkpoint for the given channels."""
-    checkpoint = {
-        CHECKPOINT_KEY_VERSION: 1,
-        CHECKPOINT_KEY_TS: datetime.now(timezone.utc).isoformat(),
-    }
+    checkpoint = Checkpoint(checkpoint, v=1, ts=datetime.now(timezone.utc).isoformat())
     for k, v in channels.items():
         try:
-            checkpoint[k] = v.checkpoint()
+            checkpoint["values"][k] = v.checkpoint()
         except EmptyChannelError:
             pass
     return checkpoint

permchain/checkpoint/base.py

@@ -1,14 +1,32 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, Mapping
+from collections import defaultdict
+from typing import Any, Mapping, TypedDict
 
 from langchain.load.serializable import Serializable
+from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import RunnableConfig
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
 from permchain.utils import StrEnum
 
 
+class Checkpoint(TypedDict):
+    v: int
+    ts: str
+    values: dict[str, Any]
+    versions: defaultdict[str, int]
+    seen: defaultdict[str, defaultdict[str, int]]
+
+
+def empty_checkpoint() -> Checkpoint:
+    return Checkpoint(
+        values={},
+        versions=defaultdict(int),
+        seen=defaultdict(lambda: defaultdict(int)),
+    )
+
+
 class CheckpointAt(StrEnum):
     END_OF_STEP = "end_of_step"
     END_OF_RUN = "end_of_run"
@@ -22,17 +40,30 @@ class BaseCheckpointAdapter(Serializable, ABC):
         return []
 
     @abstractmethod
-    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+    def get(self, config: RunnableConfig) -> Checkpoint | None:
         ...
 
     @abstractmethod
-    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
         ...
 
-    async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+    async def aget(self, config: RunnableConfig) -> Checkpoint | None:
         return await asyncio.get_running_loop().run_in_executor(None, self.get, config)
 
-    async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
         return await asyncio.get_running_loop().run_in_executor(
             None, self.put, config, checkpoint
         )
+
+
+class StepState(Serializable):
+    _checkpoint: Checkpoint
+
+    _puts: set[str] = Field(default_factory=set)
+
+    def get(self, key: str, default: Any | None = None) -> Any | None:
+        return self._checkpoint["values"].get(key)
+
+    def put(self, key: str, value: Any) -> None:
+        self._puts.add(key)
+        self._checkpoint["values"][key] = value

permchain/checkpoint/memory.py

@@ -1,14 +1,12 @@
-from typing import Any, Dict, Mapping
-
 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
+from permchain.checkpoint.base import BaseCheckpointAdapter, Checkpoint
 
 
 class MemoryCheckpoint(BaseCheckpointAdapter):
-    storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict)
+    storage: dict[str, Checkpoint] = Field(default_factory=dict)
 
     @property
     def config_specs(self) -> list[ConfigurableFieldSpec]:
@@ -23,8 +21,8 @@ class MemoryCheckpoint(BaseCheckpointAdapter):
             ),
         ]
 
-    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+    def get(self, config: RunnableConfig) -> Checkpoint | None:
         return self.storage.get(config["configurable"]["thread_id"], None)
 
-    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
         return self.storage.update({config["configurable"]["thread_id"]: checkpoint})

permchain/pregel/__init__.py

@@ -32,6 +32,7 @@ from langchain.schema.runnable import (
 from langchain.schema.runnable.base import Input, Output, coerce_to_runnable
 from langchain.schema.runnable.config import (
     RunnableConfig,
+    ensure_config,
     get_executor_for_config,
     patch_config,
 )
@@ -47,7 +48,14 @@ from permchain.channels.base import (
     EmptyChannelError,
     create_checkpoint,
 )
-from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
+from permchain.checkpoint.base import (
+    BaseCheckpointAdapter,
+    Checkpoint,
+    CheckpointAt,
+    StepState,
+    empty_checkpoint,
+)
+from permchain.checkpoint.memory import MemoryCheckpoint
 from permchain.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
 from permchain.pregel.debug import print_checkpoint, print_step_start
 from permchain.pregel.io import map_input, map_output
@@ -135,7 +143,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
     debug: bool = Field(default_factory=get_debug)
 
-    checkpoint: Optional[BaseCheckpointAdapter] = None
+    saver: Optional[BaseCheckpointAdapter] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -151,7 +159,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     def config_specs(self) -> list[ConfigurableFieldSpec]:
         return get_unique_config_specs(
             [spec for chain in self.chains.values() for spec in chain.config_specs]
-            + (self.checkpoint.config_specs if self.checkpoint is not None else [])
+            + (self.saver.config_specs if self.saver is not None else [])
         )
 
     @property
@@ -194,27 +202,27 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         input: Iterator[dict[str, Any] | Any],
         run_manager: CallbackManagerForChainRun,
         config: RunnableConfig,
+        *,
+        saver: Optional[BaseCheckpointAdapter] = None,
     ) -> Iterator[dict[str, Any] | Any]:
         if config["recursion_limit"] < 1:
             raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
+        saver = saver or self.saver
         checkpoint = (
-            self.checkpoint.get(config) if self.checkpoint is not None else None
-        )
+            self.saver.get(config) if self.saver else None
+        ) or empty_checkpoint()
         with ChannelsManager(
             self.channels, checkpoint
         ) as channels, get_executor_for_config(config) as executor:
-            next_tasks = _apply_writes_and_prepare_next_tasks(
-                processes,
+            _apply_writes(
+                checkpoint,
                 channels,
                 deque(w for c in input for w in map_input(self.input, c)),
                 config,
                 0,
             )
 
-            if not next_tasks:
-                return
-
             read = partial(_read_channel, channels)
 
             # Similarly to Bulk Synchronous Parallel / Pregel model
@@ -223,6 +231,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             # channels are guaranteed to be immutable for the duration of the step,
             # with channel updates applied only at the transition between steps
             for step in range(config["recursion_limit"]):
+                next_tasks = _prepare_next_tasks(checkpoint, processes, channels)
+
+                # if no more tasks, we're done
+                if not next_tasks:
+                    break
+
                 if self.debug:
                     print_step_start(step, next_tasks)
 
@@ -255,11 +269,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 # interrupt on failure or timeout
                 _interrupt_or_proceed(done, inflight, step)
 
-                # apply writes to channels, decide on next step
-
-                next_tasks = _apply_writes_and_prepare_next_tasks(
-                    processes, channels, pending_writes, config, step + 1
-                )
+                # apply writes to channels
+                _apply_writes(checkpoint, channels, pending_writes, config, step + 1)
 
                 if self.debug:
                     print_checkpoint(step, channels)
@@ -269,39 +280,32 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                     yield output
 
                 # save end of step checkpoint
-                if (
-                    self.checkpoint is not None
-                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
-                ):
-                    checkpoint = create_checkpoint(channels)
-                    self.checkpoint.put(config, checkpoint)
-
-                # if no more tasks, we're done
-                if not next_tasks:
-                    break
+                if self.saver is not None and self.saver.at == CheckpointAt.END_OF_STEP:
+                    checkpoint = create_checkpoint(checkpoint, channels)
+                    self.saver.put(config, checkpoint)
 
             # save end of run checkpoint
-            if (
-                self.checkpoint is not None
-                and self.checkpoint.at == CheckpointAt.END_OF_RUN
-            ):
-                checkpoint = create_checkpoint(channels)
-                self.checkpoint.put(config, checkpoint)
+            if self.saver is not None and self.saver.at == CheckpointAt.END_OF_RUN:
+                checkpoint = create_checkpoint(checkpoint, channels)
+                self.saver.put(config, checkpoint)
 
     async def _atransform(
         self,
         input: AsyncIterator[dict[str, Any] | Any],
         run_manager: AsyncCallbackManagerForChainRun,
         config: RunnableConfig,
+        *,
+        saver: Optional[BaseCheckpointAdapter] = None,
     ) -> AsyncIterator[dict[str, Any] | Any]:
         if config["recursion_limit"] < 1:
             raise ValueError("recursion_limit must be at least 1")
         processes = {**self.chains}
+        saver = saver or self.saver
         checkpoint = (
-            await self.checkpoint.aget(config) if self.checkpoint is not None else None
-        )
+            await self.saver.aget(config) if self.saver else None
+        ) or empty_checkpoint()
         async with AsyncChannelsManager(self.channels, checkpoint) as channels:
-            next_tasks = _apply_writes_and_prepare_next_tasks(
+            next_tasks = _apply_writes(
                 processes,
                 channels,
                 deque([w async for c in input for w in map_input(self.input, c)]),
@@ -356,7 +360,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 _interrupt_or_proceed(done, inflight, step)
 
                 # apply writes to channels, decide on next step
-                next_tasks = _apply_writes_and_prepare_next_tasks(
+                next_tasks = _apply_writes(
                     processes, channels, pending_writes, config, step + 1
                 )
 
@@ -369,23 +373,20 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
                 # save end of step checkpoint
                 if (
-                    self.checkpoint is not None
-                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                    checkpointer is not None
+                    and checkpointer.at == CheckpointAt.END_OF_STEP
                 ):
                     checkpoint = create_checkpoint(channels)
-                    await self.checkpoint.aput(config, checkpoint)
+                    await checkpointer.aput(config, checkpoint)
 
                 # if no more tasks, we're done
                 if not next_tasks:
                     break
 
             # save end of run checkpoint
-            if (
-                self.checkpoint is not None
-                and self.checkpoint.at == CheckpointAt.END_OF_RUN
-            ):
+            if checkpointer is not None and checkpointer.at == CheckpointAt.END_OF_RUN:
                 checkpoint = create_checkpoint(channels)
-                await self.checkpoint.aput(config, checkpoint)
+                await checkpointer.aput(config, checkpoint)
 
     def invoke(
         self,
@@ -416,6 +417,24 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             input, self._transform, config, **kwargs
         )
 
+    def step(
+        self,
+        input: dict[str, Any] | Any,
+        config: RunnableConfig | None = None,
+        **kwargs: Any,
+    ) -> Iterator[StepState]:
+        config = ensure_config(config)
+        recursion_limit = config["recursion_limit"]
+        config = patch_config(config, recursion_limit=1)
+        step_checkpointer = MemoryCheckpoint()
+        global_checkpointer = self.saver
+        for i in range(recursion_limit):
+            for chunk in self.stream(
+                input, config, checkpointer=step_checkpointer, **kwargs
+            ):
+                checkpoint = step_checkpointer.get(config)
+                yield StepState(checkpoint=checkpoint)
+
     async def ainvoke(
         self,
         input: dict[str, Any] | Any,
@@ -484,13 +503,13 @@ def _read_channel(
         return None
 
 
-def _apply_writes_and_prepare_next_tasks(
-    processes: Mapping[str, ChannelInvoke | ChannelBatch],
+def _apply_writes(
+    checkpoint: Checkpoint,
     channels: Mapping[str, BaseChannel],
     pending_writes: Sequence[tuple[str, Any]],
     config: RunnableConfig,
     for_step: int,
-) -> list[tuple[Runnable, Any, str]]:
+) -> None:
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
     # Group writes by channel
     for chan, val in pending_writes:
@@ -508,6 +527,7 @@ def _apply_writes_and_prepare_next_tasks(
     for chan, vals in pending_writes_by_channel.items():
         if chan in channels:
             channels[chan].update(vals)
+            checkpoint["versions"][chan] += 1
             updated_channels.add(chan)
         else:
             logger.warning(f"Skipping write for channel {chan} which has no readers")
@@ -516,13 +536,20 @@ def _apply_writes_and_prepare_next_tasks(
         if chan not in updated_channels:
             channels[chan].update([])
 
+
+def _prepare_next_tasks(
+    checkpoint: Checkpoint,
+    processes: Mapping[str, ChannelInvoke | ChannelBatch],
+    channels: Mapping[str, BaseChannel],
+) -> list[tuple[Runnable, Any, str]]:
     tasks: list[tuple[Runnable, Any, str]] = []
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
     for name, proc in processes.items():
+        seen = checkpoint["seen"][name]
         if isinstance(proc, ChannelInvoke):
             # If any of the channels read by this process were updated
-            if any(chan in updated_channels for chan in proc.triggers):
+            if any(checkpoint["versions"][chan] > seen[chan] for chan in proc.triggers):
                 # If all channels subscribed by this process have been initialized
                 try:
                     val = {
@@ -540,9 +567,12 @@ def _apply_writes_and_prepare_next_tasks(
                     val = val[None]
 
                 tasks.append((proc, val, name))
+                seen.update(
+                    {chan: checkpoint["versions"][chan] for chan in proc.triggers}
+                )
         elif isinstance(proc, ChannelBatch):
             # If the channel read by this process was updated
-            if proc.channel in updated_channels:
+            if checkpoint["versions"][proc.channel] > seen[proc.channel]:
                 # Here we don't catch EmptyChannelError because the channel
                 # must be intialized if the previous `if` condition is true
                 val = channels[proc.channel].get()
@@ -550,5 +580,6 @@ def _apply_writes_and_prepare_next_tasks(
                     val = [{proc.key: v} for v in val]
 
                 tasks.append((proc, val, name))
+                seen[proc.channel] = checkpoint["versions"][proc.channel]
 
     return tasks

tests/test_pregel.py

@@ -286,34 +286,34 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=memory,
+        saver=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
     assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
     checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
-    assert checkpoint.get("total") == 2
+    assert checkpoint["values"].get("total") == 2
     # total is now 2, so output is 2+3=5
     assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
     checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
-    assert checkpoint.get("total") == 7
+    assert checkpoint["values"].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
     checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
-    assert checkpoint.get("total") == 7
+    assert checkpoint["values"].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
     checkpoint = memory.get({"configurable": {"thread_id": "1"}})
     assert checkpoint is not None
-    assert checkpoint.get("total") == 7
+    assert checkpoint["values"].get("total") == 7
     checkpoint = memory.get({"configurable": {"thread_id": "2"}})
     assert checkpoint is not None
-    assert checkpoint.get("total") == 5
+    assert checkpoint["values"].get("total") == 5
 
 
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:

tests/test_pregel_async.py

@@ -299,7 +299,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     app = Pregel(
         chains={"chain_one": chain_one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        checkpoint=memory,
+        saver=memory,
     )
 
     # total starts out as 0, so output is 0+2=2