client: Make request timeout configurable

Signed-off-by: Phoevos Kalemkeris <phoevos.kalemkeris@ucl.ac.uk>

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -14,25 +14,30 @@ def __init__(
         self,
         base_url: str,
         default_model: str = None,
+        request_timeout: float = 300.0,
         polling_interval: float = 2.0,
-        timeout: float | None = None,
+        polling_timeout: float | None = None,
     ):
         """Initialize the GatewayClient with the base Gateway URL and optional parameters.
 
         Args:
             base_url (str): The base URL of the Gateway service.
             default_model (str, optional): The default model to use for tasks. Defaults to None.
+            request_timeout (float, optional): The HTTP request timeout in seconds for individual
+                requests to the Gateway. Defaults to 300.0 seconds to accommodate slower requests,
+                e.g. the ones triggering model auto-deployment, but should be adjusted as needed.
             polling_interval (float, optional): The interval in seconds to poll for task completion.
                 Defaults to 2.0 seconds, with a minimum of 0.5 and maximum of 3.0 seconds.
-            timeout (float, optional): The client polling timeout while waiting for task completion.
-                Defaults to None (no timeout). When set to a float value, a TimeoutError will be
-                raised if the task takes longer than the specified number of seconds. When None,
-                the client will wait indefinitely for task completion.
+            polling_timeout (float, optional): The client polling timeout while waiting for task
+                completion. Defaults to None (no timeout). When set to a float value, a TimeoutError
+                will be raised if the task takes longer than the specified number of seconds. When
+                None, the client will wait indefinitely for task completion.
         """
         self.base_url = base_url.rstrip("/")
         self.default_model = default_model
+        self.request_timeout = request_timeout
         self.polling_interval = polling_interval
-        self.timeout = timeout
+        self.polling_timeout = polling_timeout
         self._client = None
 
     @property
@@ -44,7 +49,7 @@ def polling_interval(self, value: float):
         self._polling_interval = max(0.5, min(value, 3.0))
 
     async def __aenter__(self):
-        self._client = httpx.AsyncClient()
+        self._client = httpx.AsyncClient(timeout=self.request_timeout)
         return self
 
     async def __aexit__(self, exc_type, exc, tb):
@@ -274,7 +279,10 @@ async def wait_for_task(
                     error_message = task.get("error_message", "Unknown error")
                     raise TaskFailedError(task_uuid, error_message, task)
                 return task
-            if self.timeout is not None and asyncio.get_event_loop().time() - start > self.timeout:
+            if (
+                self.polling_timeout is not None
+                and asyncio.get_event_loop().time() - start > self.polling_timeout
+            ):
                 raise TimeoutError(f"Timed out waiting for task '{task_uuid}' to complete")
             await asyncio.sleep(self.polling_interval)
 
@@ -469,6 +477,10 @@ def __del__(self):
     def base_url(self):
         return self._client.base_url
 
+    @base_url.setter
+    def base_url(self, value: str):
+        self._client.base_url = value
+
     @property
     def default_model(self):
         return self._client.default_model
@@ -477,6 +489,14 @@ def default_model(self):
     def default_model(self, value: str):
         self._client.default_model = value
 
+    @property
+    def request_timeout(self):
+        return self._client.request_timeout
+
+    @request_timeout.setter
+    def request_timeout(self, value: float | None):
+        self._client.request_timeout = value
+
     @property
     def polling_interval(self):
         return self._client.polling_interval
@@ -486,12 +506,12 @@ def polling_interval(self, value: float):
         self._client.polling_interval = value
 
     @property
-    def timeout(self):
-        return self._client.timeout
+    def polling_timeout(self):
+        return self._client.polling_timeout
 
-    @timeout.setter
-    def timeout(self, value: float | None):
-        self._client.timeout = value
+    @polling_timeout.setter
+    def polling_timeout(self, value: float | None):
+        self._client.polling_timeout = value
 
     def submit_task(
         self,

tests/unit/client/test_client.py

@@ -30,20 +30,23 @@ async def test_gateway_client_init():
         base_url="http://localhost:8888/",
         default_model="test-model",
         polling_interval=0.5,
-        timeout=0.1,
+        polling_timeout=0.1,
+        request_timeout=120.0,
     )
     assert client.base_url == "http://localhost:8888"
     assert client.default_model == "test-model"
     assert client.polling_interval == 0.5
-    assert client.timeout == 0.1
+    assert client.polling_timeout == 0.1
+    assert client.request_timeout == 120.0
     assert client._client is None
 
     client = GatewayClient(
         base_url="http://localhost:8888/",
         polling_interval=10,
     )
     assert client.polling_interval == 3.0  # Maximum 3.0 seconds
-    assert client.timeout is None  # Default timeout should be None
+    assert client.polling_timeout is None  # Default polling_timeout should be None
+    assert client.request_timeout == 300.0  # Default request_timeout
 
     client.polling_interval = 0.05
     assert client.polling_interval == 0.5  # Minimum is 0.5 seconds
@@ -377,16 +380,15 @@ async def test_wait_for_task_timeout(mock_httpx_async_client, mocker):
     mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
-    async with GatewayClient(base_url="http://test-gateway.com") as client:
-        client.timeout = 0.05
-        client.polling_interval = 0.5
-
+    async with GatewayClient(
+        base_url="http://test-gateway.com", polling_timeout=0.05, polling_interval=0.5
+    ) as client:
         with pytest.raises(
             TimeoutError, match="Timed out waiting for task 'task-polling' to complete"
         ):
             await client.wait_for_task("task-polling")
 
-        assert mock_get_task.await_count >= (client.timeout / client.polling_interval)
+        assert mock_get_task.await_count >= (client.polling_timeout / client.polling_interval)
 
 
 @pytest.mark.asyncio
@@ -406,10 +408,9 @@ async def mock_get_task_side_effect(*args, **kwargs):
     mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
-    async with GatewayClient(base_url="http://test-gateway.com") as client:
-        assert client.timeout is None
-        client.polling_interval = 0.01
-
+    async with GatewayClient(
+        base_url="http://test-gateway.com", polling_timeout=None, polling_interval=0.01
+    ) as client:
         result = await client.wait_for_task("task-polling")
 
         assert result["status"] == "succeeded"
@@ -937,7 +938,7 @@ def test_sync_client_is_healthy(mock_httpx_async_client):
 
 
 def test_sync_client_timeout_handling(mock_httpx_async_client):
-    """Test that timeouts work correctly in the sync client."""
+    """Test that polling_timeout works correctly in the sync client."""
     _, mock_client_instance = mock_httpx_async_client
 
     async def slow_response(*args, **kwargs):