feature: improve async / executor functionality (#2070)

- refactor: move is_event_loop_running and as_completed to async_utils
This helps separate async features from executor and engine modules, and
will prevent circular imports in these modules.
- feat: add ProgressBarManager for enhanced progress tracking
Introduces a new ProgressBarManager class to manage progress bars for
both batch and non-batch execution. This includes methods for creating
single and nested progress bars, as well as updating batch progress
bars, improving user experience during long-running tasks.
- feat: refactor async utilities for improved functionality
- Refactored `run` function from repeated async, executor, and engine
logic.
- Changed `as_completed` from async to a regular function for better
compatibility.
- Added `process_futures` to handle futures with optional progress
tracking.
- feat: enhance async task management with ProgressBarManager
- Refactored Executor to utilize new async_utils and ProgressBarManager
for improved composition
- Refactored run_async_tasks to utilize ProgressBarManager for improved
progress tracking.
  - Modified tests to validate the new behavior functions
- feat: enhance Executor with job indexing and exception handling
- Added `_jobs_processed` attribute to maintain consistent job indexing
across multiple runs.
  - Implemented `clear_jobs` method to reset job indices.
- Updated tests to verify exception handling and job indexing after
clearing jobs.
- Adjusted Jupyter notebook tests for consistency in results retrieval.
- feat: update engine's apply_transformations to recursively process
nested Transformation or Parallel procedures
- The apply_transforms() function has been updated to handle different
types of transformations recursively:
If transforms is a list, it recursively applies each transform in the
list.
If transforms is a Parallel instance, it recursively applies the
transformations contained within the Parallel object.
If transforms is a BaseGraphTransformation, it generates an execution
plan (a list of coroutines), gets a description, and then runs the
coroutines asynchronously using run_async_tasks().
If transforms is none of the above, it raises a ValueError indicating an
invalid type.
  - Move apply_nest_asyncio to async_utils for better organization
- Updated the Parallel class to support transformations with improved
type hints.
  - Added unit tests
- chore: use Sequence over List for type checking improvements; add
debug logging

---------

Co-authored-by: Ani <[email protected]>

Author: ahgraber

docs/references/executor.md

@@ -3,4 +3,3 @@
         members:
             - Executor
             - run_async_batch
-            - is_event_loop_running

src/ragas/async_utils.py

@@ -1,89 +1,171 @@
 """Async utils."""
 
 import asyncio
-from typing import Any, Coroutine, List, Optional
+import logging
+import typing as t
 
-from tqdm.auto import tqdm
+logger = logging.getLogger(__name__)
 
-from ragas.executor import is_event_loop_running
-from ragas.utils import batched
+
+def is_event_loop_running() -> bool:
+    """
+    Check if an event loop is currently running.
+    """
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return False
+    else:
+        return loop.is_running()
+
+
+def apply_nest_asyncio():
+    """Apply nest_asyncio if an event loop is running."""
+    if is_event_loop_running():
+        # an event loop is running so call nested_asyncio to fix this
+        try:
+            import nest_asyncio
+        except ImportError:
+            raise ImportError(
+                "It seems like your running this in a jupyter-like environment. Please install nest_asyncio with `pip install nest_asyncio` to make it work."
+            )
+
+        nest_asyncio.apply()
+
+
+def as_completed(
+    coroutines: t.Sequence[t.Coroutine],
+    max_workers: int = -1,
+    *,
+    cancel_check: t.Optional[t.Callable[[], bool]] = None,
+    cancel_pending: bool = True,
+) -> t.Iterator[asyncio.Future]:
+    """
+    Wrap coroutines with a semaphore if max_workers is specified.
+
+    Returns an iterator of futures that completes as tasks finish.
+    """
+    if max_workers == -1:
+        tasks = [asyncio.create_task(coro) for coro in coroutines]
+    else:
+        semaphore = asyncio.Semaphore(max_workers)
+
+        async def sema_coro(coro):
+            async with semaphore:
+                return await coro
+
+        tasks = [asyncio.create_task(sema_coro(coro)) for coro in coroutines]
+
+    ac_iter = asyncio.as_completed(tasks)
+
+    if cancel_check is None:
+        return ac_iter
+
+    def _iter_with_cancel():
+        for future in ac_iter:
+            if cancel_check():
+                if cancel_pending:
+                    for t in tasks:
+                        if not t.done():
+                            t.cancel()
+                break
+            yield future
+
+    return _iter_with_cancel()
+
+
+async def process_futures(
+    futures: t.Iterator[asyncio.Future],
+) -> t.AsyncGenerator[t.Any, None]:
+    """
+    Process futures with optional progress tracking.
+
+    Args:
+        futures: Iterator of asyncio futures to process (e.g., from asyncio.as_completed)
+
+    Yields:
+        Results from completed futures as they finish
+    """
+    # Process completed futures as they finish
+    for future in futures:
+        try:
+            result = await future
+        except asyncio.CancelledError:
+            raise  # Re-raise CancelledError to ensure proper cancellation
+        except Exception as e:
+            result = e
+
+        yield result
+
+
+def run(
+    async_func: t.Union[
+        t.Callable[[], t.Coroutine[t.Any, t.Any, t.Any]],
+        t.Coroutine[t.Any, t.Any, t.Any],
+    ],
+) -> t.Any:
+    """
+    Run an async function in the current event loop or a new one if not running.
+    """
+    # Ensure nest_asyncio is applied if we're in a running loop
+    # This is common in environments like Jupyter notebooks
+    apply_nest_asyncio()
+
+    # Create the coroutine if it's a callable, otherwise use directly
+    coro = async_func() if callable(async_func) else async_func
+    return asyncio.run(coro)
 
 
 def run_async_tasks(
-    tasks: List[Coroutine],
-    batch_size: Optional[int] = None,
+    tasks: t.Sequence[t.Coroutine],
+    batch_size: t.Optional[int] = None,
     show_progress: bool = True,
     progress_bar_desc: str = "Running async tasks",
-) -> List[Any]:
+    max_workers: int = -1,
+    *,
+    cancel_check: t.Optional[t.Callable[[], bool]] = None,
+) -> t.List[t.Any]:
     """
     Execute async tasks with optional batching and progress tracking.
 
     NOTE: Order of results is not guaranteed!
 
     Args:
-        tasks: List of coroutines to execute
+        tasks: Sequence of coroutines to execute
         batch_size: Optional size for batching tasks. If None, runs all concurrently
         show_progress: Whether to display progress bars
+        max_workers: Maximum number of concurrent tasks (-1 for unlimited)
     """
+    from ragas.utils import ProgressBarManager, batched
 
     async def _run():
         total_tasks = len(tasks)
         results = []
+        pbm = ProgressBarManager(progress_bar_desc, show_progress)
 
-        # If no batching, run all tasks concurrently with single progress bar
         if not batch_size:
-            with tqdm(
-                total=total_tasks,
-                desc=progress_bar_desc,
-                disable=not show_progress,
-            ) as pbar:
-                for future in asyncio.as_completed(tasks):
-                    result = await future
+            with pbm.create_single_bar(total_tasks) as pbar:
+                async for result in process_futures(
+                    as_completed(tasks, max_workers, cancel_check=cancel_check)
+                ):
                     results.append(result)
                     pbar.update(1)
-            return results
-
-        # With batching, show nested progress bars
-        batches = batched(tasks, batch_size)  # generator
-        n_batches = (total_tasks + batch_size - 1) // batch_size
-        with (
-            tqdm(
-                total=total_tasks,
-                desc=progress_bar_desc,
-                disable=not show_progress,
-                position=0,
-                leave=True,
-            ) as overall_pbar,
-            tqdm(
-                total=batch_size,
-                desc=f"Batch 1/{n_batches}",
-                disable=not show_progress,
-                position=1,
-                leave=False,
-            ) as batch_pbar,
-        ):
-            for i, batch in enumerate(batches, 1):
-                batch_pbar.reset(total=len(batch))
-                batch_pbar.set_description(f"Batch {i}/{n_batches}")
-                for future in asyncio.as_completed(batch):
-                    result = await future
-                    results.append(result)
-                    overall_pbar.update(1)
-                    batch_pbar.update(1)
+        else:
+            total_tasks = len(tasks)
+            batches = batched(tasks, batch_size)
+            overall_pbar, batch_pbar, n_batches = pbm.create_nested_bars(
+                total_tasks, batch_size
+            )
+            with overall_pbar, batch_pbar:
+                for i, batch in enumerate(batches, 1):
+                    pbm.update_batch_bar(batch_pbar, i, n_batches, len(batch))
+                    async for result in process_futures(
+                        as_completed(batch, max_workers, cancel_check=cancel_check)
+                    ):
+                        results.append(result)
+                        batch_pbar.update(1)
+                    overall_pbar.update(len(batch))
 
         return results
 
-    if is_event_loop_running():
-        # an event loop is running so call nested_asyncio to fix this
-        try:
-            import nest_asyncio
-        except ImportError:
-            raise ImportError(
-                "It seems like your running this in a jupyter-like environment. "
-                "Please install nest_asyncio with `pip install nest_asyncio` to make it work."
-            )
-        else:
-            nest_asyncio.apply()
-
-    results = asyncio.run(_run())
-    return results
+    return run(_run)