diff --git a/src/bespokelabs/curator/request_processor/event_loop.py b/src/bespokelabs/curator/request_processor/event_loop.py index e2ee40c6..c0541e0f 100644 --- a/src/bespokelabs/curator/request_processor/event_loop.py +++ b/src/bespokelabs/curator/request_processor/event_loop.py @@ -27,18 +27,17 @@ def run_in_event_loop(coroutine): try: # This call will raise an RuntimError if there is no event loop running. asyncio.get_running_loop() - + except RuntimeError: + # If no event loop is running, asyncio will + # return a RuntimeError (https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.get_running_loop). + # In that case, we can just use asyncio.run. + pass + else: # If there is an event loop running (the call # above doesn't raise an exception), we can # use nest_asyncio to patch the event loop. nest_asyncio.apply() - return asyncio.run(coroutine) - except RuntimeError: - # Explicitly pass, since we want to fallback to asyncio.run - pass - - # If no event loop is running, asyncio will - # return a RuntimeError (https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.get_running_loop). - # In that case, we can just use asyncio.run. + # Run the coroutine once, outside the except block, so its own errors (including + # RuntimeErrors) propagate unchanged and aren't chained to the error above. return asyncio.run(coroutine) diff --git a/tests/unittests/test_event_loop.py b/tests/unittests/test_event_loop.py new file mode 100644 index 00000000..46a4e121 --- /dev/null +++ b/tests/unittests/test_event_loop.py @@ -0,0 +1,45 @@ +import asyncio + +import pytest + +from bespokelabs.curator.request_processor.event_loop import run_in_event_loop + + +async def _fail(calls): + calls.append(1) + raise RuntimeError("original error") + + +async def _answer(): + return 42 + + +def test_coroutine_error_is_raised_without_a_running_loop(): + calls = [] + + with pytest.raises(RuntimeError, match="original error") as exc_info: + run_in_event_loop(_fail(calls)) + + assert calls == [1] + # Not chained to the "no running event loop" RuntimeError used to detect the loop. + assert exc_info.value.__context__ is None + + +def test_coroutine_error_is_raised_inside_a_running_loop(): + """In Jupyter or Colab an event loop is already running when curator is called.""" + calls = [] + + async def notebook_cell(): + with pytest.raises(RuntimeError, match="original error"): + run_in_event_loop(_fail(calls)) + + asyncio.run(notebook_cell()) + + assert calls == [1] + + +def test_result_is_returned_inside_a_running_loop(): + async def notebook_cell(): + return run_in_event_loop(_answer()) + + assert asyncio.run(notebook_cell()) == 42