From db21eab713a194d404f949bd69c3f276abd0f517 Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Mon, 12 May 2025 05:49:59 -0400 Subject: [PATCH] fix: catch TimeoutError in place of asyncio.TimeoutError (#2131) # What does this PR do? As per docs [1], since python 3.11 wait_for() raises TimeoutError. Since we currently support python 3.10+, we have to catch both. [1]: https://docs.python.org/3.12/library/asyncio-task.html#asyncio.wait_for [//]: # (If resolving an issue, uncomment and update the line below) [//]: # (Closes #[issue-number]) ## Test Plan No explicit testing; just code hardening to reflect docs. [//]: # (## Documentation) Signed-off-by: Ihar Hrachyshka --- llama_stack/distribution/providers.py | 2 +- llama_stack/distribution/routers/routers.py | 2 +- llama_stack/distribution/server/server.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/llama_stack/distribution/providers.py b/llama_stack/distribution/providers.py index 157fd1e0e..29b7109dd 100644 --- a/llama_stack/distribution/providers.py +++ b/llama_stack/distribution/providers.py @@ -99,7 +99,7 @@ class ProviderImpl(Providers): try: health = await asyncio.wait_for(impl.health(), timeout=timeout) return api_name, health - except asyncio.TimeoutError: + except (asyncio.TimeoutError, TimeoutError): return ( api_name, HealthResponse( diff --git a/llama_stack/distribution/routers/routers.py b/llama_stack/distribution/routers/routers.py index 7e80a067f..371d34904 100644 --- a/llama_stack/distribution/routers/routers.py +++ b/llama_stack/distribution/routers/routers.py @@ -630,7 +630,7 @@ class InferenceRouter(Inference): continue health = await asyncio.wait_for(impl.health(), timeout=timeout) health_statuses[provider_id] = health - except asyncio.TimeoutError: + except (asyncio.TimeoutError, TimeoutError): health_statuses[provider_id] = HealthResponse( status=HealthStatus.ERROR, message=f"Health check timed out after {timeout} seconds", diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index 85c576753..e34a62b00 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -114,7 +114,7 @@ def translate_exception(exc: Exception) -> HTTPException | RequestValidationErro return HTTPException(status_code=400, detail=str(exc)) elif isinstance(exc, PermissionError): return HTTPException(status_code=403, detail=f"Permission denied: {str(exc)}") - elif isinstance(exc, TimeoutError): + elif isinstance(exc, asyncio.TimeoutError | TimeoutError): return HTTPException(status_code=504, detail=f"Operation timed out: {str(exc)}") elif isinstance(exc, NotImplementedError): return HTTPException(status_code=501, detail=f"Not implemented: {str(exc)}") @@ -139,7 +139,7 @@ async def shutdown(app): await asyncio.wait_for(impl.shutdown(), timeout=5) else: logger.warning("No shutdown method for %s", impl_name) - except asyncio.TimeoutError: + except (asyncio.TimeoutError, TimeoutError): logger.exception("Shutdown timeout for %s ", impl_name, exc_info=True) except (Exception, asyncio.CancelledError) as e: logger.exception("Failed to shutdown %s: %s", impl_name, {e})