From 59f9976a99ea0f1c65c123fd69239c4cb16ab85c Mon Sep 17 00:00:00 2001 From: Waleed Khalid Date: Fri, 2 Oct 2026 03:16:38 +0500 Subject: [PATCH] fix: Run LiteLLM/WML async inference without relying on an implicit event loop asyncio.get_event_loop() raises when the thread has no current event loop: in worker threads, after asyncio.run() returned, and on Python 3.14+. Reuse or create the thread's loop instead, raise an actionable UnitxtError when called inside a running loop (#1647), and recreate the LiteLLM semaphore/rate limiter when inference runs in a different loop. Signed-off-by: Waleed Khalid --- src/unitxt/inference.py | 62 +++++++++++++++++++++++---- tests/library/test_inference_utils.py | 59 +++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 8 deletions(-) create mode 100644 tests/library/test_inference_utils.py diff --git a/src/unitxt/inference.py b/src/unitxt/inference.py index cd2f19ed9f..1dc133ac3a 100644 --- a/src/unitxt/inference.py +++ b/src/unitxt/inference.py @@ -64,6 +64,44 @@ def batched(lst, n): yield batch +def get_or_create_event_loop() -> asyncio.AbstractEventLoop: + """Return the event loop of the current thread, creating and setting one if needed. + + ``asyncio.get_event_loop()`` raises when the thread has no current event loop, which + happens in worker threads, after ``asyncio.run()`` has returned, and on Python 3.14+ + whenever no loop was set explicitly. + """ + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = None + if loop is None or loop.is_closed(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + return loop + + +def run_coroutine_synchronously(coroutine): + """Run a coroutine to completion from synchronous code and return its result. + + The current thread's event loop is reused across calls, so asyncio primitives + (e.g. semaphores) created by an inference engine stay bound to the same loop. + """ + loop = get_or_create_event_loop() + # nest_asyncio marks the loops it patches to allow re-entrant run_until_complete + if loop.is_running() and not getattr(loop, "_nest_patched", False): + coroutine.close() + raise UnitxtError( + "Synchronous inference cannot run inside an already running asyncio event loop " + "(e.g. in a Jupyter notebook or in an async framework such as FastAPI). " + "Either run the inference in a separate thread, " + "e.g. 'await asyncio.to_thread(engine.infer, dataset)', " + "or allow nested event loops by installing nest_asyncio ('pip install nest_asyncio') " + "and calling 'import nest_asyncio; nest_asyncio.apply()' before running the inference." + ) + return loop.run_until_complete(coroutine) + + class StandardAPIParamsMixin(Artifact): model: str frequency_penalty: Optional[float] = None @@ -2739,11 +2777,10 @@ async def handle_async_requests(start_idx, end_idx): batch_results = await asyncio.gather(*coroutines) return list(batch_results) - loop = asyncio.get_event_loop() results = [] for batch_idx in range(0, len(data), self.concurrency_limit): - batch_results = loop.run_until_complete( + batch_results = run_coroutine_synchronously( handle_async_requests( batch_idx, min(batch_idx + self.concurrency_limit, len(data)) ) @@ -3164,17 +3201,24 @@ def get_engine_id(self): def prepare_engine(self): if self.credentials is None: self.credentials = {} + self.inference_type = "litellm" + from litellm import acompletion + + self._completion = acompletion + self._prepare_async_primitives(loop_id=None) + + def _prepare_async_primitives(self, loop_id: Optional[int]): + # asyncio primitives get bound to the event loop they are first awaited in, so + # they are recreated whenever inference runs in a different loop. Comparing loop + # ids is safe: a primitive bound to a loop keeps that loop (and its id) alive. # Initialize the token bucket rate limiter self._rate_limiter = AsyncTokenBucket( rate=self.max_requests_per_second, capacity=self.max_requests_per_second, ) - self.inference_type = "litellm" - from litellm import acompletion - - self._completion = acompletion # Initialize a semaphore to limit concurrency self._semaphore = asyncio.Semaphore(round(self.max_requests_per_second)) + self._async_primitives_loop_id = loop_id async def _infer_instance( self, index: int, instance: Dict[str, Any] @@ -3228,6 +3272,9 @@ async def _infer_async( self, dataset: List[Dict[str, Any]] ) -> List[TextGenerationInferenceOutput]: """Process multiple inference requests concurrently with a progress bar.""" + loop_id = id(asyncio.get_running_loop()) + if loop_id != self._async_primitives_loop_id: + self._prepare_async_primitives(loop_id) tasks = ( self._infer_instance(i, instance) for i, instance in enumerate(dataset) ) @@ -3242,8 +3289,7 @@ def _infer( return_meta_data: bool = False, ) -> Union[List[str], List[TextGenerationInferenceOutput]]: """Main inference entry point.""" - loop = asyncio.get_event_loop() - responses = loop.run_until_complete(self._infer_async(dataset)) + responses = run_coroutine_synchronously(self._infer_async(dataset)) return self.get_return_object(responses, return_meta_data) def get_return_object(self, responses, return_meta_data): diff --git a/tests/library/test_inference_utils.py b/tests/library/test_inference_utils.py new file mode 100644 index 0000000000..e27c820e7e --- /dev/null +++ b/tests/library/test_inference_utils.py @@ -0,0 +1,59 @@ +import asyncio +import threading + +from unitxt.error_utils import UnitxtError +from unitxt.inference import run_coroutine_synchronously + +from tests.utils import UnitxtTestCase + + +async def add_one(x): + await asyncio.sleep(0) + return x + 1 + + +class TestRunCoroutineSynchronously(UnitxtTestCase): + def test_returns_coroutine_result(self): + self.assertEqual(run_coroutine_synchronously(add_one(1)), 2) + + def test_works_after_asyncio_run(self): + # asyncio.run() leaves the main thread without a current event loop + asyncio.run(add_one(0)) + self.assertEqual(run_coroutine_synchronously(add_one(1)), 2) + + def test_works_in_worker_thread(self): + results = [] + thread = threading.Thread( + target=lambda: results.append(run_coroutine_synchronously(add_one(1))) + ) + thread.start() + thread.join() + self.assertEqual(results, [2]) + + def test_reuses_event_loop_across_calls(self): + # Like the semaphore of LiteLLMInferenceEngine, which is created once and + # binds to the event loop it is first contended on. + semaphore = asyncio.Semaphore(1) + + async def guarded(x): + async with semaphore: + return await add_one(x) + + async def contend(): + return await asyncio.gather(guarded(1), guarded(2)) + + self.assertEqual(run_coroutine_synchronously(contend()), [2, 3]) + self.assertEqual(run_coroutine_synchronously(contend()), [2, 3]) + + def test_raises_informative_error_inside_running_loop(self): + coroutine = add_one(1) + + async def main(): + run_coroutine_synchronously(coroutine) + + with self.assertRaises(UnitxtError) as cm: + asyncio.run(main()) + self.assertIn("nest_asyncio", str(cm.exception)) + self.assertIn("asyncio.to_thread", str(cm.exception)) + # the coroutine is closed rather than left un-awaited + self.assertIsNone(coroutine.cr_frame)