client: Retry network errors

Add retry decorator for network-related exceptions in the client module,
including `RemoteProtocolError`, `ConnectError`, `TimeoutException`, and
`NetworkError`. Abstract HTTP request logic in the client module into a
single method and apply the retry decorator to it.

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

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -5,6 +5,8 @@
 
 import httpx
 
+from cogstack_model_gateway_client.exceptions import retry_if_network_error
+
 
 class GatewayClient:
     def __init__(
@@ -41,6 +43,33 @@ async def wrapper(self, *args, **kwargs):
 
         return wrapper
 
+    @require_client
+    @retry_if_network_error
+    async def _request(
+        self,
+        method: str,
+        url: str,
+        *,
+        data=None,
+        json=None,
+        files=None,
+        params=None,
+        headers=None,
+        **kwargs,
+    ) -> httpx.Response:
+        """Make HTTP requests with retry logic."""
+        resp = await self._client.request(
+            method=method,
+            url=url,
+            data=data,
+            json=json,
+            files=files,
+            params=params,
+            headers=headers,
+            **kwargs,
+        )
+        return resp.raise_for_status()
+
     @require_client
     async def submit_task(
         self,
@@ -59,11 +88,12 @@ async def submit_task(
         if not model_name:
             raise ValueError("Please provide a model name or set a default model for the client.")
         url = f"{self.base_url}/models/{model_name}/tasks/{task}"
-        resp = await self._client.post(
-            url, data=data, json=json, files=files, params=params, headers=headers
+
+        resp = await self._request(
+            "POST", url, data=data, json=json, files=files, params=params, headers=headers
         )
-        resp.raise_for_status()
         task_info = resp.json()
+
         if wait_for_completion:
             task_uuid = task_info["uuid"]
             task_info = await self.wait_for_task(task_uuid)
@@ -143,8 +173,7 @@ async def _get_task(self, task_uuid: str, detail: bool = True, download: bool =
         """Get a Gateway task."""
         url = f"{self.base_url}/tasks/{task_uuid}"
         params = {"detail": detail, "download": download}
-        resp = await self._client.get(url, params=params)
-        return resp.raise_for_status()
+        return await self._request("GET", url, params=params)
 
     @require_client
     async def get_task(self, task_uuid: str, detail: bool = True):
@@ -207,8 +236,7 @@ async def wait_for_task(
     async def get_models(self, verbose: bool = False):
         """Get the list of available models from the Gateway."""
         url = f"{self.base_url}/models/"
-        resp = await self._client.get(url, params={"verbose": verbose})
-        resp.raise_for_status()
+        resp = await self._request("GET", url, params={"verbose": verbose})
         return resp.json()
 
     @require_client
@@ -218,8 +246,7 @@ async def get_model(self, model_name: str = None):
         if not model_name:
             raise ValueError("Please provide a model name or set a default model for the client.")
         url = f"{self.base_url}/models/{model_name}/info"
-        resp = await self._client.get(url)
-        resp.raise_for_status()
+        resp = await self._request("GET", url)
         return resp.json()
 
     @require_client
@@ -236,8 +263,7 @@ async def deploy_model(
             raise ValueError("Please provide a model name or set a default model for the client.")
         url = f"{self.base_url}/models/{model_name}"
         data = {"tracking_id": tracking_id, "model_uri": model_uri, "ttl": ttl}
-        resp = await self._client.post(url, json=data)
-        resp.raise_for_status()
+        resp = await self._request("POST", url, json=data)
         return resp.json()
 
 

client/cogstack_model_gateway_client/exceptions.py

@@ -0,0 +1,31 @@
+import logging
+
+import httpx
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_fixed,
+)
+
+log = logging.getLogger("cmg.client")
+
+
+def is_network_error(exception: Exception):
+    """Check if the exception is a network-related error."""
+    return isinstance(
+        exception,
+        httpx.RemoteProtocolError
+        | httpx.ConnectError
+        | httpx.TimeoutException
+        | httpx.NetworkError,
+    )
+
+
+retry_if_network_error = retry(
+    stop=stop_after_attempt(3),
+    wait=wait_fixed(2),
+    retry=retry_if_exception(is_network_error),
+    before_sleep=before_sleep_log(log, logging.DEBUG),
+)

client/poetry.lock

@@ -35,6 +35,7 @@ pre-commit = "^4.2.0"
 pytest = "^8.4.1"
 pytest-asyncio = "^1.0.0"
 pytest-mock = "^3.14.1"
+tenacity = "^9.1.2"
 
 [tool.poetry.requires-plugins]
 poetry-dynamic-versioning = { version = ">=1.0.0,<2.0.0", extras = ["plugin"] }