wool 0.1rc20__py3-none-any.whl

wool/core/worker/process.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import asyncio
+import signal
+from contextlib import contextmanager
+from functools import partial
+from multiprocessing import Pipe
+from multiprocessing import Process
+from multiprocessing.connection import Connection
+from typing import TYPE_CHECKING
+
+import grpc.aio
+
+import wool
+from wool._resource_pool import ResourcePool
+from wool.core import protobuf as pb
+from wool.core.worker.service import WorkerService
+
+if TYPE_CHECKING:
+    from wool.core.worker.proxy import WorkerProxy
+
+
+class WorkerProcess(Process):
+    """Subprocess hosting a gRPC worker server.
+
+    Isolated Python process running a gRPC server for task execution.
+    Maintains its own event loop and serves as an independent worker node.
+
+    Communicates the bound port back to the parent process via pipe after
+    startup. Handles SIGTERM and SIGINT for graceful shutdown.
+
+    :param host:
+        Host address to bind.
+    :param port:
+        Port to bind. 0 for random available port.
+    :param shutdown_grace_period:
+        Graceful shutdown timeout in seconds.
+    :param proxy_pool_ttl:
+        Proxy pool TTL in seconds.
+    :param args:
+        Additional args for :class:`multiprocessing.Process`.
+    :param kwargs:
+        Additional kwargs for :class:`multiprocessing.Process`.
+    """
+
+    _port: int | None
+    _get_port: Connection
+    _set_port: Connection
+    _shutdown_grace_period: float
+    _proxy_pool_ttl: float
+
+    def __init__(
+        self,
+        *args,
+        host: str = "127.0.0.1",
+        port: int = 0,
+        shutdown_grace_period: float = 60.0,
+        proxy_pool_ttl: float = 60.0,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        if not host:
+            raise ValueError("Host must be a non-blank string")
+        self._host = host
+        if port < 0 or port > 65535:
+            raise ValueError("Port must be a positive integer")
+        self._port = port
+        if shutdown_grace_period <= 0:
+            raise ValueError("Shutdown grace period must be positive")
+        self._shutdown_grace_period = shutdown_grace_period
+        if proxy_pool_ttl <= 0:
+            raise ValueError("Proxy pool TTL must be positive")
+        self._proxy_pool_ttl = proxy_pool_ttl
+        self._get_port, self._set_port = Pipe(duplex=False)
+
+    @property
+    def address(self) -> str | None:
+        """The network address where the gRPC server is listening.
+
+        :returns:
+            The address in "host:port" format, or None if not started.
+        """
+        return self._address(self._host, self._port) if self._port else None
+
+    @property
+    def host(self) -> str | None:
+        """The host where the gRPC server is listening.
+
+        :returns:
+            The host address, or None if not started.
+        """
+        return self._host
+
+    @property
+    def port(self) -> int | None:
+        """The port where the gRPC server is listening.
+
+        :returns:
+            The port number, or None if not started.
+        """
+        return self._port or None
+
+    def start(self, *, timeout: float | None = None):
+        """Start the worker process.
+
+        Launches the worker process and waits until it has started
+        listening on a port. After starting, the :attr:`address`
+        property will contain the actual network address.
+
+        :param timeout:
+            Maximum time in seconds to wait for worker process startup.
+        :raises RuntimeError:
+            If the worker process fails to start within the timeout.
+        :raises ValueError:
+            If the timeout is not positive.
+        """
+        if timeout is not None and timeout <= 0:
+            raise ValueError("Timeout must be positive")
+        super().start()
+        if self._get_port.poll(timeout=timeout):
+            self._port = self._get_port.recv()
+        else:
+            self.terminate()
+            self.join()
+            raise RuntimeError(
+                f"Worker process failed to start within {timeout} seconds"
+            )
+        self._get_port.close()
+
+    def run(self) -> None:
+        """Run the worker process.
+
+        Sets the event loop for this process and starts the gRPC server,
+        blocking until the server is stopped.
+        """
+        wool.__proxy_pool__.set(
+            ResourcePool(
+                factory=_proxy_factory,
+                finalizer=_proxy_finalizer,
+                ttl=self._proxy_pool_ttl,
+            )
+        )
+        asyncio.run(self._serve())
+
+    async def _serve(self):
+        """Start the gRPC server in this worker process.
+
+        This method is called by the event loop to start serving
+        requests. It creates a gRPC server, adds the worker service, and
+        starts listening for incoming connections.
+        """
+        server = grpc.aio.server()
+        port = server.add_insecure_port(self._address(self._host, self._port))
+        service = WorkerService()
+        pb.add_to_server[pb.worker.WorkerServicer](service, server)
+
+        with _signal_handlers(service):
+            try:
+                await server.start()
+                try:
+                    self._set_port.send(port)
+                finally:
+                    self._set_port.close()
+                await service.stopped.wait()
+            finally:
+                await server.stop(grace=self._shutdown_grace_period)
+
+    def _address(self, host, port) -> str:
+        """Format network address for the given port.
+
+        :param port:
+            Port number to include in the address.
+        :returns:
+            Address string in "host:port" format.
+        """
+        return f"{host}:{port}"
+
+
+@contextmanager
+def _signal_handlers(service: WorkerService):
+    """Context manager for setting up signal handlers for graceful shutdown.
+
+    Installs SIGTERM and SIGINT handlers that gracefully shut down the worker
+    service when the process receives termination signals.
+
+    :param service:
+        The :class:`WorkerService` instance to shut down on signal receipt.
+    :yields:
+        Control to the calling context with signal handlers installed.
+    """
+    loop = asyncio.get_running_loop()
+
+    old_sigterm = signal.signal(signal.SIGTERM, partial(_sigterm_handler, loop, service))
+    old_sigint = signal.signal(signal.SIGINT, partial(_sigint_handler, loop, service))
+    try:
+        yield
+    finally:
+        signal.signal(signal.SIGTERM, old_sigterm)
+        signal.signal(signal.SIGINT, old_sigint)
+
+
+def _sigterm_handler(loop, service, signum, frame):
+    if loop.is_running():
+        loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(
+                service.stop(pb.worker.StopRequest(timeout=0), None)
+            )
+        )
+
+
+def _sigint_handler(loop, service, signum, frame):
+    if loop.is_running():
+        loop.call_soon_threadsafe(
+            lambda: asyncio.create_task(
+                service.stop(pb.worker.StopRequest(timeout=None), None)
+            )
+        )
+
+
+async def _proxy_factory(proxy: WorkerProxy):
+    """Factory function for WorkerProxy instances in ResourcePool.
+
+    Starts the proxy if not already started and returns it.
+    The proxy object itself is used as the cache key.
+
+    :param proxy:
+        The WorkerProxy instance to start (passed as key from
+        ResourcePool).
+    :returns:
+        The started WorkerProxy instance.
+    """
+    if not proxy.started:
+        await proxy.start()
+    return proxy
+
+
+async def _proxy_finalizer(proxy: WorkerProxy):
+    """Finalizer function for WorkerProxy instances in ResourcePool.
+
+    Stops the proxy when it's being cleaned up from the resource pool.
+    Based on the cleanup logic from WorkerProxyCache._delayed_cleanup.
+
+    :param proxy:
+        The WorkerProxy instance to clean up.
+    """
+    try:
+        await proxy.stop()
+    except Exception:
+        pass

wool/core/worker/proxy.py

@@ -0,0 +1,427 @@
+from __future__ import annotations
+
+import asyncio
+import uuid
+from typing import TYPE_CHECKING
+from typing import AsyncContextManager
+from typing import AsyncIterator
+from typing import Awaitable
+from typing import ContextManager
+from typing import Generic
+from typing import Sequence
+from typing import TypeAlias
+from typing import TypeVar
+from typing import overload
+
+import wool
+from wool._resource_pool import ResourcePool
+from wool.core.discovery.base import DiscoveryEvent
+from wool.core.discovery.base import DiscoverySubscriberLike
+from wool.core.discovery.base import WorkerInfo
+from wool.core.discovery.local import LocalDiscovery
+from wool.core.loadbalancer.base import LoadBalancerContext
+from wool.core.loadbalancer.base import LoadBalancerLike
+from wool.core.loadbalancer.roundrobin import RoundRobinLoadBalancer
+from wool.core.typing import Factory
+from wool.core.worker.connection import WorkerConnection
+
+if TYPE_CHECKING:
+    from wool._work import WoolTask
+
+T = TypeVar("T")
+
+
+class ReducibleAsyncIterator(Generic[T]):
+    """An async iterator that can be pickled via __reduce__.
+
+    Converts a sequence into an async iterator while maintaining
+    picklability for distributed task execution contexts.
+
+    :param items:
+        Sequence of items to convert to async iterator.
+    """
+
+    def __init__(self, items: Sequence[T]):
+        self._items = items
+        self._index = 0
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        return self
+
+    async def __anext__(self) -> T:
+        if self._index >= len(self._items):
+            raise StopAsyncIteration
+        item = self._items[self._index]
+        self._index += 1
+        return item
+
+    def __reduce__(self) -> tuple:
+        """Return constructor args for unpickling."""
+        return (self.__class__, (self._items,))
+
+
+async def connection_factory(target: str) -> WorkerConnection:
+    """Factory function for creating worker connections.
+
+    Creates a connection to the specified worker target.
+    The target is passed as the key from ResourcePool.
+
+    :param target:
+        The network target (host:port) to create a channel for.
+    :returns:
+        A new connection to the target.
+    """
+    return WorkerConnection(target)
+
+
+async def connection_finalizer(connection: WorkerConnection) -> None:
+    """Finalizer function for gRPC channels.
+
+    Closes the gRPC connection when it's being cleaned up from the resource pool.
+
+    :param connection:
+        The gRPC connection to close.
+    """
+    try:
+        await connection.close()
+    except Exception:
+        pass
+
+
+WorkerUri: TypeAlias = str
+
+
+# public
+class WorkerProxy:
+    """Client-side proxy for dispatching tasks to distributed workers.
+
+    Manages worker discovery, connection pooling, and load-balanced task
+    routing. The bridge between :func:`@wool.work <wool.work>` decorated
+    functions and the worker pool.
+
+    Connects to workers through discovery services, pool URIs, or static
+    worker lists. Handles connection lifecycle and fault tolerance
+    automatically.
+
+    **Connect via pool URI:**
+
+    .. code-block:: python
+
+        async with WorkerProxy("pool-abc123") as proxy:
+            result = await task()
+
+    **Connect via discovery:**
+
+    .. code-block:: python
+
+        from wool.core.discovery.lan import LanDiscovery
+
+        discovery = LanDiscovery().subscribe()
+        async with WorkerProxy(discovery=discovery) as proxy:
+            result = await task()
+
+    **Connect to static workers:**
+
+    .. code-block:: python
+
+        workers = [
+            WorkerInfo(host="10.0.0.1", port=50051, ...),
+            WorkerInfo(host="10.0.0.2", port=50051, ...),
+        ]
+        async with WorkerProxy(workers=workers) as proxy:
+            result = await task()
+
+    **Custom load balancer:**
+
+    .. code-block:: python
+
+        from wool.core.loadbalancer.roundrobin import RoundRobinLoadBalancer
+
+
+        class CustomBalancer(RoundRobinLoadBalancer):
+            async def dispatch(self, task, context, timeout=None):
+                # Custom routing strategy
+                ...
+
+
+        async with WorkerProxy(
+            discovery=discovery,
+            loadbalancer=CustomBalancer(),
+        ) as proxy:
+            result = await task()
+
+    :param pool_uri:
+        Pool identifier for discovery-based connection.
+    :param tags:
+        Additional tags for filtering discovered workers.
+    :param discovery:
+        Discovery service or event stream.
+    :param workers:
+        Static worker list for direct connection.
+    :param loadbalancer:
+        Load balancer instance, factory, or context manager.
+    """
+
+    _discovery: DiscoverySubscriberLike | Factory[DiscoverySubscriberLike]
+    _discovery_manager: (
+        AsyncContextManager[DiscoverySubscriberLike]
+        | ContextManager[DiscoverySubscriberLike]
+    )
+
+    _loadbalancer = LoadBalancerLike | Factory[LoadBalancerLike]
+    _loadbalancer_manager: (
+        AsyncContextManager[LoadBalancerLike] | ContextManager[LoadBalancerLike]
+    )
+
+    @overload
+    def __init__(
+        self,
+        *,
+        discovery: DiscoverySubscriberLike | Factory[DiscoverySubscriberLike],
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        workers: Sequence[WorkerInfo],
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        pool_uri: str,
+        *tags: str,
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
+    ): ...
+
+    def __init__(
+        self,
+        pool_uri: str | None = None,
+        *tags: str,
+        discovery: (
+            DiscoverySubscriberLike | Factory[DiscoverySubscriberLike] | None
+        ) = None,
+        workers: Sequence[WorkerInfo] | None = None,
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
+    ):
+        if not (pool_uri or discovery or workers):
+            raise ValueError(
+                "Must specify either a workerpool URI, discovery event stream, or a "
+                "sequence of workers"
+            )
+
+        self._id: uuid.UUID = uuid.uuid4()
+        self._started = False
+        self._loadbalancer = loadbalancer
+
+        match (pool_uri, discovery, workers):
+            case (pool_uri, None, None) if pool_uri is not None:
+                self._discovery = LocalDiscovery(pool_uri).subscribe(
+                    filter=lambda w: bool({pool_uri, *tags} & w.tags)
+                )
+            case (None, discovery, None) if discovery is not None:
+                self._discovery = discovery
+            case (None, None, workers) if workers is not None:
+                self._discovery = ReducibleAsyncIterator(
+                    [DiscoveryEvent(type="worker-added", worker_info=w) for w in workers]
+                )
+            case _:
+                raise ValueError(
+                    "Must specify exactly one of: "
+                    "pool_uri, discovery_event_stream, or workers"
+                )
+        self._sentinel_task: asyncio.Task[None] | None = None
+        self._loadbalancer_context: LoadBalancerContext | None = None
+
+    async def __aenter__(self):
+        """Starts the proxy and sets it as the active context."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, *args):
+        """Stops the proxy and resets the active context."""
+        await self.stop(*args)
+
+    def __hash__(self) -> int:
+        return hash(str(self.id))
+
+    def __eq__(self, value: object) -> bool:
+        return isinstance(value, WorkerProxy) and hash(self) == hash(value)
+
+    def __reduce__(self) -> tuple:
+        """Return constructor args for unpickling with proxy ID preserved.
+
+        Creates a new WorkerProxy instance with the same discovery stream and
+        load balancer type, then sets the preserved proxy ID on the new object.
+        Workers will be re-discovered on the new instance.
+
+        :returns:
+            Tuple of (callable, args, state) for unpickling.
+        """
+
+        def _restore_proxy(discovery, loadbalancer, proxy_id):
+            proxy = WorkerProxy(discovery=discovery, loadbalancer=loadbalancer)
+            proxy._id = proxy_id
+            return proxy
+
+        return (
+            _restore_proxy,
+            (self._discovery, self._loadbalancer, self._id),
+        )
+
+    @property
+    def id(self) -> uuid.UUID:
+        return self._id
+
+    @property
+    def started(self) -> bool:
+        return self._started
+
+    @property
+    def workers(self) -> list[WorkerInfo]:
+        """A list of the currently discovered worker gRPC stubs."""
+        if self._loadbalancer_context:
+            return list(self._loadbalancer_context.workers.keys())
+        else:
+            return []
+
+    async def start(self) -> None:
+        """Starts the proxy by initiating the worker discovery process.
+
+        :raises RuntimeError:
+            If the proxy has already been started.
+        """
+        if self._started:
+            raise RuntimeError("Proxy already started")
+
+        (
+            self._loadbalancer_service,
+            self._loadbalancer_context_manager,
+        ) = await self._enter_context(self._loadbalancer)
+        if not isinstance(self._loadbalancer_service, LoadBalancerLike):
+            raise ValueError
+
+        (
+            self._discovery_stream,
+            self._discovery_context_manager,
+        ) = await self._enter_context(self._discovery)
+        if not isinstance(self._discovery_stream, DiscoverySubscriberLike):
+            raise ValueError
+
+        self._proxy_token = wool.__proxy__.set(self)
+        self._connection_pool = ResourcePool(
+            factory=connection_factory, finalizer=connection_finalizer, ttl=60
+        )
+        self._loadbalancer_context = LoadBalancerContext()
+        self._sentinel_task = asyncio.create_task(self._worker_sentinel())
+        self._started = True
+
+    async def stop(self, *args) -> None:
+        """Stops the proxy, terminating discovery and clearing connections.
+
+        :raises RuntimeError:
+            If the proxy was not started first.
+        """
+        if not self._started:
+            raise RuntimeError("Proxy not started - call start() first")
+
+        await self._exit_context(self._discovery_context_manager, *args)
+        await self._exit_context(self._loadbalancer_context_manager, *args)
+
+        wool.__proxy__.reset(self._proxy_token)
+        if self._sentinel_task:
+            self._sentinel_task.cancel()
+            try:
+                await self._sentinel_task
+            except asyncio.CancelledError:
+                pass
+            self._sentinel_task = None
+        await self._connection_pool.clear()
+        self._loadbalancer_context = None
+        self._started = False
+
+    async def dispatch(self, task: WoolTask, *, timeout: float | None = None):
+        """Dispatches a task to an available worker in the pool.
+
+        This method selects a worker using a round-robin strategy. If no
+        workers are available within the timeout period, it raises an
+        exception.
+
+        :param task:
+            The :class:`WoolTask` object to be dispatched.
+        :param timeout:
+            Timeout in seconds for getting a worker.
+        :returns:
+            A protobuf result object from the worker.
+        :raises RuntimeError:
+            If the proxy is not started.
+        :raises asyncio.TimeoutError:
+            If no worker is available within the timeout period.
+        """
+        if not self._started:
+            raise RuntimeError("Proxy not started - call start() first")
+
+        await asyncio.wait_for(self._await_workers(), 60)
+
+        assert isinstance(self._loadbalancer_service, LoadBalancerLike)
+        assert self._loadbalancer_context
+        return await self._loadbalancer_service.dispatch(
+            task, context=self._loadbalancer_context, timeout=timeout
+        )
+
+    async def _enter_context(self, factory):
+        ctx = None
+        if isinstance(factory, ContextManager):
+            ctx = factory
+            obj = ctx.__enter__()
+        elif isinstance(factory, AsyncContextManager):
+            ctx = factory
+            obj = await ctx.__aenter__()
+        elif callable(factory):
+            return await self._enter_context(factory())
+        elif isinstance(factory, Awaitable):
+            obj = await factory
+        else:
+            obj = factory
+        return obj, ctx
+
+    async def _exit_context(
+        self, ctx: AsyncContextManager | ContextManager | None, *args
+    ):
+        if isinstance(ctx, AsyncContextManager):
+            await ctx.__aexit__(*args)
+        elif isinstance(ctx, ContextManager):
+            ctx.__exit__(*args)
+
+    async def _await_workers(self):
+        while not self._loadbalancer_context or not self._loadbalancer_context.workers:
+            await asyncio.sleep(0)
+
+    async def _worker_sentinel(self):
+        assert self._loadbalancer_context
+        async for event in self._discovery_stream:
+            match event.type:
+                case "worker-added":
+                    self._loadbalancer_context.add_worker(
+                        event.worker_info,
+                        lambda: self._connection_pool.get(
+                            f"{event.worker_info.host}:{event.worker_info.port}",
+                        ),
+                    )
+                case "worker-updated":
+                    self._loadbalancer_context.update_worker(
+                        event.worker_info,
+                        lambda: self._connection_pool.get(
+                            f"{event.worker_info.host}:{event.worker_info.port}",
+                        ),
+                    )
+                case "worker-dropped":
+                    self._loadbalancer_context.remove_worker(event.worker_info)