diff --git a/src/guidellm/backend/backend.py b/src/guidellm/backend/backend.py index 12e30fa3..ff625432 100644 --- a/src/guidellm/backend/backend.py +++ b/src/guidellm/backend/backend.py @@ -110,6 +110,14 @@ def info(self) -> dict[str, Any]: """ ... + @abstractmethod + async def reset(self) -> None: + """ + Reset the connection object. This is useful for backends that + reuse connections or have state that needs to be cleared. + """ + ... + async def validate(self): """ Handle final setup and validate the backend is ready for use. @@ -126,6 +134,8 @@ async def validate(self): ): # type: ignore[attr-defined] pass + await self.reset() + @abstractmethod async def check_setup(self): """ diff --git a/src/guidellm/backend/openai.py b/src/guidellm/backend/openai.py index e3f23963..b5efc847 100644 --- a/src/guidellm/backend/openai.py +++ b/src/guidellm/backend/openai.py @@ -157,6 +157,15 @@ def info(self) -> dict[str, Any]: "chat_completions_path": CHAT_COMPLETIONS_PATH, } + async def reset(self) -> None: + """ + Reset the connection object. This is useful for backends that + reuse connections or have state that needs to be cleared. + For this backend, it closes the async client if it exists. + """ + if self._async_client is not None: + await self._async_client.aclose() + async def check_setup(self): """ Check if the backend is setup correctly and can be used for requests. @@ -361,7 +370,7 @@ def _get_async_client(self) -> httpx.AsyncClient: :return: The async HTTP client. """ - if self._async_client is None: + if self._async_client is None or self._async_client.is_closed: client = httpx.AsyncClient( http2=self.http2, timeout=self.timeout, diff --git a/tests/unit/mock_backend.py b/tests/unit/mock_backend.py index baf7bda8..27bfe382 100644 --- a/tests/unit/mock_backend.py +++ b/tests/unit/mock_backend.py @@ -41,6 +41,9 @@ def model(self) -> Optional[str]: def info(self) -> dict[str, Any]: return {} + async def reset(self) -> None: + pass + async def prepare_multiprocessing(self): pass