Add register_validation_exporter plugin hook

osprey_async_worker/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "osprey-async-worker"
+version = "0.1.0"
+description = "Asyncio-based Osprey worker"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "osprey_rpc",
+    "osprey_worker",
+]
+
+[project.scripts]
+osprey-async-cli = "osprey.async_worker.cli.main:cli"
+
+[project.entry-points."osprey_async_plugin"]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+"*" = ["*.json", "*.yaml", "*.yml"]

osprey_async_worker/src/osprey/async_worker/adaptor/constants.py

@@ -0,0 +1 @@
+OSPREY_ASYNC_ADAPTOR = 'osprey_async_plugin'

osprey_async_worker/src/osprey/async_worker/adaptor/hookspecs.py

@@ -0,0 +1,55 @@
+"""Hook specifications for the async worker plugin system.
+
+Mirrors osprey.worker.adaptor.hookspecs but uses the 'osprey_async_plugin'
+entry_point group. Plugins register async output sinks and UDFs here.
+
+UDFs are shared with the sync worker (they're registered via the existing
+'osprey_plugin' hooks and wrapped with SyncUDFAdapter). Async-native UDFs
+can also be registered here.
+
+Output sinks MUST be async (AsyncBaseOutputSink) since the async worker
+doesn't use gevent.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence, Type
+
+import pluggy
+from osprey.engine.ast_validator.base_validator import BaseValidator
+from osprey.engine.udf.base import UDFBase
+
+from osprey.async_worker.adaptor.constants import OSPREY_ASYNC_ADAPTOR
+from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink
+
+if TYPE_CHECKING:
+    from osprey.worker.lib.config import Config
+
+hookspec: pluggy.HookspecMarker = pluggy.HookspecMarker(OSPREY_ASYNC_ADAPTOR)
+
+
+@hookspec
+def register_async_output_sinks(config: Config) -> Sequence[AsyncBaseOutputSink]:
+    """Register async output sinks for the async worker.
+
+    These must be AsyncBaseOutputSink instances (not sync BaseOutputSink).
+    The async worker will call `await sink.push(result)` for each result.
+    """
+    raise NotImplementedError
+
+
+@hookspec
+def register_udfs() -> Sequence[Type[UDFBase[Any, Any]]]:
+    """Register UDFs for the async worker.
+
+    These are the same UDFBase types as the sync worker. The async executor
+    runs them in a thread pool via run_in_executor. Async-native UDFs can
+    also be registered here in the future.
+    """
+    raise NotImplementedError
+
+
+@hookspec
+def register_ast_validators() -> Sequence[Type[BaseValidator]]:
+    """Register AST validators. Same interface as the sync worker."""
+    raise NotImplementedError

osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py

@@ -0,0 +1,160 @@
+"""Async plugin interfaces for the osprey async worker.
+
+AsyncUDFBase extends UDFBase so it works with the existing engine machinery
+(UDFRegistry, CallExecutor, validation, type checking, argument resolution).
+
+All I/O UDFs in the async worker MUST be AsyncUDFBase subclasses — existing
+sync UDFs that use gevent primitives will not work without monkey patching.
+Pure-computation UDFs (no I/O) can remain as regular UDFBase and run inline.
+"""
+
+import abc
+import asyncio
+from typing import Any, ClassVar, Generic, Sequence, Tuple, TypeVar
+
+from osprey.engine.executor.execution_context import ExecutionContext, ExecutionResult
+from osprey.engine.udf.base import (
+    Arguments,
+    BatchableArguments,
+    BatchableUDFBase,
+    RValue,
+    UDFBase,
+)
+from result import Result
+
+_T = TypeVar('_T')
+
+
+class AsyncUDFBase(UDFBase[Arguments, RValue]):
+    """Native async UDF base class.
+
+    Extends UDFBase so it integrates with UDFRegistry, CallExecutor, and
+    the full validation/type-checking pipeline. The async executor detects
+    AsyncUDFBase instances via isinstance() and awaits async_execute()
+    directly on the event loop — no thread pool.
+
+    The sync execute() raises so it can't accidentally be called in the
+    async executor's sync path.
+    """
+
+    execute_async: ClassVar[bool] = True
+    is_native_async: ClassVar[bool] = True
+
+    def __init__(self, validation_context, arguments):
+        super().__init__(validation_context, arguments)
+
+    @classmethod
+    def _get_udf_base_args(cls):
+        """Override to include AsyncUDFBase in the generic origin check."""
+        import typing_inspect
+
+        for base in cls.__mro__:
+            for generic_base in typing_inspect.get_generic_bases(base):
+                origin = typing_inspect.get_origin(generic_base)
+                if origin in (UDFBase, AsyncUDFBase) or (
+                    hasattr(origin, '__mro__') and UDFBase in origin.__mro__
+                ):
+                    args = typing_inspect.get_args(generic_base)
+                    # Only return if args are concrete (not TypeVars)
+                    if args and not any(isinstance(a, TypeVar) for a in args):
+                        return args
+
+        # Fallback to parent
+        return super()._get_udf_base_args()
+
+    def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue:
+        raise RuntimeError(
+            f'{self.__class__.__name__} is a native async UDF. '
+            f'Use async_execute() instead of execute().'
+        )
+
+    @abc.abstractmethod
+    async def async_execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue:
+        """Override this to implement the UDF's async execution logic."""
+        raise NotImplementedError
+
+
+class AsyncBatchableUDFBase(BatchableUDFBase[Arguments, RValue, BatchableArguments]):
+    """Native async batchable UDF base class.
+
+    Same as AsyncUDFBase but for batchable UDFs. The async executor detects
+    these and awaits async_execute_batch() directly.
+    """
+
+    is_native_async: ClassVar[bool] = True
+
+    def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue:
+        raise RuntimeError(
+            f'{self.__class__.__name__} is a native async UDF. '
+            f'Use async_execute() instead of execute().'
+        )
+
+    def execute_batch(
+        self,
+        execution_context: ExecutionContext,
+        udfs: Sequence[UDFBase[Any, Any]],
+        arguments: Sequence[BatchableArguments],
+    ) -> Sequence[Result[RValue, Exception]]:
+        raise RuntimeError(
+            f'{self.__class__.__name__} is a native async UDF. '
+            f'Use async_execute_batch() instead of execute_batch().'
+        )
+
+    @abc.abstractmethod
+    async def async_execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def async_execute_batch(
+        self,
+        execution_context: ExecutionContext,
+        udfs: Sequence[UDFBase[Any, Any]],
+        arguments: Sequence[BatchableArguments],
+    ) -> Sequence[Result[RValue, Exception]]:
+        raise NotImplementedError
+
+
+# --- Output sinks and input streams (unchanged) ---
+
+
+class AsyncBaseOutputSink(abc.ABC):
+    """Async output sink."""
+
+    timeout: float = 2.0
+    max_retries: int = 0
+
+    @abc.abstractmethod
+    def will_do_work(self, result: ExecutionResult) -> bool:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def push(self, result: ExecutionResult) -> None:
+        raise NotImplementedError
+
+    async def stop(self) -> None:
+        pass
+
+
+class AsyncMultiOutputSink(AsyncBaseOutputSink):
+    """Tees execution results to multiple async output sinks."""
+
+    def __init__(self, sinks: Sequence[AsyncBaseOutputSink]):
+        self._sinks = sinks
+
+    def will_do_work(self, result: ExecutionResult) -> bool:
+        return any(sink.will_do_work(result) for sink in self._sinks)
+
+    async def push(self, result: ExecutionResult) -> None:
+        for sink in self._sinks:
+            if sink.will_do_work(result):
+                try:
+                    async with asyncio.timeout(sink.timeout):
+                        await sink.push(result)
+                except TimeoutError:
+                    pass
+                except Exception:
+                    pass
+
+    async def stop(self) -> None:
+        for sink in self._sinks:
+            await sink.stop()

osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py

@@ -0,0 +1,118 @@
+"""Plugin manager for the async worker.
+
+Discovers plugins via the 'osprey_async_plugin' setuptools entry_point group.
+Also loads sync plugins from 'osprey_plugin' for UDFs (wrapped in adapters).
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, List, Sequence, Type
+
+import pluggy
+from osprey.engine.ast_validator import ValidatorRegistry
+from osprey.engine.executor.udf_execution_helpers import HasHelper, UDFHelpers
+from osprey.engine.udf.base import UDFBase
+from osprey.engine.udf.registry import UDFRegistry
+
+from osprey.async_worker.adaptor import hookspecs as async_hookspecs
+from osprey.async_worker.adaptor.constants import OSPREY_ASYNC_ADAPTOR
+from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink, AsyncMultiOutputSink
+
+if TYPE_CHECKING:
+    from osprey.worker.lib.config import Config
+
+hookimpl_osprey_async: pluggy.HookimplMarker = pluggy.HookimplMarker(OSPREY_ASYNC_ADAPTOR)
+
+plugin_manager = pluggy.PluginManager(OSPREY_ASYNC_ADAPTOR)
+plugin_manager.add_hookspecs(async_hookspecs)
+
+
+def _flatten(seq: List[List[Any]]) -> List[Any]:
+    return sum(seq, [])
+
+
+@lru_cache(maxsize=1)
+def load_all_async_plugins() -> None:
+    """Load all plugins registered under the 'osprey_async_plugin' entry_point group."""
+    plugin_manager.load_setuptools_entrypoints(OSPREY_ASYNC_ADAPTOR)
+    plugin_manager.check_pending()
+
+
+def bootstrap_async_udfs(load_sync_plugins: bool = False) -> tuple[UDFRegistry, UDFHelpers]:
+    """Bootstrap UDFs from async plugins + stdlib.
+
+    Always loads stdlib UDFs (JsonData, StringLength, Rule, etc.) since
+    they're needed for basic rule compilation. Optionally loads from
+    osprey_plugin too (triggers gevent side effects).
+    """
+    # Always load stdlib UDFs — they don't trigger gevent side effects
+    from osprey.worker._stdlibplugin.udf_register import register_udfs as stdlib_register_udfs
+
+    load_all_async_plugins()
+    udf_helpers = UDFHelpers()
+
+    # Load stdlib + async plugin UDFs
+    all_udfs: List[Type[UDFBase[Any, Any]]] = list(stdlib_register_udfs()) + _flatten(plugin_manager.hook.register_udfs())
+
+    if load_sync_plugins:
+        from osprey.worker.adaptor.plugin_manager import load_all_osprey_plugins, plugin_manager as sync_plugin_manager
+
+        load_all_osprey_plugins()
+        sync_udfs: List[Type[UDFBase[Any, Any]]] = _flatten(sync_plugin_manager.hook.register_udfs())
+
+        seen = {udf for udf in all_udfs}
+        for udf in sync_udfs:
+            if udf not in seen:
+                seen.add(udf)
+                all_udfs.append(udf)
+
+    for udf in all_udfs:
+        if issubclass(udf, HasHelper):
+            try:
+                udf_helpers.set_udf_helper(udf, udf.create_provider())
+            except Exception:
+                # Skip helper creation for UDFs that fail (e.g., etcd not available).
+                # These UDFs will fail at execution time via the legacy fallback,
+                # which is expected — errors get captured in error_infos.
+                pass
+
+    udf_registry = UDFRegistry.with_udfs(*all_udfs)
+    return udf_registry, udf_helpers
+
+
+def bootstrap_async_output_sinks(config: Config) -> AsyncMultiOutputSink:
+    """Bootstrap async output sinks from async plugins only.
+
+    Does NOT load sync output sinks — the async worker uses only async sinks.
+    """
+    load_all_async_plugins()
+    sinks: List[AsyncBaseOutputSink] = _flatten(plugin_manager.hook.register_async_output_sinks(config=config))
+    return AsyncMultiOutputSink(sinks)
+
+
+def bootstrap_async_ast_validators(load_sync_plugins: bool = False) -> None:
+    """Bootstrap AST validators from async plugins + stdlib.
+
+    Always loads stdlib validators (ValidateCallKwargs, etc.) since they're
+    needed for rule compilation. Optionally loads from osprey_plugin too.
+    """
+    # Always load stdlib validators — they don't trigger gevent side effects
+    from osprey.worker._stdlibplugin.validator_regsiter import register_ast_validators as stdlib_register_validators
+
+    load_all_async_plugins()
+    validators = list(stdlib_register_validators()) + _flatten(plugin_manager.hook.register_ast_validators())
+
+    if load_sync_plugins:
+        from osprey.worker.adaptor.plugin_manager import load_all_osprey_plugins, plugin_manager as sync_plugin_manager
+
+        load_all_osprey_plugins()
+        sync_validators = _flatten(sync_plugin_manager.hook.register_ast_validators())
+        validators = validators + sync_validators
+
+    registry = ValidatorRegistry.get_instance()
+    seen = set()
+    for validator in validators:
+        if validator not in seen:
+            seen.add(validator)
+            registry.register_to_instance(validator)