Implement checkpoints

permchain/channels/base.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
 from typing import (
     Any,
     AsyncGenerator,
@@ -13,6 +14,8 @@ from typing import (
 
 from typing_extensions import Self
 
+from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 Checkpoint = TypeVar("Checkpoint")
@@ -72,10 +75,10 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
     @abstractmethod
     def checkpoint(self) -> Checkpoint | None:
-        """Return a string representation of the channel's current state,
-        or None if the channel doesn't support checkpoints.
+        """Return a string representation of the channel's current state.
 
-        Raises EmptyChannelError if the channel is empty (never updated yet)."""
+        Raises EmptyChannelError if the channel is empty (never updated yet),
+        or doesn't supportcheckpoints."""
 
 
 @contextmanager
@@ -111,5 +114,13 @@ async def AsyncChannelsManager(
 
 def create_checkpoint(channels: Mapping[str, BaseChannel]) -> Mapping[str, Any]:
     """Create a checkpoint for the given channels."""
-    checkpoint = {k: v.checkpoint() for k, v in channels.items()}
-    return {k: v for k, v in checkpoint.items() if v is not None}
+    checkpoint = {
+        CHECKPOINT_KEY_VERSION: 1,
+        CHECKPOINT_KEY_TS: datetime.utcnow().isoformat(),
+    }
+    for k, v in channels.items():
+        try:
+            checkpoint[k] = v.checkpoint()
+        except EmptyChannelError:
+            pass
+    return checkpoint

permchain/channels/context.py

@@ -107,4 +107,4 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> None:
-        return None
+        raise EmptyChannelError()

permchain/checkpoint/base.py

@@ -0,0 +1,44 @@
+import asyncio
+import enum
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Sequence
+
+from langchain.load.serializable import Serializable
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+
+# Before Python 3.11 native StrEnum is not available
+class StrEnum(str, enum.Enum):
+    """A string enum."""
+
+    pass
+
+
+class CheckpointAt(StrEnum):
+    END_OF_STEP = "end_of_step"
+    END_OF_RUN = "end_of_run"
+
+
+class BaseCheckpointAdapter(Serializable, ABC):
+    at: CheckpointAt = CheckpointAt.END_OF_RUN
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return []
+
+    @abstractmethod
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        ...
+
+    @abstractmethod
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        ...
+
+    async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return 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(
+            None, self.put, config, checkpoint
+        )

permchain/constants.py

@@ -0,0 +1,4 @@
+CONFIG_KEY_SEND = "__pregel_send"
+CONFIG_KEY_READ = "__pregel_read"
+CHECKPOINT_KEY_VERSION = "__pregel_version"
+CHECKPOINT_KEY_TS = "__pregel_ts"

permchain/pregel/__init__.py

@@ -42,8 +42,10 @@ from permchain.channels.base import (
     BaseChannel,
     ChannelsManager,
     EmptyChannelError,
+    create_checkpoint,
 )
-from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
+from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
+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
 from permchain.pregel.log import logger
@@ -115,6 +117,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
     debug: bool = Field(default_factory=get_debug)
 
+    checkpoint: Optional[BaseCheckpointAdapter] = None
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -128,7 +132,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     @property
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
         return get_unique_config_specs(
-            spec for chain in self.chains.values() for spec in chain.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 []
         )
 
     @property
@@ -173,10 +180,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         config: RunnableConfig,
     ) -> Iterator[dict[str, Any] | Any]:
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        with ChannelsManager(self.channels, None) as channels, get_executor_for_config(
-            config
-        ) as executor:
+        checkpoint = (
+            self.checkpoint.get(config) if self.checkpoint is not None else None
+        )
+        with ChannelsManager(
+            self.channels, checkpoint
+        ) as channels, get_executor_for_config(config) as executor:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
@@ -239,12 +248,26 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # 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
 
+            # 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)
+
     async def _atransform(
         self,
         input: AsyncIterator[dict[str, Any] | Any],
@@ -252,8 +275,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         config: RunnableConfig,
     ) -> AsyncIterator[dict[str, Any] | Any]:
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        async with AsyncChannelsManager(self.channels, None) as channels:
+        checkpoint = (
+            await self.checkpoint.aget(config) if self.checkpoint is not None else None
+        )
+        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
@@ -319,12 +344,26 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # save end of step checkpoint
+                if (
+                    self.checkpoint is not None
+                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                ):
+                    checkpoint = create_checkpoint(channels)
+                    await self.checkpoint.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
+            ):
+                checkpoint = create_checkpoint(channels)
+                await self.checkpoint.aput(config, checkpoint)
+
     def invoke(
         self,
         input: dict[str, Any] | Any,

permchain/pregel/constants.py

@@ -1,2 +0,0 @@
-CONFIG_KEY_SEND = "__pregel_send"
-CONFIG_KEY_READ = "__pregel_read"

permchain/pregel/read.py

@@ -18,7 +18,7 @@ from langchain.schema.runnable.base import (
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
 from permchain.channels.base import BaseChannel
-from permchain.pregel.constants import CONFIG_KEY_READ
+from permchain.constants import CONFIG_KEY_READ
 
 
 class ChannelRead(RunnableLambda):

permchain/pregel/write.py

@@ -9,7 +9,7 @@ from langchain.schema.runnable import (
 )
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
-from permchain.pregel.constants import CONFIG_KEY_SEND
+from permchain.constants import CONFIG_KEY_SEND
 
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 

tests/test_channels.py

@@ -273,17 +273,8 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
         with pytest.raises(InvalidUpdateError):
             channel.update([5])  # type: ignore
 
-        checkpoint = channel.checkpoint()
-        assert checkpoint is None
-    with Context(httpx.Client).empty(checkpoint) as channel:
-        assert channel.ValueType is httpx.Client
-        with pytest.raises(InvalidUpdateError):
-            assert channel.UpdateType is None
-
-        assert isinstance(channel.get(), httpx.Client)
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])
+        with pytest.raises(EmptyChannelError):
+            channel.checkpoint()
 
 
 async def test_ctx_manager_async(mocker: MockerFixture) -> None: