wool 0.1rc11__tar.gz → 0.1rc12__tar.gz

{wool-0.1rc11 → wool-0.1rc12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wool
-Version: 0.1rc11
+Version: 0.1rc12
 Summary: A Python framework for distributed multiprocessing.
 Author-email: Conrad Bzura <conrad@wool.io>
 Maintainer-email: maintainers@wool.io

{wool-0.1rc11 → wool-0.1rc12}/wool/__init__.py

@@ -19,10 +19,10 @@ from wool._work import routine
 from wool._work import work
 from wool._worker import Worker
 from wool._worker import WorkerService
-from wool._worker_discovery import DiscoveryService
-from wool._worker_discovery import LanDiscoveryService
-from wool._worker_discovery import LanRegistryService
-from wool._worker_discovery import RegistryService
+from wool._worker_discovery import Discovery
+from wool._worker_discovery import LanDiscovery
+from wool._worker_discovery import LanRegistrar
+from wool._worker_discovery import Registrar
 from wool._worker_pool import WorkerPool
 from wool._worker_proxy import WorkerProxy
 
@@ -54,20 +54,22 @@ try:
 except PackageNotFoundError:
     __version__ = "unknown"
 
-__proxy__: Final[ContextVar[WorkerProxy | None]] = ContextVar("__proxy__", default=None)
+__proxy__: Final[ContextVar[WorkerProxy | None]] = ContextVar(
+    "__proxy__", default=None
+)
 
 __proxy_pool__: Final[ContextVar[ResourcePool[WorkerProxy] | None]] = ContextVar(
     "__proxy_pool__", default=None
 )
 
 __all__ = [
-    "LanDiscoveryService",
-    "LanRegistryService",
+    "LanDiscovery",
+    "LanRegistrar",
     "Worker",
-    "DiscoveryService",
+    "Discovery",
     "WorkerPool",
     "WorkerProxy",
-    "RegistryService",
+    "Registrar",
     "WorkerService",
     "WoolTask",
     "WoolTaskEvent",

{wool-0.1rc11 → wool-0.1rc12}/wool/_worker.py

@@ -13,7 +13,10 @@ from multiprocessing import Process
 from multiprocessing.connection import Connection
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import AsyncContextManager
 from typing import AsyncIterator
+from typing import Awaitable
+from typing import ContextManager
 from typing import Final
 from typing import Generic
 from typing import Protocol
@@ -31,15 +34,13 @@ from wool import _protobuf as pb
 from wool._resource_pool import ResourcePool
 from wool._work import WoolTask
 from wool._work import WoolTaskEvent
-from wool._worker_discovery import RegistryServiceLike
+from wool._worker_discovery import Factory
+from wool._worker_discovery import RegistrarLike
 from wool._worker_discovery import WorkerInfo
 
 if TYPE_CHECKING:
     from wool._worker_proxy import WorkerProxy
 
-_ip_address: str | None = None
-_EXTERNAL_DNS_SERVER: Final[str] = "8.8.8.8"  # Google DNS for IP detection
-
 
 @contextmanager
 def _signal_handlers(service: "WorkerService"):
@@ -80,13 +81,8 @@ def _signal_handlers(service: "WorkerService"):
         signal.signal(signal.SIGINT, old_sigint)
 
 
-_T_RegistryService = TypeVar(
-    "_T_RegistryService", bound=RegistryServiceLike, covariant=True
-)
-
-
 # public
-class Worker(ABC, Generic[_T_RegistryService]):
+class Worker(ABC):
     """Abstract base class for worker implementations in the wool framework.
 
     Workers are individual processes that execute distributed tasks within
@@ -94,36 +90,40 @@ class Worker(ABC, Generic[_T_RegistryService]):
     a discovery service to be found by client sessions.
 
     This class defines the core interface that all worker implementations
-    must provide, including lifecycle management and registry service
+    must provide, including lifecycle management and registrar service
     integration for peer-to-peer discovery.
 
     :param tags:
         Capability tags associated with this worker for filtering and
         selection by client sessions.
-    :param registry_service:
-        Service instance for worker registration and discovery within
-        the distributed pool.
+    :param registrar:
+        Service instance or factory for worker registration and discovery
+        within the distributed pool. Can be provided as:
+
+        - **Instance**: Direct registrar service object
+        - **Factory function**: Function returning a registrar service instance
+        - **Context manager factory**: Function returning a context manager
+            that yields a registrar service
     :param extra:
         Additional arbitrary metadata as key-value pairs.
     """
 
     _info: WorkerInfo | None = None
     _started: bool = False
-    _registry_service: RegistryServiceLike
+    _registrar: RegistrarLike | Factory[RegistrarLike]
+    _registrar_service: RegistrarLike | None = None
+    _registrar_context: Any | None = None
     _uid: Final[str]
     _tags: Final[set[str]]
     _extra: Final[dict[str, Any]]
 
     def __init__(
-        self,
-        *tags: str,
-        registry_service: _T_RegistryService,
-        **extra: Any,
+        self, *tags: str, registrar: RegistrarLike | Factory[RegistrarLike], **extra: Any
     ):
         self._uid = f"worker-{uuid.uuid4().hex}"
         self._tags = set(tags)
         self._extra = extra
-        self._registry_service = registry_service
+        self._registrar = registrar
 
     @property
     def uid(self) -> str:
@@ -167,14 +167,21 @@ class Worker(ABC, Generic[_T_RegistryService]):
 
         This method is a final implementation that calls the abstract
         `_start` method to initialize the worker process and register
-        it with the registry service.
+        it with the registrar service.
         """
         if self._started:
             raise RuntimeError("Worker has already been started")
-        if self._registry_service:
-            await self._registry_service.start()
+
+        self._registrar_service, self._registrar_context = await self._enter_context(
+            self._registrar
+        )
+        if not isinstance(self._registrar_service, RegistrarLike):
+            raise ValueError("Registrar factory must return a RegistrarLike instance")
+
         await self._start()
         self._started = True
+        assert self._info
+        await self._registrar_service.register(self._info)
 
     @final
     async def stop(self):
@@ -182,13 +189,23 @@ class Worker(ABC, Generic[_T_RegistryService]):
 
         This method is a final implementation that calls the abstract
         `_stop` method to gracefully shut down the worker process and
-        unregister it from the registry service.
+        unregister it from the registrar service.
         """
         if not self._started:
             raise RuntimeError("Worker has not been started")
-        await self._stop()
-        if self._registry_service:
-            await self._registry_service.stop()
+        try:
+            if not self._info:
+                raise RuntimeError("Cannot unregister - worker has no info")
+            assert self._registrar_service is not None
+            await self._registrar_service.unregister(self._info)
+        finally:
+            try:
+                await self._stop()
+            finally:
+                await self._exit_context(self._registrar_context)
+                self._registrar_service = None
+                self._registrar_context = None
+                self._started = False
 
     @abstractmethod
     async def _start(self):
@@ -208,10 +225,36 @@ class Worker(ABC, Generic[_T_RegistryService]):
         """
         ...
 
+    async def _enter_context(self, factory):
+        """Enter context for factory objects, handling different factory types."""
+        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
+    ):
+        """Exit context for context managers."""
+        if isinstance(ctx, AsyncContextManager):
+            await ctx.__aexit__(*args)
+        elif isinstance(ctx, ContextManager):
+            ctx.__exit__(*args)
+
 
 # public
-class WorkerFactory(Generic[_T_RegistryService], Protocol):
-    """Protocol for creating worker instances with registry integration.
+class WorkerFactory(Protocol):
+    """Protocol for creating worker instances with registrar integration.
 
     Defines the callable interface for worker factory implementations
     that can create :py:class:`Worker` instances configured with specific
@@ -221,7 +264,7 @@ class WorkerFactory(Generic[_T_RegistryService], Protocol):
     worker processes with consistent configuration.
     """
 
-    def __call__(self, *tags: str, **_) -> Worker[_T_RegistryService]:
+    def __call__(self, *tags: str, **_) -> Worker:
         """Create a new worker instance.
 
         :param tags:
@@ -235,12 +278,12 @@ class WorkerFactory(Generic[_T_RegistryService], Protocol):
 
 
 # public
-class LocalWorker(Worker[_T_RegistryService]):
+class LocalWorker(Worker):
     """Local worker implementation that runs tasks in a separate process.
 
     :py:class:`LocalWorker` creates and manages a dedicated worker process
     that hosts a gRPC server for executing distributed wool tasks. Each
-    worker automatically registers itself with the provided registry service
+    worker automatically registers itself with the provided registrar service
     for discovery by client sessions.
 
     The worker process runs independently and can handle multiple concurrent
@@ -250,8 +293,8 @@ class LocalWorker(Worker[_T_RegistryService]):
     :param tags:
         Capability tags to associate with this worker for filtering
         and selection by client sessions.
-    :param registry_service:
-        Service instance for worker registration and discovery.
+    :param registrar:
+        Service instance or factory for worker registration and discovery.
     :param extra:
         Additional arbitrary metadata as key-value pairs.
     """
@@ -263,10 +306,10 @@ class LocalWorker(Worker[_T_RegistryService]):
         *tags: str,
         host: str = "127.0.0.1",
         port: int = 0,
-        registry_service: _T_RegistryService,
+        registrar: RegistrarLike | Factory[RegistrarLike],
         **extra: Any,
     ):
-        super().__init__(*tags, registry_service=registry_service, **extra)
+        super().__init__(*tags, registrar=registrar, **extra)
         self._worker_process = WorkerProcess(host=host, port=port)
 
     @property
@@ -299,9 +342,9 @@ class LocalWorker(Worker[_T_RegistryService]):
     async def _start(self):
         """Start the worker process and register it with the pool.
 
-        Initializes the registry service, starts the worker process
+        Initializes the registrar service, starts the worker process
         with its gRPC server, and registers the worker's network
-        address with the registry for discovery by client sessions.
+        address with the registrar for discovery by client sessions.
         """
         loop = asyncio.get_running_loop()
         await loop.run_in_executor(None, self._worker_process.start)
@@ -324,30 +367,24 @@ class LocalWorker(Worker[_T_RegistryService]):
             tags=self._tags,
             extra=self._extra,
         )
-        await self._registry_service.register(self._info)
 
     async def _stop(self):
         """Stop the worker process and unregister it from the pool.
 
-        Unregisters the worker from the registry service, gracefully
+        Unregisters the worker from the registrar service, gracefully
         shuts down the worker process using SIGINT, and cleans up
-        the registry service. If graceful shutdown fails, the process
+        the registrar service. If graceful shutdown fails, the process
         is forcefully terminated.
         """
+        if not self._worker_process.is_alive():
+            return
         try:
-            if not self._info:
-                raise RuntimeError("Cannot unregister - worker has no info")
-            await self._registry_service.unregister(self._info)
-        finally:
-            if not self._worker_process.is_alive():
-                return
-            try:
-                if self._worker_process.pid:
-                    os.kill(self._worker_process.pid, signal.SIGINT)
-                    self._worker_process.join()
-            except OSError:
-                if self._worker_process.is_alive():
-                    self._worker_process.kill()
+            if self._worker_process.pid:
+                os.kill(self._worker_process.pid, signal.SIGINT)
+                self._worker_process.join()
+        except OSError:
+            if self._worker_process.is_alive():
+                self._worker_process.kill()
 
 
 class WorkerProcess(Process):

{wool-0.1rc11 → wool-0.1rc12}/wool/_worker_discovery.py

@@ -11,7 +11,6 @@ from abc import abstractmethod
 from collections import deque
 from dataclasses import dataclass
 from dataclasses import field
-from typing import TYPE_CHECKING
 from typing import Any
 from typing import AsyncContextManager
 from typing import AsyncIterator
@@ -23,15 +22,13 @@ from typing import Dict
 from typing import Generic
 from typing import Literal
 from typing import Protocol
+from typing import Self
 from typing import Tuple
 from typing import TypeAlias
 from typing import TypeVar
 from typing import final
 from typing import runtime_checkable
 
-if TYPE_CHECKING:
-    pass
-
 from zeroconf import IPVersion
 from zeroconf import ServiceInfo
 from zeroconf import ServiceListener
@@ -39,9 +36,6 @@ from zeroconf import Zeroconf
 from zeroconf.asyncio import AsyncServiceBrowser
 from zeroconf.asyncio import AsyncZeroconf
 
-if TYPE_CHECKING:
-    pass
-
 
 # public
 @dataclass
@@ -314,7 +308,6 @@ class DiscoveryEvent:
 _T_co = TypeVar("_T_co", covariant=True)
 
 
-# public
 class Reducible(Protocol):
     """Protocol for objects that support pickling via __reduce__."""
 
@@ -322,21 +315,22 @@ class Reducible(Protocol):
 
 
 # public
-class ReducibleAsyncIteratorLike(Reducible, Protocol, Generic[_T_co]):
+@runtime_checkable
+class DiscoveryLike(Protocol):
     """Protocol for async iterators that yield discovery events.
 
     Implementations must be pickleable via __reduce__ to support
     task-specific session contexts in distributed environments.
     """
 
-    def __aiter__(self) -> ReducibleAsyncIteratorLike[_T_co]: ...
+    def __aiter__(self) -> DiscoveryLike: ...
 
-    def __anext__(self) -> Awaitable[_T_co]: ...
+    def __anext__(self) -> Awaitable[DiscoveryEvent]: ...
 
 
 # public
 @runtime_checkable
-class Factory(Protocol, Generic[_T_co]):
+class _Factory(Protocol, Generic[_T_co]):
     def __call__(
         self,
     ) -> (
@@ -344,8 +338,16 @@ class Factory(Protocol, Generic[_T_co]):
     ): ...
 
 
+Factory: TypeAlias = (
+    Awaitable[_T_co]
+    | AsyncContextManager[_T_co]
+    | ContextManager[_T_co]
+    | _Factory[_T_co]
+)
+
+
 # public
-class DiscoveryService(ABC):
+class Discovery(ABC):
     """Abstract base class for discovering worker services.
 
     When started, implementations should discover all existing services that
@@ -358,7 +360,7 @@ class DiscoveryService(ABC):
         - Workers updated to satisfy the filter should trigger worker-added events
         - Workers updated to no longer satisfy the filter should trigger
           worker-removed events
-        - Tracked workers removed from the registry entirely should always
+        - Tracked workers removed from the registrar entirely should always
           trigger worker-removed
 
     :param filter:
@@ -436,7 +438,7 @@ class DiscoveryService(ABC):
         """Yields discovery events as they occur.
 
         Returns an asynchronous iterator that yields discovery events for
-        workers being added, updated, or removed from the registry. Events
+        workers being added, updated, or removed from the registrar. Events
         are filtered according to the filter function provided during
         initialization.
 
@@ -457,21 +459,18 @@ class DiscoveryService(ABC):
         ...
 
 
-_T_DiscoveryServiceLike = TypeVar("_T_DiscoveryServiceLike", bound=DiscoveryService)
+_T_DiscoveryLike = TypeVar("_T_DiscoveryLike", bound=Discovery)
 
 
 # public
-class RegistryServiceLike(Protocol):
+@runtime_checkable
+class RegistrarLike(Protocol):
     """Abstract base class for a service where workers can register themselves.
 
     Provides the interface for worker registration, unregistration, and
     property updates within a distributed worker pool system.
     """
 
-    async def start(self) -> None: ...
-
-    async def stop(self) -> None: ...
-
     async def register(self, worker_info: WorkerInfo) -> None: ...
 
     async def unregister(self, worker_info: WorkerInfo) -> None: ...
@@ -480,7 +479,7 @@ class RegistryServiceLike(Protocol):
 
 
 # public
-class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
+class Registrar(Generic[_T_DiscoveryLike], ABC):
     """Abstract base class for a service where workers can register themselves.
 
     Provides the interface for worker registration, unregistration, and
@@ -494,19 +493,26 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
         self._started = False
         self._stopped = False
 
+    async def __aenter__(self) -> Self:
+        await self.start()
+        return self
+
+    async def __aexit__(self, *_):
+        await self.stop()
+
     async def start(self) -> None:
-        """Starts the registry service, making it ready to accept registrations.
+        """Starts the registrar service, making it ready to accept registrations.
 
         :raises RuntimeError:
             If the service has already been started.
         """
         if self._started:
-            raise RuntimeError("Registry service already started")
+            raise RuntimeError("Registrar service already started")
         await asyncio.wait_for(self._start(), timeout=60)
         self._started = True
 
     async def stop(self) -> None:
-        """Stops the registry service and cleans up any resources.
+        """Stops the registrar service and cleans up any resources.
 
         :raises RuntimeError:
             If the service has not been started.
@@ -514,18 +520,18 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
         if self._stopped:
             return
         if not self._started:
-            raise RuntimeError("Registry service not started")
+            raise RuntimeError("Registrar service not started")
         await self._stop()
         self._stopped = True
 
     @abstractmethod
     async def _start(self) -> None:
-        """Starts the registry service, making it ready to accept registrations."""
+        """Starts the registrar service, making it ready to accept registrations."""
         ...
 
     @abstractmethod
     async def _stop(self) -> None:
-        """Stops the registry service and cleans up any resources."""
+        """Stops the registrar service and cleans up any resources."""
         ...
 
     async def register(
@@ -538,12 +544,12 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
             The :class:`~wool._worker_discovery.WorkerInfo` instance containing all
             worker details.
         :raises RuntimeError:
-            If the registry service is not running.
+            If the registrar service is not running.
         """
         if not self._started:
-            raise RuntimeError("Registry service not started - call start() first")
+            raise RuntimeError("Registrar service not started - call start() first")
         if self._stopped:
-            raise RuntimeError("Registry service already stopped")
+            raise RuntimeError("Registrar service already stopped")
         await self._register(worker_info)
 
     @abstractmethod
@@ -566,12 +572,12 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
             The :class:`~wool._worker_discovery.WorkerInfo` instance of the worker to
             unregister.
         :raises RuntimeError:
-            If the registry service is not running.
+            If the registrar service is not running.
         """
         if not self._started:
-            raise RuntimeError("Registry service not started - call start() first")
+            raise RuntimeError("Registrar service not started - call start() first")
         if self._stopped:
-            raise RuntimeError("Registry service already stopped")
+            raise RuntimeError("Registrar service already stopped")
         await self._unregister(worker_info)
 
     @abstractmethod
@@ -590,12 +596,12 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
         :param worker_info:
             The updated :class:`~wool._worker_discovery.WorkerInfo` instance.
         :raises RuntimeError:
-            If the registry service is not running.
+            If the registrar service is not running.
         """
         if not self._started:
-            raise RuntimeError("Registry service not started - call start() first")
+            raise RuntimeError("Registrar service not started - call start() first")
         if self._stopped:
-            raise RuntimeError("Registry service already stopped")
+            raise RuntimeError("Registrar service already stopped")
         await self._update(worker_info)
 
     @abstractmethod
@@ -609,7 +615,7 @@ class RegistryService(Generic[_T_DiscoveryServiceLike], ABC):
 
 
 # public
-class LanDiscoveryService(DiscoveryService):
+class LanDiscovery(Discovery):
     """Implements worker discovery on the local network using Zeroconf.
 
     This service browses the local network for DNS-SD services and delivers
@@ -705,12 +711,12 @@ class LanDiscoveryService(DiscoveryService):
 
         def add_service(self, zc: Zeroconf, type_: str, name: str):
             """Called by Zeroconf when a service is added."""
-            if type_ == LanRegistryService.service_type:
+            if type_ == LanRegistrar.service_type:
                 asyncio.create_task(self._handle_add_service(type_, name))
 
         def remove_service(self, zc: Zeroconf, type_: str, name: str):
             """Called by Zeroconf when a service is removed."""
-            if type_ == LanRegistryService.service_type:
+            if type_ == LanRegistrar.service_type:
                 if worker := self._service_cache.pop(name, None):
                     asyncio.create_task(
                         self._event_queue.put(
@@ -720,7 +726,7 @@ class LanDiscoveryService(DiscoveryService):
 
         def update_service(self, zc: Zeroconf, type_, name):
             """Called by Zeroconf when a service is updated."""
-            if type_ == LanRegistryService.service_type:
+            if type_ == LanRegistrar.service_type:
                 asyncio.create_task(self._handle_update_service(type_, name))
 
         async def _handle_add_service(self, type_: str, name: str):
@@ -787,11 +793,11 @@ class LanDiscoveryService(DiscoveryService):
 
 
 # public
-class LanRegistryService(RegistryService[LanDiscoveryService]):
-    """Implements a worker registry using Zeroconf to advertise on the LAN.
+class LanRegistrar(Registrar[LanDiscovery]):
+    """Implements a worker registrar using Zeroconf to advertise on the LAN.
 
     This service registers workers by publishing a DNS-SD service record on
-    the local network, allowing :class:`LanDiscoveryService` to find them.
+    the local network, allowing :class:`LanDiscovery` to find them.
     """
 
     aiozc: AsyncZeroconf | None
@@ -824,10 +830,10 @@ class LanRegistryService(RegistryService[LanDiscoveryService]):
             The :class:`~wool._worker_discovery.WorkerInfo` instance containing all
             worker details.
         :raises RuntimeError:
-            If the registry service is not properly initialized.
+            If the registrar service is not properly initialized.
         """
         if self.aiozc is None:
-            raise RuntimeError("Registry service not properly initialized")
+            raise RuntimeError("Registrar service not properly initialized")
         address = f"{worker_info.host}:{worker_info.port}"
         ip_address, port = self._resolve_address(address)
         service_name = f"{worker_info.uid}.{self.service_type}"
@@ -848,10 +854,10 @@ class LanRegistryService(RegistryService[LanDiscoveryService]):
             The :class:`~wool._worker_discovery.WorkerInfo` instance of the worker to
             unregister.
         :raises RuntimeError:
-            If the registry service is not properly initialized.
+            If the registrar service is not properly initialized.
         """
         if self.aiozc is None:
-            raise RuntimeError("Registry service not properly initialized")
+            raise RuntimeError("Registrar service not properly initialized")
         service = self.services[worker_info.uid]
         await self.aiozc.async_unregister_service(service)
         del self.services[worker_info.uid]
@@ -866,12 +872,12 @@ class LanRegistryService(RegistryService[LanDiscoveryService]):
         :param worker_info:
             The updated :class:`~wool._worker_discovery.WorkerInfo` instance.
         :raises RuntimeError:
-            If the registry service is not properly initialized.
+            If the registrar service is not properly initialized.
         :raises Exception:
             If the Zeroconf service update fails.
         """
         if self.aiozc is None:
-            raise RuntimeError("Registry service not properly initialized")
+            raise RuntimeError("Registrar service not properly initialized")
 
         service = self.services[worker_info.uid]
         new_properties = _serialize_worker_info(worker_info)
@@ -971,12 +977,12 @@ def _deserialize_worker_info(info: ServiceInfo) -> WorkerInfo:
 
 
 # public
-class LocalRegistryService(RegistryService):
-    """Implements a worker registry using shared memory for local pools.
+class LocalRegistrar(Registrar):
+    """Implements a worker registrar using shared memory for local pools.
 
     This service registers workers by writing their information to a shared memory
-    block, allowing LocalDiscoveryService instances to find them efficiently.
-    The registry stores worker ports as integers in a simple array format,
+    block, allowing LocalDiscovery instances to find them efficiently.
+    The registrar stores worker ports as integers in a simple array format,
     providing fast local discovery without network overhead.
 
     :param uri:
@@ -1018,7 +1024,7 @@ class LocalRegistryService(RegistryService):
         if self._shared_memory:
             try:
                 self._shared_memory.close()
-                # Unlink the shared memory if this registry created it
+                # Unlink the shared memory if this registrar created it
                 if self._created_shared_memory:
                     self._shared_memory.unlink()
             except Exception:
@@ -1033,10 +1039,10 @@ class LocalRegistryService(RegistryService):
             The :class:`~wool._worker_discovery.WorkerInfo` instance containing all
             worker details. Only the port is stored in shared memory.
         :raises RuntimeError:
-            If the registry service is not properly initialized.
+            If the registrar service is not properly initialized.
         """
         if self._shared_memory is None:
-            raise RuntimeError("Registry service not properly initialized")
+            raise RuntimeError("Registrar service not properly initialized")
 
         if worker_info.port is None:
             raise ValueError("Worker port must be specified")
@@ -1048,7 +1054,7 @@ class LocalRegistryService(RegistryService):
                 struct.pack_into("I", self._shared_memory.buf, i, worker_info.port)
                 break
         else:
-            raise RuntimeError("No available slots in shared memory registry")
+            raise RuntimeError("No available slots in shared memory registrar")
 
     async def _unregister(self, worker_info: WorkerInfo) -> None:
         """Unregister a worker by removing its port from shared memory.
@@ -1057,10 +1063,10 @@ class LocalRegistryService(RegistryService):
             The :class:`~wool._worker_discovery.WorkerInfo` instance of the worker to
             unregister.
         :raises RuntimeError:
-            If the registry service is not properly initialized.
+            If the registrar service is not properly initialized.
         """
         if self._shared_memory is None:
-            raise RuntimeError("Registry service not properly initialized")
+            raise RuntimeError("Registrar service not properly initialized")
 
         if worker_info.port is None:
             return
@@ -1075,7 +1081,7 @@ class LocalRegistryService(RegistryService):
     async def _update(self, worker_info: WorkerInfo) -> None:
         """Update a worker's properties in shared memory.
 
-        For the simple port-based registry, update is the same as register.
+        For the simple port-based registrar, update is the same as register.
 
         :param worker_info:
             The updated :class:`~wool._worker_discovery.WorkerInfo` instance.
@@ -1084,7 +1090,7 @@ class LocalRegistryService(RegistryService):
 
 
 # public
-class LocalDiscoveryService(DiscoveryService):
+class LocalDiscovery(Discovery):
     """Implements worker discovery using shared memory for local pools.
 
     This service reads worker ports from a shared memory block and
@@ -1172,8 +1178,8 @@ class LocalDiscoveryService(DiscoveryService):
                                 uid=f"worker-{port}",
                                 host="localhost",
                                 port=port,
-                                pid=0,  # Not available in simple registry
-                                version="unknown",  # Not available in simple registry
+                                pid=0,  # Not available in simple registrar
+                                version="unknown",  # Not available in simple registrar
                                 tags=set(),
                                 extra={},
                             )
@@ -1211,7 +1217,7 @@ class LocalDiscoveryService(DiscoveryService):
                 event = DiscoveryEvent(type="worker_removed", worker_info=worker_info)
                 await self._event_queue.put(event)
 
-        # Find updated workers (minimal for port-only registry)
+        # Find updated workers (minimal for port-only registrar)
         for uid, worker_info in current_workers.items():
             if uid in self._service_cache:
                 old_worker = self._service_cache[uid]

{wool-0.1rc11 → wool-0.1rc12}/wool/_worker_pool.py

@@ -4,22 +4,17 @@ import asyncio
 import hashlib
 import os
 import uuid
-from functools import partial
 from multiprocessing.shared_memory import SharedMemory
-from typing import AsyncIterator
 from typing import Final
 from typing import overload
 
 from wool._worker import LocalWorker
 from wool._worker import Worker
 from wool._worker import WorkerFactory
-from wool._worker_discovery import DiscoveryEvent
+from wool._worker_discovery import DiscoveryLike
 from wool._worker_discovery import Factory
-from wool._worker_discovery import LocalDiscoveryService
-from wool._worker_discovery import LocalRegistryService
-from wool._worker_discovery import ReducibleAsyncIteratorLike
-from wool._worker_discovery import RegistryServiceLike
-from wool._worker_proxy import LoadBalancerFactory
+from wool._worker_discovery import LocalDiscovery
+from wool._worker_discovery import LocalRegistrar
 from wool._worker_proxy import LoadBalancerLike
 from wool._worker_proxy import RoundRobinLoadBalancer
 from wool._worker_proxy import WorkerProxy
@@ -65,13 +60,11 @@ class WorkerPool:
     .. code-block:: python
 
         from wool import WorkerPool, LocalWorker
-        from wool._worker_discovery import LocalRegistryService
+        from wool._worker_discovery import LocalRegistrar
         from functools import partial
 
         # Custom worker factory with specific tags
-        worker_factory = partial(
-            LocalWorker, registry_service=LocalRegistryService("my-pool")
-        )
+        worker_factory = partial(LocalWorker, registrar=LocalRegistrar("my-pool"))
 
         async with WorkerPool(
             "gpu-capable",
@@ -86,10 +79,10 @@ class WorkerPool:
     .. code-block:: python
 
         from wool import WorkerPool
-        from wool._worker_discovery import LanDiscoveryService
+        from wool._worker_discovery import LanDiscovery
 
         # Connect to existing workers on the network
-        discovery = LanDiscoveryService(filter=lambda w: "production" in w.tags)
+        discovery = LanDiscovery(filter=lambda w: "production" in w.tags)
 
         async with WorkerPool(discovery=discovery) as pool:
             results = await gather_metrics()
@@ -158,15 +151,15 @@ class WorkerPool:
         Examples::
 
             # Direct instance
-            discovery=LanDiscoveryService(filter=lambda w: "prod" in w.tags)
+            discovery=LanDiscovery(filter=lambda w: "prod" in w.tags)
 
             # Instance factory
-            discovery=lambda: LocalDiscoveryService("pool-123")
+            discovery=lambda: LocalDiscovery("pool-123")
 
             # Context manager factory
             @asynccontextmanager
             async def discovery():
-                service = await DatabaseDiscoveryService.create(connection_string)
+                service = await DatabaseDiscovery.create(connection_string)
                 try:
                     ...
                     yield service
@@ -186,8 +179,10 @@ class WorkerPool:
         self,
         *tags: str,
         size: int = 0,
-        worker: WorkerFactory[RegistryServiceLike] = LocalWorker[LocalRegistryService],
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        worker: WorkerFactory = LocalWorker,
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
     ):
         """
         Create an ephemeral pool of workers, spawning the specified quantity of workers
@@ -199,11 +194,10 @@ class WorkerPool:
     def __init__(
         self,
         *,
-        discovery: (
-            ReducibleAsyncIteratorLike[DiscoveryEvent]
-            | Factory[AsyncIterator[DiscoveryEvent]]
-        ),
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        discovery: DiscoveryLike | Factory[DiscoveryLike],
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
     ):
         """
         Connect to an existing pool of workers discovered by the specified discovery
@@ -216,12 +210,10 @@ class WorkerPool:
         *tags: str,
         size: int | None = None,
         worker: WorkerFactory | None = None,
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
-        discovery: (
-            ReducibleAsyncIteratorLike[DiscoveryEvent]
-            | Factory[AsyncIterator[DiscoveryEvent]]
-            | None
-        ) = None,
+        discovery: DiscoveryLike | Factory[DiscoveryLike] | None = None,
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
     ):
         self._workers = []
 
@@ -244,7 +236,7 @@ class WorkerPool:
                         self._shared_memory.buf[i] = 0
                     await self._spawn_workers(uri, *tags, size=size, factory=worker)
                     return WorkerProxy(
-                        discovery=LocalDiscoveryService(uri),
+                        discovery=LocalDiscovery(uri),
                         loadbalancer=loadbalancer,
                     )
 
@@ -269,7 +261,7 @@ class WorkerPool:
                         self._shared_memory.buf[i] = 0
                     await self._spawn_workers(uri, *tags, size=size, factory=worker)
                     return WorkerProxy(
-                        discovery=LocalDiscoveryService(uri),
+                        discovery=LocalDiscovery(uri),
                         loadbalancer=loadbalancer,
                     )
 
@@ -289,7 +281,7 @@ class WorkerPool:
     async def __aenter__(self) -> WorkerPool:
         """Starts the worker pool and its services, returning a session.
 
-        This method starts the worker registry, creates a client session,
+        This method starts the worker registrar, creates a client session,
         launches all worker processes, and registers them.
 
         :returns:
@@ -332,6 +324,6 @@ class WorkerPool:
 
     def _default_worker_factory(self, uri):
         def factory(*tags, **_):
-            return LocalWorker(*tags, registry_service=LocalRegistryService(uri))
+            return LocalWorker(*tags, registrar=LocalRegistrar(uri))
 
         return factory

{wool-0.1rc11 → wool-0.1rc12}/wool/_worker_proxy.py

@@ -27,9 +27,9 @@ from wool._resource_pool import Resource
 from wool._resource_pool import ResourcePool
 from wool._worker import WorkerClient
 from wool._worker_discovery import DiscoveryEvent
+from wool._worker_discovery import DiscoveryLike
 from wool._worker_discovery import Factory
-from wool._worker_discovery import LocalDiscoveryService
-from wool._worker_discovery import ReducibleAsyncIteratorLike
+from wool._worker_discovery import LocalDiscovery
 from wool._worker_discovery import WorkerInfo
 
 if TYPE_CHECKING:
@@ -106,6 +106,7 @@ class NoWorkersAvailable(Exception):
     """
 
 
+# public
 @runtime_checkable
 class LoadBalancerLike(Protocol):
     """Protocol for load balancer v2 that directly dispatches tasks.
@@ -122,19 +123,16 @@ class LoadBalancerLike(Protocol):
     def dispatch(self, task: WoolTask) -> AsyncIterator: ...
 
     def worker_added_callback(
-        self, client: Resource[WorkerClient], info: WorkerInfo
+        self, client: Callable[[], Resource[WorkerClient]], info: WorkerInfo
     ): ...
 
     def worker_updated_callback(
-        self, client: Resource[WorkerClient], info: WorkerInfo
+        self, client: Callable[[], Resource[WorkerClient]], info: WorkerInfo
     ): ...
 
     def worker_removed_callback(self, info: WorkerInfo): ...
 
 
-LoadBalancerFactory: TypeAlias = Factory[LoadBalancerLike]
-
-
 DispatchCall: TypeAlias = grpc.aio.UnaryStreamCall[pb.task.Task, pb.worker.Response]
 
 
@@ -213,10 +211,14 @@ class RoundRobinLoadBalancer:
             f"All {len(self._workers)} workers failed with transient errors"
         )
 
-    def worker_added_callback(self, client: Resource[WorkerClient], info: WorkerInfo):
+    def worker_added_callback(
+        self, client: Callable[[], Resource[WorkerClient]], info: WorkerInfo
+    ):
         self._workers[info] = client
 
-    def worker_updated_callback(self, client: Resource[WorkerClient], info: WorkerInfo):
+    def worker_updated_callback(
+        self, client: Callable[[], Resource[WorkerClient]], info: WorkerInfo
+    ):
         self._workers[info] = client
 
     def worker_removed_callback(self, info: WorkerInfo):
@@ -251,15 +253,12 @@ class WorkerProxy:
         Load balancer implementation or factory for task distribution.
     """
 
-    _discovery: (
-        ReducibleAsyncIteratorLike[DiscoveryEvent]
-        | Factory[AsyncIterator[DiscoveryEvent]]
-    )
+    _discovery: DiscoveryLike | Factory[DiscoveryLike]
     _discovery_manager: (
-        AsyncContextManager[AsyncIterator[DiscoveryEvent]]
-        | ContextManager[AsyncIterator[DiscoveryEvent]]
+        AsyncContextManager[DiscoveryLike] | ContextManager[DiscoveryLike]
     )
-    _loadbalancer = LoadBalancerLike | LoadBalancerFactory
+
+    _loadbalancer = LoadBalancerLike | Factory[LoadBalancerLike]
     _loadbalancer_manager: (
         AsyncContextManager[LoadBalancerLike] | ContextManager[LoadBalancerLike]
     )
@@ -268,11 +267,10 @@ class WorkerProxy:
     def __init__(
         self,
         *,
-        discovery: (
-            ReducibleAsyncIteratorLike[DiscoveryEvent]
-            | Factory[AsyncIterator[DiscoveryEvent]]
-        ),
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        discovery: DiscoveryLike | Factory[DiscoveryLike],
+        loadbalancer: (
+            LoadBalancerLike | Factory[LoadBalancerLike]
+        ) = RoundRobinLoadBalancer,
     ): ...
 
     @overload
@@ -280,7 +278,8 @@ class WorkerProxy:
         self,
         *,
         workers: Sequence[WorkerInfo],
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
     ): ...
 
     @overload
@@ -288,20 +287,18 @@ class WorkerProxy:
         self,
         pool_uri: str,
         *tags: str,
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
     ): ...
 
     def __init__(
         self,
         pool_uri: str | None = None,
         *tags: str,
-        discovery: (
-            ReducibleAsyncIteratorLike[DiscoveryEvent]
-            | Factory[AsyncIterator[DiscoveryEvent]]
-            | None
-        ) = None,
+        discovery: (DiscoveryLike | Factory[DiscoveryLike] | None) = None,
         workers: Sequence[WorkerInfo] | None = None,
-        loadbalancer: LoadBalancerLike | LoadBalancerFactory = RoundRobinLoadBalancer,
+        loadbalancer: LoadBalancerLike
+        | Factory[LoadBalancerLike] = RoundRobinLoadBalancer,
     ):
         if not (pool_uri or discovery or workers):
             raise ValueError(
@@ -316,7 +313,7 @@ class WorkerProxy:
 
         match (pool_uri, discovery, workers):
             case (pool_uri, None, None) if pool_uri is not None:
-                self._discovery = LocalDiscoveryService(
+                self._discovery = LocalDiscovery(
                     pool_uri, filter=lambda w: bool({pool_uri, *tags} & w.tags)
                 )
             case (None, discovery, None) if discovery is not None:
@@ -399,7 +396,7 @@ class WorkerProxy:
         self._discovery_service, self._discovery_ctx = await self._enter_context(
             self._discovery
         )
-        if not isinstance(self._discovery_service, AsyncIterator):
+        if not isinstance(self._discovery_service, DiscoveryLike):
             raise ValueError
 
         self._proxy_token = wool.__proxy__.set(self)
@@ -463,16 +460,16 @@ class WorkerProxy:
 
     async def _enter_context(self, factory):
         ctx = None
-        if callable(factory):
-            obj = factory()
-            if isinstance(obj, ContextManager):
-                ctx = obj
-                obj = obj.__enter__()
-            elif isinstance(obj, AsyncContextManager):
-                ctx = obj
-                obj = await obj.__aenter__()
-            elif isinstance(obj, Awaitable):
-                obj = await obj
+        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