wool 0.1rc20__py3-none-any.whl

wool/core/worker/local.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import asyncio
+from types import MappingProxyType
+from typing import Any
+
+import grpc.aio
+
+import wool
+from wool.core import protobuf as pb
+from wool.core.discovery.base import WorkerInfo
+from wool.core.worker.base import Worker
+from wool.core.worker.process import WorkerProcess
+
+
+# public
+class LocalWorker(Worker):
+    """Worker running in a local subprocess.
+
+    Spawns a dedicated process hosting a gRPC server for task execution.
+    Handles multiple concurrent tasks in an isolated asyncio event loop.
+
+    **Basic usage:**
+
+    .. code-block:: python
+
+        worker = LocalWorker("gpu-capable")
+        await worker.start()
+        # Worker is now accepting tasks
+        await worker.stop()
+
+    **Custom configuration:**
+
+    .. code-block:: python
+
+        worker = LocalWorker(
+            "production",
+            "high-memory",
+            host="0.0.0.0",  # Listen on all interfaces
+            port=50051,  # Fixed port
+            shutdown_grace_period=30.0,
+        )
+
+    :param tags:
+        Capability tags for filtering and selection.
+    :param host:
+        Host address to bind. Defaults to localhost.
+    :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 extra:
+        Additional metadata as key-value pairs.
+    """
+
+    _worker_process: WorkerProcess
+
+    def __init__(
+        self,
+        *tags: str,
+        host: str = "127.0.0.1",
+        port: int = 0,
+        shutdown_grace_period: float = 60.0,
+        proxy_pool_ttl: float = 60.0,
+        **extra: Any,
+    ):
+        super().__init__(*tags, **extra)
+        self._worker_process = WorkerProcess(
+            host=host,
+            port=port,
+            shutdown_grace_period=shutdown_grace_period,
+            proxy_pool_ttl=proxy_pool_ttl,
+        )
+
+    @property
+    def address(self) -> str | None:
+        """The network address where the worker is listening.
+
+        :returns:
+            The address in "host:port" format, or None if not started.
+        """
+        return self._worker_process.address
+
+    @property
+    def host(self) -> str | None:
+        """The host where the worker is listening.
+
+        :returns:
+            The host address, or None if not started.
+        """
+        return self._info.host if self._info else None
+
+    @property
+    def port(self) -> int | None:
+        """The port where the worker is listening.
+
+        :returns:
+            The port number, or None if not started.
+        """
+        return self._info.port if self._info else None
+
+    async def _start(self, timeout: float | None):
+        """Start the worker process and register it with the pool.
+
+        Initializes the registrar service, starts the worker process
+        with its gRPC server, and registers the worker's network
+        address with the registrar for discovery by client sessions.
+
+        :param timeout:
+            Maximum time in seconds to wait for worker process startup.
+        """
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None, lambda t: self._worker_process.start(timeout=t), timeout
+        )
+        if not self._worker_process.address:
+            raise RuntimeError("Worker process failed to start - no address")
+        if not self._worker_process.pid:
+            raise RuntimeError("Worker process failed to start - no PID")
+
+        host, port_str = self._worker_process.address.split(":")
+        port = int(port_str)
+
+        self._info = WorkerInfo(
+            uid=self._uid,
+            host=host,
+            port=port,
+            pid=self._worker_process.pid,
+            version=wool.__version__,
+            tags=frozenset(self._tags),
+            extra=MappingProxyType(self._extra),
+        )
+
+    async def _stop(self, timeout: float | None):
+        """Stop the worker process and unregister it from the pool.
+
+        Unregisters the worker from the registrar service, gracefully
+        shuts down the worker process using SIGINT, and cleans up
+        the registrar service. If graceful shutdown fails, the process
+        is forcefully terminated.
+        """
+        if self._worker_process.is_alive():
+            assert self.address
+            channel = grpc.aio.insecure_channel(self.address)
+            stub = pb.worker.WorkerStub(channel)
+            await stub.stop(pb.worker.StopRequest(timeout=timeout))

wool/core/worker/pool.py

@@ -0,0 +1,386 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+import uuid
+from contextlib import asynccontextmanager
+from typing import AsyncContextManager
+from typing import Awaitable
+from typing import ContextManager
+from typing import Coroutine
+from typing import Final
+from typing import overload
+
+from wool.core.discovery.base import DiscoveryLike
+from wool.core.discovery.base import DiscoveryPublisherLike
+from wool.core.discovery.local import LocalDiscovery
+from wool.core.typing import Factory
+from wool.core.worker.base import WorkerFactory
+from wool.core.worker.base import WorkerLike
+from wool.core.worker.local import LocalWorker
+from wool.core.worker.proxy import LoadBalancerLike
+from wool.core.worker.proxy import RoundRobinLoadBalancer
+from wool.core.worker.proxy import WorkerProxy
+
+
+# public
+class WorkerPool:
+    """Orchestrates distributed workers for task execution.
+
+    The core of wool's distributed runtime. Manages worker lifecycle,
+    discovery, and load balancing across two modes:
+
+    **Ephemeral pools** spawn local workers automatically managed within the
+    pool's lifecycle. Perfect for development and single-machine deployments.
+
+    **Durable pools** connect to existing remote workers through discovery
+    services. Workers run independently, serving multiple clients across
+    distributed deployments.
+
+    **Basic ephemeral pool:**
+
+    .. code-block:: python
+
+        import wool
+
+
+        @wool.work
+        async def fibonacci(n: int) -> int:
+            if n <= 1:
+                return n
+            a = await fibonacci(n - 1)
+            b = await fibonacci(n - 2)
+            return a + b
+
+
+        async with wool.WorkerPool() as pool:
+            result = await fibonacci(10)
+
+    **Ephemeral with tags:**
+
+    .. code-block:: python
+
+        async with WorkerPool("gpu-capable", size=4) as pool:
+            result = await gpu_task()
+
+    **Custom worker factory:**
+
+    .. code-block:: python
+
+        from functools import partial
+
+        worker_factory = partial(LocalWorker, host="0.0.0.0")
+
+        async with WorkerPool(size=8, worker=worker_factory) as pool:
+            result = await task()
+
+    **Durable pool:**
+
+    .. code-block:: python
+
+        from wool.core.discovery.lan import LanDiscovery
+
+        async with WorkerPool(discovery=LanDiscovery()) as pool:
+            result = await task()
+
+    **Filtered discovery:**
+
+    .. code-block:: python
+
+        discovery = LanDiscovery().subscribe(filter=lambda w: "production" in w.tags)
+        async with WorkerPool(discovery=discovery) as pool:
+            result = await task()
+
+    **Hybrid pool:**
+
+    .. code-block:: python
+
+        # Spawn local workers AND discover remote workers
+        async with WorkerPool(size=4, discovery=LanDiscovery()) as pool:
+            result = await task()
+
+    **Custom load balancer:**
+
+    .. code-block:: python
+
+        from wool.core.loadbalancer.roundrobin import RoundRobinLoadBalancer
+
+
+        class PriorityBalancer(RoundRobinLoadBalancer):
+            async def dispatch(self, task, context, timeout=None):
+                # Custom routing logic
+                ...
+
+
+        async with WorkerPool(loadbalancer=PriorityBalancer()) as pool:
+            result = await task()
+
+    **Custom discovery:**
+
+    .. code-block:: python
+
+        from contextlib import asynccontextmanager
+
+
+        @asynccontextmanager
+        async def custom_discovery():
+            svc = await DatabaseDiscovery.connect()
+            try:
+                yield svc.subscribe()
+            finally:
+                await svc.close()
+
+
+        async with WorkerPool(discovery=custom_discovery) as pool:
+            result = await task()
+
+    :param tags:
+        Capability tags for spawned workers.
+    :param size:
+        Number of workers to spawn (0 = CPU count).
+    :param worker:
+        Worker factory callable. Defaults to :class:`LocalWorker`.
+    :param loadbalancer:
+        Load balancer instance, factory, or context manager.
+    :param discovery:
+        Discovery service instance, factory, or context manager.
+    :raises ValueError:
+        If configuration is invalid or CPU count unavailable.
+    """
+
+    _workers: Final[dict[WorkerLike, Coroutine]]
+
+    @overload
+    def __init__(
+        self,
+        *tags: str,
+        size: int = 0,
+        worker: WorkerFactory = LocalWorker,
+        discovery: DiscoveryLike | Factory[DiscoveryLike] | None = None,
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
+    ):
+        """
+        Create an ephemeral pool of workers, spawning the specified quantity of workers
+        using the specified worker factory.
+        """
+        ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        discovery: DiscoveryLike | Factory[DiscoveryLike],
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
+    ):
+        """
+        Connect to an existing pool of workers discovered by the specified discovery
+        protocol.
+        """
+        ...
+
+    def __init__(
+        self,
+        *tags: str,
+        size: int | None = None,
+        worker: WorkerFactory | None = None,
+        discovery: DiscoveryLike | Factory[DiscoveryLike] | None = None,
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
+    ):
+        self._workers = {}
+
+        match (size, discovery):
+            case (size, discovery) if size is not None and discovery is not None:
+                if size == 0:
+                    cpu_count = os.cpu_count()
+                    if cpu_count is None:
+                        raise ValueError("Unable to determine CPU count")
+                    size = cpu_count
+                elif size < 0:
+                    raise ValueError("Size must be non-negative")
+
+                @asynccontextmanager
+                async def create_proxy():
+                    discovery_svc, discovery_ctx = await self._enter_context(discovery)
+                    if not isinstance(discovery_svc, DiscoveryLike):
+                        raise TypeError(
+                            f"Expected DiscoveryLike, got: {type(discovery_svc)}"
+                        )
+
+                    try:
+                        async with self._worker_context(
+                            *tags,
+                            size=size,
+                            factory=worker,
+                            publisher=discovery_svc.publisher,
+                        ):
+                            async with WorkerProxy(
+                                discovery=discovery_svc.subscribe(_predicate(tags)),
+                                loadbalancer=loadbalancer,
+                            ):
+                                yield
+                    finally:
+                        await self._exit_context(discovery_ctx)
+
+            case (size, None) if size is not None:
+                if size == 0:
+                    cpu_count = os.cpu_count()
+                    if cpu_count is None:
+                        raise ValueError("Unable to determine CPU count")
+                    size = cpu_count
+                elif size < 0:
+                    raise ValueError("Size must be non-negative")
+
+                namespace = f"pool-{uuid.uuid4().hex}"
+
+                @asynccontextmanager
+                async def create_proxy():
+                    discovery = LocalDiscovery(namespace)
+                    async with self._worker_context(
+                        *tags,
+                        size=size,
+                        factory=worker,
+                        publisher=discovery.publisher,
+                    ):
+                        async with WorkerProxy(
+                            discovery=discovery.subscribe(_predicate(tags)),
+                            loadbalancer=loadbalancer,
+                        ):
+                            yield
+
+            case (None, discovery) if discovery is not None:
+
+                @asynccontextmanager
+                async def create_proxy():
+                    discovery_svc, discovery_ctx = await self._enter_context(discovery)
+                    if not isinstance(discovery_svc, DiscoveryLike):
+                        raise ValueError
+                    try:
+                        async with WorkerProxy(
+                            discovery=discovery_svc.subscriber,
+                            loadbalancer=loadbalancer,
+                        ):
+                            yield
+                    finally:
+                        await self._exit_context(discovery_ctx)
+
+            case (None, None):
+                cpu_count = os.cpu_count()
+                if cpu_count is None:
+                    raise ValueError("Unable to determine CPU count")
+                size = cpu_count
+
+                namespace = f"pool-{uuid.uuid4().hex}"
+
+                @asynccontextmanager
+                async def create_proxy():
+                    discovery = LocalDiscovery(namespace)
+                    async with self._worker_context(
+                        *tags,
+                        size=size,
+                        factory=worker,
+                        publisher=discovery.publisher,
+                    ):
+                        async with WorkerProxy(
+                            discovery=discovery.subscriber,
+                            loadbalancer=loadbalancer,
+                        ):
+                            yield
+
+            case _:
+                raise RuntimeError
+
+        self._proxy_factory = create_proxy
+
+    async def __aenter__(self) -> WorkerPool:
+        """Starts the worker pool and its services, returning a session.
+
+        This method starts the worker registrar, creates a connection,
+        launches all worker processes, and registers them.
+
+        :returns:
+            The :class:`WorkerPool` instance itself for method chaining.
+        """
+        self._proxy_context = self._proxy_factory()
+        await self._proxy_context.__aenter__()
+        return self
+
+    async def __aexit__(self, *args):
+        """Stops all workers and tears down the pool and its services."""
+        await self._proxy_context.__aexit__(*args)
+
+    @asynccontextmanager
+    async def _worker_context(
+        self,
+        *tags: str,
+        size: int,
+        factory: WorkerFactory | None,
+        publisher: DiscoveryPublisherLike,
+    ):
+        if factory is None:
+            factory = self._default_worker_factory()
+        publisher_svc, publisher_ctx = await self._enter_context(publisher)
+        if not isinstance(publisher_svc, DiscoveryPublisherLike):
+            raise ValueError
+
+        tasks = []
+        for _ in range(size):
+            worker = factory(*tags)
+
+            async def start(worker):
+                await worker.start()
+                await publisher.publish("worker-added", worker.info)
+
+            async def stop(worker):
+                await publisher.publish("worker-dropped", worker.info)
+                await worker.stop()
+
+            task = asyncio.create_task(start(worker))
+            tasks.append(task)
+            self._workers[worker] = stop(worker)
+
+        try:
+            await asyncio.gather(*tasks, return_exceptions=True)
+            yield [w.info for w in self._workers if w.info]
+        finally:
+            tasks = [asyncio.create_task(stop) for stop in self._workers.values()]
+            await asyncio.gather(*tasks, return_exceptions=True)
+            await self._exit_context(publisher_ctx)
+
+    def _default_worker_factory(self):
+        def factory(*tags, **_):
+            return LocalWorker(*tags)
+
+        return factory
+
+    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):
+        if isinstance(ctx, AsyncContextManager):
+            await ctx.__aexit__(*sys.exc_info())
+        elif isinstance(ctx, ContextManager):
+            ctx.__exit__(*sys.exc_info())
+
+
+def _predicate(tags):
+    return lambda w: bool(w.tags & set(tags)) if tags else True