fix(executor): simplify executor by removing threading (explodinggradients#1093)

jjmachan · web-flow · commit f2470c578844 · 2024-07-12T20:31:52.000+05:30
ragas `0.1.10` broke how the executor was functioning and there were a few hard to reproduce bugs overall. This fixes all those. - uses `nest_asyncio` in jupyter notebooks - added tests in jupyter notebook with `nbmake` - removed the Runner. this means all the `Tasks` will be run in the same event loop as main. for jupyter notebooks we will use `nest_async` lib takes inspiration from explodinggradients#689 fixes: explodinggradients#681
diff --git a/Makefile b/Makefile
@@ -27,7 +27,7 @@ clean: ## Clean all generated files
 run-ci: format lint type ## Running all CI checks
 test: ## Run tests
 	@echo "Running tests..."
-	@pytest tests/unit $(shell if [ -n "$(k)" ]; then echo "-k $(k)"; fi)
+	@pytest --nbmake tests/unit $(shell if [ -n "$(k)" ]; then echo "-k $(k)"; fi)
 test-e2e: ## Run end2end tests
 	echo "running end2end tests..."
 	@pytest tests/e2e -s
diff --git a/requirements/test.txt b/requirements/test.txt
@@ -2,3 +2,4 @@ pytest
 pytest-xdist[psutil]
 pytest-asyncio
 llama_index
+nbmake
diff --git a/src/ragas/executor.py b/src/ragas/executor.py
@@ -2,8 +2,6 @@
 
 import asyncio
 import logging
-import sys
-import threading
 import typing as t
 from dataclasses import dataclass, field
 
@@ -16,77 +14,28 @@
 logger = logging.getLogger(__name__)
 
 
-def runner_exception_hook(args: threading.ExceptHookArgs):
-    raise args.exc_type
+def is_event_loop_running() -> bool:
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return False
+    else:
+        return loop.is_running()
 
 
-# set a custom exception hook
-# threading.excepthook = runner_exception_hook
-
-
-def as_completed(loop, coros, max_workers):
-    loop_arg_dict = {"loop": loop} if sys.version_info[:2] < (3, 10) else {}
+def as_completed(coros, max_workers):
     if max_workers == -1:
-        return asyncio.as_completed(coros, **loop_arg_dict)
+        return asyncio.as_completed(coros)
 
-    # loop argument is removed since Python 3.10
-    semaphore = asyncio.Semaphore(max_workers, **loop_arg_dict)
+    semaphore = asyncio.Semaphore(max_workers)
 
     async def sema_coro(coro):
         async with semaphore:
             return await coro
 
     sema_coros = [sema_coro(c) for c in coros]
-    return asyncio.as_completed(sema_coros, **loop_arg_dict)
-
-
-class Runner(threading.Thread):
-    def __init__(
-        self,
-        jobs: t.List[t.Tuple[t.Coroutine, str]],
-        desc: str,
-        keep_progress_bar: bool = True,
-        raise_exceptions: bool = True,
-        run_config: t.Optional[RunConfig] = None,
-    ):
-        super().__init__()
-        self.jobs = jobs
-        self.desc = desc
-        self.keep_progress_bar = keep_progress_bar
-        self.raise_exceptions = raise_exceptions
-        self.run_config = run_config or RunConfig()
-
-        # create task
-        try:
-            self.loop = asyncio.get_event_loop()
-        except RuntimeError:
-            self.loop = asyncio.new_event_loop()
-        self.futures = as_completed(
-            loop=self.loop,
-            coros=[coro for coro, _ in self.jobs],
-            max_workers=self.run_config.max_workers,
-        )
-
-    async def _aresults(self) -> t.List[t.Any]:
-        results = []
-        for future in tqdm(
-            self.futures,
-            desc=self.desc,
-            total=len(self.jobs),
-            # whether you want to keep the progress bar after completion
-            leave=self.keep_progress_bar,
-        ):
-            r = await future
-            results.append(r)
 
-        return results
-
-    def run(self):
-        results = []
-        try:
-            results = self.loop.run_until_complete(self._aresults())
-        finally:
-            self.results = results
+    return asyncio.as_completed(sema_coros)
 
 
 @dataclass
@@ -95,21 +44,22 @@ class Executor:
     keep_progress_bar: bool = True
     jobs: t.List[t.Any] = field(default_factory=list, repr=False)
     raise_exceptions: bool = False
-    run_config: t.Optional[RunConfig] = field(default_factory=RunConfig, repr=False)
+    run_config: t.Optional[RunConfig] = field(default=None, repr=False)
 
     def wrap_callable_with_index(self, callable: t.Callable, counter):
         async def wrapped_callable_async(*args, **kwargs):
             result = np.nan
             try:
                 result = await callable(*args, **kwargs)
             except MaxRetriesExceeded as e:
+                # this only for testset generation v2
                 logger.warning(f"max retries exceeded for {e.evolution}")
             except Exception as e:
                 if self.raise_exceptions:
                     raise e
                 else:
                     logger.error(
-                        "Runner in Executor raised an exception", exc_info=True
+                        "Runner in Executor raised an exception", exc_info=False
                     )
 
             return counter, result
@@ -120,29 +70,40 @@ def submit(
         self, callable: t.Callable, *args, name: t.Optional[str] = None, **kwargs
     ):
         callable_with_index = self.wrap_callable_with_index(callable, len(self.jobs))
-        self.jobs.append((callable_with_index(*args, **kwargs), name))
+        self.jobs.append((callable_with_index, args, kwargs, name))
 
     def results(self) -> t.List[t.Any]:
-        executor_job = Runner(
-            jobs=self.jobs,
-            desc=self.desc,
-            keep_progress_bar=self.keep_progress_bar,
-            raise_exceptions=self.raise_exceptions,
-            run_config=self.run_config,
-        )
-        executor_job.start()
-        try:
-            executor_job.join()
-        finally:
-            ...
-
-        if executor_job.results is None:
-            if self.raise_exceptions:
-                raise RuntimeError(
-                    "Executor failed to complete. Please check logs above for full info."
+        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:
-                logger.error("Executor failed to complete. Please check logs above.")
-                return []
-        sorted_results = sorted(executor_job.results, key=lambda x: x[0])
+
+            nest_asyncio.apply()
+
+        # create a generator for which returns tasks as they finish
+        futures_as_they_finish = as_completed(
+            coros=[afunc(*args, **kwargs) for afunc, args, kwargs, _ in self.jobs],
+            max_workers=(self.run_config or RunConfig()).max_workers,
+        )
+
+        async def _aresults() -> t.List[t.Any]:
+            results = []
+            for future in tqdm(
+                futures_as_they_finish,
+                desc=self.desc,
+                total=len(self.jobs),
+                # whether you want to keep the progress bar after completion
+                leave=self.keep_progress_bar,
+            ):
+                r = await future
+                results.append(r)
+
+            return results
+
+        results = asyncio.run(_aresults())
+        sorted_results = sorted(results, key=lambda x: x[0])
         return [r[1] for r in sorted_results]
diff --git a/tests/unit/test_executor.py b/tests/unit/test_executor.py
@@ -1,7 +1,12 @@
-def test_order_of_execution():
+import asyncio
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_order_of_execution():
     from ragas.executor import Executor
 
-    async def echo_order(index):
+    async def echo_order(index: int):
         return index
 
     # Arrange
@@ -14,3 +19,72 @@ async def echo_order(index):
     results = executor.results()
     # Assert
     assert results == list(range(1, 11))
+
+
+@pytest.mark.asyncio
+async def test_executor_in_script():
+    from ragas.executor import Executor
+
+    async def echo_order(index: int):
+        await asyncio.sleep(0.1)
+        return index
+
+    # Arrange
+    executor = Executor()
+
+    # Act
+    # add 10 jobs to the executor
+    for i in range(1, 4):
+        executor.submit(echo_order, i, name=f"echo_order_{i}")
+    results = executor.results()
+    # Assert
+    assert results == list(range(1, 4))
+
+
+@pytest.mark.asyncio
+async def test_executor_with_running_loop():
+    import asyncio
+    from ragas.executor import Executor
+
+    loop = asyncio.new_event_loop()
+    loop.run_until_complete(asyncio.sleep(0.1))
+
+    async def echo_order(index: int):
+        await asyncio.sleep(0.1)
+        return index
+
+    # Arrange
+    executor = Executor()
+
+    # Act
+    # add 10 jobs to the executor
+    for i in range(1, 4):
+        executor.submit(echo_order, i, name=f"echo_order_{i}")
+    results = executor.results()
+    # Assert
+    assert results == list(range(1, 4))
+
+
+def test_is_event_loop_running_in_script():
+    from ragas.executor import is_event_loop_running
+
+    assert is_event_loop_running() is False
+
+
+def test_as_completed_in_script():
+    from ragas.executor import as_completed
+
+    async def echo_order(index: int):
+        await asyncio.sleep(index)
+        return index
+
+    async def _run():
+        results = []
+        for t in as_completed([echo_order(1), echo_order(2), echo_order(3)], 3):
+            r = await t
+            results.append(r)
+        return results
+
+    results = asyncio.run(_run())
+
+    assert results == [1, 2, 3]
diff --git a/tests/unit/test_executor_in_jupyter.ipynb b/tests/unit/test_executor_in_jupyter.ipynb
@@ -0,0 +1,80 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e89bde88d0054ede891c3659b439a02f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Evaluating:   0%|          | 0/10 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# from ragas.executor import Executor\n",
+    "# from asyncio import sleep\n",
+    "\n",
+    "# exec = Executor(raise_exceptions=True)\n",
+    "# for i in range(10):\n",
+    "#     exec.submit(sleep, i)\n",
+    "    \n",
+    "# try:\n",
+    "#     exec.results()\n",
+    "# except Exception:\n",
+    "#     print(\"error\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ragas.executor import is_event_loop_running\n",
+    "\n",
+    "assert is_event_loop_running() is True"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ragas",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
