Improve retrying downloads code and testing

dandi/consts.py

@@ -141,7 +141,7 @@ def urls(self) -> Iterator[str]:
         "linc-staging",
         "https://staging.lincbrain.org",
         "https://staging-api.lincbrain.org/api",
-    )
+    ),
 }
 # to map back url: name
 known_instances_rev = {
@@ -160,7 +160,7 @@ def urls(self) -> Iterator[str]:
 
 #: HTTP response status codes that should always be retried (until we run out
 #: of retries)
-RETRY_STATUSES = (500, 502, 503, 504)
+RETRY_STATUSES = (429, 500, 502, 503, 504)
 
 VIDEO_FILE_EXTENSIONS = [".mp4", ".avi", ".wmv", ".mov", ".flv", ".mkv"]
 VIDEO_FILE_MODULES = ["processing", "acquisition"]

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,10 +1086,10 @@
     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 = random.random() * 5 * attempt
+    sleep_amount: float = -1.0
     if os.environ.get("DANDI_DOWNLOAD_AGGRESSIVE_RETRY"):
         # in such a case if we downloaded a little more --
         # consider it a successful attempt
@@ -1122,19 +1123,36 @@
             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",
@@ -1143,17 +1161,9 @@
                 sleep_amount,
                 exc,
             )
-        else:
-            lgr.debug(
-                "%s - download failed on attempt #%d: %s, will sleep a bit and retry",
-                path,
-                attempt,
-                exc,
-            )
-    # if is_access_denied(exc) or attempt >= 2:
-    #     raise
-    # sleep a little and retry
-    else:
+    if sleep_amount < 0:
+        # it was not Retry-after set, so we come up with random duration to sleep
+        sleep_amount = random.random() * 5 * attempt
         lgr.debug(
             "%s - download failed on attempt #%d: %s, will sleep a bit and retry",
             path,

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