Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Parse correctly the Retry-After header when it is of type http-date #118

Merged
merged 2 commits into from
Aug 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions firecrest/AsyncClient.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@
import firecrest.FirecrestException as fe
import firecrest.types as t
from firecrest.AsyncExternalStorage import AsyncExternalUpload, AsyncExternalDownload
from firecrest.utilities import time_block, slurm_state_completed
from firecrest.utilities import (
parse_retry_after, slurm_state_completed, time_block
)


if sys.version_info >= (3, 8):
Expand Down Expand Up @@ -127,9 +129,11 @@ async def wrapper(*args, **kwargs):
else:
reset = resp.headers.get(
"Retry-After",
default=resp.headers.get("RateLimit-Reset", default=10),
default=resp.headers.get(
"RateLimit-Reset", default=10
),
)
reset = int(reset)
reset = parse_retry_after(reset, client.log)
try:
f = kwargs["files"]["file"]
client.log(
Expand Down
8 changes: 5 additions & 3 deletions firecrest/BasicClient.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
import firecrest.FirecrestException as fe
import firecrest.types as t
from firecrest.ExternalStorage import ExternalUpload, ExternalDownload
from firecrest.utilities import time_block, slurm_state_completed
from firecrest.utilities import (parse_retry_after, slurm_state_completed, time_block)

if sys.version_info >= (3, 8):
from typing import Literal
Expand Down Expand Up @@ -89,14 +89,16 @@ def wrapper(*args, **kwargs):
else:
reset = resp.headers.get(
"Retry-After",
default=resp.headers.get("RateLimit-Reset", default=10),
default=resp.headers.get(
"RateLimit-Reset", default=10
),
)
reset = parse_retry_after(reset, client.log)
client.log(
logging.INFO,
f"Rate limit is reached, will sleep for "
f"{reset} seconds and try again"
)
reset = int(reset)
try:
f = kwargs["files"]["file"]
client.log(
Expand Down
30 changes: 30 additions & 0 deletions firecrest/utilities.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import email.utils as eut
import logging
import time
from contextlib import contextmanager

Expand Down Expand Up @@ -28,3 +30,31 @@ def slurm_state_completed(state):
return all(s in completion_states for s in state.split(','))

return False


def parse_retry_after(retry_after_header, log_func):
"""
Parse the Retry-After header.

:param retry_after_header: The value of the Retry-After header.
:return: A non-negative floating point number representing when the retry
should occur.
"""
try:
# Try to parse it as a delta-seconds
delta_seconds = int(retry_after_header)
return max(delta_seconds, 0)
except ValueError:
pass

try:
# Try to parse it as an HTTP-date
retry_after_date = eut.parsedate_to_datetime(retry_after_header)
delta_seconds = retry_after_date.timestamp() - time.time()
return max(delta_seconds, 0)
except Exception:
log_func(
logging.WARNING,
f"Could not parse Retry-After header: {retry_after_header}"
)
return 10
Loading