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

wool/__init__.py

@@ -1,23 +1,30 @@
 import logging
 from contextvars import ContextVar
 from importlib.metadata import entry_points
-from typing import Final, Literal
+from typing import Final
 
-from wool._cli import WoolPoolCommand, cli
-from wool._client import NullClient, WoolClient
-from wool._future import WoolFuture
+from tblib import pickling_support
+
+from wool._cli import WorkerPoolCommand
+from wool._cli import cli
+from wool._future import Future
 from wool._logging import __log_format__
-from wool._pool import WoolPool
-from wool._task import (
-    WoolTask,
-    WoolTaskEvent,
-    WoolTaskEventCallback,
-    WoolTaskException,
-    current_task,
-    task,
-)
+from wool._pool import WorkerPool
+from wool._pool import pool
+from wool._session import LocalSession
+from wool._session import WorkerPoolSession
+from wool._session import session
+from wool._task import Task
+from wool._task import TaskEvent
+from wool._task import TaskEventCallback
+from wool._task import TaskException
+from wool._task import current_task
+from wool._task import task
+from wool._worker import Scheduler
 from wool._worker import Worker
 
+pickling_support.install()
+
 # PUBLIC
 __log_format__: str = __log_format__
 
@@ -25,26 +32,29 @@ __log_format__: str = __log_format__
 __log_level__: int = logging.INFO
 
 # PUBLIC
-__wool_client__: Final[ContextVar[WoolClient]] = ContextVar(
-    "__wool_client__", default=NullClient()
+__wool_session__: Final[ContextVar[WorkerPoolSession]] = ContextVar(
+    "__wool_session__", default=LocalSession()
 )
 
-__wool_worker__: Worker | Literal[True] | None = None
+__wool_worker__: Worker | None = None
 
 __all__ = [
-    "WoolTaskException",
-    "WoolFuture",
-    "WoolTask",
-    "WoolTaskEvent",
-    "WoolTaskEventCallback",
-    "WoolPool",
-    "WoolClient",
-    "WoolPoolCommand",
+    "TaskException",
+    "Future",
+    "Task",
+    "TaskEvent",
+    "TaskEventCallback",
+    "WorkerPool",
+    "WorkerPoolSession",
+    "WorkerPoolCommand",
+    "Scheduler",
     "__log_format__",
     "__log_level__",
-    "__wool_client__",
+    "__wool_session__",
     "cli",
     "current_task",
+    "pool",
+    "session",
     "task",
 ]
 
@@ -57,10 +67,10 @@ for symbol in __all__:
     except AttributeError:
         continue
 
-for plugin in entry_points(group="wool.cli.plugins"):
+for plugin in entry_points(group="wool_cli_plugins"):
     try:
         plugin.load()
         logging.info(f"Loaded CLI plugin {plugin.name}")
     except Exception as e:
         logging.error(f"Failed to load CLI plugin {plugin.name}: {e}")
-        continue
+        raise

wool/_cli.py

@@ -9,15 +9,23 @@ from time import perf_counter
 import click
 
 import wool
-from wool._client import WoolClient
-from wool._pool import WoolPool
+from wool._pool import WorkerPool
+from wool._session import WorkerPoolSession
 from wool._task import task
 
 DEFAULT_PORT = 48800
 
 
 # PUBLIC
-class WoolPoolCommand(click.core.Command):
+class WorkerPoolCommand(click.core.Command):
+    """
+    Custom Click command class for worker pool commands.
+
+    :param default_host: Default host address.
+    :param default_port: Default port number.
+    :param default_authkey: Default authentication key.
+    """
+
     def __init__(
         self,
         *args,
@@ -43,7 +51,11 @@ class WoolPoolCommand(click.core.Command):
 
 @contextmanager
 def timer():
-    """Context manager to measure the execution time of a code block."""
+    """
+    Context manager to measure the execution time of a code block.
+
+    :return: A function to retrieve the elapsed time.
+    """
     start = end = perf_counter()
     yield lambda: end - start
     end = perf_counter()
@@ -53,13 +65,10 @@ def to_bytes(context: click.Context, parameter: click.Parameter, value: str):
     """
     Convert the given value to bytes.
 
-    Args:
-        context (click.Context): Click context.
-        parameter (click.Parameter): Click parameter.
-        value (str): Value to convert.
-
-    Returns:
-        bytes: The converted value in bytes.
+    :param context: Click context.
+    :param parameter: Click parameter.
+    :param value: Value to convert.
+    :return: The converted value in bytes.
     """
     if value is None:
         return b""
@@ -72,13 +81,10 @@ def assert_nonzero(
     """
     Assert that the given value is non-zero.
 
-    Args:
-        context (click.Context): Click context.
-        parameter (click.Parameter): Click parameter.
-        value (int): Value to check.
-
-    Returns:
-        int: The original value if it is non-zero.
+    :param context: Click context.
+    :param parameter: Click parameter.
+    :param value: Value to check.
+    :return: The original value if it is non-zero.
     """
     if value is None:
         return value
@@ -88,12 +94,11 @@ def assert_nonzero(
 
 def debug(ctx, param, value):
     """
-    Enable debugging with debugpy.
+    Enable debugging mode with a specified port.
 
-    Args:
-        ctx (click.Context): Click context.
-        param (click.Parameter): Click parameter.
-        value (bool): Flag value indicating whether to enable debugging.
+    :param ctx: The Click context object.
+    :param param: The parameter being handled.
+    :param value: The port number for the debugger.
     """
     if not value or ctx.resilient_parsing:
         return
@@ -112,7 +117,10 @@ def debug(ctx, param, value):
     "-d",
     callback=debug,
     expose_value=False,
-    help="Run with debugger listening on the specified port. Execution will block until the debugger is attached.",
+    help=(
+        "Run with debugger listening on the specified port. Execution will "
+        "block until the debugger is attached."
+    ),
     is_eager=True,
     type=int,
 )
@@ -128,10 +136,7 @@ def cli(verbosity: int):
     """
     CLI command group with options for verbosity, debugging, and version.
 
-    Args:
-        verbosity (int): Verbosity level for logging.
-        debug (bool): Flag to enable debugging.
-        version (bool): Flag to display the version and exit.
+    :param verbosity: Verbosity level for logging.
     """
     match verbosity:
         case 4:
@@ -150,10 +155,14 @@ def cli(verbosity: int):
 
 
 @cli.group()
-def pool(): ...
+def pool():
+    """
+    CLI command group for managing worker pools.
+    """
+    pass
 
 
-@pool.command(cls=partial(WoolPoolCommand, default_port=DEFAULT_PORT))
+@pool.command(cls=partial(WorkerPoolCommand, default_port=DEFAULT_PORT))
 @click.option(
     "--breadth", "-b", type=int, default=cpu_count(), callback=assert_nonzero
 )
@@ -163,14 +172,26 @@ def pool(): ...
     "-m",
     multiple=True,
     type=str,
-    help="Python module containing workerpool task definitions to be executed by this pool.",
+    help=(
+        "Python module containing workerpool task definitions to be executed "
+        "by this pool."
+    ),
 )
 def up(host, port, authkey, breadth, modules):
+    """
+    Start a worker pool with the specified configuration.
+
+    :param host: The host address for the worker pool.
+    :param port: The port number for the worker pool.
+    :param authkey: The authentication key for the worker pool.
+    :param breadth: The number of worker processes in the pool.
+    :param modules: Python modules containing task definitions.
+    """
     for module in modules:
         importlib.import_module(module)
     if not authkey:
         logging.warning("No authkey specified")
-    workerpool = WoolPool(
+    workerpool = WorkerPool(
         address=(host, port),
         breadth=breadth,
         authkey=authkey,
@@ -180,7 +201,7 @@ def up(host, port, authkey, breadth, modules):
     workerpool.join()
 
 
-@pool.command(cls=partial(WoolPoolCommand, default_port=DEFAULT_PORT))
+@pool.command(cls=partial(WorkerPoolCommand, default_port=DEFAULT_PORT))
 @click.option(
     "--wait",
     "-w",
@@ -189,17 +210,32 @@ def up(host, port, authkey, breadth, modules):
     help="Wait for in-flight tasks to complete before shutting down.",
 )
 def down(host, port, authkey, wait):
+    """
+    Shut down the worker pool.
+
+    :param host: The host address of the worker pool.
+    :param port: The port number of the worker pool.
+    :param authkey: The authentication key for the worker pool.
+    :param wait: Whether to wait for in-flight tasks to complete.
+    """
     assert port
     if not host:
         host = "localhost"
     if not authkey:
         authkey = b""
-    client = WoolClient(address=(host, port), authkey=authkey).connect()
-    client.stop(wait=wait)
+    with WorkerPoolSession(address=(host, port), authkey=authkey) as client:
+        client.stop(wait=wait)
 
 
-@cli.command(cls=partial(WoolPoolCommand, default_port=DEFAULT_PORT))
+@cli.command(cls=partial(WorkerPoolCommand, default_port=DEFAULT_PORT))
 def ping(host, port, authkey):
+    """
+    Ping the worker pool to check connectivity.
+
+    :param host: The host address of the worker pool.
+    :param port: The port number of the worker pool.
+    :param authkey: The authentication key for the worker pool.
+    """
     assert port
     if not host:
         host = "localhost"
@@ -208,14 +244,19 @@ def ping(host, port, authkey):
 
     async def _():
         with timer() as t:
-            await asyncio.create_task(_ping())
+            await _ping()
 
         print(f"Ping: {int((t() * 1000) + 0.5)} ms")
 
-    with WoolClient(address=(host, port), authkey=authkey):
+    with WorkerPoolSession(address=(host, port), authkey=authkey):
         asyncio.get_event_loop().run_until_complete(_())
 
 
 @task
 async def _ping():
+    """
+    Asynchronous task to log a ping message.
+
+    :return: None
+    """
     logging.debug("Ping!")

wool/_event.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import logging
+from time import perf_counter_ns
+from typing import TYPE_CHECKING
+from typing import Literal
+
+from wool._typing import PassthroughDecorator
+
+if TYPE_CHECKING:
+    from wool._task import Task
+    from wool._task import TaskEventCallback
+
+
+# PUBLIC
+class TaskEvent:
+    """
+    Represents an event related to a Wool task, such as creation, queuing,
+    scheduling, starting, stopping, or completion.
+
+    Task events are emitted during the lifecycle of a task. These events can
+    be used to track task execution and measure performance, such as CPU
+    utilization.
+
+    :param type: The type of the task event.
+    :param task: The task associated with the event.
+    """
+
+    type: TaskEventType
+    task: Task
+
+    _handlers: dict[str, list[TaskEventCallback]] = {}
+
+    def __init__(self, type: TaskEventType, /, task: Task) -> None:
+        """
+        Initialize a WoolTaskEvent instance.
+
+        :param type: The type of the task event.
+        :param task: The task associated with the event.
+        """
+        self.type = type
+        self.task = task
+
+    @classmethod
+    def handler(
+        cls, *event_types: TaskEventType
+    ) -> PassthroughDecorator[TaskEventCallback]:
+        """
+        Register a handler function for specific task event types.
+
+        :param event_types: The event types to handle.
+        :return: A decorator to register the handler function.
+        """
+
+        def _handler(
+            fn: TaskEventCallback,
+        ) -> TaskEventCallback:
+            for event_type in event_types:
+                cls._handlers.setdefault(event_type, []).append(fn)
+            return fn
+
+        return _handler
+
+    def emit(self):
+        """
+        Emit the task event, invoking all registered handlers for the event
+        type.
+
+        Handlers are called with the event instance and a timestamp.
+
+        :raises Exception: If any handler raises an exception.
+        """
+        logging.debug(
+            f"Emitting {self.type} event for "
+            f"task {self.task.id} "
+            f"({self.task.callable.__qualname__})"
+        )
+        if handlers := self._handlers.get(self.type):
+            timestamp = perf_counter_ns()
+            for handler in handlers:
+                handler(self, timestamp)
+
+
+# PUBLIC
+TaskEventType = Literal[
+    "task-created",
+    "task-queued",
+    "task-scheduled",
+    "task-started",
+    "task-stopped",
+    "task-completed",
+]
+"""
+Defines the types of events that can occur during the lifecycle of a Wool 
+task.
+
+- "task-created": 
+    Emitted when a task is created.
+- "task-queued": 
+    Emitted when a task is added to the queue.
+- "task-scheduled": 
+    Emitted when a task is scheduled for execution in a worker's event loop.
+- "task-started": 
+    Emitted when a task starts execution.
+- "task-stopped": 
+    Emitted when a task stops execution.
+- "task-completed": 
+    Emitted when a task completes execution.
+"""

wool/_future.py

@@ -3,16 +3,34 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 import logging
-from typing import Any, Generator, Generic, TypeVar, cast
+from typing import Any
+from typing import Generator
+from typing import Generic
+from typing import TypeVar
+from typing import cast
 
-from wool._utils import Undefined, UndefinedType
+from wool._utils import Undefined
+from wool._utils import UndefinedType
 
 T = TypeVar("T")
 
 
 # PUBLIC
-class WoolFuture(Generic[T]):
+class Future(Generic[T]):
+    """
+    A future object representing the result of an asynchronous operation.
+
+    WoolFuture provides methods to retrieve the result or exception of an
+    asynchronous operation, set the result or exception, and await its
+    completion.
+
+    :param T: The type of the result.
+    """
+
     def __init__(self) -> None:
+        """
+        Initialize a WoolFuture instance.
+        """
         self._result: T | UndefinedType = Undefined
         self._exception: (
             BaseException | type[BaseException] | UndefinedType
@@ -21,6 +39,12 @@ class WoolFuture(Generic[T]):
         self._cancelled: bool = False
 
     def __await__(self) -> Generator[Any, None, T]:
+        """
+        Await the completion of the future.
+
+        :return: The result of the future.
+        """
+
         async def _():
             while not self.done():
                 await asyncio.sleep(0)
@@ -30,6 +54,13 @@ class WoolFuture(Generic[T]):
         return _().__await__()
 
     def result(self) -> T:
+        """
+        Retrieve the result of the future.
+
+        :return: The result of the future.
+        :raises BaseException: If the future completed with an exception.
+        :raises asyncio.InvalidStateError: If the future is not yet completed.
+        """
         if self._exception is not Undefined:
             assert (
                 isinstance(self._exception, BaseException)
@@ -43,6 +74,12 @@ class WoolFuture(Generic[T]):
             raise asyncio.InvalidStateError
 
     def set_result(self, result: T) -> None:
+        """
+        Set the result of the future.
+
+        :param result: The result to set.
+        :raises asyncio.InvalidStateError: If the future is already completed.
+        """
         if self.done():
             raise asyncio.InvalidStateError
         else:
@@ -50,7 +87,14 @@ class WoolFuture(Generic[T]):
             self._done = True
 
     def exception(self) -> BaseException | type[BaseException]:
-        if self._exception is not Undefined:
+        """
+        Retrieve the exception of the future, if any.
+
+        :return: The exception of the future.
+        :raises asyncio.InvalidStateError: If the future is not yet completed
+            or has no exception.
+        """
+        if self.done() and self._exception is not Undefined:
             return cast(BaseException | type[BaseException], self._exception)
         else:
             raise asyncio.InvalidStateError
@@ -58,6 +102,12 @@ class WoolFuture(Generic[T]):
     def set_exception(
         self, exception: BaseException | type[BaseException]
     ) -> None:
+        """
+        Set the exception of the future.
+
+        :param exception: The exception to set.
+        :raises asyncio.InvalidStateError: If the future is already completed.
+        """
         if self.done():
             raise asyncio.InvalidStateError
         else:
@@ -65,17 +115,35 @@ class WoolFuture(Generic[T]):
             self._done = True
 
     def done(self) -> bool:
-        return self._done or self._cancelled
+        """
+        Check if the future is completed.
+
+        :return: True if the future is completed, False otherwise.
+        """
+        return self._done
 
     def cancel(self) -> None:
-        self.set_exception(asyncio.CancelledError)
-        self._cancelled = True
+        """
+        Cancel the future.
+
+        :raises asyncio.InvalidStateError: If the future is already completed.
+        """
+        if self.done():
+            raise asyncio.InvalidStateError
+        else:
+            self._cancelled = True
+            self._done = True
 
     def cancelled(self) -> bool:
+        """
+        Check if the future was cancelled.
+
+        :return: True if the future was cancelled, False otherwise.
+        """
         return self._cancelled
 
 
-async def poll(future: WoolFuture, task: concurrent.futures.Future) -> None:
+async def poll(future: Future, task: concurrent.futures.Future) -> None:
     while True:
         if future.cancelled():
             task.cancel()
@@ -86,15 +154,18 @@ async def poll(future: WoolFuture, task: concurrent.futures.Future) -> None:
             await asyncio.sleep(0)
 
 
-def fulfill(future: WoolFuture):
+def fulfill(future: Future):
     def callback(task: concurrent.futures.Future):
         try:
-            future.set_result(task.result())
+            result = task.result()
         except concurrent.futures.CancelledError:
-            if not future.cancelled():
+            if not future.done():
                 future.cancel()
         except BaseException as e:
             logging.exception(e)
             future.set_exception(e)
+        else:
+            if not future.done():
+                future.set_result(result)
 
     return callback

wool/_logging.py

@@ -12,6 +12,13 @@ def italic(text: str) -> str:
 
 class WoolLogFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
+        """
+        Modify the log record to include a reference to the source file and
+        line number.
+
+        :param record: The log record to modify.
+        :return: True to indicate the record should be logged.
+        """
         pathname: str = record.pathname
         cwd: str = os.getcwd()
         if pathname.startswith(cwd):