Repository: ibm-granite/granite-io
Pull request #30: Re-enable asyncio support in backends
Author: login: frreiss
Created: 2025-02-28T03:26:03Z | Merged: 2025-03-11T20:47:22Z

This PR re-enables `asyncio` compatibility, which had been temporarily lost in recent refactoring. Both backend implementations now support multiple simultaneous requests via `asyncio`.

Note that some modifications to #29 will be required to merge that change with this one; specifically replacing LiteLLM's `completion` with [`acompletion`](https://docs.litellm.ai/docs/completion/stream#async-completion).


Diff:
diff --git a/notebooks/.gitignore b/notebooks/.gitignore
new file mode 100644
index 00000000..84812d18
--- /dev/null
+++ b/notebooks/.gitignore
@@ -0,0 +1,3 @@
+# Temporary notebooks for testing
+temp_*.ipynb
+
diff --git a/notebooks/io.ipynb b/notebooks/io.ipynb
index e310905f..72ece865 100644
--- a/notebooks/io.ipynb
+++ b/notebooks/io.ipynb
@@ -175,7 +175,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -189,7 +189,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,
diff --git a/src/granite_io/backend/base.py b/src/granite_io/backend/base.py
index f227c0af..b6e2433f 100644
--- a/src/granite_io/backend/base.py
+++ b/src/granite_io/backend/base.py
@@ -24,7 +24,7 @@ class Backend(FactoryConstructible):
     """
 
     @abc.abstractmethod
-    def generate(
+    async def generate(
         self, input_str: str, num_return_sequences: int = 1
     ) -> GenerateResults:
         """
@@ -43,7 +43,7 @@ class ChatCompletionBackend(FactoryConstructible):
     """
 
     @abc.abstractmethod
-    def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
         """
         Callback to invoke the model to generate a response.
         """
diff --git a/src/granite_io/backend/litellm.py b/src/granite_io/backend/litellm.py
index 88edf8c9..8493cbd1 100644
--- a/src/granite_io/backend/litellm.py
+++ b/src/granite_io/backend/litellm.py
@@ -34,7 +34,7 @@ class LiteLLMBackend(Backend):
     def __init__(self, config: aconfig.Config):
         self._model_str = config.model_name
 
-    def generate(
+    async def generate(
         self, input_str: str, num_return_sequences: int = 1
     ) -> GenerateResults:
         """Run a direct /completions call"""
@@ -48,7 +48,7 @@ def generate(
             # Third Party
             import litellm
 
-        result = litellm.text_completion(
+        result = await litellm.atext_completion(
             # model strings start with provider/
             # model="watsonx/ibm/granite-3-2-8b-instruct",
             # model="ollama/llama3.1:latest",
diff --git a/src/granite_io/backend/openai.py b/src/granite_io/backend/openai.py
index a6001272..b772f151 100644
--- a/src/granite_io/backend/openai.py
+++ b/src/granite_io/backend/openai.py
@@ -29,26 +29,23 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
-
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
         )
 
-    def generate(
+    async def generate(
         self, input_str: str, num_return_sequences: int = 1
     ) -> GenerateResults:
         """Run a direct /completions call"""
@@ -58,7 +55,7 @@ def generate(
                 f"Invalid value for num_return_sequences ({num_return_sequences})"
             )
 
-        result = self.openai_client.completions.create(
+        result = await self._openai_client.completions.create(
             model=self._model_str,
             prompt=input_str,
             best_of=num_return_sequences,
diff --git a/src/granite_io/backend/transformers.py b/src/granite_io/backend/transformers.py
index 0ba1d203..0430a1ae 100644
--- a/src/granite_io/backend/transformers.py
+++ b/src/granite_io/backend/transformers.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Standard
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+import asyncio
+import concurrent.futures
 import dataclasses
 
 # Third Party
@@ -15,10 +17,24 @@
 
 if TYPE_CHECKING:
     # Third Party
-    import torch  # noqa: F401
     import transformers
 
 
+def _detect_hw_accel():
+    """:returns: Pytorch's name for the best available hardware accelerator on the
+    current machine."""
+    with import_optional("transformers"):
+        # Third Party
+        import torch
+    if torch.cuda.is_available():
+        result = "cuda"
+    elif torch.backends.mps.is_available():
+        result = "mps"
+    else:
+        result = "cpu"
+    return result
+
+
 @dataclasses.dataclass
 class _GenerationInputs:
     """Dataclass to encapsulate inputs for calling the Transformers generate() method"""
@@ -41,8 +57,12 @@ class _GenerationInputs:
 )
 class TransformersBackend(Backend):
     _model_str: str
+    _torch_device_name: str
+    "Device used for hardware acceleration by this backend."
     _model: "transformers.AutoModelForCausalLM"
     _tokenizer: "transformers.AutoTokenizer"
+    _executor: concurrent.futures.ThreadPoolExecutor
+    "Single background thread, wrapped in an executor for queueing"
 
     def __init__(self, config: aconfig.Config):
         """
@@ -59,42 +79,33 @@ def __init__(self, config: aconfig.Config):
             import transformers
 
         self._model_str = config.model_name
-        self._torch_device_name = config.device
-
-        if not self._torch_device_name:
-            if torch.cuda.is_available():
-                self._torch_device_name = "cuda"
-            # Re-enable once we have a smaller model that can run in a laptop GPU
-            elif torch.backends.mps.is_available():
-                self._torch_device_name = "mps"
-            else:
-                self._torch_device_name = "cpu"
-                # CPU mode; prevent thrashing
-                torch.set_num_threads(4)
-
+        self._torch_device_name = config.device if config.device else _detect_hw_accel()
+        if self._torch_device_name == "cpu":
+            # CPU mode; prevent thrashing
+            torch.set_num_threads(4)
         self._model = transformers.AutoModelForCausalLM.from_pretrained(
-            self._model_str
+            config.model_name
         ).to(self._torch_device_name)
-        self._tokenizer = transformers.AutoTokenizer.from_pretrained(self._model_str)
+        self._tokenizer = transformers.AutoTokenizer.from_pretrained(config.model_name)
+        self._executor = concurrent.futures.ThreadPoolExecutor()
 
-    def generate(
+    async def generate(
         self, input_str: str, num_return_sequences: int = 1
     ) -> GenerateResults:
         if num_return_sequences < 1:
             raise ValueError(
                 f"Invalid value for num_return_sequences ({num_return_sequences})"
             )
-
         generation_inputs = self._prepare_for_generate(
             input_str, num_return_sequences=num_return_sequences
         )
 
-        model_output = self._model.generate(
-            **(generation_inputs.model_input),
-            generation_config=generation_inputs.generation_config,
-            # Re-enable to use constrained decoding
-            # prefix_allowed_tokens_fn=generation_inputs.prefix_allowed_tokens_fn
+        # Wrap the call to the non-async Transformers library in a thread pool
+        concurrent_futures_future = self._executor.submit(
+            self._generate_callback, generation_inputs
         )
+        async_future = asyncio.wrap_future(concurrent_futures_future)
+        model_output = await async_future
 
         # The result of generate() is of course the prompt concatenated with the
         # additional tokens generated. Strip off the prompt.
@@ -141,8 +152,6 @@ def _prepare_for_generate(
             import torch  # noqa: F401
             import transformers
 
-        # Call model.generate(). This is much harder than it should be.
-
         # Turn the conversation prefix into tokens for input to the model.
         model_input = self._tokenizer(
             # The conversation up to this point
@@ -241,3 +250,31 @@ def _prepare_for_generate(
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        def do_generate():
+            return self._model.generate(
+                **(generation_inputs.model_input),
+                generation_config=generation_inputs.generation_config,
+            )
+
+        if self._torch_device_name == "cuda":
+            # Make sure computations for this thread will happen in a separate CUDA
+            # context.
+            stream = torch.cuda.Stream()
+            with torch.cuda.stream(stream):
+                return do_generate()
+        else:
+            return do_generate()
diff --git a/src/granite_io/io/base.py b/src/granite_io/io/base.py
index 74e83563..334830e5 100644
--- a/src/granite_io/io/base.py
+++ b/src/granite_io/io/base.py
@@ -6,6 +6,8 @@
 
 # Standard
 import abc
+import asyncio
+import threading
 
 # Third Party
 import aconfig
@@ -21,6 +23,32 @@
 )
 
 
+def _workaround_for_horrible_design_flaw_in_asyncio(coroutine_to_run):
+    """
+    Horrible hack that seems to be the only way to call a coroutine from a non-async
+    function that is being called from an async function.
+
+    Creates a background thread with its own event loop. The background thread then
+    runs a single task in that event loop.
+    """
+    result_holder = []
+
+    def _callback():
+        # Python's thread library requires you to write your own code to return results
+        # or raise exceptions from a background thread.
+        try:
+            result_holder.append(asyncio.run(coroutine_to_run))
+        except Exception as e:  # pylint: disable=W0718
+            result_holder.append(e)
+
+    thread = threading.Thread(target=_callback)
+    thread.start()
+    thread.join()
+    if isinstance(result_holder[0], Exception):
+        raise result_holder[0]
+    return result_holder[0]
+
+
 class InputOutputProcessor(FactoryConstructible):
     """
     Interface for generic input-output processors. An input-output processor exposes an
@@ -33,10 +61,24 @@ def __init__(self, config: aconfig.Config | None = None):  # pylint: disable=unu
         """By default an IO processor doesn't require config"""
 
     @abc.abstractmethod
+    async def acreate_chat_completion(
+        self, inputs: ChatCompletionInputs
+    ) -> ChatCompletionResult:
+        """
+        :param inputs: Structured representation of the inputs to a chat completion
+            request, possibly including additional fields that only this input-output
+            processor can consume
+
+        :returns: The next message that the model produces when fed the specified
+            inputs, plus additional information about the low-level request.
+        """
+
     def create_chat_completion(
         self, inputs: ChatCompletionInputs
     ) -> ChatCompletionResult:
         """
+        Non-async version of :func:`acreate_chat_completion()`
+
         :param inputs: Structured representation of the inputs to a chat completion
             request, possibly including additional fields that only this input-output
             processor can consume
@@ -44,6 +86,17 @@ def create_chat_completion(
         :returns: The next message that the model produces when fed the specified
             inputs, plus additional information about the low-level request.
         """
+        # Fall back on async version of this method by default.  Subclasses may override
+        # this method if they have a more efficient way of doing non-async operation.
+        coroutine_to_run = self.acreate_chat_completion(inputs)
+        try:  # Exceptions as control flow. Sorry, asyncio forces this design on us.
+            asyncio.get_running_loop()
+        except RuntimeError:
+            # If we get here, this code is not running inside an async function.
+            return asyncio.run(coroutine_to_run)
+
+        # If we get here, this code is running inside an async function.
+        return _workaround_for_horrible_design_flaw_in_asyncio(coroutine_to_run)
 
 
 class ModelDirectInputOutputProcessor(InputOutputProcessor):
@@ -68,7 +121,7 @@ def __init__(
         super().__init__(config)
         self._backend = backend
 
-    def create_chat_completion(
+    async def acreate_chat_completion(
         self, inputs: ChatCompletionInputs
     ) -> ChatCompletionResults:
         if self._backend is None:
@@ -77,7 +130,7 @@ def create_chat_completion(
                 "configuring an inference backend."
             )
         input_string = self.inputs_to_string(inputs)
-        generation_results = self._backend.generate(input_string)
+        generation_results = await self._backend.generate(input_string)
         return self.output_to_result(generation_results, inputs)
 
     @abc.abstractmethod
diff --git a/tests/io/test_granite_3_2.py b/tests/io/test_granite_3_2.py
index 84732fa6..807a55d7 100644
--- a/tests/io/test_granite_3_2.py
+++ b/tests/io/test_granite_3_2.py
@@ -26,6 +26,7 @@
 from granite_io.types import (
     AssistantMessage,
     ChatCompletionInputs,
+    ChatCompletionResults,
     GenerateResult,
     GenerateResults,
     UserMessage,
@@ -200,10 +201,11 @@ def test_basic_inputs_to_string():
 def test_run_processor(backend_x: Backend, input_json_str: str):
     inputs = ChatCompletionInputs.model_validate_json(input_json_str)
     io_processor = make_io_processor(_MODEL_NAME, backend=backend_x)
-    _ = io_processor.create_chat_completion(inputs)
+    results: ChatCompletionResults = io_processor.create_chat_completion(inputs)
 
-    # TODO: Once the prerelease model has settled down and we have implemented
-    # temperature controls, verify outputs
+    assert isinstance(results, ChatCompletionResults)
+    assert len(results.results) == 1
+    # TODO: Verify outputs in greater detail
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_backends.py b/tests/test_backends.py
index e1185e97..229258f7 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -12,27 +12,30 @@
 from granite_io.types import GenerateResults
 
 
+@pytest.mark.asyncio
 @pytest.mark.vcr
 @pytest.mark.block_network
-def test_simple(backend_x):
-    ret = backend_x.generate("hello")
+async def test_simple(backend_x):
+    ret = await backend_x.generate("hello")
     assert isinstance(ret, GenerateResults)
     assert len(ret.results) == 1
 
 
+@pytest.mark.asyncio
 @pytest.mark.vcr
 @pytest.mark.block_network
-def test_num_return_sequences_1(backend_x):
-    ret = backend_x.generate("hello", num_return_sequences=1)
+async def test_num_return_sequences_1(backend_x):
+    ret = await backend_x.generate("hello", num_return_sequences=1)
     assert isinstance(ret, GenerateResults)
     assert len(ret.results) == 1
 
 
+@pytest.mark.asyncio
 @pytest.mark.vcr
 @pytest.mark.block_network
-def test_num_return_sequences_3(backend_x):
+async def test_num_return_sequences_3(backend_x):
     try:
-        ret = backend_x.generate("what is up?", num_return_sequences=3)
+        ret = await backend_x.generate("what is up?", num_return_sequences=3)
     except UnsupportedParamsError:
         # Specific exception from LiteLLMBackend
         # xfail because LiteLLM is telling us that ollama does not support
@@ -53,12 +56,13 @@ def test_num_return_sequences_3(backend_x):
     assert num_returned == 3
 
 
+@pytest.mark.asyncio
 @pytest.mark.vcr
 @pytest.mark.block_network
-@pytest.mark.flaky(retries=3, delay=5)  # VCR recording flakey
+@pytest.mark.flaky(retries=3, delay=5)  # VCR recording flaky
 @pytest.mark.parametrize("n", [-1, 0])
-def test_num_return_sequences_invalid(backend_x, n):
+async def test_num_return_sequences_invalid(backend_x, n):
     with pytest.raises(
         ValueError, match=re.escape(f"Invalid value for num_return_sequences ({n})")
     ):
-        _ = backend_x.generate("hello", num_return_sequences=n)
+        _ = await backend_x.generate("hello", num_return_sequences=n)


Reviews:
user: login: hickeyma
body: Thanks @frreiss  for updating. Waiting until #13 and #29 are merged before reviewing.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T17:06:53Z
user: login: markstur
body: Will need example and notebook fixing.  I'm trying to fix the notebook now which is how I see problems like torch.cuda.Stream.

See inline comments too
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-02-28T21:37:16Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:14:51Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:20:46Z
user: login: markstur
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:25:24Z
user: login: markstur
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:30:22Z
user: login: markstur
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:31:01Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-01T04:45:44Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-01T04:46:19Z
user: login: hickeyma
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T10:47:56Z
user: login: hickeyma
body: @frreiss It is looking ok apart from some small nits. I think it would be worth also updating `LiteLLM` also.

Overall, the question still remains if we need asynchronous APIs for now or not.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T11:20:27Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T19:39:16Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T20:07:17Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T20:07:36Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T20:07:57Z
user: login: hickeyma
body: Looking food @frreiss. Open questions around `create_chat_completion()` async/sync APIs.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T11:54:13Z
user: login: hickeyma
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T12:04:01Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T21:06:48Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T17:44:08Z
user: login: frreiss
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T17:54:33Z
user: login: markstur
body: lgtm
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-11T19:42:39Z
user: login: hickeyma
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-11T20:43:56Z

Review comments:
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: markstur
body: Do we have any use for chat/completion?   I thought you said we always do completion (generate).

So this was removed and the backend is just a Backend and not a ChatCompletionBackend

but please update me if we have any use for this
created_at: 2025-02-28T21:31:36Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
path: src/granite_io/backend/transformers.py
user: login: markstur
body: this does not seem to work on Mac
created_at: 2025-02-28T21:34:31Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
+
+        with torch.cuda.stream(stream):
+            return self._model.generate(
+                **(generation_inputs.model_input),
+                generation_config=generation_inputs.generation_config,
+                # Re-enable to use constrained decoding
path: src/granite_io/backend/transformers.py
user: login: markstur
body: Is this a todo or ?

Probably delete to cleanup, but probably todo comments are okay at this stage
created_at: 2025-02-28T21:35:41Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: frreiss
body: The chat completion backend provides basic, lowest-common-denominator chat completions on models for which we do not currently have an IO processor. This backend type was not my idea, and I'm happy to remove this functionality entirely, up to and including the base class `backend.ChatCompletionBackend`, if you, @hickeyma, and @gabe-l-hart agree.
created_at: 2025-02-28T22:14:51Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
+
+        with torch.cuda.stream(stream):
+            return self._model.generate(
+                **(generation_inputs.model_input),
+                generation_config=generation_inputs.generation_config,
+                # Re-enable to use constrained decoding
path: src/granite_io/backend/transformers.py
user: login: frreiss
body: Deleting this for now.
created_at: 2025-02-28T22:20:46Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: markstur
body: I agree and want to hear from others.  So yeah I had only deleted it here (not used not tested) but I didn't remove the base class. 
created_at: 2025-02-28T22:25:23Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
+
+        with torch.cuda.stream(stream):
path: src/granite_io/backend/transformers.py
user: login: markstur
body: ```suggestion
        from contextlib import nullcontext
        with torch.cuda.stream(torch.cuda.Stream()) if self._torch_device_name == "cuda" else nullcontext():
```
created_at: 2025-02-28T22:30:22Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
+
+        with torch.cuda.stream(stream):
path: src/granite_io/backend/transformers.py
user: login: markstur
body: suggestion works on mac (import should be moved to top)
created_at: 2025-02-28T22:31:01Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
path: src/granite_io/backend/transformers.py
user: login: frreiss
body: Oops! Updated this part so that it will run without cuda.
created_at: 2025-03-01T04:45:44Z
author_association: COLLABORATOR
diff_hunk: @@ -211,3 +214,27 @@ def _prepare_for_generate(self, prompt: str) -> _GenerationInputs:
             model_input=model_input,
             # prefix_allowed_tokens_fn=prefix_allowed_tokens_fn
         )
+
+    def _generate_callback(self, generation_inputs: _GenerationInputs) -> Any:
+        """
+        :param generation_inputs: Wrapper for the required inputs to a Transformers
+            :func:`generate()` call, prepared in the event loop thread.
+
+        :returns: Whatever type this model's `generate()` method happens to return on
+            this day of the week, which changes without warning from one version to the
+            next of the library.
+        """
+        with import_optional("transformers"):
+            # Third Party
+            import torch
+
+        # Make sure computations for this thread will happen in a separate CUDA context.
+        stream = torch.cuda.Stream()
+
+        with torch.cuda.stream(stream):
path: src/granite_io/backend/transformers.py
user: login: frreiss
body: Did something similar earlier today, just seeing this now. Thanks for the suggestion!
created_at: 2025-03-01T04:46:19Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: hickeyma
body: @frreiss  Where did the idea come from and how would you use it for a non supported model?
created_at: 2025-03-03T10:47:55Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: hickeyma
body: Doing a bit more digging, it would seem that "completions" API maybe legacy or unsupported for some models. There is a good explanation how OpenAI platform handles both APIs: https://platform.openai.com/docs/guides/completions
created_at: 2025-03-03T11:12:48Z
author_association: COLLABORATOR
diff_hunk: @@ -217,25 +217,29 @@ def test_basic_inputs_to_string():
     assert chatRequest.endswith("")
 
 
-def test_run_transformers(
+# Session scope for asyncio because the tests in this class all share the same vLLM
+# backend
path: tests/io/test_granite_3_2.py
user: login: hickeyma
body: Comment looks incorrect with mention to vLLM?
created_at: 2025-03-03T11:17:22Z
author_association: COLLABORATOR
diff_hunk: @@ -217,25 +217,29 @@ def test_basic_inputs_to_string():
     assert chatRequest.endswith("")
 
 
-def test_run_transformers(
+# Session scope for asyncio because the tests in this class all share the same vLLM
+# backend
path: tests/io/test_granite_3_2.py
user: login: frreiss
body: Good point, will remove the mention of vLLM.
created_at: 2025-03-03T19:39:16Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: frreiss
body: Implementations of the OpenAI protocol that do not have a `create_completion` endpoint aren't going to work with I/O processors that apply a chat template, because the API doesn't allow callers to send raw prompts and receive raw results. We may want to implement a separate I/O processor for such endpoints at some point in the future.
created_at: 2025-03-07T20:07:17Z
author_association: COLLABORATOR
diff_hunk: @@ -29,29 +29,53 @@
 )
 class OpenAIBackend(Backend):
     _model_str: str
-    _openai_client: "openai.OpenAI"
+    _openai_client: "openai.AsyncOpenAI"
 
     def __init__(self, config: aconfig.Config):
+        with import_optional("openai"):
+            # Third Party
+            import openai
+
         self._model_str = config.model_name
         api_key = config.get("openai_api_key", "ollama")
         base_url = config.get("openai_base_url", "http://localhost:11434/v1")
-
         default_headers = {"RITS_API_KEY": api_key} if api_key else None
 
-        with import_optional("openai"):
-            # Third Party
-            import openai
+        self._openai_client = openai.AsyncOpenAI(
+            base_url=base_url, api_key=api_key, default_headers=default_headers
+        )
 
-        self.openai_client = openai.OpenAI(
-            base_url=base_url,
-            api_key=api_key,
-            default_headers=default_headers,
+    async def create_chat_completion(self, input_chat: ChatCompletionInputs) -> str:
path: src/granite_io/backend/openai.py
user: login: frreiss
body: I've changed this code to remove `create_chate_completion()` for now.
created_at: 2025-03-07T20:07:36Z
author_association: COLLABORATOR
diff_hunk: @@ -217,25 +217,29 @@ def test_basic_inputs_to_string():
     assert chatRequest.endswith("")
 
 
-def test_run_transformers(
+# Session scope for asyncio because the tests in this class all share the same vLLM
+# backend
path: tests/io/test_granite_3_2.py
user: login: frreiss
body: Removed now.
created_at: 2025-03-07T20:07:57Z
author_association: COLLABORATOR
diff_hunk: @@ -33,17 +61,42 @@ def __init__(self, config: aconfig.Config | None = None):  # pylint: disable=unu
         """By default an IO processor doesn't require config"""
 
     @abc.abstractmethod
+    async def acreate_chat_completion(
+        self, inputs: ChatCompletionInputs
+    ) -> ChatCompletionResult:
+        """
+        :param inputs: Structured representation of the inputs to a chat completion
+            request, possibly including additional fields that only this input-output
+            processor can consume
+
+        :returns: The next message that the model produces when fed the specified
+            inputs, plus additional information about the low-level request.
+        """
+
     def create_chat_completion(
path: src/granite_io/io/base.py
user: login: hickeyma
body: Why are we including a synchronous API when all it is doing is calling the asynchronous API?  It would make sense if it was implemented as natively synchronous. Suggest dropping it and just having `create_chat_completion()` as asynchronous.
created_at: 2025-03-10T11:52:41Z
author_association: COLLABORATOR
diff_hunk: @@ -41,8 +57,12 @@ class _GenerationInputs:
 )
 class TransformersBackend(Backend):
     _model_str: str
+    _torch_device_name: str
+    """Device used for hardware acceleration by this backend."""
     _model: "transformers.AutoModelForCausalLM"
     _tokenizer: "transformers.AutoTokenizer"
+    _executor: concurrent.futures.ThreadPoolExecutor
+    """Single background thread, wrapped in an executor for queueing"""
 
path: src/granite_io/backend/transformers.py
user: login: hickeyma
body: Comments look skewed .. maybe also using sinle line comment?
created_at: 2025-03-10T12:04:00Z
author_association: COLLABORATOR
diff_hunk: @@ -41,8 +57,12 @@ class _GenerationInputs:
 )
 class TransformersBackend(Backend):
     _model_str: str
+    _torch_device_name: str
+    """Device used for hardware acceleration by this backend."""
     _model: "transformers.AutoModelForCausalLM"
     _tokenizer: "transformers.AutoTokenizer"
+    _executor: concurrent.futures.ThreadPoolExecutor
+    """Single background thread, wrapped in an executor for queueing"""
 
path: src/granite_io/backend/transformers.py
user: login: frreiss
body: These strings are docstrings. They don't show up in `__doc__`, but some Python development tools can use them. For example, hover the cursor over `self._torch_device_name` in VS Code and you get this popup: 
![Screenshot 2025-03-10 at 2 05 37 PM](https://github.com/user-attachments/assets/7497ac10-afe7-4374-9ee3-d54249c1a56b)

created_at: 2025-03-10T21:06:47Z
author_association: COLLABORATOR
diff_hunk: @@ -33,17 +61,42 @@ def __init__(self, config: aconfig.Config | None = None):  # pylint: disable=unu
         """By default an IO processor doesn't require config"""
 
     @abc.abstractmethod
+    async def acreate_chat_completion(
+        self, inputs: ChatCompletionInputs
+    ) -> ChatCompletionResult:
+        """
+        :param inputs: Structured representation of the inputs to a chat completion
+            request, possibly including additional fields that only this input-output
+            processor can consume
+
+        :returns: The next message that the model produces when fed the specified
+            inputs, plus additional information about the low-level request.
+        """
+
     def create_chat_completion(
path: src/granite_io/io/base.py
user: login: frreiss
body: The intent here is that some I/O processors will override this base-class implementation with their own natively synchronous implementation. Other I/O processors will only implement the asynchronous version of `create_chat_completion()` and will rely on this method in the base class for compatibility with non-async code.

I'd also be fine with only having async APIs, but there were some strong objections the last time I recommended that approach.
created_at: 2025-03-11T17:44:08Z
author_association: COLLABORATOR
diff_hunk: @@ -41,8 +57,12 @@ class _GenerationInputs:
 )
 class TransformersBackend(Backend):
     _model_str: str
+    _torch_device_name: str
+    """Device used for hardware acceleration by this backend."""
     _model: "transformers.AutoModelForCausalLM"
     _tokenizer: "transformers.AutoTokenizer"
+    _executor: concurrent.futures.ThreadPoolExecutor
+    """Single background thread, wrapped in an executor for queueing"""
 
path: src/granite_io/backend/transformers.py
user: login: frreiss
body: Changed these docstrings to single-quoted strings to reduce the visual clutter.
created_at: 2025-03-11T17:54:32Z
author_association: COLLABORATOR

PR comments:
user: login: frreiss
created_at: 2025-03-01T05:12:42Z
body: Notebooks and litellm are working now on this branch.
author_association: COLLABORATOR
user: login: frreiss
created_at: 2025-03-07T22:56:44Z
body: @hickeyma @markstur @gabe-l-hart This code should be ready for final review now.
author_association: COLLABORATOR