From 99b712ca5fc0540eac235b8531900e8e399d51eb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 28 Nov 2023 11:26:46 +0000 Subject: [PATCH] Change .join() implementation to have non-subscribed channels as input arg --- permchain/pregel/__init__.py | 48 +++++++++++++++++++++++------------- permchain/pregel/read.py | 42 +++++++++++++++++++++++-------- 2 files changed, 63 insertions(+), 27 deletions(-) diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index a1b65eec1..8545adfec 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import concurrent.futures from collections import defaultdict, deque +from functools import partial from typing import ( Any, AsyncIterator, @@ -98,7 +99,8 @@ class Channel: {key: channels} if isinstance(channels, str) else {chan: chan for chan in channels}, - ) + ), + triggers=[channels] if isinstance(channels, str) else channels, ) @classmethod @@ -211,11 +213,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 0, ) - def read(chan: str) -> Any: - try: - return channels[chan].get() - except EmptyChannelError: - return None + if not next_tasks: + return + + read = partial(_read_channel, channels) # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates @@ -309,11 +310,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 0, ) - def read(chan: str) -> Any: - try: - return channels[chan].get() - except EmptyChannelError: - return None + if not next_tasks: + return + + read = partial(_read_channel, channels) # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates @@ -476,6 +476,15 @@ def _interrupt_or_proceed( raise TimeoutError(f"Timed out at step {step}") +def _read_channel( + channels: Mapping[str, BaseChannel], chan: str, catch: bool = True +) -> Any: + try: + return channels[chan].get() + except EmptyChannelError: + return None + + def _apply_writes_and_prepare_next_tasks( processes: Mapping[str, ChannelInvoke | ChannelBatch], channels: Mapping[str, BaseChannel], @@ -514,19 +523,24 @@ def _apply_writes_and_prepare_next_tasks( for name, proc in processes.items(): 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.channels.values()): - # If all channels read by this process have been initialized + if any(chan in updated_channels for chan in proc.triggers): + # If all channels subscribed by this process have been initialized try: - val = {k: channels[chan].get() for k, chan in proc.channels.items()} + val = { + k: _read_channel( + channels, chan, catch=chan not in proc.triggers + ) + for k, chan in proc.channels.items() + } except EmptyChannelError: continue # Processes that subscribe to a single keyless channel get # the value directly, instead of a dict if list(proc.channels.keys()) == [None]: - tasks.append((proc, val[None], name)) - else: - tasks.append((proc, val, name)) + val = val[None] + + tasks.append((proc, val, name)) elif isinstance(proc, ChannelBatch): # If the channel read by this process was updated if proc.channel in updated_channels: diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py index 66a20367c..4084229b6 100644 --- a/permchain/pregel/read.py +++ b/permchain/pregel/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence +from typing import Any, Callable, List, Mapping, Optional, Sequence from langchain.pydantic_v1 import Field from langchain.schema.runnable import ( @@ -67,6 +67,10 @@ default_bound = RunnablePassthrough() class ChannelInvoke(RunnableBindingBase): channels: Mapping[None, str] | Mapping[str, str] + triggers: List[str] = Field(default_factory=list) + + skip: Optional[Callable[[Any], bool]] = None + bound: Runnable[Any, Any] = Field(default=default_bound) kwargs: Mapping[str, Any] = Field(default_factory=dict) @@ -74,6 +78,7 @@ class ChannelInvoke(RunnableBindingBase): def __init__( self, channels: Mapping[None, str] | Mapping[str, str], + triggers: Sequence[str], *, bound: Optional[Runnable[Any, Any]] = None, kwargs: Optional[Mapping[str, Any]] = None, @@ -82,6 +87,7 @@ class ChannelInvoke(RunnableBindingBase): ) -> None: super().__init__( channels=channels, + triggers=triggers, bound=bound or default_bound, kwargs=kwargs or {}, config=config, @@ -92,13 +98,19 @@ class ChannelInvoke(RunnableBindingBase): 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} + assert all( + k is not None for k in self.channels.keys() + ), "all channels must be named when using .join()" + return ChannelInvoke( + channels={ + **self.channels, + **{chan: chan for chan in channels}, + }, + triggers=self.triggers, + bound=self.bound, + kwargs=self.kwargs, + config=self.config, ) - if self.bound is default_bound: - return ChannelInvoke(channels=self.channels, bound=joiner) - else: - return ChannelInvoke(channels=self.channels, bound=self.bound | joiner) def __or__( self, @@ -108,11 +120,21 @@ class ChannelInvoke(RunnableBindingBase): ) -> ChannelInvoke: if self.bound is default_bound: return ChannelInvoke( - channels=self.channels, bound=coerce_to_runnable(other) + channels=self.channels, + triggers=self.triggers, + bound=coerce_to_runnable(other), + kwargs=self.kwargs, + config=self.config, ) else: - # delegate to __or__ in self.bound - return ChannelInvoke(channels=self.channels, bound=self.bound | other) + return ChannelInvoke( + channels=self.channels, + triggers=self.triggers, + # delegate to __or__ in self.bound + bound=self.bound | other, + kwargs=self.kwargs, + config=self.config, + ) def __ror__( self,