[Iguazio] Add retry on session verification (#5600)

mlrun/utils/async_http.py

@@ -24,7 +24,7 @@
 from aiohttp_retry.client import _RequestContext
 
 from mlrun.config import config
-from mlrun.errors import err_to_str
+from mlrun.errors import err_to_str, raise_for_status
 
 from .helpers import logger as mlrun_logger
 
@@ -46,12 +46,21 @@ def __init__(
         *args,
         **kwargs,
     ):
+        # do not retry on PUT / PATCH as they might have side effects (not truly idempotent)
+        blacklisted_methods = (
+            blacklisted_methods
+            if blacklisted_methods is not None
+            else [
+                "POST",
+                "PUT",
+                "PATCH",
+            ]
+        )
         super().__init__(
             *args,
             retry_options=ExponentialRetryOverride(
                 retry_on_exception=retry_on_exception,
-                # do not retry on PUT / PATCH as they might have side effects (not truly idempotent)
-                blacklisted_methods=blacklisted_methods or ["POST", "PUT", "PATCH"],
+                blacklisted_methods=blacklisted_methods,
                 attempts=max_retries,
                 statuses=retry_on_status_codes,
                 factor=retry_backoff_factor,
@@ -63,6 +72,12 @@ def __init__(
             **kwargs,
         )
 
+    def methods_blacklist_update_required(self, new_blacklist: str):
+        self._retry_options: ExponentialRetryOverride
+        return set(self._retry_options.blacklisted_methods).difference(
+            set(new_blacklist)
+        )
+
     def _make_requests(
         self,
         params_list: list[RequestParams],
@@ -173,7 +188,7 @@ async def _do_request(self) -> aiohttp.ClientResponse:
                 last_attempt = current_attempt == self._retry_options.attempts
                 if self._is_status_code_ok(response.status) or last_attempt:
                     if self._raise_for_status:
-                        response.raise_for_status()
+                        raise_for_status(response)
 
                     self._response = response
                     return response
@@ -275,6 +290,11 @@ def verify_exception_type(self, exc):
                 if isinstance(exc.os_error, exc_type):
                     return
         if exc.__cause__:
-            return self.verify_exception_type(exc.__cause__)
+            # If the cause exception is retriable, return, otherwise, raise the original exception
+            try:
+                self.verify_exception_type(exc.__cause__)
+            except Exception:
+                raise exc
+            return
         else:
             raise exc

server/api/utils/clients/iguazio.py

@@ -974,23 +974,32 @@ async def verify_request_session(
             "cookie": request.headers.get("cookie"),
             "x-request-id": request.state.request_id,
         }
-        async with self._send_request_to_api_async(
-            "POST",
-            mlrun.mlconf.httpdb.authentication.iguazio.session_verification_endpoint,
-            "Failed verifying iguazio session",
-            headers=headers,
-        ) as response:
+        async with (
+            self._send_request_to_api_async(
+                "POST",
+                mlrun.mlconf.httpdb.authentication.iguazio.session_verification_endpoint,
+                "Failed verifying iguazio session",
+                blacklisted_methods=[],  # iguazio session verification endpoint is idempotent
+                headers=headers,
+            ) as response
+        ):
             return self._generate_auth_info_from_session_verification_response(
                 response.headers, await response.json()
             )
 
     @contextlib.asynccontextmanager
     async def _send_request_to_api_async(
-        self, method, path: str, error_message: str, session=None, **kwargs
+        self,
+        method,
+        path: str,
+        error_message: str,
+        session=None,
+        blacklisted_methods=None,
+        **kwargs,
     ) -> aiohttp.ClientResponse:
         url = f"{self._api_url}/api/{path}"
         self._prepare_request_kwargs(session, path, kwargs=kwargs)
-        await self._ensure_async_session()
+        await self._ensure_async_session(blacklisted_methods)
         response = None
         try:
             response = await self._async_session.request(
@@ -1009,10 +1018,16 @@ async def _send_request_to_api_async(
             if response:
                 response.release()
 
-    async def _ensure_async_session(self):
-        if not self._async_session:
+    async def _ensure_async_session(self, blacklisted_methods=None):
+        if (
+            not self._async_session
+            or self._async_session.methods_blacklist_update_required(
+                blacklisted_methods
+            )
+        ):
             self._async_session = mlrun.utils.AsyncClientWithRetry(
                 retry_on_exception=mlrun.mlconf.httpdb.projects.retry_leader_request_on_exception
                 == mlrun.common.schemas.HTTPSessionRetryMode.enabled.value,
                 logger=logger,
+                blacklisted_methods=blacklisted_methods,
             )