Repository: caikit/caikit (language: python)

Issue:
Wrap sync in async for fastapi handler
## Is your feature request related to a problem? Please describe.

The fastapi handlers are async, and our module inference apis are sync.

## Describe the solution you'd like

We'll need to do something about that, a generic little sync-over-async wrapper should be fine. I bet @gabe-l-hart already has a POC

## Describe alternatives you've considered

We could see if there's a sync handler api, but word on the street is that SSE support (server-side events, how we want to do "json streaming") generally needs asnyc

We could also change our generic module inference code to be async, but that would probably just punt the problem down the stack because the `caikit.module` api is synchronous. We probably don't want to ask everybody to write async modules to run their ML models.

## Additional context


Optional deps 228
## Description

This PR sets up the framework for optional dependencies in `caikit` and refactors the existing dependency management into the following sets:

Closes: #228 

### User Deps
* `caikit`: Dependencies for `caikit.core`
* `caikit[runtime-grpc]`: Dependencies for the `grpc` runtime
* `caikit[all]`: All dependencies for user functionality in `caikit` (currently core + `runtime-grpc`)

### Developer Deps
* `dev-test`: Dependencies for running unit tests
* `dev-docs`: Dependencies for building documentation
* `dev-fmt`: Dependencies for formatting, linting, and other stylistic checks
* `dev-build`: Dependencies for building the wheel
* `all-dev`: All developer dependencies (including all user dependencies)

**NOTE**: The choice of `-` as a delimiter seems to be required when using `flit` as our build tool. I had originally used `.` which seemed a little cleaner to me, but the `flit` build commands used to build the temporary wheel that `tox` then used to create the venvs was not being built correctly.


Discussion / hints:

NOTE: I was a bit concerned about using `caikit[...]` inside the `all` deps as references, so I went digging. When the wheel gets built locally, the extras metadata looks like the following:

```
Requires-Dist: alchemy-config>=1.1.1,<2.0.0
Requires-Dist: alchemy-logging>=1.0.4,<2.0.0
Requires-Dist: anytree>=2.7.0,<3.0
Requires-Dist: docstring-parser>=0.14.1,<0.16.0
Requires-Dist: grpcio>=1.35.0,<2.0,!=1.55.0
Requires-Dist: ijson>=3.1.4,<3.3.0
Requires-Dist: munch>=2.5.0,<4.0
Requires-Dist: numpy>=1.20,<2
Requires-Dist: protobuf>=3.19.0,<5
Requires-Dist: py-to-proto>=0.3.0,<0.4.0,!=0.2.1
Requires-Dist: PyYAML>=6.0,<7.0
Requires-Dist: semver>=2.13.0,<4.0
Requires-Dist: six>=1.16.0,<2.0.0
Requires-Dist: tqdm>=4.59.0,<5.0.0
Requires-Dist: caikit[runtime-grpc] ; extra == "all"
Requires-Dist: caikit[all, dev-test, dev-docs, dev-fmt, dev-build] ; extra == "all-dev"
Requires-Dist: flit==3.8 ; extra == "dev-build"
Requires-Dist: sphinx>=4.0.2,<8.0 ; extra == "dev-docs"
Requires-Dist: sphinx-autoapi>=2.1.0 ; extra == "dev-docs"
Requires-Dist: sphinx-rtd-theme~=1.2.1 ; extra == "dev-docs"
Requires-Dist: pre-commit>=3.0.4,<4.0 ; extra == "dev-fmt"
Requires-Dist: pylint>=2.16.2,<3.0 ; extra == "dev-fmt"
Requires-Dist: pydeps==1.12.3 ; extra == "dev-fmt"
Requires-Dist: pytest>=6.2.5,<7.0 ; extra == "dev-test"
Requires-Dist: pytest-cov>=2.10.1,<3.0 ; extra == "dev-test"
Requires-Dist: pytest-html>=3.1.1,<4.0 ; extra == "dev-test"
Requires-Dist: tls_test_tools>=0.1.1 ; extra == "dev-test"
Requires-Dist: wheel>=0.38.4 ; extra == "dev-test"
Requires-Dist: grpcio-health-checking>=1.35.0,<2.0 ; extra == "runtime-grpc"
Requires-Dist: grpcio-reflection>=1.35.0,<2.0 ; extra == "runtime-grpc"
Requires-Dist: prometheus_client>=0.12.0,<1.0 ; extra == "runtime-grpc"
Requires-Dist: py-grpc-prometheus>=0.7.0,<0.8 ; extra == "runtime-grpc"
Project-URL: Source, https://github.com/caikit/caikit
Provides-Extra: all
Provides-Extra: all-dev
Provides-Extra: dev-build
Provides-Extra: dev-docs
Provides-Extra: dev-fmt
Provides-Extra: dev-test
Provides-Extra: runtime-grpc
```

In particular, the fact that you see `Requires-Dist: caikit[all, dev-test, dev-docs, dev-fmt, dev-build] ; extra == "all-dev"` worried me that a `pip install caikit[all-dev]` would go back out to the internet to resolve `caikit[all]`, thus breaking installs from local wheel files.

According to [the internet](https://discuss.python.org/t/pyproject-toml-optional-dependencies-redundancy-aka-dry-extras/8428/5), however, this is actually safe. Specifically:

>Since packagename is already in the set of dependencies to be resolved, mentioning it in its own dependencies doesn’t cause a cycle.

I confirmed that this does work with a local wheel, and since these optional dependency sets don't exist in public pypi yet, it must be getting them from the wheel itself.

Fix patch:
diff --git a/caikit/core/toolkit/sync_to_async.py b/caikit/core/toolkit/sync_to_async.py
new file mode 100644
index 000000000..3deab20ed
--- /dev/null
+++ b/caikit/core/toolkit/sync_to_async.py
@@ -0,0 +1,62 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module holds utilities for wrapping synchronous functionality into async
+wrapper functions
+"""
+# Standard
+from concurrent.futures import ThreadPoolExecutor
+from typing import AsyncGenerator, Iterable, Optional
+import asyncio
+
+
+def async_wrap_iter(
+    it: Iterable,
+    pool: Optional[ThreadPoolExecutor] = None,
+) -> AsyncGenerator:
+    """Wrap blocking iterable into an asynchronous one
+
+    CITE: https://stackoverflow.com/a/62297994
+
+    This function manages a single thread for the iteration that shuttles the
+    synchronous outputs back to the async generator using a queue
+    """
+    loop = asyncio.get_event_loop()
+    q = asyncio.Queue(1)
+    exception = None
+    _END = object()
+
+    async def yield_queue_items():
+        while True:
+            next_item = await q.get()
+            if next_item is _END:
+                break
+            yield next_item
+        if exception is not None:
+            # the iterator has raised, propagate the exception
+            raise exception
+
+    def iter_to_queue():
+        nonlocal exception
+        try:
+            for item in it:
+                # This runs outside the event loop thread, so we
+                # must use thread-safe API to talk to the queue.
+                asyncio.run_coroutine_threadsafe(q.put(item), loop).result()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            exception = e
+        finally:
+            asyncio.run_coroutine_threadsafe(q.put(_END), loop).result()
+
+    loop.run_in_executor(pool, iter_to_queue)
+    return yield_queue_items()
diff --git a/pyproject.toml b/pyproject.toml
index 015a2c659..5bda08be0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,9 +50,10 @@ all = [
 ## Dev Extra Sets ##
 
 dev-test = [
-    "pytest>=6.2.5,<8.0",
+    "pytest-asyncio>=0.21.0,<1",
     "pytest-cov>=2.10.1,<5.0",
     "pytest-html>=3.1.1,<4.0",
+    "pytest>=6.2.5,<8.0",
     "tls_test_tools>=0.1.1",
     "wheel>=0.38.4",
 ]


Test patch:
diff --git a/tests/core/toolkit/test_sync_to_async.py b/tests/core/toolkit/test_sync_to_async.py
new file mode 100644
index 000000000..f4e589565
--- /dev/null
+++ b/tests/core/toolkit/test_sync_to_async.py
@@ -0,0 +1,133 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for sync -> async wrappers"""
+# Standard
+from concurrent.futures import ThreadPoolExecutor
+import asyncio
+import threading
+import time
+
+# Third Party
+import pytest
+
+# Local
+from caikit.core.toolkit.sync_to_async import async_wrap_iter
+
+## Helpers #####################################################################
+
+
+class IterableHelper:
+    def __init__(self, iterable, delay=0, raise_element=None):
+        self.iterable = iterable
+        self.delay = delay
+        self.raise_element = raise_element
+        self.idx = 0
+        self.iter = iter(self.iterable)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.delay:
+            time.sleep(self.delay)
+        next_element = next(self.iter)
+        if self.idx == self.raise_element:
+            raise RuntimeError("Raising")
+        self.idx += 1
+        return next_element
+
+
+async def azip(*aiterables):
+    aiterators = [aiterable.__aiter__() for aiterable in aiterables]
+    while aiterators:
+        results = await asyncio.gather(
+            *[ait.__anext__() for ait in aiterators],
+            return_exceptions=True,
+        )
+        if any(isinstance(result, StopAsyncIteration) for result in results):
+            return
+        yield tuple(results)
+
+
+## Tests #######################################################################
+
+
+@pytest.mark.asyncio
+async def test_async_wrap_iter_happy():
+    """Make sure a simple iterable can be wrapped as async and round trip
+    cleanly
+    """
+    lst = [1, 2, 3, 4]
+    async_round_trip = [x async for x in async_wrap_iter(lst)]
+    assert lst == async_round_trip
+
+
+@pytest.mark.asyncio
+async def test_async_wrap_iter_concurrent():
+    """Make sure that multiple wrapped synchronous iterators can interleve in
+    the event loop
+    """
+    lst1 = [1, 2, 3]
+    lst2 = [4, 5, 6]
+    iter1 = IterableHelper(lst1)
+    iter2 = IterableHelper(lst2)
+    aiter1 = async_wrap_iter(iter1)
+    aiter2 = async_wrap_iter(iter2)
+    zipped = [elt async for elt in azip(aiter1, aiter2)]
+    assert zipped == list(zip(lst1, lst2))
+
+
+@pytest.mark.asyncio
+async def test_async_wrap_iter_exception_propagation():
+    """Make sure exceptions thrown during iteration are raised in the parent
+    async context
+    """
+    with pytest.raises(RuntimeError):
+        async for _ in async_wrap_iter(IterableHelper([1, 2, 3], raise_element=1)):
+            pass
+
+
+@pytest.mark.asyncio
+async def test_async_wrap_iter_shared_executor():
+    """Make sure that a shared ThreadPoolExecutor can be used to manage the work
+    without exhausting thread counts
+    """
+    lst = [1, 2, 3, 4]
+    pool = ThreadPoolExecutor(max_workers=1)
+
+    start_event = threading.Event()
+    finish_event = threading.Event()
+
+    def slow_task():
+        start_event.set()
+        finish_event.wait()
+
+    # Start the slow task and make sure it's started
+    slow_future = pool.submit(slow_task)
+    start_event.wait()
+
+    # Use the same pool to start the async iteration
+    agen = async_wrap_iter(lst, pool)
+
+    # Make sure the async iteration is blocked
+    assert pool._work_queue.qsize() == 1
+
+    # Wait for the slow task to finish
+    finish_event.set()
+    slow_future.result()
+
+    # Make sure the async iteration is unblocked
+    assert pool._work_queue.qsize() == 0
+    async_round_trip = [x async for x in agen]
+    assert lst == async_round_trip