client: Set default timeout to None (no timeout)

Set the default timeout for the GatewayClient to None, indicating no
timeout while waiting for task completion. A float value can still be
explicitly set to enforce a timeout.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -14,7 +14,7 @@ def __init__(
         base_url: str,
         default_model: str = None,
         polling_interval: float = 2.0,
-        timeout: float = 300.0,
+        timeout: float | None = None,
     ):
         """Initialize the GatewayClient with the base Gateway URL and optional parameters.
 
@@ -24,9 +24,9 @@ def __init__(
             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 300.0 seconds. A TimeoutError in this case should rarely indicate that
-                something went wrong, but rather that the task is taking longer than expected (e.g.
-                common with long running tasks like training).
+                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
@@ -248,7 +248,7 @@ async def wait_for_task(
                     error_message = task.get("error_message", "Unknown error")
                     raise RuntimeError(f"Task '{task_uuid}' failed: {error_message}")
                 return task
-            if asyncio.get_event_loop().time() - start > self.timeout:
+            if self.timeout is not None and asyncio.get_event_loop().time() - start > self.timeout:
                 raise TimeoutError(f"Timed out waiting for task '{task_uuid}' to complete")
             await asyncio.sleep(self.polling_interval)
 
@@ -350,7 +350,7 @@ def timeout(self):
         return self._client.timeout
 
     @timeout.setter
-    def timeout(self, value: float):
+    def timeout(self, value: float | None):
         self._client.timeout = value
 
     def submit_task(

tests/unit/client/test_client.py

@@ -42,6 +42,7 @@ async def test_gateway_client_init():
         polling_interval=10,
     )
     assert client.polling_interval == 3.0  # Maximum 3.0 seconds
+    assert client.timeout is None  # Default timeout should be None
 
     client.polling_interval = 0.05
     assert client.polling_interval == 0.5  # Minimum is 0.5 seconds
@@ -368,6 +369,35 @@ async def test_wait_for_task_timeout(mock_httpx_async_client, mocker):
         assert mock_get_task.await_count >= (client.timeout / client.polling_interval)
 
 
+@pytest.mark.asyncio
+async def test_wait_for_task_no_timeout(mock_httpx_async_client, mocker):
+    """Test wait_for_task doesn't timeout when timeout is None."""
+    call_count = 0
+
+    async def mock_get_task_side_effect(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count < 5:  # Return pending for first 4 calls
+            return {"uuid": "task-polling", "status": "pending"}
+        else:  # Return succeeded on 5th call
+            return {"uuid": "task-polling", "status": "succeeded"}
+
+    mock_get_task = AsyncMock(side_effect=mock_get_task_side_effect)
+    mocker.patch(
+        "client.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
+
+        result = await client.wait_for_task("task-polling")
+
+        assert result["status"] == "succeeded"
+        assert mock_get_task.await_count == 5
+
+
 @pytest.mark.asyncio
 async def test_wait_for_task_failed_raise_on_error(mock_httpx_async_client, mocker):
     """Test wait_for_task raises RuntimeError on task failure with raise_on_error."""