wool 0.1rc6__py3-none-any.whl → 0.1rc7__py3-none-any.whl

wool/_task.py

@@ -4,57 +4,115 @@ import asyncio
 import logging
 import traceback
 from collections.abc import Callable
-from contextvars import Context, ContextVar
+from contextvars import Context
+from contextvars import ContextVar
 from dataclasses import dataclass
 from functools import wraps
 from sys import modules
-from time import perf_counter_ns
+from types import ModuleType
 from types import TracebackType
-from typing import (
-    Any,
-    Coroutine,
-    Literal,
-    ParamSpec,
-    Protocol,
-    TypeVar,
-    cast,
-)
-from uuid import UUID, uuid4
-
-import wool
-from wool._future import WoolFuture
-from wool._pool import WoolClient
-from wool._typing import PassthroughDecorator
+from typing import Coroutine
+from typing import Dict
+from typing import ParamSpec
+from typing import Protocol
+from typing import SupportsInt
+from typing import Tuple
+from typing import Type
+from typing import TypeVar
+from typing import cast
+from uuid import UUID
+from uuid import uuid4
+
+from wool._event import TaskEvent
+from wool._future import Future
+from wool._pool import WorkerPoolSession
+from wool._session import current_session
 
 AsyncCallable = Callable[..., Coroutine]
 C = TypeVar("C", bound=AsyncCallable)
 
-Args = tuple
-Kwargs = dict
-Timeout = int
-Timestamp = int
+Args = Tuple
+Kwargs = Dict
+Timeout = SupportsInt
+Timestamp = SupportsInt
 
 
+# PUBLIC
 def task(fn: C) -> C:
     """
-    A decorator to declare a function as remotely executed using a worker pool.
-    This decorator allows a function to be executed either locally or remotely
-    using a worker pool. If a worker pool is provided, the function will be
-    submitted to the pool for remote execution. If no pool is provided, the
-    function will be executed locally.
-    The decorator also handles the case where the function is being executed
-    within a worker, ensuring that it does not resubmit itself to the pool.
+    A decorator to declare an asynchronous function as remotely executable by a
+    worker pool. When the wrapped function is invoked, it is dispatched to the
+    worker pool associated with the current worker pool session context.
+
+    Tasks behave like coroutines, meaning they can be awaited as well as
+    cancelled.
+
+    :param fn: The task function.
+    :return: A Wool task declaration.
+
+    **Best practices and considerations for designing tasks:**
+
+    1. **Picklability**:
+        - Task arguments and return values must be picklable, as they are
+            serialized and transferred between processes. Avoid passing
+            unpicklable objects such as open file handles, database
+            connections, or lambda functions.
+        - Ensure that any custom objects used as arguments or return values
+            implement the necessary methods for pickling (e.g.,
+            ``__getstate__`` and ``__setstate__``).
+
+    2. **Synchronization**:
+        - Tasks are not guaranteed to execute on the same process between
+            invocations. Each invocation may run on a different worker process.
+        - Standard ``asyncio`` synchronization primitives (e.g.,
+            ``asyncio.Lock``) will not behave as expected in a multi-process
+            environment, as they are designed for single-process applications.
+            Use the specialized ``wool.locking`` synchronization primitives to
+            achieve inter-worker and inter-pool synchronization.
+
+    3. **Statelessness and idempotency**:
+        - Design tasks to be stateless and idemptoent. Avoid relying on global
+            variables or shared mutable state. This ensures predictable
+            behavior, avoids race conditions, and enables safe retries.
+
+    4. **Cancellation**:
+        - Task cancellation and propagation thereof mimics that of standard
+            Python coroutines.
+
+    5. **Error propagation**:
+        - Wool makes every effort to execute tasks transparently to the user,
+            and this includes error propagation. Unhandled exceptions raised
+            within a task will be propagated to the caller as they would
+            normally.
+
+    6. **Performance**:
+        - Minimize the size of arguments and return values to reduce
+            serialization overhead.
+        - For large datasets, consider using shared memory or passing
+            references (e.g., file paths) instead of transferring the entire
+            data.
+
+    **Usage**::
+
+    .. code-block:: python
+
+        import wool
+
+
+        @wool.task
+        async def foo(...):
+            ...
     """
 
     @wraps(fn)
     def wrapper(
         *args,
         __wool_remote__: bool = False,
-        __wool_client__: WoolClient | None = None,
+        __wool_session__: WorkerPoolSession | None = None,
         **kwargs,
     ) -> Coroutine:
         # Handle static and class methods in a picklable way.
-        parent, function = resolve(fn)
+        parent, function = _resolve(fn)
         assert parent is not None
         assert callable(function)
 
@@ -64,7 +122,7 @@ def task(fn: C) -> C:
         else:
             # Otherwise, submit the task to the pool.
             return _put(
-                __wool_client__ or wool.__wool_client__.get(),
+                __wool_session__ or current_session(),
                 wrapper.__module__,
                 wrapper.__qualname__,
                 function,
@@ -76,14 +134,15 @@ def task(fn: C) -> C:
 
 
 def _put(
-    client: WoolClient,
+    session: WorkerPoolSession,
     module: str,
     qualname: str,
     function: AsyncCallable,
     *args,
     **kwargs,
 ) -> Coroutine:
-    assert client.connected
+    if not session.connected:
+        session.connect()
 
     # Skip self argument if function is a method.
     _args = args[1:] if hasattr(function, "__self__") else args
@@ -98,28 +157,29 @@ def _put(
     # pool, so we set the `__wool_remote__` flag to true.
     kwargs["__wool_remote__"] = True
 
-    task = WoolTask(
+    task = Task(
         id=uuid4(),
         callable=function,
         args=args,
         kwargs=kwargs,
         tag=f"{module}.{qualname}({signature})",
     )
-    assert isinstance(client, WoolClient)
-    future: WoolFuture = client.put(task)
+    assert isinstance(session, WorkerPoolSession)
+    future: Future = session.put(task)
 
-    async def coroutine(future):
+    @wraps(function)
+    async def coroutine(future: Future):
         try:
             while not future.done():
                 await asyncio.sleep(0)
             else:
-                try:
-                    return future.result()
-                except Exception as e:
-                    raise Exception().with_traceback(e.__traceback__) from e
+                return future.result()
+        except ConnectionResetError as e:
+            raise asyncio.CancelledError from e
         except asyncio.CancelledError:
-            logging.debug("Cancelling...")
-            future.cancel()
+            if not future.done():
+                logging.debug("Cancelling...")
+                future.cancel()
             raise
 
     return coroutine(future)
@@ -133,50 +193,78 @@ def _execute(fn: AsyncCallable, parent, *args, **kwargs):
 
 
 # PUBLIC
-def current_task() -> WoolTask | None:
+def current_task() -> Task | None:
     """
     Get the current task from the context variable if we are inside a task
     context, otherwise return None.
+
+    :return: The current task or None if no task is active.
     """
     return _current_task.get()
 
 
 # PUBLIC
 @dataclass
-class WoolTask:
+class Task:
+    """
+    Represents a task to be executed in the worker pool.
+
+    :param id: The unique identifier for the task.
+    :param callable: The asynchronous function to execute.
+    :param args: Positional arguments for the function.
+    :param kwargs: Keyword arguments for the function.
+    :param timeout: The timeout for the task in seconds.
+        Defaults to 0 (no timeout).
+    :param caller: The ID of the calling task, if any.
+    :param exception: The exception raised during task execution, if any.
+    :param filename: The filename where the task was defined.
+    :param function: The name of the function being executed.
+    :param line_no: The line number where the task was defined.
+    :param tag: An optional tag for the task.
+    """
+
     id: UUID
     callable: AsyncCallable
     args: Args
     kwargs: Kwargs
     timeout: Timeout = 0
     caller: UUID | None = None
-    exception: WoolTaskException | None = None
+    exception: TaskException | None = None
     filename: str | None = None
     function: str | None = None
     line_no: int | None = None
     tag: str | None = None
 
     def __post_init__(self, **kwargs):
+        """
+        Initialize the task and emit a "task-created" event.
+
+        :param kwargs: Additional keyword arguments.
+        """
         if caller := _current_task.get():
             self.caller = caller.id
-        WoolTaskEvent("task-created", task=self).emit()
+        TaskEvent("task-created", task=self).emit()
 
     def __enter__(self) -> Callable[[], Coroutine]:
-        logging.info(f"Entering {self.__class__.__name__} with ID {self.id}")
-        WoolTaskEvent("task-queued", task=self).emit()
+        """
+        Enter the context of the task.
+
+        :return: The task's run method.
+        """
+        logging.debug(f"Entering {self.__class__.__name__} with ID {self.id}")
         return self.run
 
     def __exit__(
         self,
-        exception_type: type,
-        exception_value: Exception,
-        exception_traceback: TracebackType,
+        exception_type: type[BaseException] | None,
+        exception_value: BaseException | None,
+        exception_traceback: TracebackType | None,
     ):
-        logging.info(f"Exiting {self.__class__.__name__} with ID {self.id}")
+        logging.debug(f"Exiting {self.__class__.__name__} with ID {self.id}")
         if exception_value:
             this = asyncio.current_task()
             assert this
-            self.exception = WoolTaskException(
+            self.exception = TaskException(
                 exception_type.__qualname__,
                 traceback=[
                     y
@@ -189,9 +277,14 @@ class WoolTask:
             this.add_done_callback(self._finish, context=Context())
 
     def _finish(self, _):
-        WoolTaskEvent("task-completed", task=self).emit()
+        TaskEvent("task-completed", task=self).emit()
 
     def run(self) -> Coroutine:
+        """
+        Execute the task's callable with its arguments.
+
+        :return: A coroutine representing the task execution.
+        """
         work = self._with_task(self.callable)
         return work(*self.args, **self.kwargs)
 
@@ -210,67 +303,29 @@ class WoolTask:
 
 
 # PUBLIC
-class WoolTaskEvent:
-    """
-    Task events are emitted when a task is created, queued, started, stopped,
-    and completed. Tasks can be started and stopped multiple times by the event
-    loop. The cumulative time between start and stop events can be used to
-    determine a task's CPU utilization.
+@dataclass
+class TaskException:
     """
+    Represents an exception raised during task execution.
 
-    type: WoolTaskEventType
-    task: WoolTask
-
-    _handlers: dict[str, list[WoolTaskEventCallback]] = {}
-
-    def __init__(self, type: WoolTaskEventType, /, task: WoolTask) -> None:
-        self.type = type
-        self.task = task
-
-    @classmethod
-    def handler(
-        cls, *event_types: WoolTaskEventType
-    ) -> PassthroughDecorator[WoolTaskEventCallback]:
-        def _handler(
-            fn: WoolTaskEventCallback,
-        ) -> WoolTaskEventCallback:
-            for event_type in event_types:
-                cls._handlers.setdefault(event_type, []).append(fn)
-            return fn
-
-        return _handler
-
-    def emit(self):
-        logging.debug(f"Emitting {self.type} event for task {self.task.id}")
-        if handlers := self._handlers.get(self.type):
-            timestamp = perf_counter_ns()
-            for handler in handlers:
-                handler(self, timestamp)
-
+    :param type: The type of the exception.
+    :param traceback: The traceback of the exception.
+    """
 
-# PUBLIC
-WoolTaskEventType = Literal[
-    "task-created",
-    "task-queued",
-    "task-started",
-    "task-stopped",
-    "task-completed",
-]
+    type: str
+    traceback: list[str]
 
 
 # PUBLIC
-class WoolTaskEventCallback(Protocol):
-    def __call__(self, event: WoolTaskEvent, timestamp: Timestamp) -> None: ...
-
+class TaskEventCallback(Protocol):
+    """
+    Protocol for WoolTaskEvent callback functions.
+    """
 
-# PUBLIC
-@dataclass
-class WoolTaskException:
-    type: str
-    traceback: list[str]
+    def __call__(self, event: TaskEvent, timestamp: Timestamp) -> None: ...
 
 
-_current_task: ContextVar[WoolTask | None] = ContextVar(
+_current_task: ContextVar[Task | None] = ContextVar(
     "_current_task", default=None
 )
 
@@ -279,11 +334,11 @@ def _run(fn):
     @wraps(fn)
     def wrapper(self, *args, **kwargs):
         if current_task := self._context.get(_current_task):
-            WoolTaskEvent("task-started", task=current_task).emit()
+            TaskEvent("task-started", task=current_task).emit()
             try:
                 result = fn(self, *args, **kwargs)
             finally:
-                WoolTaskEvent("task-stopped", task=current_task).emit()
+                TaskEvent("task-stopped", task=current_task).emit()
             return result
         else:
             return fn(self, *args, **kwargs)
@@ -298,13 +353,14 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
-def resolve(method: Callable[P, R]) -> tuple[Any, Callable[P, R]]:
-    """
-    Make static and class methods picklable from within their decorators.
-    """
+def _resolve(
+    method: Callable[P, R],
+) -> Tuple[Type | ModuleType | None, Callable[P, R]]:
     scope = modules[method.__module__]
     parent = None
     for name in method.__qualname__.split("."):
         parent = scope
         scope = getattr(scope, name)
+        assert scope
+    assert isinstance(parent, (Type, ModuleType))
     return parent, cast(Callable[P, R], scope)

wool/_typing.py

@@ -1,4 +1,8 @@
-from typing import Annotated, Callable, Literal, SupportsInt, TypeVar
+from typing import Annotated
+from typing import Callable
+from typing import Literal
+from typing import SupportsInt
+from typing import TypeVar
 
 from annotated_types import Gt
 

wool/_utils.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import threading
-from typing import Callable, Final, Generic, TypeVar
+from typing import Callable
+from typing import Final
+from typing import Generic
+from typing import TypeVar
 
 
 class UndefinedType:
@@ -23,12 +26,6 @@ Undefined: Final = UndefinedType()
 
 
 class PredicatedEvent(threading.Event):
-    """
-    This class extends the threading.Event class and adds functionality to
-    automatically set the event if the given predicate function returns True
-    when the event's state is checked.
-    """
-
     def __init__(self, predicate: Callable[[], bool]):
         self._predicate = predicate
         super().__init__()
@@ -53,18 +50,14 @@ class Property(Generic[T]):
         self._default = default
 
     def get(self) -> T:
-        if self._value is Undefined and self._default is Undefined:
-            raise LookupError
-        elif self._value is Undefined:
-            assert not isinstance(self._default, UndefinedType)
+        if isinstance(self._value, UndefinedType):
+            if isinstance(self._default, UndefinedType):
+                raise ValueError("Property value is undefined")
             return self._default
-        else:
-            assert not isinstance(self._value, UndefinedType)
-            return self._value
+        return self._value
 
     def set(self, value: T) -> None:
-        if self._value is Undefined:
-            self._value = value
+        self._value = value
 
-    def unset(self) -> None:
+    def reset(self) -> None:
         self._value = Undefined