From 6b5737953cdd6a5c5c1afe85438104c59e1eff20 Mon Sep 17 00:00:00 2001 From: Yaroslav Halchenko Date: Fri, 20 Dec 2024 16:50:36 -0500 Subject: [PATCH] Handle Retry-After as a date, and add tests for all logic of checking attempts --- dandi/download.py | 48 ++++++++++++------ dandi/tests/test_download.py | 94 ++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 15 deletions(-) diff --git a/dandi/download.py b/dandi/download.py index 71c79a1c5..dda2703b5 100644 --- a/dandi/download.py +++ b/dandi/download.py @@ -3,7 +3,8 @@ from collections import Counter, deque from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import InitVar, dataclass, field -from datetime import datetime +from datetime import UTC, datetime +from email.utils import parsedate_to_datetime from enum import Enum from functools import partial import hashlib @@ -1085,7 +1086,7 @@ def _check_if_more_attempts_allowed( exc: requests.RequestException, attempt: int, attempts_allowed: int, - downloaded_in_attempt: int, + downloaded_in_attempt: int = 0, ) -> int | None: """Check if we should retry the download, return potentially adjusted 'attempts_allowed'""" sleep_amount: float = -1.0 @@ -1122,19 +1123,36 @@ def _check_if_more_attempts_allowed( return None elif retry_after := exc.response.headers.get("Retry-After"): # playing safe - if not str(retry_after).isdigit(): - # our code is wrong, do not crash but issue warning so - # we might get report/fix it up - lgr.warning( - "%s - download failed due to response %d with non-integer" - " Retry-After=%r: %s", - path, - exc.response.status_code, - retry_after, - exc, - ) - return None - sleep_amount = int(retry_after) + retry_after = str(retry_after) + if retry_after.isdigit(): + sleep_amount = int(retry_after) + else: + # else if it is a datestamp like "Wed, 21 Oct 2015 07:28:00 GMT" + # we could parse it and calculate how long to sleep + try: + retry_after_date = parsedate_to_datetime(retry_after) + current_date = datetime.now(UTC) + difference = retry_after_date - current_date + sleep_amount = int(difference.total_seconds()) + if sleep_amount < 0: + lgr.warning( + "%s - date in Retry-After=%r is in the past (current is %r)", + path, + retry_after, + current_date, + ) + except ValueError as exc_ve: + # our code or response is wrong, do not crash but issue warning + # and continue with "default" sleep logic + lgr.warning( + "%s - download failed due to response %d with non-integer or future date" + " Retry-After=%r: %s with %s upon converting assuming it is a date", + path, + exc.response.status_code, + retry_after, + exc, + exc_ve, + ) lgr.debug( "%s - download failed due to response %d with " "Retry-After=%d: %s, will sleep and retry", diff --git a/dandi/tests/test_download.py b/dandi/tests/test_download.py index fb9626d36..3d65e19a7 100644 --- a/dandi/tests/test_download.py +++ b/dandi/tests/test_download.py @@ -2,6 +2,8 @@ from collections.abc import Callable from contextlib import nullcontext +from email.utils import parsedate_to_datetime +from functools import partial from glob import glob import json import logging @@ -12,10 +14,13 @@ import re from shutil import rmtree import time +from unittest import mock import numpy as np import pytest from pytest_mock import MockerFixture +import requests +from requests.exceptions import HTTPError import responses import zarr @@ -32,6 +37,7 @@ PathType, ProgressCombiner, PYOUTHelper, + _check_if_more_attempts_allowed, download, ) from ..exceptions import NotFoundError @@ -1169,3 +1175,91 @@ def test_DownloadDirectory_exc( assert dl.dirpath.exists() assert dl.fp is None assert dl.writefile.read_bytes() == b"456" + + +def test__check_if_more_attempts_allowed(): + f = partial(_check_if_more_attempts_allowed, "some/path") + + response403 = requests.Response() + response403.status_code = 403 # no retry + + response500 = requests.Response() + response500.status_code = 500 + + # we do retry if cause is unknown (no response) + with mock.patch("time.sleep") as mock_sleep: + assert f(HTTPError(), attempt=1, attempts_allowed=2) == 2 + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] >= 0 + + # or if some 500 + with mock.patch("time.sleep") as mock_sleep: + assert f(HTTPError(response=response500), attempt=1, attempts_allowed=2) == 2 + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] >= 0 + + # do not bother if already at limit + with mock.patch("time.sleep") as mock_sleep: + assert f(HTTPError(), attempt=2, attempts_allowed=2) is None + mock_sleep.assert_not_called() + + # do not bother if 403 + with mock.patch("time.sleep") as mock_sleep: + assert f(HTTPError(response=response403), attempt=1, attempts_allowed=2) is None + mock_sleep.assert_not_called() + + # And in case of "Aggressive setting" when DANDI_DOWNLOAD_AGGRESSIVE_RETRY + # env var is set to 1, we retry if there was extra content downloaded + # patch env var DANDI_DOWNLOAD_AGGRESSIVE_RETRY + with mock.patch.dict(os.environ, {"DANDI_DOWNLOAD_AGGRESSIVE_RETRY": "1"}): + with mock.patch("time.sleep") as mock_sleep: + assert ( + f(HTTPError(), attempt=2, attempts_allowed=2, downloaded_in_attempt=0) + is None + ) + mock_sleep.assert_not_called() + + assert ( + f(HTTPError(), attempt=2, attempts_allowed=2, downloaded_in_attempt=1) + == 3 + ) + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] >= 0 + + +@pytest.mark.parametrize("status_code", [429, 503]) +def test__check_if_more_attempts_allowed_retries(status_code): + f = partial(_check_if_more_attempts_allowed, "some/path") + + response = requests.Response() + response.status_code = status_code + + response.headers["Retry-After"] = "10" + with mock.patch("time.sleep") as mock_sleep: + assert f(HTTPError(response=response), attempt=1, attempts_allowed=2) == 2 + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] == 10 + + response.headers["Retry-After"] = "Wed, 21 Oct 2015 07:28:00 GMT" + with mock.patch("time.sleep") as mock_sleep, mock.patch( + "dandi.download.datetime" + ) as mock_datetime: + # shifted by 2 minutes + mock_datetime.now.return_value = parsedate_to_datetime( + "Wed, 21 Oct 2015 07:26:00 GMT" + ) + assert f(HTTPError(response=response), attempt=1, attempts_allowed=2) == 2 + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] == 120 + + # we would still sleep some time if Retry-After is not decypherable + response.headers["Retry-After"] = "nondecypherable" + with mock.patch("time.sleep") as mock_sleep: + assert ( + _check_if_more_attempts_allowed( + "some/path", HTTPError(response=response), attempt=1, attempts_allowed=2 + ) + == 2 + ) + mock_sleep.assert_called_once() + assert mock_sleep.call_args.args[0] > 0