Skip to content

Commit

Permalink
[Iguazio] Retry on session verification (#5673)
Browse files Browse the repository at this point in the history
  • Loading branch information
alonmr committed Jun 2, 2024
1 parent 49822c8 commit f79df3c
Showing 1 changed file with 25 additions and 4 deletions.
29 changes: 25 additions & 4 deletions server/api/utils/clients/iguazio.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,11 +139,22 @@ class Client(
):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
retry_on_exception = (
mlrun.mlconf.httpdb.projects.retry_leader_request_on_exception
== mlrun.common.schemas.HTTPSessionRetryMode.enabled.value
)
self._session = mlrun.utils.HTTPSessionWithRetry(
retry_on_exception=mlrun.mlconf.httpdb.projects.retry_leader_request_on_exception
== mlrun.common.schemas.HTTPSessionRetryMode.enabled.value,
retry_on_exception=retry_on_exception,
verbose=True,
)
self._retry_on_post_session = None
if retry_on_exception:
self._retry_on_post_session = mlrun.utils.HTTPSessionWithRetry(
retry_on_exception=mlrun.mlconf.httpdb.projects.retry_leader_request_on_exception
== mlrun.common.schemas.HTTPSessionRetryMode.enabled.value,
retry_on_post=True,
verbose=True,
)
self._api_url = mlrun.mlconf.iguazio_api_url
# The job is expected to be completed in less than 5 seconds. If 10 seconds have passed and the job
# has not been completed, increase the interval to retry every 5 seconds
Expand Down Expand Up @@ -172,6 +183,7 @@ def verify_request_session(
"authorization": request.headers.get("authorization"),
"cookie": request.headers.get("cookie"),
},
retry_on_post=True,
)
return self._generate_auth_info_from_session_verification_response(
response.headers, response.json()
Expand Down Expand Up @@ -618,11 +630,20 @@ def _get_job_from_iguazio(self, session: str, job_id: str) -> dict:
return response.json()

def _send_request_to_api(
self, method, path, error_message: str, session=None, **kwargs
self,
method,
path,
error_message: str,
session=None,
retry_on_post=False,
**kwargs,
):
url = f"{self._api_url}/api/{path}"
self._prepare_request_kwargs(session, path, kwargs=kwargs)
response = self._session.request(
http_session = self._session
if retry_on_post and self._retry_on_post_session:
http_session = self._retry_on_post_session
response = http_session.request(
method, url, verify=mlrun.mlconf.httpdb.http.verify, **kwargs
)
if not response.ok:
Expand Down

0 comments on commit f79df3c

Please sign in to comment.