Make the remote inference engine runnable in jupyter notebooks.

Author: taenin

src/oumi/core/async_utils.py

@@ -0,0 +1,22 @@
+import asyncio
+from multiprocessing.pool import ThreadPool
+from typing import Any, Awaitable
+
+
+def safe_asyncio_run(main: Awaitable[Any]) -> Any:
+    """Run a series of Awaitables in a new thread. Blocks until the thread is finished.
+
+    This circumvents the issue of running async functions in the main thread when
+    an event loop is already running (Jupyter notebooks, for example).
+
+    Prefer using `safe_asyncio_run` over `asyncio.run` to allow upstream callers to
+    ignore our dependency on asyncio.
+
+    Args:
+        main: The awaitable to resolve.
+
+    Returns:
+        The result of the awaitable.
+    """
+    pool = ThreadPool(processes=1)
+    return pool.apply(asyncio.run, (main,))

src/oumi/inference/remote_inference_engine.py

@@ -5,6 +5,7 @@
 
 import aiohttp
 
+from oumi.core.async_utils import safe_asyncio_run
 from oumi.core.configs import GenerationConfig, ModelParams, RemoteParams
 from oumi.core.inference import BaseInferenceEngine
 from oumi.core.types.turn import Conversation, Message, Role, Type
@@ -244,7 +245,7 @@ def infer_online(
         """
         if not generation_config.remote_params:
             raise ValueError("Remote params must be provided in generation_config.")
-        conversations = asyncio.run(
+        conversations = safe_asyncio_run(
             self._infer(input, generation_config, generation_config.remote_params)
         )
         if generation_config.output_filepath:
@@ -271,7 +272,7 @@ def infer_from_file(
         if not generation_config.remote_params:
             raise ValueError("Remote params must be provided in generation_config.")
         input = self._read_conversations(input_filepath)
-        conversations = asyncio.run(
+        conversations = safe_asyncio_run(
             self._infer(input, generation_config, generation_config.remote_params)
         )
         if generation_config.output_filepath:

tests/core/test_async_utils.py

@@ -0,0 +1,35 @@
+import asyncio
+import re
+
+import pytest
+
+from oumi.core.async_utils import safe_asyncio_run
+
+
+def test_safe_asyncio_run_nested():
+    async def nested():
+        return 1
+
+    def method_using_asyncio():
+        return asyncio.run(nested())
+
+    def method_using_safe_asyncio_run():
+        return safe_asyncio_run(nested())
+
+    with pytest.raises(
+        RuntimeError,
+        match=re.escape("asyncio.run() cannot be called from a running event loop"),
+    ):
+
+        async def main_async():
+            return method_using_asyncio()
+
+        # This will raise a RuntimeError because we are trying to run an async function
+        # inside a running event loop.
+        asyncio.run(main_async())
+
+    async def safe_main():
+        return method_using_safe_asyncio_run()
+
+    result = safe_asyncio_run(safe_main())
+    assert result == 1