[Iguazio] Add retry on session verification [1.6.x] (#5609)

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
 
@@ -48,12 +48,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,
@@ -65,6 +74,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],
@@ -175,7 +190,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
@@ -277,6 +292,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

@@ -975,37 +975,49 @@ async def verify_request_session(
         """
         Proxy the request to one of the session verification endpoints (which will verify the session of the request)
         """
-        async with self._send_request_to_api_async(
-            "POST",
-            mlrun.mlconf.httpdb.authentication.iguazio.session_verification_endpoint,
-            "Failed verifying iguazio session",
-            headers={
-                "authorization": request.headers.get("authorization"),
-                "cookie": request.headers.get("cookie"),
-            },
-        ) 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={
+                    "authorization": request.headers.get("authorization"),
+                    "cookie": request.headers.get("cookie"),
+                },
+            ) as response
+        ):
             return self._generate_auth_info_from_session_verification_response(
                 response.headers, await response.json()
             )
 
     async def verify_session(self, session: str) -> mlrun.common.schemas.AuthInfo:
-        async with self._send_request_to_api_async(
-            "POST",
-            mlrun.mlconf.httpdb.authentication.iguazio.session_verification_endpoint,
-            "Failed verifying iguazio session",
-            session,
-        ) as response:
+        async with (
+            self._send_request_to_api_async(
+                "POST",
+                mlrun.mlconf.httpdb.authentication.iguazio.session_verification_endpoint,
+                "Failed verifying iguazio session",
+                session,
+                blacklisted_methods=[],  # iguazio session verification endpoint is idempotent
+            ) 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, 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(
@@ -1024,10 +1036,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,
             )

tests/api/utils/clients/test_async_nuclio.py

@@ -14,7 +14,6 @@
 #
 import http
 
-import aiohttp
 import pytest
 from aioresponses import aioresponses as aioresponses_
 
@@ -89,5 +88,5 @@ async def test_nuclio_list_api_gateways(
     assert r == response_body
 
     mock_aioresponse.get(request_url, status=http.HTTPStatus.UNAUTHORIZED)
-    with pytest.raises(aiohttp.client_exceptions.ClientResponseError):
+    with pytest.raises(mlrun.errors.MLRunUnauthorizedError):
         await nuclio_client.list_api_gateways()