Alternative proposal for concurrent.futures support in effectful

effectful/handlers/futures/__init__.py

@@ -0,0 +1,255 @@
+"""
+Futures handler for effectful - provides integration with concurrent.futures.
+
+This module provides operations for working with concurrent.futures, allowing
+effectful operations to be executed asynchronously in thread pools with
+automatic preservation of interpretation context.
+"""
+
+import concurrent.futures as futures
+import functools
+from collections.abc import Callable, Iterable
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Literal
+
+from effectful.ops.semantics import defop
+from effectful.ops.syntax import ObjectInterpretation, defdata, implements
+from effectful.ops.types import NotHandled, Term
+
+
+class Executor:
+    """Namespace for executor-related operations."""
+
+    @staticmethod
+    @defop  # type: ignore
+    def submit[**P, T](
+        task: Callable[P, T], *args: P.args, **kwargs: P.kwargs
+    ) -> Future[T]:
+        """
+        Submit a task for asynchronous execution.
+
+        This operation should be handled by providing a FuturesInterpretation
+        which automatically preserves the interpretation context across thread boundaries.
+
+        :param task: The callable to execute asynchronously
+        :param args: Positional arguments for the task
+        :param kwargs: Keyword arguments for the task
+        :return: A Future representing the asynchronous computation
+
+        Example:
+            >>> from concurrent.futures import ThreadPoolExecutor
+            >>> from effectful.handlers.futures import ThreadPoolFuturesInterpretation
+            >>> from effectful.ops.semantics import handler
+            >>>
+            >>> pool = ThreadPoolExecutor()
+            >>> with handler(ThreadPoolFuturesInterpretation(pool)):
+            >>>     future = Executor.submit(my_function, arg1, arg2)
+        """
+        raise NotHandled
+
+    @staticmethod
+    @defop
+    def map[T, R](
+        func: Callable[[T], R],
+        *iterables: Iterable[T],
+        timeout: float | None = None,
+        chunksize: int = 1,
+    ) -> Iterable[R]:
+        """
+        Map a function over iterables, executing asynchronously.
+
+        Returns an iterator yielding results as they complete. Equivalent to
+        map(func, *iterables) but executes asynchronously.
+
+        This operation should be handled by providing a FuturesInterpretation
+        which automatically preserves the interpretation context across thread boundaries.
+
+        :param func: The function to map over the iterables
+        :param iterables: One or more iterables to map over
+        :param timeout: Maximum time to wait for a result (default: None)
+        :param chunksize: Size of chunks for ProcessPoolExecutor (default: 1)
+        :return: An iterator yielding results
+
+        Example:
+            >>> from effectful.handlers.futures import ThreadPoolFuturesInterpretation
+            >>> from effectful.ops.semantics import handler
+            >>>
+            >>> def square(x):
+            >>>     return x ** 2
+            >>>
+            >>> with handler(ThreadPoolFuturesInterpretation()):
+            >>>     results = list(Executor.map(square, range(10)))
+            >>>     print(results)  # [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
+        """
+        raise NotHandled
+
+
+class FuturesInterpretation(ObjectInterpretation):
+    """
+    Base interpretation for concurrent.futures executors.
+
+    This interpretation automatically preserves the effectful interpretation context
+    when submitting tasks to worker threads, ensuring that effectful operations
+    work correctly across thread boundaries.
+    """
+
+    def __init__(self, executor: futures.Executor):
+        """
+        Initialize the futures interpretation.
+
+        :param executor: The executor to use (ThreadPoolExecutor or ProcessPoolExecutor)
+        """
+        super().__init__()
+        self.executor: futures.Executor = executor
+
+    def shutdown(self, *args, **kwargs):
+        self.executor.shutdown(*args, **kwargs)
+
+    @implements(Executor.submit)
+    def submit(self, task: Callable, *args, **kwargs) -> Future:
+        """
+        Submit a task to the executor with automatic context preservation.
+
+        Captures the current interpretation context and ensures it is restored
+        in the worker thread before executing the task.
+        """
+        from effectful.internals.runtime import get_interpretation, interpreter
+
+        # Capture the current interpretation context
+        context = get_interpretation()
+
+        @functools.wraps(task)
+        def wrapped_task(*task_args, **task_kwargs):
+            # Restore the interpretation context in the worker thread
+            with interpreter(context):
+                return task(*task_args, **task_kwargs)
+
+        # Submit the wrapped task to the underlying executor
+        return self.executor.submit(wrapped_task, *args, **kwargs)
+
+    @implements(Executor.map)
+    def map(self, func: Callable, *iterables, timeout=None, chunksize=1):
+        """
+        Map a function over iterables with automatic context preservation.
+
+        Captures the current interpretation context and ensures it is restored
+        in each worker thread before executing the function.
+        """
+        from effectful.internals.runtime import get_interpretation, interpreter
+
+        # Capture the current interpretation context
+        context = get_interpretation()
+
+        @functools.wraps(func)
+        def wrapped_func(*args, **kwargs):
+            # Restore the interpretation context in the worker thread
+            with interpreter(context):
+                return func(*args, **kwargs)
+
+        # Call the executor's map with the wrapped function
+        return self.executor.map(
+            wrapped_func, *iterables, timeout=timeout, chunksize=chunksize
+        )
+
+
+class ThreadPoolFuturesInterpretation(FuturesInterpretation):
+    """
+    Interpretation for ThreadPoolExecutor with automatic context preservation.
+
+    Example:
+        >>> from concurrent.futures import ThreadPoolExecutor, Future
+        >>> from effectful.ops.semantics import defop, handler
+        >>> from effectful.handlers.futures import Executor, ThreadPoolFuturesInterpretation
+        >>>
+        >>> @defop
+        >>> def async_pow(n: int, k: int) -> Future[int]:
+        >>>     return Executor.submit(pow, n, k)
+        >>>
+        >>> pool = ThreadPoolExecutor()
+        >>> with handler(ThreadPoolFuturesInterpretation(pool)):
+        >>>     result = async_pow(2, 10).result()
+        >>>     print(result)  # 1024
+    """
+
+    def __init__(self, max_workers=None):
+        """
+        Initialize with a ThreadPoolExecutor.
+
+        :param max_workers: Maximum number of worker threads (default: None, uses default from ThreadPoolExecutor)
+        """
+        super().__init__(ThreadPoolExecutor(max_workers=max_workers))
+
+
+type ReturnOptions = Literal["All_COMPLETED", "FIRST_COMPLETED", "FIRST_EXCEPTION"]
+
+
+@dataclass(frozen=True)
+class DoneAndNotDoneFutures[T]:
+    done: set[Future[T]]
+    not_done: set[Future[T]]
+
+
+@defdata.register(DoneAndNotDoneFutures)
+class _DoneAndNotDoneFuturesTerm[T](Term[DoneAndNotDoneFutures[T]]):
+    """Term representing a DoneAndNotDoneFutures result."""
+
+    def __init__(self, op, *args, **kwargs):
+        self._op = op
+        self._args = args
+        self._kwargs = kwargs
+
+    @property
+    def op(self):
+        return self._op
+
+    @property
+    def args(self):
+        return self._args
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+    @defop  # type: ignore[prop-decorator]
+    @property
+    def done(self) -> set[Future[T]]:
+        """Get the set of done futures."""
+        if not isinstance(self, Term):
+            return self.done
+        else:
+            raise NotHandled
+
+    @defop  # type: ignore[prop-decorator]
+    @property
+    def not_done(self) -> set[Future[T]]:
+        """Get the set of not done futures."""
+        if not isinstance(self, Term):
+            return self.not_done
+        else:
+            raise NotHandled
+
+
+@defop
+def wait[T](
+    fs: Iterable[Future[T]],
+    timeout: int | None = None,
+    return_when: ReturnOptions = futures.ALL_COMPLETED,  # type: ignore
+) -> DoneAndNotDoneFutures[T]:
+    if (
+        isinstance(timeout, Term)
+        or isinstance(return_when, Term)
+        or any(not isinstance(t, Future) for t in fs)
+    ):
+        raise NotHandled
+    return futures.wait(fs, timeout, return_when)  # type: ignore
+
+
+@defop
+def as_completed[T](
+    fs: Iterable[Future[T]],
+    timeout: int | None = None,
+) -> Iterable[Future[T]]:
+    if isinstance(timeout, Term) or any(isinstance(t, Term) for t in fs):
+        raise NotHandled
+    return futures.as_completed(fs, timeout)

effectful/handlers/llm/providers.py

@@ -6,6 +6,7 @@
 import string
 import typing
 from collections.abc import Hashable, Iterable, Mapping, Sequence
+from concurrent.futures import Future
 from typing import Any, get_type_hints
 
 import pydantic
@@ -22,6 +23,7 @@
 
 from openai.types.responses import FunctionToolParam
 
+from effectful.handlers.futures import Executor
 from effectful.handlers.llm import Template
 from effectful.ops.semantics import fwd
 from effectful.ops.syntax import ObjectInterpretation, defop, implements
@@ -282,13 +284,13 @@ def __init__(self, client: openai.OpenAI, model_name: str = "gpt-4o"):
         self._client = client
         self._model_name = model_name
 
-    @implements(Template.__call__)
-    def _call[**P, T](
-        self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs
-    ) -> T:
-        ret_type = template.__signature__.return_annotation
-        bound_args = template.__signature__.bind(*args, **kwargs)
-        bound_args.apply_defaults()
+    def _openai_api_call[**P, T, RT](
+        self,
+        template: Template[P, T],
+        bound_args: inspect.BoundArguments,
+        ret_type: type[RT],
+    ) -> RT:
+        """Execute the actual OpenAI API call and decode the response."""
         prompt = _OpenAIPromptFormatter().format_as_messages(
             template.__prompt_template__, **bound_args.arguments
         )
@@ -368,3 +370,24 @@ def _call[**P, T](
         result = Result.model_validate_json(result_str)
         assert isinstance(result, Result)
         return result.value  # type: ignore[attr-defined]
+
+    @implements(Template.__call__)
+    def _call[**P, T](
+        self, template: Template[P, T], *args: P.args, **kwargs: P.kwargs
+    ) -> T:
+        ret_type = template.__signature__.return_annotation
+        bound_args = template.__signature__.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        # Check if return type is Future[T]
+        origin = typing.get_origin(ret_type)
+        if origin is Future:
+            inner_type = typing.get_args(ret_type)[0]
+            return Executor.submit(
+                self._openai_api_call,  # type: ignore
+                template,  # type: ignore
+                bound_args,  # type: ignore
+                inner_type,
+            )
+
+        return self._openai_api_call(template, bound_args, ret_type)

effectful/internals/runtime.py

@@ -1,20 +1,75 @@
 import contextlib
-import dataclasses
 import functools
+import threading
 from collections.abc import Callable, Mapping
 
 from effectful.ops.syntax import defop
 from effectful.ops.types import Interpretation, Operation
 
+# Global reentrant lock for mutual exclusion of handler execution
+_handler_lock = threading.RLock()
+
+
+class Runtime[S, T](threading.local):
+    """Thread-local runtime for effectful interpretations."""
 
-@dataclasses.dataclass
-class Runtime[S, T]:
     interpretation: "Interpretation[S, T]"
 
+    def __init__(self):
+        super().__init__()
+        self.interpretation = {}
+
+
+def get_handler_lock() -> threading.RLock:
+    """Get the global handler execution lock.
+
+    This lock ensures mutual exclusion for handler execution by default.
+    Use release_handler_lock() to temporarily release it for concurrent operations.
+    """
+    return _handler_lock
+
+
+@contextlib.contextmanager
+def release_handler_lock():
+    """Context manager to temporarily release the handler lock.
+
+    Use this when performing I/O operations or other blocking calls
+    that should allow other threads to execute handlers concurrently.
+
+    Example:
+        @defop
+        def llm_call(prompt: str) -> str:
+            with release_handler_lock():
+                # HTTP request can run while other threads execute handlers
+                response = requests.post(url, json={"prompt": prompt})
+            return response.json()["result"]
+    """
+    lock = get_handler_lock()
+    lock.release()
+    try:
+        yield
+    finally:
+        lock.acquire()
+
+
+@contextlib.contextmanager
+def acquire_handler_lock():
+    """Context manager to acquire the handler lock.
+
+    This is called automatically by apply() for handler execution.
+    Most users won't need to call this directly.
+    """
+    lock = get_handler_lock()
+    lock.acquire()
+    try:
+        yield
+    finally:
+        lock.release()
+
 
 @functools.lru_cache(maxsize=1)
 def get_runtime() -> Runtime:
-    return Runtime(interpretation={})
+    return Runtime()
 
 
 def get_interpretation():