feat: optimize async execution

dynamiq/cache/utils.py

@@ -1,3 +1,4 @@
+import asyncio
 from functools import wraps
 from typing import Any, Callable
 
@@ -82,3 +83,65 @@ def wrapper(*args: Any, **kwargs: Any) -> tuple[Any, bool]:
         return wrapper
 
     return _cache
+
+
+def cache_wf_entity_async(
+    entity_id: str,
+    cache_enabled: bool = False,
+    cache_manager_cls: type[WorkflowCacheManager] = WorkflowCacheManager,
+    cache_config: CacheConfig | None = None,
+    func_kwargs_to_remove: tuple[str] = FUNC_KWARGS_TO_REMOVE,
+) -> Callable:
+    """Async decorator to cache workflow entity outputs.
+
+    Like cache_wf_entity but wraps an async function. Cache I/O (Redis get/set)
+    is offloaded to threads via asyncio.to_thread to avoid blocking the event loop.
+
+    Args:
+        entity_id (str): Identifier for the entity.
+        cache_enabled (bool): Flag to enable caching.
+        cache_manager_cls (type[WorkflowCacheManager]): Cache manager class.
+        cache_config (CacheConfig | None): Cache configuration.
+        func_kwargs_to_remove (tuple[str]): List of params to remove from callable function kwargs.
+
+    Returns:
+        Callable: Wrapped async function with caching.
+    """
+    def _cache(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> tuple[Any, bool]:
+            cache_manager = None
+            from_cache = False
+            input_data = kwargs.pop("input_data", args[0] if args else {})
+            input_data = dict(input_data) if isinstance(input_data, BaseModel) else input_data
+
+            cleaned_kwargs = {k: v for k, v in kwargs.items() if k not in func_kwargs_to_remove}
+            if cache_enabled and cache_config:
+                logger.debug(f"Entity_id {entity_id}: async cache used")
+                cache_manager = cache_manager_cls(config=cache_config)
+                output = await asyncio.to_thread(
+                    cache_manager.get_entity_output,
+                    entity_id=entity_id,
+                    input_data=input_data,
+                    **cleaned_kwargs,
+                )
+                if output is not None:
+                    from_cache = True
+                    return output, from_cache
+
+            output = await func(*args, **kwargs)
+
+            if cache_manager:
+                await asyncio.to_thread(
+                    cache_manager.set_entity_output,
+                    entity_id=entity_id,
+                    input_data=input_data,
+                    output_data=output,
+                    **cleaned_kwargs,
+                )
+
+            return output, from_cache
+
+        return wrapper
+
+    return _cache

dynamiq/flows/flow.py

@@ -11,6 +11,7 @@
 
 from dynamiq.connections.managers import ConnectionManager
 from dynamiq.executors.base import BaseExecutor
+from dynamiq.executors.context import ContextAwareThreadPoolExecutor
 from dynamiq.executors.pool import ThreadExecutor
 from dynamiq.flows.base import BaseFlow
 from dynamiq.nodes.node import Node, NodeReadyToRun
@@ -360,13 +361,8 @@ async def run_async(self, input_data: Any, config: RunnableConfig = None, **kwar
         """
         Run the flow asynchronously with the given input data and configuration.
 
-        Args:
-            input_data (Any): Input data for the flow.
-            config (RunnableConfig, optional): Configuration for the run. Defaults to None.
-            **kwargs: Additional keyword arguments.
-
-        Returns:
-            RunnableResult: Result of the flow execution.
+        Creates a dedicated ContextAwareThreadPoolExecutor for this flow run,
+        isolating it from other concurrent flow executions.
         """
         self.reset_run_state()
         run_id = uuid4()
@@ -375,6 +371,9 @@ async def run_async(self, input_data: Any, config: RunnableConfig = None, **kwar
             "parent_run_id": kwargs.get("parent_run_id", run_id),
         }
 
+        max_workers = (config.max_node_workers if config else None) or self.max_node_workers
+        executor = ContextAwareThreadPoolExecutor(max_workers=max_workers)
+
         logger.info(f"Flow {self.id}: execution started.")
         self.run_on_flow_start(input_data, config, **merged_kwargs)
         time_start = datetime.now()
@@ -391,6 +390,7 @@ async def run_async(self, input_data: Any, config: RunnableConfig = None, **kwar
                                 input_data=node.input_data,
                                 depends_result=node.depends_result,
                                 config=config,
+                                executor=executor,
                                 **(merged_kwargs | {"parent_run_id": run_id}),
                             )
                             for node in nodes_to_run
@@ -435,6 +435,8 @@ async def run_async(self, input_data: Any, config: RunnableConfig = None, **kwar
                 error=RunnableResultError.from_exception(e, failed_nodes=failed_nodes),
             )
         finally:
+            # wait=False is safe: all node tasks have been awaited via asyncio.gather()
+            executor.shutdown(wait=False)
             try:
                 await self._cleanup_dry_run_async(config)
             except Exception as e: