move async handling to parallelizer

Author: chenmoneygithub

dspy/evaluate/evaluate.py

@@ -6,7 +6,6 @@
 if TYPE_CHECKING:
     import pandas as pd
 
-import asyncio
 
 import tqdm
 
@@ -161,7 +160,7 @@ async def acall(
         if callback_metadata:
             logger.debug(f"Evaluate.acall is called with callback metadata: {callback_metadata}")
         tqdm.tqdm._instances.clear()
-        results = await self._execute_with_event_loop(program, metric, devset, num_threads)
+        results = await self._execute_with_event_loop(program, metric, devset, num_threads, display_progress)
         return self._process_evaluate_result(devset, results, metric, display_table)
 
     def _resolve_call_args(self, metric, devset, num_threads, display_progress, display_table):
@@ -220,41 +219,22 @@ async def _execute_with_event_loop(
         metric: Callable,
         devset: list["dspy.Example"],
         num_threads: int,
+        disable_progress_bar: bool,
     ):
-        queue = asyncio.Queue()
-        results = [None] * len(devset)
-        for i, example in enumerate(devset):
-            await queue.put((i, example))
-
-        for _ in range(num_threads):
-            # Add a sentinel value to indicate that the worker should exit
-            await queue.put((-1, None))
-
-        # Create tqdm progress bar
-        pbar = tqdm.tqdm(total=len(devset), dynamic_ncols=True)
-
-        async def worker():
-            while True:
-                index, example = await queue.get()
-                if index == -1:
-                    break
-                prediction = await program.acall(**example.inputs())
-                score = metric(example, prediction)
-                results[index] = (prediction, score)
-
-                vals = [r[-1] for r in results if r is not None]
-                nresults = sum(vals)
-                ntotal = len(vals)
-                pct = round(100 * nresults / ntotal, 1) if ntotal else 0
-                pbar.set_description(f"Average Metric: {nresults:.2f} / {ntotal} ({pct}%)")
-                pbar.update(1)
-                queue.task_done()
-
-        workers = [asyncio.create_task(worker()) for _ in range(num_threads)]
-        await asyncio.gather(*workers)
-        pbar.close()
-
-        return results
+        executor = ParallelExecutor(
+            num_threads=num_threads,
+            disable_progress_bar=disable_progress_bar,
+            max_errors=(self.max_errors or dspy.settings.max_errors),
+            provide_traceback=self.provide_traceback,
+            compare_results=True,
+        )
+
+        async def process_item(example):
+            prediction = await program.acall(**example.inputs())
+            score = metric(example, prediction)
+            return prediction, score
+
+        return await executor.aexecute(process_item, devset)
 
     def _construct_result_table(
         self, results: list[tuple["dspy.Example", "dspy.Example", Any]], metric_name: str

dspy/utils/parallelizer.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 import copy
 import logging
@@ -42,29 +43,60 @@ def __init__(
         self.error_lock = threading.Lock()
         self.cancel_jobs = threading.Event()
 
+        self.error_lock_async = asyncio.Lock()
+        self.cancel_jobs_async = asyncio.Event()
+
     def execute(self, function, data):
         tqdm.tqdm._instances.clear()
-        wrapped = self._wrap_function(function)
+        wrapped = self._wrap_function(function, async_mode=False)
         return self._execute_parallel(wrapped, data)
 
-    def _wrap_function(self, user_function):
-        def safe_func(item):
+    async def aexecute(self, function, data):
+        tqdm.tqdm._instances.clear()
+        wrapped = self._wrap_function(function, async_mode=True)
+        return await self._execute_parallel_async(wrapped, data)
+
+    def _handle_error(self, item, e):
+        with self.error_lock:
+            self.error_count += 1
+            if self.error_count >= self.max_errors:
+                self.cancel_jobs.set()
+        if self.provide_traceback:
+            logger.error(f"Error for {item}: {e}\n{traceback.format_exc()}")
+        else:
+            logger.error(f"Error for {item}: {e}. Set `provide_traceback=True` for traceback.")
+
+    async def _handle_error_async(self, item, e):
+        async with self.error_lock_async:
+            self.error_count += 1
+            if self.error_count >= self.max_errors:
+                self.cancel_jobs_async.set()
+        if self.provide_traceback:
+            logger.error(f"Error for {item}: {e}\n{traceback.format_exc()}")
+
+    def _wrap_function(self, user_function, async_mode=False):
+        async def _async_safe_func(item):
+            if self.cancel_jobs.is_set():
+                return None
+            try:
+                return await user_function(item)
+            except Exception as e:
+                await self._handle_error_async(item, e)
+                return None
+
+        def _sync_safe_func(item):
             if self.cancel_jobs.is_set():
                 return None
             try:
                 return user_function(item)
             except Exception as e:
-                with self.error_lock:
-                    self.error_count += 1
-                    if self.error_count >= self.max_errors:
-                        self.cancel_jobs.set()
-                if self.provide_traceback:
-                    logger.error(f"Error for {item}: {e}\n{traceback.format_exc()}")
-                else:
-                    logger.error(f"Error for {item}: {e}. Set `provide_traceback=True` for traceback.")
+                self._handle_error(item, e)
                 return None
 
-        return safe_func
+        if async_mode:
+            return _async_safe_func
+        else:
+            return _sync_safe_func
 
     def _execute_parallel(self, function, data):
         results = [None] * len(data)
@@ -204,6 +236,50 @@ def all_done():
 
         return results
 
+    async def _execute_parallel_async(self, function, data):
+        queue = asyncio.Queue()
+        results = [None] * len(data)
+        for i, example in enumerate(data):
+            await queue.put((i, example))
+
+        for _ in range(self.num_threads):
+            # Add a sentinel value to indicate that the worker should exit
+            await queue.put((-1, None))
+
+        # Create tqdm progress bar
+        pbar = tqdm.tqdm(total=len(data), dynamic_ncols=True)
+
+        async def worker():
+            while True:
+                if self.cancel_jobs_async.is_set():
+                    break
+                index, example = await queue.get()
+                if index == -1:
+                    break
+                function_outputs = await function(example)
+                results[index] = function_outputs
+
+                if self.compare_results:
+                    vals = [r[-1] for r in results if r is not None]
+                    self._update_progress(pbar, sum(vals), len(vals))
+                else:
+                    self._update_progress(
+                        pbar,
+                        len([r for r in results if r is not None]),
+                        len(data),
+                    )
+
+                queue.task_done()
+
+        workers = [asyncio.create_task(worker()) for _ in range(self.num_threads)]
+        await asyncio.gather(*workers)
+        pbar.close()
+        if self.cancel_jobs_async.is_set():
+            logger.warning("Execution cancelled due to errors or interruption.")
+            raise Exception("Execution cancelled due to errors or interruption.")
+
+        return results
+
     def _update_progress(self, pbar, nresults, ntotal):
         if self.compare_results:
             pct = round(100 * nresults / ntotal, 1) if ntotal else 0

tests/utils/test_parallelizer.py

@@ -59,3 +59,63 @@ def task(item):
 
     # Verify that the results exclude the failed task
     assert results == [1, 2, None, 4, 5]
+
+
+@pytest.mark.asyncio
+async def test_worker_threads_independence_async():
+    async def task(item):
+        # Each thread maintains its own state by appending to a thread-local list
+        return item * 2
+
+    data = [1, 2, 3, 4, 5]
+    executor = ParallelExecutor(num_threads=3)
+    results = await executor.aexecute(task, data)
+
+    assert results == [2, 4, 6, 8, 10]
+
+
+@pytest.mark.asyncio
+async def test_parallel_execution_speed_async():
+    async def task(item):
+        time.sleep(0.1)  # Simulate a time-consuming task
+        return item
+
+    data = [1, 2, 3, 4, 5]
+    executor = ParallelExecutor(num_threads=5)
+
+    start_time = time.time()
+    await executor.aexecute(task, data)
+    end_time = time.time()
+
+    assert end_time - start_time < len(data)
+
+
+@pytest.mark.asyncio
+async def test_max_errors_handling_async():
+    async def task(item):
+        if item == 3:
+            raise ValueError("Intentional error")
+        return item
+
+    data = [1, 2, 3, 4, 5]
+    executor = ParallelExecutor(num_threads=3, max_errors=1)
+
+    with pytest.raises(Exception, match="Execution cancelled due to errors or interruption."):
+        await executor.aexecute(task, data)
+
+
+@pytest.mark.asyncio
+async def test_max_errors_not_met_async():
+    async def task(item):
+        if item == 3:
+            raise ValueError("Intentional error")
+        return item
+
+    data = [1, 2, 3, 4, 5]
+    executor = ParallelExecutor(num_threads=3, max_errors=2)
+
+    # Ensure that the execution completes without crashing when max_errors is not met
+    results = await executor.aexecute(task, data)
+
+    # Verify that the results exclude the failed task
+    assert results == [1, 2, None, 4, 5]