From b1b68cde9473847b153e907bdfe659931e65e90c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 26 Dec 2023 19:05:21 -0800 Subject: [PATCH] WIP --- permchain/channels/base.py | 27 +++---- permchain/checkpoint/base.py | 41 +++++++++-- permchain/checkpoint/memory.py | 10 ++- permchain/pregel/__init__.py | 127 ++++++++++++++++++++------------- tests/test_pregel.py | 12 ++-- tests/test_pregel_async.py | 2 +- 6 files changed, 137 insertions(+), 82 deletions(-) diff --git a/permchain/channels/base.py b/permchain/channels/base.py index bd53e37b4..eda3e13e7 100644 --- a/permchain/channels/base.py +++ b/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 diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py index 7c4b96384..ba777d066 100644 --- a/permchain/checkpoint/base.py +++ b/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 diff --git a/permchain/checkpoint/memory.py b/permchain/checkpoint/memory.py index e036c877b..f58144078 100644 --- a/permchain/checkpoint/memory.py +++ b/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}) diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 3edcc2f14..b400405f3 100644 --- a/permchain/pregel/__init__.py +++ b/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 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 4d5c32f74..b0d1411d9 100644 --- a/tests/test_pregel.py +++ b/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: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 8a59fa4a0..2a7efcaf7 100644 --- a/tests/test_pregel_async.py +++ b/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