diff --git a/docs/integration/integrations.md b/docs/integration/integrations.md index 4d906558..7c37af31 100644 --- a/docs/integration/integrations.md +++ b/docs/integration/integrations.md @@ -77,6 +77,8 @@ UDFs can also produce **effects**: structured outputs that downstream systems ac UDFs that perform I/O (e.g. network calls or database reads) should subclass `AsyncUDFBase` when used in the async worker; see [`osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py`](https://github.com/roostorg/osprey/blob/main/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py) for an example. Pure-computation UDFs like `TextContains` can be reused in both workers without modification. +Native async UDFs time out after two seconds by default. Set `OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT` to change the process default. A UDF can set its own numeric `timeout` class attribute in seconds, which takes precedence over the process default. Values must be greater than zero and can't be NaN or infinity. Async UDFs must not suppress `asyncio.CancelledError`; catch it only to clean up, then re-raise it. + ### Registering UDFs Return your UDF classes from the `register_udfs` hook; for example: diff --git a/osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py b/osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py index ca961239..f33f3104 100644 --- a/osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py +++ b/osprey_async_worker/src/osprey/async_worker/adaptor/interfaces.py @@ -22,6 +22,7 @@ from result import Result _T = TypeVar('_T') +DEFAULT_ASYNC_UDF_TIMEOUT = 2.0 class AsyncUDFBase(UDFBase[Arguments, RValue]): @@ -34,10 +35,15 @@ class AsyncUDFBase(UDFBase[Arguments, RValue]): The sync execute() raises so it can't accidentally be called in the async executor's sync path. + + `timeout` sets the maximum runtime in seconds. Plugins may override it for + a UDF. `async_execute()` must not suppress `asyncio.CancelledError`; catch it + only to clean up, then re-raise it. """ execute_async: ClassVar[bool] = True is_native_async: ClassVar[bool] = True + timeout: ClassVar[float] = DEFAULT_ASYNC_UDF_TIMEOUT def __init__(self, validation_context, arguments): super().__init__(validation_context, arguments) @@ -74,10 +80,13 @@ class AsyncBatchableUDFBase(BatchableUDFBase[Arguments, RValue, BatchableArgumen """Native async batchable UDF base class. Same as AsyncUDFBase but for batchable UDFs. The async executor detects - these and awaits async_execute_batch() directly. + these and awaits async_execute_batch() directly. `timeout` sets the maximum + batch runtime in seconds. `async_execute_batch()` must not suppress + `asyncio.CancelledError`; catch it only to clean up, then re-raise it. """ is_native_async: ClassVar[bool] = True + timeout: ClassVar[float] = DEFAULT_ASYNC_UDF_TIMEOUT def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> RValue: raise RuntimeError( diff --git a/osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py b/osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py index 38c41bf1..03cbdbea 100644 --- a/osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py +++ b/osprey_async_worker/src/osprey/async_worker/adaptor/plugin_manager.py @@ -7,13 +7,19 @@ from __future__ import annotations import logging +import math from functools import lru_cache from typing import TYPE_CHECKING, Any, List, Type, cast import pluggy 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 +from osprey.async_worker.adaptor.interfaces import ( + DEFAULT_ASYNC_UDF_TIMEOUT, + AsyncBaseOutputSink, + AsyncBatchableUDFBase, + AsyncUDFBase, +) from osprey.async_worker.sinks.sink.output_sink import AsyncMultiOutputSink from osprey.engine.ast_validator import ValidatorRegistry from osprey.engine.executor.udf_execution_helpers import HasHelper, UDFHelpers @@ -24,6 +30,8 @@ if TYPE_CHECKING: from osprey.worker.lib.config import Config +OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT = 'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT' + hookimpl_osprey_async: pluggy.HookimplMarker = pluggy.HookimplMarker(OSPREY_ASYNC_ADAPTOR) plugin_manager = pluggy.PluginManager(OSPREY_ASYNC_ADAPTOR) @@ -66,6 +74,17 @@ def _deduplicate_udfs( return deduplicated +def _validate_udf_timeout(value: object, source: str) -> None: + """Validate a timeout value. + + Raises ValueError naming the source and value when not finite or non-positive. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f'{source} must be a positive finite number, got {value!r}') + if not math.isfinite(value) or value <= 0: + raise ValueError(f'{source} must be a positive finite number, got {value}') + + def bootstrap_async_udfs(config: 'Config | None' = None) -> tuple[UDFRegistry, UDFHelpers]: """Bootstrap UDFs from async plugins + stdlib. @@ -76,9 +95,21 @@ def bootstrap_async_udfs(config: 'Config | None' = None) -> tuple[UDFRegistry, U `_async_stdlib_plugin`, which registers through the same `register_udfs` hook as third-party plugins. No sync fallbacks — all I/O UDFs must be native async. + + Resolves the default UDF timeout from OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT config key, + defaulting to 2.0. Each call replaces the process-wide inherited default and + must finish before UDF execution starts. Validation completes before mutation. """ from osprey.worker._stdlibplugin.udf_register import register_udfs as stdlib_register_udfs + # Resolve and validate the configured timeout before any mutations + resolved_timeout = ( + DEFAULT_ASYNC_UDF_TIMEOUT + if config is None + else config.get_float(OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT, DEFAULT_ASYNC_UDF_TIMEOUT) + ) + _validate_udf_timeout(resolved_timeout, f'config[{OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT!r}]') + load_all_async_plugins() udf_helpers = UDFHelpers() @@ -86,6 +117,13 @@ def bootstrap_async_udfs(config: 'Config | None' = None) -> tuple[UDFRegistry, U plugin_udfs = _flatten(plugin_manager.hook.register_udfs()) all_udfs = _deduplicate_udfs(stdlib_udfs, plugin_udfs) + for udf in all_udfs: + if issubclass(udf, (AsyncUDFBase, AsyncBatchableUDFBase)): + _validate_udf_timeout(udf.timeout, f'{udf.__name__}.timeout') + + AsyncUDFBase.timeout = resolved_timeout + AsyncBatchableUDFBase.timeout = resolved_timeout + # Auto-register helpers for UDFs that extend HasHelper for udf in all_udfs: if issubclass(udf, HasHelper): diff --git a/osprey_async_worker/src/osprey/async_worker/cli/main.py b/osprey_async_worker/src/osprey/async_worker/cli/main.py index 313f3735..376dfa7b 100644 --- a/osprey_async_worker/src/osprey/async_worker/cli/main.py +++ b/osprey_async_worker/src/osprey/async_worker/cli/main.py @@ -50,6 +50,7 @@ def bootstrap_stdlib_engine(rules_path: str) -> tuple[AsyncOspreyEngine, UDFHelp """Bootstrap engine with only stdlib UDFs — no external plugins, no Postgres, no labels. This avoids loading example_plugins or any third-party plugins that require database connections. + It registers no native async UDFs, so native timeout configuration does not apply. """ from osprey.engine.ast_validator import ValidatorRegistry from osprey.worker._stdlibplugin.udf_register import register_udfs as stdlib_register_udfs diff --git a/osprey_async_worker/src/osprey/async_worker/executor.py b/osprey_async_worker/src/osprey/async_worker/executor.py index 19d6c92f..c6b10e3c 100644 --- a/osprey_async_worker/src/osprey/async_worker/executor.py +++ b/osprey_async_worker/src/osprey/async_worker/executor.py @@ -237,8 +237,12 @@ async def _execute_async_udf( """Execute a native async UDF. Awaited directly on the event loop. `pre_resolved_arguments`, when given, is the Arguments already computed for this same - chain (e.g. by `_enqueue_batches` while checking whether a batch would form), so this - skips a redundant `resolve_arguments` call for the same message. + chain while checking whether a batch would form, so execution does not resolve it twice. + + Timeout enforcement: asyncio.timeout() enforces udf.timeout only around + async_execute() after semaphore admission. TimeoutError is converted to + NodeErrorInfo through the existing exception handler, preserving current + failure handling and Sentry capture behavior without modification. """ async with semaphore: call_executor: CallExecutor = chain.executor # type: ignore @@ -254,7 +258,8 @@ async def _execute_async_udf( else udf.resolve_arguments(context, call_executor) ) with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01): - result = await udf.async_execute(context, resolved_arguments) + async with asyncio.timeout(type(udf).timeout): + result = await udf.async_execute(context, resolved_arguments) execution_result = Ok(udf.check_result_type(result)) except Exception as e: if not isinstance(e, NodeFailurePropagationException): @@ -274,15 +279,26 @@ async def _execute_async_batch( context: ExecutionContext, error_info_: list[NodeErrorInfo], ) -> Sequence[NodeResult]: - """Execute a batch of native async batchable UDFs.""" + """Execute a batch of native async batchable UDFs. + + Timeout enforcement: the shared batch deadline is computed as + max(udf.timeout for udf in udfs). asyncio.timeout() enforces this + deadline only around async_execute_batch() after semaphore admission. + TimeoutError is converted to NodeErrorInfo through the existing exception + handler, preserving current failure handling and Sentry capture behavior. + """ async with semaphore: assert len(udfs) == len(nodes) == len(batchable_args) num_executions = len(udfs) metric_tags = _get_metric_tags(context, udfs[0]) + # Compute the shared deadline using the maximum timeout from all UDFs + batch_timeout = max(type(udf).timeout for udf in udfs) + try: with metrics.timed('udf_execution_batch_duration', tags=metric_tags, sample_rate=0.01): - results = await udfs[0].async_execute_batch(context, udfs, batchable_args) + async with asyncio.timeout(batch_timeout): + results = await udfs[0].async_execute_batch(context, udfs, batchable_args) assert len(results) == num_executions except Exception as e: if not isinstance(e, NodeFailurePropagationException): @@ -343,7 +359,7 @@ async def _enqueue_batches( A native async chain can include arguments that this function resolved. Legacy chains include `None`. """ - batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, ArgumentsBase, Any]]] = defaultdict(list) + batch_chains: dict[tuple[type, str, bool], list[tuple[DependencyChain, ArgumentsBase, Any]]] = defaultdict(list) chains_to_remove: list[DependencyChain] = [] for async_chain in ready_async: @@ -355,11 +371,14 @@ async def _enqueue_batches( udf = call_executor._udf batch_type = udf.get_batchable_arguments_type() + is_native = isinstance(udf, AsyncBatchableUDFBase) try: resolved_arguments = udf.resolve_arguments(context, call_executor) batchable_arguments = udf.get_batchable_arguments(resolved_arguments) routing_key = udf.get_batch_routing_key(batchable_arguments) - batch_chains[(batch_type, routing_key)].append((async_chain, resolved_arguments, batchable_arguments)) + batch_chains[(batch_type, routing_key, is_native)].append( + (async_chain, resolved_arguments, batchable_arguments) + ) except Exception as e: if not isinstance(e, NodeFailurePropagationException): error_infos.append(NodeErrorInfo(e, call_executor.node)) diff --git a/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py b/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py index 3ff9ca88..7bef654b 100644 --- a/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py +++ b/osprey_async_worker/src/osprey/async_worker/stdlib_udfs/async_mx_lookup.py @@ -17,6 +17,8 @@ from osprey.engine.stdlib.udfs.mx_lookup import MXLookup as SyncMXLookup _DNS_TIMEOUT = 5.0 +_DNS_TRIES = 3 +_MX_LOOKUP_TIMEOUT = 31.0 _resolver: aiodns.DNSResolver | None = None @@ -25,7 +27,7 @@ def _get_resolver() -> aiodns.DNSResolver: global _resolver loop = asyncio.get_running_loop() if _resolver is None or _resolver.loop is not loop: - _resolver = aiodns.DNSResolver(timeout=_DNS_TIMEOUT, loop=loop) + _resolver = aiodns.DNSResolver(timeout=_DNS_TIMEOUT, tries=_DNS_TRIES, loop=loop) return _resolver @@ -33,6 +35,9 @@ class MXLookup(AsyncUDFBase[Arguments, str]): # type: ignore[misc] """Async MXLookup — uses aiodns for non-blocking DNS resolution.""" category = SyncMXLookup.category + # One nameserver can consume 5s for each of 3 tries on both sequential + # MX and A queries; reserve one additional second for scheduling slack + timeout = _MX_LOOKUP_TIMEOUT @classmethod def _get_udf_base_args(cls): diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py b/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py index c28f3af7..73db2d63 100644 --- a/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_async_executor.py @@ -13,7 +13,11 @@ import pytest from osprey.async_worker import executor as async_executor -from osprey.async_worker.adaptor.interfaces import AsyncBatchableUDFBase, AsyncUDFBase +from osprey.async_worker.adaptor.interfaces import ( + DEFAULT_ASYNC_UDF_TIMEOUT, + AsyncBatchableUDFBase, + AsyncUDFBase, +) from osprey.async_worker.executor import execute from osprey.engine.ast.grammar import Source from osprey.engine.ast.sources import Sources @@ -25,11 +29,21 @@ from osprey.engine.executor.udf_execution_helpers import UDFHelpers from osprey.engine.stdlib import get_config_registry from osprey.engine.udf.arguments import ArgumentsBase -from osprey.engine.udf.base import UDFBase +from osprey.engine.udf.base import BatchableUDFBase, UDFBase from osprey.engine.udf.registry import UDFRegistry from result import Ok, Result +@pytest.fixture(autouse=True) +def reset_udf_timeouts(): + """Reset both async UDF base timeouts before and after each test.""" + AsyncUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + AsyncBatchableUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + yield + AsyncUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + AsyncBatchableUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + + class CountingBatchableArguments(ArgumentsBase): key: str value: str @@ -381,6 +395,492 @@ async def test_parity_complex_graph(async_execute_fn): assert result['RuleB'] is False +# Test argument type for timeout testing +class TimeoutTestArguments(ArgumentsBase): + """Simple test argument type.""" + + value: str + + +# Timeout behavior + + +@pytest.mark.asyncio +async def test_native_udf_base_exposes_timeout(): + """Native async UDF bases expose a two-second timeout.""" + assert hasattr(AsyncUDFBase, 'timeout') + assert AsyncUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + assert hasattr(AsyncBatchableUDFBase, 'timeout') + assert AsyncBatchableUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + + # Legacy bases should NOT have timeout + assert not hasattr(UDFBase, 'timeout') + assert not hasattr(BatchableUDFBase, 'timeout') + + +@pytest.mark.asyncio +async def test_subclass_timeout_override_used(async_execute_with_result): + """The executor enforces a positive subclass timeout override.""" + + class CustomTimeoutUDF(AsyncUDFBase[TimeoutTestArguments, str]): + timeout: ClassVar[float] = 0.1 # Short override + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + # This will timeout because 0.1s is very short + await asyncio.sleep(0.5) + return arguments.value + + registry = UDFRegistry.with_udfs(CustomTimeoutUDF) + result = await async_execute_with_result( + 'Result = CustomTimeoutUDF(value="test")', + data={}, + udf_registry=registry, + ) + + # Should have failed with timeout + assert len(result.error_infos) > 0 + assert isinstance(result.error_infos[0].error, TimeoutError) + + +@pytest.mark.asyncio +async def test_instance_timeout_shadow_does_not_bypass_class_deadline(async_execute_with_result): + """The validated class deadline wins over an instance attribute shadow.""" + + class InstanceShadowUDF(AsyncUDFBase[TimeoutTestArguments, str]): + timeout: ClassVar[float] = 0.1 + + def __init__(self, validation_context, arguments): + super().__init__(validation_context, arguments) + setattr(self, 'timeout', float('inf')) + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + await asyncio.sleep(0.5) + return arguments.value + + result = await async_execute_with_result( + 'Result = InstanceShadowUDF(value="test")', + data={}, + udf_registry=UDFRegistry.with_udfs(InstanceShadowUDF), + ) + + assert isinstance(result.error_infos[0].error, TimeoutError) + + +@pytest.mark.asyncio +async def test_udf_completes_within_deadline(async_execute_with_result): + """A UDF completing before its deadline retains its result without errors.""" + + class ControlledAsyncUDF(AsyncUDFBase[TimeoutTestArguments, str]): + """Async UDF controlled by an event for testing.""" + + _release_event: asyncio.Event | None = None + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + """Wait for release event before completing.""" + if self._release_event: + await self._release_event.wait() + return arguments.value + + release_event = asyncio.Event() + ControlledAsyncUDF._release_event = release_event + + registry = UDFRegistry.with_udfs(ControlledAsyncUDF) + + # Start execution and immediately release + task = asyncio.create_task( + async_execute_with_result( + 'Result = ControlledAsyncUDF(value="success")', + data={}, + udf_registry=registry, + ) + ) + + # Give it a moment to start + await asyncio.sleep(0.01) + + # Release the UDF + release_event.set() + + result = await task + + # Should have completed successfully + assert result.extracted_features['Result'] == 'success' + assert len(result.error_infos) == 0 + + +@pytest.mark.asyncio +async def test_udf_timeout_recorded_as_error(async_execute_with_result): + """A UDF exceeding its deadline records a TimeoutError.""" + release_event = asyncio.Event() + + class ShortTimeoutUDF(AsyncUDFBase[TimeoutTestArguments, str]): + timeout: ClassVar[float] = 0.1 # 100ms timeout + _release_event: asyncio.Event | None = None + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + # Wait for release event; timeout will trigger before it fires + if ShortTimeoutUDF._release_event: + await ShortTimeoutUDF._release_event.wait() + return arguments.value + + ShortTimeoutUDF._release_event = release_event + + registry = UDFRegistry.with_udfs(ShortTimeoutUDF) + + result = await async_execute_with_result( + 'Result = ShortTimeoutUDF(value="timeout")', + data={}, + udf_registry=registry, + ) + + # Should have a timeout error + assert len(result.error_infos) > 0 + assert isinstance(result.error_infos[0].error, TimeoutError) + + +@pytest.mark.asyncio +async def test_semaphore_wait_not_counted_against_deadline(async_execute_with_result): + """Semaphore wait time does not consume the UDF execution deadline. + + Coordination: + 1. First (blocking) acquires semaphore and signals block_event + 2. Main test waits for block_event to confirm First owns semaphore + 3. Main test holds First in semaphore for 0.6s (longer than Second's 0.5s timeout) + 4. Main test releases First; Second can now acquire semaphore + 5. Second completes within its 0.5s timeout (because semaphore wait doesn't count) + """ + + class QuickAsyncUDF(AsyncUDFBase[TimeoutTestArguments, str]): + """Async UDF with a short timeout for testing semaphore isolation.""" + + timeout: ClassVar[float] = 0.5 # Short timeout to test semaphore isolation + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + """Complete immediately.""" + # Return immediately to avoid blocking semaphore + await asyncio.sleep(0) + return arguments.value + + class BlockingAsyncUDF(AsyncUDFBase[TimeoutTestArguments, str]): + """Async UDF that blocks the semaphore for testing.""" + + _block_event: asyncio.Event | None = None + _release_event: asyncio.Event | None = None + + @classmethod + def _get_udf_base_args(cls): + return (TimeoutTestArguments, str) + + async def async_execute(self, execution_context: ExecutionContext, arguments: TimeoutTestArguments) -> str: + """Signal that we own the semaphore, then wait for release.""" + if BlockingAsyncUDF._block_event: + BlockingAsyncUDF._block_event.set() + if BlockingAsyncUDF._release_event: + await BlockingAsyncUDF._release_event.wait() + return arguments.value + + block_event = asyncio.Event() + release_event = asyncio.Event() + + # Set test coordination events + BlockingAsyncUDF._block_event = block_event + BlockingAsyncUDF._release_event = release_event + + registry = UDFRegistry.with_udfs(BlockingAsyncUDF, QuickAsyncUDF) + + async def run_execution(): + return await async_execute_with_result( + """ + First = BlockingAsyncUDF(value="first") + Second = QuickAsyncUDF(value="second") + """, + data={}, + udf_registry=registry, + max_concurrent=1, # Only one concurrent execution + ) + + # Start execution + result_task = asyncio.create_task(run_execution()) + + # Wait for First (blocking) to own the semaphore and signal block_event + try: + await asyncio.wait_for(block_event.wait(), timeout=1.0) + except asyncio.TimeoutError: + raise AssertionError('BlockingAsyncUDF did not acquire semaphore') + + # Hold First in semaphore for 0.6s (longer than Second's 0.5s execution timeout) + # Second is waiting for semaphore; the wait must not count against its deadline + await asyncio.sleep(0.6) + + # Release First; Second acquires semaphore and must complete in its 0.5s timeout + release_event.set() + + # Get the result + result = await result_task + + # Both should succeed - QuickAsyncUDF should not have timed out + # even though it waited 0.6s for the semaphore + assert result.extracted_features['First'] == 'first' + assert result.extracted_features['Second'] == 'second' + assert len(result.error_infos) == 0 + + +# Test argument type for batch testing +class BatchTestArgs(ArgumentsBase): + """Test argument type for batch UDFs.""" + + id: str + routing_key: str + + +@pytest.mark.asyncio +async def test_batch_mixed_timeout_uses_highest(async_execute_with_result): + """A mixed-timeout batch uses the highest timeout without changing grouping. + + Tests that a batch with mixed timeouts uses the MAX timeout: + - Both UDFs batch together (grouping unchanged) + - Batch sleeps 0.3s which would timeout at 0.2s but succeeds at 1.0s + - Verifies executor uses max(0.2, 1.0) = 1.0s timeout for the batch + """ + # Isolated capture state for this test execution + captured_timeouts: list[float] = [] + + class MixedTimeoutBatchUDF(AsyncBatchableUDFBase[BatchTestArgs, str, BatchTestArgs]): + timeout: ClassVar[float] = 2.0 + + @classmethod + def _get_udf_base_args(cls): + return (BatchTestArgs, str, BatchTestArgs) + + def get_batchable_arguments(self, arguments: BatchTestArgs) -> BatchTestArgs: + return arguments + + async def async_execute(self, execution_context: ExecutionContext, arguments: BatchTestArgs) -> str: + return f'result-{arguments.id}' + + async def async_execute_batch( + self, + execution_context: ExecutionContext, + udfs: Sequence, + arguments: Sequence[BatchTestArgs], + ) -> Sequence[Result[str, Exception]]: + # Record the actual timeouts in local capture state + nonlocal captured_timeouts + captured_timeouts = [u.timeout for u in udfs] + # Sleep 0.3s: would fail with 0.2s timeout but pass with 1.0s + await asyncio.sleep(0.3) + return [Ok(f'result-{arg.id}') for arg in arguments] + + class ShortTimeoutBatchUDF(MixedTimeoutBatchUDF): + timeout: ClassVar[float] = 0.2 + + class LongTimeoutBatchUDF(MixedTimeoutBatchUDF): + timeout: ClassVar[float] = 1.0 + + registry = UDFRegistry.with_udfs(ShortTimeoutBatchUDF, LongTimeoutBatchUDF) + + result = await async_execute_with_result( + """ + Short = ShortTimeoutBatchUDF(id="short", routing_key="shared") + Long = LongTimeoutBatchUDF(id="long", routing_key="shared") + """, + data={}, + udf_registry=registry, + ) + + # Verify both UDFs were in the same batch + timeouts = sorted(captured_timeouts) + assert timeouts == [0.2, 1.0] + + # Both should have succeeded (batch used max timeout of 1.0s, not min of 0.2s) + assert result.extracted_features['Short'] == 'result-short' + assert result.extracted_features['Long'] == 'result-long' + assert len(result.error_infos) == 0 + + +@pytest.mark.asyncio +async def test_native_and_legacy_batches_with_same_route_execute_separately(async_execute_with_result): + """Native and legacy batch implementations never share one execution batch.""" + + class NativeBatchUDF(AsyncBatchableUDFBase[BatchTestArgs, str, BatchTestArgs]): + @classmethod + def _get_udf_base_args(cls): + return (BatchTestArgs, str, BatchTestArgs) + + def get_batchable_arguments(self, arguments: BatchTestArgs) -> BatchTestArgs: + return arguments + + async def async_execute(self, execution_context: ExecutionContext, arguments: BatchTestArgs) -> str: + return f'native-{arguments.id}' + + async def async_execute_batch( + self, + execution_context: ExecutionContext, + udfs: Sequence, + arguments: Sequence[BatchTestArgs], + ) -> Sequence[Result[str, Exception]]: + return [Ok(f'native-{arg.id}') for arg in arguments] + + class LegacyBatchUDF(BatchableUDFBase[BatchTestArgs, str, BatchTestArgs]): + @classmethod + def _get_udf_base_args(cls): + return (BatchTestArgs, str, BatchTestArgs) + + def get_batchable_arguments(self, arguments: BatchTestArgs) -> BatchTestArgs: + return arguments + + def execute(self, execution_context: ExecutionContext, arguments: BatchTestArgs) -> str: + return f'legacy-{arguments.id}' + + def execute_batch( + self, + execution_context: ExecutionContext, + udfs: Sequence, + arguments: Sequence[BatchTestArgs], + ) -> Sequence[Result[str, Exception]]: + return [Ok(f'legacy-{arg.id}') for arg in arguments] + + result = await async_execute_with_result( + """ + Native1 = NativeBatchUDF(id="one", routing_key="shared") + Legacy1 = LegacyBatchUDF(id="one", routing_key="shared") + Native2 = NativeBatchUDF(id="two", routing_key="shared") + Legacy2 = LegacyBatchUDF(id="two", routing_key="shared") + """, + data={}, + udf_registry=UDFRegistry.with_udfs(NativeBatchUDF, LegacyBatchUDF), + ) + + assert result.extracted_features['Native1'] == 'native-one' + assert result.extracted_features['Legacy1'] == 'legacy-one' + assert result.extracted_features['Native2'] == 'native-two' + assert result.extracted_features['Legacy2'] == 'legacy-two' + assert result.error_infos == [] + + +@pytest.mark.asyncio +async def test_batch_within_deadline_retains_results(async_execute_with_result): + """A batch within its shared deadline retains every result without errors.""" + + class FastBatchUDF(AsyncBatchableUDFBase[BatchTestArgs, str, BatchTestArgs]): + timeout: ClassVar[float] = 0.5 + + @classmethod + def _get_udf_base_args(cls): + return (BatchTestArgs, str, BatchTestArgs) + + def get_batchable_arguments(self, arguments: BatchTestArgs) -> BatchTestArgs: + return arguments + + async def async_execute(self, execution_context: ExecutionContext, arguments: BatchTestArgs) -> str: + return f'fast-{arguments.id}' + + async def async_execute_batch( + self, + execution_context: ExecutionContext, + udfs: Sequence, + arguments: Sequence[BatchTestArgs], + ) -> Sequence[Result[str, Exception]]: + # Quick execution within timeout + await asyncio.sleep(0.1) + return [Ok(f'fast-{arg.id}') for arg in arguments] + + registry = UDFRegistry.with_udfs(FastBatchUDF) + + result = await async_execute_with_result( + """ + Result1 = FastBatchUDF(id="first", routing_key="shared") + Result2 = FastBatchUDF(id="second", routing_key="shared") + """, + data={}, + udf_registry=registry, + ) + + # Both should succeed + assert result.extracted_features['Result1'] == 'fast-first' + assert result.extracted_features['Result2'] == 'fast-second' + assert len(result.error_infos) == 0 + + +@pytest.mark.asyncio +async def test_batch_over_deadline_fails_all_nodes(async_execute_with_result): + """A batch exceeding its shared deadline records timeout and fails every node. + + Direct assertion: async_execute_batch() receives asyncio.CancelledError + when the batch timeout fires, confirming the coroutine is cancelled. + """ + # Isolated flag to capture cancellation signal + batch_received_cancellation = False + + class SlowBatchUDF(AsyncBatchableUDFBase[BatchTestArgs, str, BatchTestArgs]): + timeout: ClassVar[float] = 0.1 # 100ms timeout + + @classmethod + def _get_udf_base_args(cls): + return (BatchTestArgs, str, BatchTestArgs) + + def get_batchable_arguments(self, arguments: BatchTestArgs) -> BatchTestArgs: + return arguments + + async def async_execute(self, execution_context: ExecutionContext, arguments: BatchTestArgs) -> str: + return f'slow-{arguments.id}' + + async def async_execute_batch( + self, + execution_context: ExecutionContext, + udfs: Sequence, + arguments: Sequence[BatchTestArgs], + ) -> Sequence[Result[str, Exception]]: + # Exceed the 0.1s timeout + nonlocal batch_received_cancellation + try: + await asyncio.sleep(0.5) + return [Ok(f'slow-{arg.id}') for arg in arguments] + except asyncio.CancelledError: + # Direct assertion: batch coroutine receives CancelledError + batch_received_cancellation = True + raise + + registry = UDFRegistry.with_udfs(SlowBatchUDF) + + result = await async_execute_with_result( + """ + Result1 = SlowBatchUDF(id="first", routing_key="shared") + Result2 = SlowBatchUDF(id="second", routing_key="shared") + """, + data={}, + udf_registry=registry, + ) + + # The batch coroutine receives cancellation before node errors are recorded + assert batch_received_cancellation, 'async_execute_batch() did not receive CancelledError' + + # Both results should have failed with TimeoutError in the batch + assert len(result.error_infos) >= 2 + timeout_errors = [ei for ei in result.error_infos if isinstance(ei.error, TimeoutError)] + assert len(timeout_errors) >= 2, f'Expected 2+ TimeoutErrors, got {len(timeout_errors)} from {result.error_infos}' + + @pytest.fixture() def counting_batchable_udf(): """Registers CountingBatchableUdf and resets its call-count/raise state around the test.""" diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_plugin_manager.py b/osprey_async_worker/src/osprey/async_worker/tests/test_plugin_manager.py index 0934dff2..a168e71d 100644 --- a/osprey_async_worker/src/osprey/async_worker/tests/test_plugin_manager.py +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_plugin_manager.py @@ -10,12 +10,18 @@ import pytest from osprey.async_worker.adaptor import plugin_manager as pm +from osprey.async_worker.adaptor.interfaces import ( + DEFAULT_ASYNC_UDF_TIMEOUT, + AsyncBatchableUDFBase, + AsyncUDFBase, +) from osprey.async_worker.stdlib_udfs import _async_stdlib_plugin from osprey.async_worker.stdlib_udfs.async_mx_lookup import MXLookup as AsyncMXLookup from osprey.engine.stdlib.udfs.json_data import JsonData from osprey.engine.stdlib.udfs.labels import HasLabel as SyncHasLabel from osprey.engine.stdlib.udfs.mx_lookup import MXLookup as SyncMXLookup from osprey.engine.stdlib.udfs.rules import Rule +from osprey.worker.lib.config import Config @pytest.fixture(autouse=True) @@ -24,12 +30,17 @@ def reset_plugin_manager(): plugin_manager is a module-level singleton. Without this, state from one test (e.g. a registered plugin) leaks into the next. + Restore both native base timeouts before and after each test. """ pm.load_all_async_plugins.cache_clear() + AsyncUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + AsyncBatchableUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT yield pm.load_all_async_plugins.cache_clear() if pm.plugin_manager.is_registered(_async_stdlib_plugin): pm.plugin_manager.unregister(_async_stdlib_plugin) + AsyncUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT + AsyncBatchableUDFBase.timeout = DEFAULT_ASYNC_UDF_TIMEOUT def test_async_stdlib_plugin_returns_async_mx_lookup() -> None: @@ -58,6 +69,11 @@ def test_bootstrap_resolves_mx_lookup_to_async_version() -> None: ) +def test_mx_lookup_deadline_covers_sequential_dns_queries() -> None: + """MX lookup allows every configured try for both sequential resolver calls.""" + assert AsyncMXLookup.timeout == 31.0 + + def test_bootstrap_does_not_register_sync_mx_lookup() -> None: """Sync MXLookup must not appear in the merged registry under any name.""" registry, _helpers = pm.bootstrap_async_udfs(config=None) @@ -152,9 +168,9 @@ def test_bootstrap_applies_register_udf_helpers_bindings() -> None: plugin = _UDFHelpersPlugin(_StubUDF, helper, captured) pm.plugin_manager.register(plugin) try: - fake_config = object() - _registry, helpers = pm.bootstrap_async_udfs(config=fake_config) # type: ignore[arg-type] - assert captured == [fake_config], 'register_udf_helpers must receive the config' + bound_config = Config({}) + _registry, helpers = pm.bootstrap_async_udfs(config=bound_config) + assert captured == [bound_config], 'register_udf_helpers must receive the config' # UDFHelpers.get_udf_helper expects an instance (it calls type()). # Inspect the underlying dict directly since _StubUDF is not instantiable. assert helpers._helpers[_StubUDF] is helper @@ -188,8 +204,8 @@ def register_udf_helpers(self, config): plugin = _BrokenPlugin() pm.plugin_manager.register(plugin) try: - fake_config = object() - registry, _helpers = pm.bootstrap_async_udfs(config=fake_config) # type: ignore[arg-type] + bound_config = Config({}) + registry, _helpers = pm.bootstrap_async_udfs(config=bound_config) # Standard UDFs still resolved despite the broken hook. assert registry.get('JsonData') is JsonData finally: @@ -203,3 +219,134 @@ def test_no_residual_register_labels_service_or_provider_hookspec() -> None: assert not hasattr(pm.plugin_manager.hook, 'register_labels_service_or_provider'), ( 'register_labels_service_or_provider should be removed in favor of register_udf_helpers' ) + + +# Absent configuration keeps both inherited native UDF defaults at 2.0 +def test_absent_config_keeps_native_udf_defaults_at_two() -> None: + """bootstrap with config=None and Config({}) leaves both native base defaults at 2.0.""" + # config=None case + pm.bootstrap_async_udfs(config=None) + assert AsyncUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + assert AsyncBatchableUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + + # Config({}) case (initialized config with no timeout key) + pm.bootstrap_async_udfs(config=Config({})) + assert AsyncUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + assert AsyncBatchableUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + + +# A positive finite configured timeout changes both inherited defaults +def test_positive_finite_config_changes_native_udf_defaults() -> None: + """a numeric string under OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT becomes the timeout on both native bases.""" + config = Config({'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT': '3.5'}) + pm.bootstrap_async_udfs(config=config) + assert AsyncUDFBase.timeout == 3.5 + assert AsyncBatchableUDFBase.timeout == 3.5 + + +# Direct and inherited UDF timeout overrides win over the configured default +def test_direct_and_inherited_udf_timeout_overrides_win() -> None: + """define one direct override and one concrete udf inheriting an override from an + intermediate plugin base; bootstrap with a different configured default and assert + both overrides remain unchanged.""" + + class DirectOverrideUDF(AsyncUDFBase): + timeout = 1.5 + + class PluginBaseUDF(AsyncUDFBase): + timeout = 2.5 + + class InheritedOverrideUDF(PluginBaseUDF): + pass # Inherits timeout = 2.5 + + config = Config({'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT': '5.0'}) + pm.bootstrap_async_udfs(config=config) + + # Native base defaults changed to 5.0 + assert AsyncUDFBase.timeout == 5.0 + assert AsyncBatchableUDFBase.timeout == 5.0 + + # But direct and inherited overrides remain unchanged + assert DirectOverrideUDF.timeout == 1.5 + assert InheritedOverrideUDF.timeout == 2.5 + + +@pytest.mark.parametrize( + 'udf_base,timeout_value,inherits_override', + [ + (AsyncUDFBase, 0.0, False), + (AsyncBatchableUDFBase, float('inf'), False), + (AsyncUDFBase, float('nan'), True), + (AsyncBatchableUDFBase, -1.0, True), + (AsyncUDFBase, '5', False), + (AsyncBatchableUDFBase, True, True), + ], +) +def test_invalid_registered_udf_timeout_override_fails_bootstrap( + udf_base: type, + timeout_value: object, + inherits_override: bool, +) -> None: + if inherits_override: + plugin_base = type('PluginTimeoutBase', (udf_base,), {'timeout': timeout_value}) + invalid_udf_class = type('InvalidTimeoutUDF', (plugin_base,), {}) + else: + invalid_udf_class = type('InvalidTimeoutUDF', (udf_base,), {'timeout': timeout_value}) + + class InvalidTimeoutPlugin: + @pm.hookimpl_osprey_async + def register_udfs(self): + return [invalid_udf_class] + + plugin = InvalidTimeoutPlugin() + pm.plugin_manager.register(plugin) + try: + with pytest.raises(ValueError, match='InvalidTimeoutUDF.timeout'): + pm.bootstrap_async_udfs(config=Config({})) + finally: + pm.plugin_manager.unregister(plugin) + + +# Malformed, non-positive, and non-finite timeout values fail before mutation +@pytest.mark.parametrize( + 'config_value,expected_exception', + [ + ('not_a_number', TypeError), + ('0', ValueError), + ('-1.5', ValueError), + ('nan', ValueError), + ('inf', ValueError), + ('-inf', ValueError), + ], +) +def test_invalid_timeout_config_fails_before_mutation(config_value: str, expected_exception: type) -> None: + """malformed, zero, negative, nan, inf, and -inf values fail bootstrap with appropriate errors.""" + config = Config({'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT': config_value}) + with pytest.raises(expected_exception, match='OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT'): + pm.bootstrap_async_udfs(config=config) + # Verify state unchanged + assert AsyncUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + assert AsyncBatchableUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + + +# Repeated bootstrap replaces prior timeout defaults without leaking state +def test_repeated_bootstrap_replaces_prior_timeout_default() -> None: + """bootstrap with one configured value, then another, then config=None; + assert each call replaces only the inherited base defaults solely through bootstrap behavior.""" + + # First bootstrap with 3.0 + config1 = Config({'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT': '3.0'}) + pm.bootstrap_async_udfs(config=config1) + assert AsyncUDFBase.timeout == 3.0 + assert AsyncBatchableUDFBase.timeout == 3.0 + + # Second bootstrap with 4.5 (no manual reset; bootstrap must replace 3.0 with 4.5) + config2 = Config({'OSPREY_ASYNC_UDF_DEFAULT_TIMEOUT': '4.5'}) + pm.bootstrap_async_udfs(config=config2) + assert AsyncUDFBase.timeout == 4.5 + assert AsyncBatchableUDFBase.timeout == 4.5 + + # Final bootstrap with config=None (no manual reset; bootstrap must restore 2.0) + pm.bootstrap_async_udfs(config=None) + assert AsyncUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT + assert AsyncBatchableUDFBase.timeout == DEFAULT_ASYNC_UDF_TIMEOUT