API: Fix blocking iterator execution

Run these iterators on the background thread. On startup, the API
spawns a background thread as needed to run sync code on without blocking
the event loop.

Use asyncio's run_thread function since it allows for errors to be
propegated.

Signed-off-by: kingbri <bdashore3@proton.me>

common/concurrency.py

@@ -2,12 +2,40 @@
 
 import asyncio
 import inspect
+from fastapi.concurrency import run_in_threadpool  # noqa
 from functools import partialmethod
 from typing import AsyncGenerator, Generator, Union
 
 generate_semaphore = asyncio.Semaphore(1)
 
 
+# Originally from https://github.com/encode/starlette/blob/master/starlette/concurrency.py
+# Uses generators instead of generics
+class _StopIteration(Exception):
+    """Wrapper for StopIteration because it doesn't send across threads."""
+
+    pass
+
+
+def gen_next(generator: Generator):
+    """Threaded function to get the next value in an iterator."""
+
+    try:
+        return next(generator)
+    except StopIteration as e:
+        raise _StopIteration from e
+
+
+async def iterate_in_threadpool(generator: Generator) -> AsyncGenerator:
+    """Iterates a generator within a threadpool."""
+
+    while True:
+        try:
+            yield await asyncio.to_thread(gen_next, generator)
+        except _StopIteration:
+            break
+
+
 def release_semaphore():
     generate_semaphore.release()
 
@@ -16,19 +44,18 @@ async def generate_with_semaphore(generator: Union[AsyncGenerator, Generator]):
     """Generate with a semaphore."""
 
     async with generate_semaphore:
-        if inspect.isasyncgenfunction:
-            async for result in generator():
-                yield result
-        else:
-            for result in generator():
-                yield result
+        if not inspect.isasyncgenfunction:
+            generator = iterate_in_threadpool(generator())
+
+        async for result in generator():
+            yield result
 
 
 async def call_with_semaphore(callback: partialmethod):
     """Call with a semaphore."""
 
     async with generate_semaphore:
-        if inspect.iscoroutinefunction(callback):
-            return await callback()
-        else:
-            return callback()
+        if not inspect.iscoroutinefunction:
+            callback = run_in_threadpool(callback)
+
+        return await callback()

common/model.py

@@ -72,7 +72,7 @@ async def load_model_gen(model_path: pathlib.Path, **kwargs):
 
 
 async def load_model(model_path: pathlib.Path, **kwargs):
-    async for _, _, _ in load_model_gen(model_path, **kwargs):
+    async for _ in load_model_gen(model_path, **kwargs):
         pass