feat(iorails): Use single engine with correct lifecycle (#1649)

* Initial checkin of code, no existing tests broken

* Move test data into its own file

* Clean up tests

* Rewrite logic on when IORails is used vs LLMRails

* Remove todos (will do in future PR), fix docstring for _last_content_by_role, remove unused guardrails_models.py

* Implement a single engine (IORails / LLMRails) and update tests

* Standardise start/stop for objects which need lifespan hooks (i.e. creating worker tasks, opening clients, etc)

* Fix start/stop ordering and add finally-clause to make sure self._running doesn't get into an inconsistent state

* Clean up start/stop exception handling, fix inaccurate docstrings

* Clean up docstrings and multiple-worker start/stop code in AsyncWorkQueue and ModelManager

* Use _flow_name() to extract the flow name and fix urljoin() in model_engine.py

* Fix IORails start/stop exception handling, add new tests to get line coverage back to 100%

* Revert changes to model manager start/stop

* Fix _has_only_iorails_flows set bug, update test_start_failure_allows_retry to match start/stop methods

* Clean up async work queue start() and tests

* Unpack GenerationOptions object (passed from server api.py) and pass on llm_params to the main LLM call in ModelManager

* Unpack llm parameters to top-level rather than nested under llm_params in HTTP body

* Reformat after rebasing onto develop

* Rebase-merge cleanups

Author: tgasser-nv

nemoguardrails/guardrails/async_work_queue.py

@@ -73,29 +73,42 @@ async def start(self) -> None:
         """Starts the worker pool. Call this during service startup."""
         if self._running:
             return
-        self._running = True
+
         self._busy_count = 0
         self._workers = []
-        for i in range(self._max_concurrency):
-            task = asyncio.create_task(self._worker_loop(), name=f"{self._name}_worker_id{i}")
-            self._workers.append(task)
+        # Try to start all workers, cancelling any on failure
+        try:
+            for i in range(self._max_concurrency):
+                task = asyncio.create_task(self._worker_loop(), name=f"{self._name}_worker_id{i}")
+                self._workers.append(task)
+        except Exception:
+            # Cancel any tasks that did start
+            for task in self._workers:
+                task.cancel()
+            await asyncio.gather(*self._workers, return_exceptions=True)
+            self._workers = []
+            raise
+
+        self._running = True
 
     async def stop(self, wait_for_completion: bool = True) -> None:
         """Stops the worker pool. Call this during service shutdown."""
         if not self._running:
             return
-        self._running = False
 
-        if wait_for_completion:
-            await self._queue.join()
+        try:
+            if wait_for_completion:
+                await self._queue.join()
 
-        for task in self._workers:
-            task.cancel()
+            for task in self._workers:
+                task.cancel()
 
-        # Swallow cancellations to prevent noise during shutdown
-        await asyncio.gather(*self._workers, return_exceptions=True)
-        self._workers = []
-        self._busy_count = 0
+            # Swallow cancellations to prevent noise during shutdown
+            await asyncio.gather(*self._workers, return_exceptions=True)
+            self._workers = []
+            self._busy_count = 0
+        finally:
+            self._running = False
 
     async def submit(self, func: Callable[..., Awaitable[T]], *args: Any, **kwargs: Any) -> T:
         """

nemoguardrails/guardrails/guardrails.py

@@ -16,12 +16,13 @@
 """Top-level Guardrails interface module.
 
 This module provides a simplified, user-friendly interface for interacting with
-NeMo Guardrails. The Guardrails class wraps the LLMRails functionality and provides
-a streamlined API for generating LLM responses with programmable guardrails.
+NeMo Guardrails. The Guardrails class wraps either IORails or LLMRails (chosen
+automatically based on config) and provides a streamlined API for generating
+LLM responses with programmable guardrails.
 """
 
 import logging
-from typing import AsyncIterator, Optional, Tuple, Union, overload
+from typing import AsyncIterator, Optional, Tuple, Union, cast, overload
 
 from langchain_core.language_models import BaseChatModel, BaseLLM
 
@@ -63,10 +64,8 @@ def __init__(
         self.verbose = verbose
 
         # Whether to use IORailsEngine for inference requests
-        self._use_iorails_engine: bool = use_iorails and self._has_only_iorails_flows()
-        self._iorails = IORails(config)
-        self._llmrails = LLMRails(config, llm, verbose)
-        self.rails_engine = self._iorails if self._use_iorails_engine else self._llmrails
+        use_iorails_engine = use_iorails and self._has_only_iorails_flows()
+        self._rails_engine = IORails(config) if use_iorails_engine else LLMRails(config, llm, verbose)
 
         # Async work queue for managing concurrent generate_async requests
         self._generate_async_queue: AsyncWorkQueue = AsyncWorkQueue(
@@ -80,22 +79,16 @@ def __init__(
         self._queues = [self._generate_async_queue]
 
     @property
-    def iorails(self) -> IORails:
-        """Get immutable IORails object"""
-        return self._iorails
-
-    @property
-    def llmrails(self) -> LLMRails:
+    def rails_engine(self) -> IORails | LLMRails:
         """Get immutable LLMRails object"""
-        return self._llmrails
+        return self._rails_engine
 
     @staticmethod
     def _convert_to_messages(prompt: str | None = None, messages: LLMMessages | None = None) -> LLMMessages:
-        """Convert prompt or simplified messages to LLMRails standard format.
+        """Return messages in standard format, converting a prompt string if needed.
 
-        Converts from Guardrails simplified format to LLMRails standard format:
-        - Simplified: [{"user": "text"}]
-        - Standard: [{"role": "user", "content": "Hello"}]
+        If messages is provided, returns it as-is.
+        If prompt is provided, wraps it as [{"role": "user", "content": prompt}].
         """
 
         # Priority: messages first, then prompt
@@ -113,7 +106,7 @@ def _has_only_iorails_flows(self):
 
         # If we have any rails outside of `input` and `output` we don't support them
         rails_set = self.config.rails.model_fields_set
-        if rails_set > IORAILS_RAILS:
+        if rails_set - IORAILS_RAILS:
             return False
 
         for flow in self.config.rails.input.flows:
@@ -133,11 +126,14 @@ def generate(
     ) -> Union[str, dict, GenerationResponse, Tuple[dict, dict]]:
         """Generate an LLM response synchronously with guardrails applied."""
 
-        if self._use_iorails_engine:
+        if isinstance(self.rails_engine, IORails):
             raise NotImplementedError("IORails doesn't support generate()")
 
         generate_messages = self._convert_to_messages(prompt, messages)
-        response = self._llmrails.generate(messages=generate_messages, **kwargs)
+
+        # self.rails_engine must be LLMRails since we raise above if we're using IORails
+        llmrails = cast(LLMRails, self.rails_engine)
+        response = llmrails.generate(messages=generate_messages, **kwargs)
         return response
 
     @overload
@@ -178,46 +174,56 @@ def stream_async(
         Only supported when using LLMRails
         """
 
-        if self._use_iorails_engine:
+        if isinstance(self.rails_engine, IORails):
             raise NotImplementedError("IORails doesn't support stream_async()")
 
         stream_messages = self._convert_to_messages(prompt, messages)
-        return self._llmrails.stream_async(messages=stream_messages, **kwargs)
+        # self.rails_engine must be LLMRails since we raise above if we're using IORails
+        llmrails = cast(LLMRails, self.rails_engine)
+        return llmrails.stream_async(messages=stream_messages, **kwargs)
 
     def explain(self) -> ExplainInfo:
         """Get the latest ExplainInfo object for debugging.
         Only supported for LLMRails
         """
 
-        if self._use_iorails_engine:
+        if isinstance(self.rails_engine, IORails):
             raise NotImplementedError("IORails doesn't support explain()")
 
-        return self._llmrails.explain()
+        # self.rails_engine must be LLMRails since we raise above if we're using IORails
+        llmrails = cast(LLMRails, self.rails_engine)
+        return llmrails.explain()
 
     def update_llm(self, llm: Union[BaseLLM, BaseChatModel]) -> None:
         """Replace the main LLM with a new one.
         Only supported for LLMRails, since IORails doesn't take LLM as argument
         """
-        if self._use_iorails_engine:
+        if isinstance(self.rails_engine, IORails):
             raise NotImplementedError("IORails doesn't support update_llm()")
 
-        self._llmrails.update_llm(llm)
+        # self.rails_engine must be LLMRails since we raise above if we're using IORails
+        llmrails = cast(LLMRails, self.rails_engine)
+        llmrails.update_llm(llm)
 
     async def startup(self) -> None:
-        """Lifecycle method to create worker threads and infrastructure"""
+        """Lifecycle method to start async worker tasks and the rails engine"""
         for queue in self._queues:
             await queue.start()
+        if isinstance(self.rails_engine, IORails):
+            await self.rails_engine.start()
 
     async def shutdown(self) -> None:
-        """Lifecycle method to cleanly shutdown worker threads and infrastructure"""
+        """Lifecycle method to stop async worker tasks and the rails engine"""
         for queue in self._queues:
             await queue.stop()
+        if isinstance(self.rails_engine, IORails):
+            await self.rails_engine.stop()
 
     async def __aenter__(self):
-        """Async context manager entry - starts the queues."""
+        """Async context manager entry."""
         await self.startup()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Async context manager exit - shuts down the queues."""
+        """Async context manager exit."""
         await self.shutdown()

nemoguardrails/guardrails/iorails.py

@@ -16,9 +16,8 @@
 """Optimized IORails Engine for specific guardrail configurations.
 
 This module provides an optimized inference path for guardrail configurations that
-only use specific supported flows (input/output content safety, topic safety,
-jailbreak detection, etc.). For configurations outside this supported set, the
-standard LLMRails engine should be used instead.
+only use specific supported flows (input/output content safety). For configurations
+outside this supported set, the standard LLMRails engine should be used instead.
 """
 
 import logging
@@ -27,6 +26,7 @@
 from nemoguardrails.guardrails.model_manager import ModelManager
 from nemoguardrails.guardrails.rails_manager import RailsManager
 from nemoguardrails.rails.llm.config import RailsConfig
+from nemoguardrails.rails.llm.options import GenerationOptions
 
 log = logging.getLogger(__name__)
 
@@ -37,12 +37,46 @@ class IORails:
     """Workflow engine for accelerated Input/Output rails inference."""
 
     def __init__(self, config: RailsConfig) -> None:
+        self._running = False
+
         # Model Manager has one or more ModelEngine inside. Each ModelEngine calls a single model or API
         self.model_manager = ModelManager(config.models)
 
         # Rails Manager is responsible for running rails by making calls to Model Manager
         self.rails_manager = RailsManager(config, self.model_manager)
 
+    async def start(self) -> None:
+        """Start the IORails engine. Call this during service startup."""
+        if self._running:
+            return
+
+        # When starting up, make sure self._running is always set to True even on exceptions.
+        # This allows the stop() method to clean up any state
+        try:
+            await self.model_manager.start()
+        finally:
+            self._running = True
+
+    async def stop(self) -> None:
+        """Stop the IORails engine. Call this during service shutdown."""
+        if not self._running:
+            return
+
+        # If any exceptions are thrown when stopping ModelManager, set the _running to False
+        try:
+            await self.model_manager.stop()
+        finally:
+            self._running = False
+
+    async def __aenter__(self):
+        """Context manager (used for testing rather than long-lived instance)"""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Context manager (used for testing rather than long-lived instance)"""
+        await self.stop()
+
     async def generate_async(self, messages: LLMMessages, **kwargs) -> LLMMessage:
         """Run input rails, generation, and output rails. Return response if safe."""
 
@@ -53,7 +87,14 @@ async def generate_async(self, messages: LLMMessages, **kwargs) -> LLMMessage:
             return {"role": "assistant", "content": REFUSAL_MESSAGE}
 
         # Step 2: Generate response from main LLM
-        response_text = await self.model_manager.generate_async("main", messages)
+        # If we got an `options=GenerationOptions`, then unpack GenerationOptions.llm_params and add
+        # that to the main LLM call
+        llm_kwargs = {}
+        if kwargs.get("options") and isinstance(kwargs["options"], GenerationOptions):
+            generation_options = kwargs["options"]
+            llm_kwargs = generation_options.llm_params if generation_options.llm_params else {}
+
+        response_text = await self.model_manager.generate_async("main", messages, **llm_kwargs)
 
         # Step 3: Check output rails
         output_result = await self.rails_manager.is_output_safe(messages, response_text)

nemoguardrails/guardrails/model_engine.py

@@ -23,7 +23,6 @@
 import logging
 import os
 from typing import Any, Optional, cast
-from urllib.parse import urljoin
 
 import aiohttp
 from aiohttp_retry import ExponentialRetry, RetryClient
@@ -83,20 +82,30 @@ def __init__(self, model_config: Model) -> None:
             exceptions={aiohttp.ClientConnectionError},
         )
         self._client: Optional[RetryClient] = None
+        self._running = False
 
     async def start(self) -> None:
-        """Create this engine's RetryClient."""
-        if self._client is None:
-            self._client = RetryClient(
-                retry_options=self._retry_options,
-                client_session=aiohttp.ClientSession(timeout=self._timeout),
-            )
+        """Create this engine's RetryClient. Call this during service startup."""
+        if self._running:
+            return
+
+        self._client = RetryClient(
+            retry_options=self._retry_options,
+            client_session=aiohttp.ClientSession(timeout=self._timeout),
+        )
+        self._running = True
 
     async def stop(self) -> None:
-        """Close this engine's RetryClient."""
-        if self._client:
-            await self._client.close()
-            self._client = None
+        """Close this engine's RetryClient. Call this during service shutdown."""
+        if not self._running:
+            return
+
+        try:
+            if self._client:
+                await self._client.close()
+                self._client = None
+        finally:
+            self._running = False
 
     def _resolve_base_url(self) -> str:
         """Resolve the base URL from model parameters or engine type."""
@@ -116,9 +125,7 @@ def _resolve_base_url(self) -> str:
     def _get_environment_variable(self, variable_name: str) -> str | None:
         """Return the value stored in environment variable `variable_name`."""
         env_value = os.environ.get(variable_name)
-        if env_value:
-            return env_value
-        return None
+        return env_value
 
     def _resolve_api_key(self, engine: str | None) -> Optional[str]:
         """Resolve the API key from model config or environment."""
@@ -159,17 +166,17 @@ async def call(
             The parsed JSON response dict from the API.
 
         Raises:
-            ModelEngineError: If the request fails after all retries or the client is not started.
+            ModelEngineError: If the request fails after all retries.
         """
 
         # Lazy-initialize client if `start()` hasn't yet been called.
-        if self._client is None:
+        if not self._running:
             await self.start()
 
         # Cast as RetryClient so type-checking knows it isn't None
         client = cast(RetryClient, self._client)
 
-        url = urljoin(self.base_url, _CHAT_COMPLETIONS_ENDPOINT)
+        url = self.base_url.rstrip("/") + _CHAT_COMPLETIONS_ENDPOINT
 
         headers: dict[str, str] = {"Content-Type": "application/json"}
         if self.api_key:
@@ -200,6 +207,15 @@ async def call(
                 model_name=self.model_name,
             ) from exc
 
+    async def __aenter__(self):
+        """Context manager (used for testing rather than long-lived instance)"""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Context manager (used for testing rather than long-lived instance)"""
+        await self.stop()
+
 
 async def _safe_read_body(response: aiohttp.ClientResponse) -> str:
     """Read response body for error messages, truncating if too large."""