diff --git a/great_ai/remote/call_remote_great_ai.py b/great_ai/remote/call_remote_great_ai.py index 7c124a3..d0277a9 100644 --- a/great_ai/remote/call_remote_great_ai.py +++ b/great_ai/remote/call_remote_great_ai.py @@ -17,7 +17,7 @@ def call_remote_great_ai( base_uri: str, data: Mapping[str, Any], retry_count: int = 4, - timeout_in_seconds: Optional[int] = 300, + timeout_in_seconds: Optional[float] = 300, model_class: Optional[Type[T]] = None, ) -> Trace[T]: """Communicate with a GreatAI object through an HTTP request. diff --git a/great_ai/remote/call_remote_great_ai_async.py b/great_ai/remote/call_remote_great_ai_async.py index 39f0fab..587cef7 100644 --- a/great_ai/remote/call_remote_great_ai_async.py +++ b/great_ai/remote/call_remote_great_ai_async.py @@ -14,7 +14,7 @@ async def call_remote_great_ai_async( base_uri: str, data: Mapping[str, Any], retry_count: int = 4, - timeout_in_seconds: Optional[int] = 300, + timeout_in_seconds: Optional[float] = 300, model_class: Optional[Type[T]] = None, ) -> Trace[T]: """Communicate with a GreatAI object through an HTTP request. @@ -47,9 +47,12 @@ async def call_remote_great_ai_async( transport = httpx.AsyncHTTPTransport(retries=0) try: - async with httpx.AsyncClient( - transport=transport, timeout=timeout_in_seconds - ) as client: + async with ( + asyncio.timeout(timeout_in_seconds), + httpx.AsyncClient( + transport=transport, timeout=timeout_in_seconds + ) as client, + ): for attempt in range(retry_count + 1): try: response = await client.post(base_uri, json=data) @@ -64,6 +67,8 @@ async def call_remote_great_ai_async( raise RemoteCallError( f"Unexpected status code, reason: {response.text}" ) + except TimeoutError as e: + raise RemoteCallError("Remote prediction exceeded its overall timeout") from e except Exception as e: raise RemoteCallError(str(e)) from e diff --git a/tests/test_remote_regressions.py b/tests/test_remote_regressions.py index b605419..078fd74 100644 --- a/tests/test_remote_regressions.py +++ b/tests/test_remote_regressions.py @@ -1,9 +1,10 @@ -from typing import Any +import asyncio +import json +from typing import Any, AsyncIterator from unittest.mock import patch import httpx import pytest - from great_ai import RemoteCallError, call_remote_great_ai_async @@ -40,6 +41,29 @@ async def test_remote_call_retries_response_read_failures() -> None: assert calls == 2 +@pytest.mark.asyncio +async def test_overall_timeout_limits_slow_streaming_responses() -> None: + class SlowStream(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + payload = json.dumps(response_trace()).encode() + for offset in range(0, len(payload), 10): + await asyncio.sleep(0.01) + yield payload[offset : offset + 10] + + transport = httpx.MockTransport( + lambda request: httpx.Response(200, stream=SlowStream()) + ) + with patch( + "great_ai.remote.call_remote_great_ai_async.httpx.AsyncHTTPTransport", + return_value=transport, + ): + with pytest.raises(RemoteCallError, match="overall timeout") as error: + await call_remote_great_ai_async( + "http://test", {}, timeout_in_seconds=0.025 + ) + assert isinstance(error.value.__cause__, TimeoutError) + + @pytest.mark.asyncio async def test_remote_retry_limit_is_respected() -> None: calls = 0