From f96ebb3357f6c102780199baf8941ed3fc3a61db Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 Nov 2023 12:14:18 +0000 Subject: [PATCH] Implement checkpoints --- permchain/channels/base.py | 21 +++++++++--- permchain/channels/context.py | 2 +- permchain/checkpoint/__init__.py | 0 permchain/checkpoint/base.py | 44 ++++++++++++++++++++++++ permchain/constants.py | 4 +++ permchain/pregel/__init__.py | 59 ++++++++++++++++++++++++++------ permchain/pregel/constants.py | 2 -- permchain/pregel/read.py | 2 +- permchain/pregel/write.py | 2 +- tests/test_channels.py | 13 ++----- 10 files changed, 118 insertions(+), 31 deletions(-) create mode 100644 permchain/checkpoint/__init__.py create mode 100644 permchain/checkpoint/base.py create mode 100644 permchain/constants.py delete mode 100644 permchain/pregel/constants.py diff --git a/permchain/channels/base.py b/permchain/channels/base.py index 1983147f5..eac51deac 100644 --- a/permchain/channels/base.py +++ b/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 diff --git a/permchain/channels/context.py b/permchain/channels/context.py index 97a95727d..13098409f 100644 --- a/permchain/channels/context.py +++ b/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() diff --git a/permchain/checkpoint/__init__.py b/permchain/checkpoint/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py new file mode 100644 index 000000000..ef6ae43e5 --- /dev/null +++ b/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 + ) diff --git a/permchain/constants.py b/permchain/constants.py new file mode 100644 index 000000000..bb96e6361 --- /dev/null +++ b/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" diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 3148e0667..49ad344fc 100644 --- a/permchain/pregel/__init__.py +++ b/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, diff --git a/permchain/pregel/constants.py b/permchain/pregel/constants.py deleted file mode 100644 index 6e5f9c37f..000000000 --- a/permchain/pregel/constants.py +++ /dev/null @@ -1,2 +0,0 @@ -CONFIG_KEY_SEND = "__pregel_send" -CONFIG_KEY_READ = "__pregel_read" diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py index 4793dee67..abb693e63 100644 --- a/permchain/pregel/read.py +++ b/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): diff --git a/permchain/pregel/write.py b/permchain/pregel/write.py index 1cf9a5ccb..27480aa33 100644 --- a/permchain/pregel/write.py +++ b/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] diff --git a/tests/test_channels.py b/tests/test_channels.py index ed4e28028..104206c10 100644 --- a/tests/test_channels.py +++ b/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: