Bound the complete remote request and its retries by an overall deadline

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.

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
 

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