[Performance] Improve engine and coordinator performance

osprey_async_worker/src/osprey/async_worker/executor.py

@@ -47,6 +47,9 @@
 
 logger = get_logger(__name__)
 
+# Shared placeholder for the "make mypy happy" default below; always overwritten before use.
+_UNSET_RESULT: NodeResult = Err(None)
+
 _DEFAULT_MAX_ASYNC_PER_EXECUTION = 12
 
 
@@ -119,7 +122,7 @@ def _execute_sync(
     error_info_: list[NodeErrorInfo],
 ) -> NodeResult:
     """Execute a sync UDF inline. For pure computation only — no I/O."""
-    execution_result: NodeResult = Err(None)
+    execution_result: NodeResult = _UNSET_RESULT
     try:
         execution_result = Ok(chain.executor.execute(execution_context=context))
     except Exception as e:
@@ -148,7 +151,7 @@ def _execute_legacy_sync(
         call_node: CallExecutor = chain.executor
         metric_tags += [f'udf:{call_node._udf.__class__.__name__}']
 
-    execution_result: NodeResult = Err(None)
+    execution_result: NodeResult = _UNSET_RESULT
     try:
         with metrics.timed('udf_execution_duration', tags=metric_tags, sample_rate=0.01):
             execution_result = Ok(chain.executor.execute(execution_context=context))
@@ -227,17 +230,27 @@ async def _execute_async_udf(
     chain: DependencyChain,
     context: ExecutionContext,
     error_info_: list[NodeErrorInfo],
+    pre_resolved_arguments: Any | None = None,
 ) -> NodeResult:
-    """Execute a native async UDF. Awaited directly on the event loop."""
+    """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.
+    """
     async with semaphore:
         call_executor: CallExecutor = chain.executor  # type: ignore
         udf: AsyncUDFBase[Any, Any] = call_executor._udf  # type: ignore
         metric_tags = _get_metric_tags(context) + [f'udf:{udf.__class__.__name__}']
 
         caught_exception: Exception | None = None
-        execution_result: NodeResult = Err(None)
+        execution_result: NodeResult = _UNSET_RESULT
         try:
-            resolved_arguments = udf.resolve_arguments(context, call_executor)
+            resolved_arguments = (
+                pre_resolved_arguments
+                if pre_resolved_arguments is not None
+                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)
             execution_result = Ok(udf.check_result_type(result))
@@ -320,12 +333,15 @@ async def _enqueue_batches(
     context: ExecutionContext,
     error_infos: list[NodeErrorInfo],
     ready_async: Sequence[DependencyChain],
-) -> tuple[Sequence[DependencyChain], dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]]]:
-    """Collect batchable async chains and launch them as tasks.
+) -> tuple[
+    Sequence[tuple[DependencyChain, Any | None]],
+    dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]],
+]:
+    """Launch batches and return the chains that remain.
 
-    Returns (remaining non-batched chains, dict of batch tasks -> chains).
+    A native async chain can include arguments that this function resolved. Legacy chains include `None`.
     """
-    batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, Any]]] = defaultdict(list)
+    batch_chains: dict[tuple[type, str], list[tuple[DependencyChain, Any, Any]]] = defaultdict(list)
     chains_to_remove: list[DependencyChain] = []
 
     for async_chain in ready_async:
@@ -341,20 +357,25 @@ async def _enqueue_batches(
             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, batchable_arguments))
+            batch_chains[(batch_type, routing_key)].append((async_chain, resolved_arguments, batchable_arguments))
         except Exception as e:
             if not isinstance(e, NodeFailurePropagationException):
                 error_infos.append(NodeErrorInfo(e, call_executor.node))
             chains_to_remove.append(async_chain)
             context.set_resolved_value(async_chain, Err(None))
 
     new_batch_tasks: dict[asyncio.Task[Sequence[NodeResult]], Sequence[DependencyChain]] = {}
+    pre_resolved_by_chain: dict[DependencyChain, Any] = {}
 
     for _, chains_and_args in batch_chains.items():
         if len(chains_and_args) < 2:
+            # Reuse native async arguments when a batch does not form.
+            for chain, resolved_arguments, _ in chains_and_args:
+                if isinstance(chain.executor, CallExecutor) and isinstance(chain.executor._udf, AsyncBatchableUDFBase):
+                    pre_resolved_by_chain[chain] = resolved_arguments
             continue
 
-        chains, args = zip(*chains_and_args)
+        chains, _resolved_args, args = zip(*chains_and_args)
         chains_to_remove.extend(chains)
 
         batch_udfs = [chain.executor._udf for chain in chains]
@@ -375,7 +396,7 @@ async def _run_legacy_batch(s, u, n, a, c, e):
             )
         new_batch_tasks[task] = chains
 
-    remaining = [chain for chain in ready_async if chain not in chains_to_remove]
+    remaining = [(chain, pre_resolved_by_chain.get(chain)) for chain in ready_async if chain not in chains_to_remove]
     return remaining, new_batch_tasks
 
 
@@ -449,12 +470,14 @@ async def execute(
                 )
             in_progress_batches.update(new_batch_tasks)
 
-            for async_chain in remaining_ready_async:
+            for async_chain, pre_resolved_arguments in remaining_ready_async:
                 # Native async UDF → await on event loop
                 if isinstance(async_chain.executor, CallExecutor) and isinstance(
                     async_chain.executor._udf, (AsyncUDFBase, AsyncBatchableUDFBase)
                 ):
-                    task = asyncio.create_task(_execute_async_udf(semaphore, async_chain, context, error_infos))
+                    task = asyncio.create_task(
+                        _execute_async_udf(semaphore, async_chain, context, error_infos, pre_resolved_arguments)
+                    )
                 else:
                     # Legacy sync UDF with execute_async=True → thread pool fallback
                     task = asyncio.create_task(

osprey_async_worker/src/osprey/async_worker/lib/external_service.py

@@ -1,7 +1,7 @@
 """Async external service utilities for the async worker.
 
 Port of osprey.engine.executor.external_service_utils with asyncio instead of gevent.
-Uses asyncio.Future instead of gevent.event.AsyncResult for cache entries.
+Uses asyncio.Task instead of gevent.event.AsyncResult for single-value cache entries.
 """
 
 import asyncio
@@ -59,8 +59,9 @@ class ExternalServiceAccessor(Generic[KeyT, ValueT]):
 
     def __init__(self, service: AsyncExternalService[KeyT, ValueT]):
         self._service = service
-        # Key -> Tuple[ Future[ValueT], Expiration datetime ]
-        self._cache: dict[KeyT, tuple[asyncio.Future[ValueT], datetime | None]] = {}
+        # Key -> (Future[ValueT], expiration datetime, count-error-once eligibility)
+        self._cache: dict[KeyT, tuple[asyncio.Future[ValueT], datetime | None, bool]] = {}
+        self._active_batch_loaders: set[asyncio.Task[None]] = set()
 
     def _is_past_cache_expiration(self, cache_expiration: datetime | None) -> bool:
         """
@@ -77,77 +78,100 @@ def _get_cache_expiration_datetime(self) -> datetime | None:
         ttl = self._service.cache_ttl()
         return datetime.now() + ttl if ttl is not None else None
 
+    def _make_task(self, key: KeyT) -> asyncio.Task[ValueT]:
+        task = asyncio.create_task(self._service.get_from_service(key))
+        task.add_done_callback(self._consume_future_exception)
+        return task
+
+    @staticmethod
+    def _consume_future_exception(future: asyncio.Future[ValueT]) -> None:
+        if not future.cancelled():
+            future.exception()
+
     def _make_future(self) -> asyncio.Future[ValueT]:
-        """Create a new Future on the running event loop."""
-        return asyncio.get_running_loop().create_future()
+        future: asyncio.Future[ValueT] = asyncio.get_running_loop().create_future()
+        future.add_done_callback(self._consume_future_exception)
+        return future
+
+    def _make_batch_loader(self, keys: Sequence[KeyT], futures: Sequence[asyncio.Future[ValueT]]) -> asyncio.Task[None]:
+        loader = asyncio.create_task(self._load_batch(keys, futures))
+        self._active_batch_loaders.add(loader)
+        loader.add_done_callback(self._active_batch_loaders.discard)
+        return loader
+
+    async def _load_batch(self, keys: Sequence[KeyT], futures: Sequence[asyncio.Future[ValueT]]) -> None:
+        try:
+            results = await self._service.batch_get_from_service(keys)
+            if len(results) != len(keys):
+                raise ValueError(f'batch service returned {len(results)} results for {len(keys)} keys')
+            for future, result in zip(futures, results):
+                if result.is_ok():
+                    future.set_result(result.unwrap())
+                else:
+                    future.set_exception(cast(BaseException, result.value))
+        except asyncio.CancelledError:
+            for key, future in zip(keys, futures):
+                cache_entry = self._cache.get(key)
+                if cache_entry is not None and cache_entry[0] is future:
+                    del self._cache[key]
+                if not future.done():
+                    future.cancel()
+            raise
+        except Exception as error:
+            for future in futures:
+                if not future.done():
+                    future.set_exception(error)
 
     async def get_without_cache(self, key: KeyT) -> ValueT:
         """
         Ignores any cached values and performs a read-through `get` to the external service.
         The new value is then used to update the cache entry for subsequent `get` calls.
         """
-        future: asyncio.Future[ValueT] = self._make_future()
-        cache_entry: tuple[asyncio.Future[ValueT], datetime | None] = (
-            future,
+        task = self._make_task(key)
+        cache_entry: tuple[asyncio.Future[ValueT], datetime | None, bool] = (
+            task,
             self._get_cache_expiration_datetime(),
+            False,
         )
         self._cache[key] = cache_entry
-        try:
-            result = await self._service.get_from_service(key)
-            future.set_result(result)
-        except Exception as e:
-            future.set_exception(e)
-
-        return await future
+        return await asyncio.shield(task)
 
     async def get(self, key: KeyT) -> ValueT:
         cache_entry = self._cache.get(key)
-        if cache_entry is not None and not self._is_past_cache_expiration(cache_entry[1]):
-            # Cache hit — await the existing future (may still be in-flight from another caller)
-            return await cache_entry[0]
-
-        future: asyncio.Future[ValueT] = self._make_future()
-        cache_entry = (future, self._get_cache_expiration_datetime())
-        self._cache[key] = cache_entry
+        is_creator = cache_entry is None or self._is_past_cache_expiration(cache_entry[1])
+        if is_creator:
+            task = self._make_task(key)
+            cache_entry = (task, self._get_cache_expiration_datetime(), True)
+            self._cache[key] = cache_entry
+        assert cache_entry is not None
         try:
-            result = await self._service.get_from_service(key)
-            future.set_result(result)
-        except Exception as e:
-            if self._service.count_error_once():
-                future.set_result(cast(ValueT, None))
-            else:
-                future.set_exception(e)
-                future.exception()
+            return await asyncio.shield(cache_entry[0])
+        except Exception:
+            if self._service.count_error_once() and cache_entry[2] and not is_creator:
+                return cast(ValueT, None)
             raise
 
-        return await future
-
     async def batch_get(self, keys: Sequence[KeyT]) -> Sequence[Result[ValueT, Exception]]:
-        cached_entries = [self._cache.get(key) for key in keys]
-        non_cached_keys = [
-            key
-            for key, cache_entry in zip(keys, cached_entries)
-            if cache_entry is None or self._is_past_cache_expiration(cache_entry[1])
-        ]
+        non_cached_keys = []
+        for key in dict.fromkeys(keys):
+            cache_entry = self._cache.get(key)
+            if cache_entry is None or self._is_past_cache_expiration(cache_entry[1]):
+                non_cached_keys.append(key)
         if non_cached_keys:
+            futures = []
             for key in non_cached_keys:
-                self._cache[key] = (self._make_future(), self._get_cache_expiration_datetime())
-            try:
-                result = await self._service.batch_get_from_service(non_cached_keys)
-                for i, key in enumerate(non_cached_keys):
-                    if result[i].is_ok():
-                        self._cache[key][0].set_result(result[i].unwrap())
-                    else:
-                        self._cache[key][0].set_exception(cast(BaseException, result[i].value))
-            except Exception as e:
-                for key in non_cached_keys:
-                    self._cache[key][0].set_exception(e)
+                future = self._make_future()
+                self._cache[key] = (future, self._get_cache_expiration_datetime(), False)
+                futures.append(future)
+        futures_by_key = {key: self._cache[key][0] for key in keys}
+        if non_cached_keys:
+            loader = self._make_batch_loader(non_cached_keys, futures)
+            await asyncio.shield(loader)
 
         results: list[Result[ValueT, Exception]] = []
         for key in keys:
-            future = self._cache[key][0]
             try:
-                value = await future
+                value = await asyncio.shield(futures_by_key[key])
                 results.append(Ok(value))
             except Exception as e:
                 results.append(Err(e))