diff --git a/permchain/connection.py b/permchain/connection.py index 116875e91..f35bf0aa2 100644 --- a/permchain/connection.py +++ b/permchain/connection.py @@ -1,8 +1,7 @@ -from abc import ABC, abstractmethod import asyncio +from abc import ABC, abstractmethod from typing import Any, Callable - PubSubListener = Callable[[Any], None] diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py index f5547bbe9..2c5bf5a99 100644 --- a/permchain/connection_inmemory.py +++ b/permchain/connection_inmemory.py @@ -1,5 +1,5 @@ -from collections import defaultdict import threading +from collections import defaultdict from typing import Any from permchain.connection import PubSubConnection, PubSubListener diff --git a/permchain/constants.py b/permchain/constants.py new file mode 100644 index 000000000..4c9c12ee7 --- /dev/null +++ b/permchain/constants.py @@ -0,0 +1,2 @@ +CONFIG_GET_KEY = "pubsub_get" +CONFIG_SEND_KEY = "pubsub_send" diff --git a/permchain/pubsub.py b/permchain/pubsub.py index 8a6c3f158..ed845c459 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -1,8 +1,9 @@ from __future__ import annotations +from ctypes import Union import queue from abc import ABC -from concurrent.futures import CancelledError, Future, ThreadPoolExecutor +from concurrent.futures import CancelledError, Future from functools import partial from itertools import filterfalse from typing import ( @@ -17,13 +18,17 @@ from typing import ( TypeVar, ) -from langchain.callbacks.manager import CallbackManager from langchain.load.dump import dumpd from langchain.load.serializable import Serializable -from langchain.schema.runnable import Runnable, RunnableConfig -from langchain.schema.runnable.base import patch_config +from langchain.schema.runnable import Runnable, RunnableConfig, patch_config +from langchain.schema.runnable.config import ( + ensure_config, + get_callback_manager_for_config, + get_executor_for_config, +) from permchain.connection import PubSubConnection +from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY from permchain.topic import INPUT_TOPIC, OUTPUT_TOPIC, RunnableSubscriber T = TypeVar("T") @@ -68,32 +73,34 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): collected.append(chunk) return collected + def batch( + self, + input: Sequence[Any], + config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None, + ) -> Any: + configs = self._get_config_list(config) + with get_executor_for_config(configs[0]) as executor: + return super().batch( + input, + [patch_config(config, executor=executor) for config in configs], + ) + def stream( self, input: Any, config: Optional[RunnableConfig] = None, - *, - max_concurrency: Optional[int] = None, ) -> Iterator[Any]: input_processes, listener_processes = partition( lambda r: r.topic.name == INPUT_TOPIC, self.processes ) # setup callbacks - config = config or {} - callback_manager = CallbackManager.configure( - inheritable_callbacks=config.get("callbacks"), - local_callbacks=None, - verbose=False, - inheritable_tags=config.get("tags"), - local_tags=None, - inheritable_metadata=config.get("metadata"), - local_metadata=None, - ) + config = ensure_config(config) + callback_manager = get_callback_manager_for_config(config) # start the root run run_manager = callback_manager.on_chain_start(dumpd(self), {"input": input}) - with ThreadPoolExecutor(max_workers=max_concurrency) as executor: + with get_executor_for_config(config) as executor: # Track inflight futures inflight: Set[Future] = set() # Track exceptions @@ -101,16 +108,20 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # Track output output = IterableQueue() + # Namespace topics for each run + def prefix_topic_name(topic_name: str) -> str: + return f"{run_manager.run_id}/{topic_name}" + def send(topic_name: str, message: Any) -> None: """Send a message to a topic. Injected into config.""" if topic_name == OUTPUT_TOPIC: output.put(message) else: - self.connection.send(topic_name, message) + self.connection.send(prefix_topic_name(topic_name), message) def cleanup_run(fut: Future) -> None: """Cleanup after a process runs.""" - inflight.remove(fut) + inflight.discard(fut) try: exc = fut.exception() @@ -146,10 +157,12 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): value, config={ **patch_config( - config, run_manager.get_child(process.topic.name) + config, + callbacks=run_manager.get_child(process.topic.name), + executor=executor, ), - "send": send, - "get": get, + CONFIG_SEND_KEY: send, + CONFIG_GET_KEY: get, }, ) @@ -159,7 +172,9 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): # Listen on all subscribed topics for process in listener_processes: - self.connection.listen(process.topic.name, partial(run_once, process)) + self.connection.listen( + prefix_topic_name(process.topic.name), partial(run_once, process) + ) # Run input processes once for process in input_processes: @@ -176,11 +191,11 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): final_output += chunk finally: # Cleanup - for fut in inflight: - fut.cancel() + while inflight: + inflight.pop().cancel() for process in listener_processes: - self.connection.disconnect(process.topic.name) + self.connection.disconnect(prefix_topic_name(process.topic.name)) # Raise exceptions if any if exceptions: diff --git a/permchain/topic.py b/permchain/topic.py index c804e5c67..32b4b7c16 100644 --- a/permchain/topic.py +++ b/permchain/topic.py @@ -1,7 +1,6 @@ from __future__ import annotations from abc import ABC - from typing import ( Any, Callable, @@ -13,6 +12,7 @@ from typing import ( ) from langchain.load.serializable import Serializable +from langchain.pydantic_v1 import Field from langchain.schema.runnable import ( Runnable, RunnableBinding, @@ -21,7 +21,8 @@ from langchain.schema.runnable import ( RunnableSequence, ) from langchain.schema.runnable.base import Other, coerce_to_runnable -from pydantic import Field + +from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY T = TypeVar("T") T_in = TypeVar("T_in") @@ -97,7 +98,7 @@ class RunnablePublisher(RunnablePassthrough[T]): topic: Topic[T] def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T: - send = config.get("send", None) + send = config.get(CONFIG_SEND_KEY, None) if send is not None: send(self.topic.name, input) return super().invoke(input, config) @@ -117,7 +118,7 @@ class RunnableCurrentValue(Serializable, Runnable[Any, T]): topic: Topic[T] def invoke(self, input: T, config: Optional[RunnableConfigForPubSub] = None) -> T: - get = config.get("get", None) + get = config.get(CONFIG_GET_KEY, None) if get is not None: return get(self.topic.name) else: