diff --git a/autotest/test_get_modflow.py b/autotest/test_get_modflow.py index 103e573d8..764fa1f3e 100644 --- a/autotest/test_get_modflow.py +++ b/autotest/test_get_modflow.py @@ -1,11 +1,13 @@ """Test get-modflow utility.""" +import io import os import sys +import urllib.request from os.path import expandvars from pathlib import Path from platform import system -from urllib.error import HTTPError +from urllib.error import HTTPError, URLError import pytest from flaky import flaky @@ -13,8 +15,8 @@ from modflow_devtools.misc import run_py_script from autotest.conftest import get_project_root_path -from flopy.utils import get_modflow -from flopy.utils.get_modflow import get_release, get_releases, select_bindir +from flopy.utils import get_modflow_module as get_modflow +from flopy.utils.get_modflow import get_release, get_releases, run_main, select_bindir rate_limit_msg = "rate limit exceeded" flopy_dir = get_project_root_path() @@ -81,6 +83,13 @@ def create_home_local_bin(): home_local.mkdir(parents=True, exist_ok=True) +@pytest.fixture(autouse=True) +def fast_retries(monkeypatch): + # make retries instant + monkeypatch.setattr(get_modflow, "http_retry_delay", 0.0) + monkeypatch.setenv("GET_MODFLOW_RETRY_DELAY", "0") + + def run_get_modflow_script(*args): return run_py_script(get_modflow_script, *args, verbose=True) @@ -100,6 +109,187 @@ def test_get_releases_bad_page_size(per_page): get_releases(repo="executables", per_page=per_page) +class FakeResponse(io.BytesIO): + def __init__(self, data=b"", headers=None): + super().__init__(data) + self.headers = headers or {} + + +def fake_urlopen(fail_times, exc, data=b"{}", headers=None): + # raises exc the first fail_times calls, then returns a fresh FakeResponse; + # call count is tracked on .calls + state = {"n": 0} + + def _fake(request, timeout=10, quiet=False): + state["n"] += 1 + if state["n"] <= fail_times: + raise exc + return FakeResponse(data, headers) + + _fake.calls = state + return _fake + + +def reset_error(): + return URLError(ConnectionResetError(104, "Connection reset by peer")) + + +def test_fetch_returns_body_and_headers(monkeypatch): + fake = fake_urlopen( + 0, None, data=b'{"a": 1}', headers={"x-ratelimit-remaining": "42"} + ) + monkeypatch.setattr(get_modflow, "urlopen", fake) + body, headers = get_modflow.fetch(object()) + assert body == b'{"a": 1}' + assert headers.get("x-ratelimit-remaining") == "42" + + +def test_fetch_retries_transient_then_succeeds(monkeypatch): + fake = fake_urlopen(2, reset_error(), data=b"OK") + monkeypatch.setattr(get_modflow, "urlopen", fake) + + body, _ = get_modflow.fetch(object(), tries=3, delay=0) + + assert body == b"OK" + assert fake.calls["n"] == 3 + + +def test_fetch_default_tries_from_module(monkeypatch): + # no tries= arg: falls back to the module-level max_http_tries + fake = fake_urlopen(99, reset_error()) + monkeypatch.setattr(get_modflow, "urlopen", fake) + monkeypatch.setattr(get_modflow, "max_http_tries", 4) + + with pytest.raises(URLError): + get_modflow.fetch(object(), delay=0) + assert fake.calls["n"] == 4 + + +def test_fetch_reraises_after_cap(monkeypatch): + fake = fake_urlopen(99, reset_error()) + monkeypatch.setattr(get_modflow, "urlopen", fake) + + with pytest.raises(URLError): + get_modflow.fetch(object(), tries=3, delay=0) + assert fake.calls["n"] == 3 + + +def test_fetch_tries_1_disables_retry(monkeypatch): + fake = fake_urlopen(99, reset_error()) + monkeypatch.setattr(get_modflow, "urlopen", fake) + + with pytest.raises(URLError): + get_modflow.fetch(object(), tries=1) + assert fake.calls["n"] == 1 + + +def test_fetch_non_transient_raises_immediately(monkeypatch): + slept = [] + monkeypatch.setattr(get_modflow.time, "sleep", slept.append) + fake = fake_urlopen(99, HTTPError("u", 401, "Unauthorized", {}, None)) + monkeypatch.setattr(get_modflow, "urlopen", fake) + + with pytest.raises(HTTPError): + get_modflow.fetch(object(), tries=3, delay=0) + assert fake.calls["n"] == 1 + assert slept == [] + + +def test_fetch_404_not_retried(monkeypatch): + slept = [] + monkeypatch.setattr(get_modflow.time, "sleep", slept.append) + fake = fake_urlopen(99, HTTPError("u", 404, "Not Found", {}, None)) + monkeypatch.setattr(get_modflow, "urlopen", fake) + + with pytest.raises(HTTPError) as exc_info: + get_modflow.fetch(object(), tries=3, delay=0) + assert exc_info.value.code == 404 + assert fake.calls["n"] == 1 + assert slept == [] + + +def test_sleep_before_retry_backoff_and_cap(monkeypatch): + slept = [] + monkeypatch.setattr(get_modflow.time, "sleep", slept.append) + err = reset_error() + + get_modflow._sleep_before_retry(1, 3, 2.0, err, quiet=True) + get_modflow._sleep_before_retry(2, 3, 2.0, err, quiet=True) + get_modflow._sleep_before_retry(3, 3, 2.0, err, quiet=True) + assert slept == [2.0, 4.0, 8.0] # exponential + + # ceiling + get_modflow._sleep_before_retry(1, 3, 100.0, err, quiet=True) + assert slept[-1] == get_modflow._max_retry_delay + + # Retry-After header wins when larger + retry_after = HTTPError("u", 503, "err", {"Retry-After": "30"}, None) + get_modflow._sleep_before_retry(1, 3, 2.0, retry_after, quiet=True) + assert slept[-1] == 30.0 + + # zero delay stays zero + get_modflow._sleep_before_retry(5, 3, 0.0, err, quiet=True) + assert slept[-1] == 0.0 + + +def test_download_atomic_success(tmp_path, monkeypatch): + monkeypatch.setattr( + get_modflow, "urlopen", fake_urlopen(1, reset_error(), data=b"PK\x03\x04zip") + ) + dest = tmp_path / "asset.zip" + + request = urllib.request.Request("http://example/asset.zip") + get_modflow.download(request, dest, quiet=True, tries=3, delay=0) + + assert dest.read_bytes() == b"PK\x03\x04zip" + assert not (tmp_path / "asset.zip.part").exists() + + +def test_download_failure_preserves_existing(tmp_path, monkeypatch): + monkeypatch.setattr(get_modflow, "urlopen", fake_urlopen(99, reset_error())) + dest = tmp_path / "asset.zip" + dest.write_bytes(b"OLD-GOOD-CACHE") + + request = urllib.request.Request("http://example/asset.zip") + with pytest.raises(URLError): + get_modflow.download(request, dest, quiet=True, tries=3, delay=0) + + assert dest.read_bytes() == b"OLD-GOOD-CACHE" + assert not (tmp_path / "asset.zip.part").exists() + + +def test_cli_forwards_retry_flags_to_run_main(monkeypatch): + captured = {} + monkeypatch.setattr( + get_modflow, "run_main", lambda **kwargs: captured.update(kwargs) + ) + monkeypatch.setattr( + sys, + "argv", + ["get_modflow.py", str(Path.home()), "--retries", "9", "--retry-delay", "0.5"], + ) + tries_before = get_modflow.max_http_tries + + get_modflow.cli_main() + + assert captured["retries"] == 9 + assert captured["retry_delay"] == 0.5 + assert get_modflow.max_http_tries == tries_before + + +def test_cli_omitted_retry_flags_are_none(monkeypatch): + captured = {} + monkeypatch.setattr( + get_modflow, "run_main", lambda **kwargs: captured.update(kwargs) + ) + monkeypatch.setattr(sys, "argv", ["get_modflow.py", str(Path.home())]) + + get_modflow.cli_main() + + assert captured["retries"] is None + assert captured["retry_delay"] is None + + @flaky @requires_github @pytest.mark.parametrize("repo", repo_options.keys()) @@ -267,7 +457,7 @@ def test_script(function_tmpdir, owner, repo, downloads_dir): def test_python_api(function_tmpdir, owner, repo, downloads_dir): bindir = str(function_tmpdir) try: - get_modflow(bindir, owner=owner, repo=repo, downloads_dir=downloads_dir) + run_main(bindir, owner=owner, repo=repo, downloads_dir=downloads_dir) except HTTPError as err: if err.code == 403: pytest.skip(f"GitHub {rate_limit_msg}") diff --git a/flopy/utils/get_modflow.py b/flopy/utils/get_modflow.py index 73e791edf..f49551ba7 100755 --- a/flopy/utils/get_modflow.py +++ b/flopy/utils/get_modflow.py @@ -8,12 +8,14 @@ See https://developer.github.com/v3/repos/releases/ for GitHub Releases API. """ +import http.client import json import os import shutil import ssl import sys import tempfile +import time import urllib import urllib.request import warnings @@ -37,6 +39,29 @@ } available_repos = list(renamed_prefix.keys()) max_http_tries = 3 +http_retry_delay = 2.0 # base seconds between retries; grows exponentially +_max_retry_delay = 60.0 # ceiling on any single backoff sleep + +# HTTP status codes worth retrying +_retry_http_codes = {429, 500, 502, 503, 504} + + +def _env(name, default): + """ + Parse an env var as the type of ``default``. + Fall back to ``default`` on any error. + """ + val = os.environ.get(name) + if val is None: + return default + try: + return type(default)(val) + except (TypeError, ValueError): + return default + + +max_http_tries = max(1, _env("GET_MODFLOW_RETRIES", max_http_tries)) +http_retry_delay = max(0.0, _env("GET_MODFLOW_RETRY_DELAY", http_retry_delay)) # Check if this is running from flopy within_flopy = False @@ -119,7 +144,77 @@ def urlopen(request, timeout=10, quiet=False): return urllib.request.urlopen(request, timeout=timeout, context=context) -def get_releases(owner=None, repo=None, quiet=False, per_page=None) -> List[str]: +def _is_transient(err) -> bool: + """Whether an exception from urlopen + read is worth retrying.""" + if isinstance(err, urllib.error.HTTPError): + return err.code in _retry_http_codes + if isinstance(err, urllib.error.URLError): + # SSLCertVerificationError is handled inside urlopen(); any other + # transport-level failure (reset, DNS blip, timeout) is retryable. + return not isinstance(err.reason, ssl.SSLCertVerificationError) + return isinstance( + err, (http.client.IncompleteRead, ConnectionResetError, TimeoutError) + ) + + +def _retry_config(tries, delay): + """Fill unset retry knobs from the module defaults and clamp to sane bounds.""" + tries = max_http_tries if tries is None else tries + delay = http_retry_delay if delay is None else delay + return max(1, tries), max(0.0, delay) + + +def _sleep_before_retry(attempt, tries, delay, err, quiet): + """Sleep with exponential backoff, capped at _max_retry_delay.""" + secs = min(delay * (2 ** (attempt - 1)), _max_retry_delay) + retry_after = getattr(err, "headers", None) and err.headers.get("Retry-After") + if retry_after and str(retry_after).isdigit(): + secs = max(secs, int(retry_after)) + if not quiet: + print(f" attempt {attempt}/{tries} failed ({err}); retrying in {secs:.0f}s") + time.sleep(secs) + + +def fetch(request, timeout=10, quiet=False, tries=None, delay=None): + """Send a request, returning (response body, headers). Retries on failure.""" + tries, delay = _retry_config(tries, delay) + for attempt in range(1, tries + 1): + try: + with urlopen(request, timeout=timeout, quiet=quiet) as resp: + return resp.read(), resp.headers + except Exception as err: + if attempt == tries or not _is_transient(err): + raise + _sleep_before_retry(attempt, tries, delay, err, quiet) + + +def download( + request, dest_path, timeout=120, quiet=False, tries=None, delay=None +) -> None: + """Download (stream) a file to dest_path (atomic via a .part file).""" + tries, delay = _retry_config(tries, delay) + dest_path = Path(dest_path) + part_path = dest_path.with_name(dest_path.name + ".part") + for attempt in range(1, tries + 1): + try: + with urlopen(request, timeout=timeout, quiet=quiet) as resp: + with open(part_path, "wb") as f: + shutil.copyfileobj(resp, f) + os.replace(part_path, dest_path) + return + except Exception as err: + try: + part_path.unlink() + except FileNotFoundError: + pass + if attempt == tries or not _is_transient(err): + raise + _sleep_before_retry(attempt, tries, delay, err, quiet) + + +def get_releases( + owner=None, repo=None, quiet=False, per_page=None, tries=None, delay=None +) -> List[str]: """Get list of available releases.""" owner = default_owner if owner is None else owner repo = default_repo if repo is None else repo @@ -132,25 +227,16 @@ def get_releases(owner=None, repo=None, quiet=False, per_page=None) -> List[str] params["per_page"] = per_page request = get_request(req_url, params=params) - num_tries = 0 - while True: - num_tries += 1 - try: - with urlopen(request, timeout=10, quiet=quiet) as resp: - result = resp.read() - break - except urllib.error.HTTPError as err: - if err.code == 401 and os.environ.get("GITHUB_TOKEN"): - raise ValueError("GITHUB_TOKEN env is invalid") from err - elif err.code == 403 and "rate limit exceeded" in err.reason: - raise ValueError( - f"use GITHUB_TOKEN env to bypass rate limit ({err})" - ) from err - elif err.code in {404, 503} and num_tries < max_http_tries: - # GitHub sometimes returns this error for valid URLs, so retry - print(f"URL request {num_tries} did not work ({err})") - continue - raise RuntimeError(f"cannot retrieve data from {req_url}") from err + try: + result, _ = fetch(request, timeout=10, quiet=quiet, tries=tries, delay=delay) + except urllib.error.HTTPError as err: + if err.code == 401 and os.environ.get("GITHUB_TOKEN"): + raise ValueError("GITHUB_TOKEN env is invalid") from err + elif err.code == 403 and "rate limit exceeded" in err.reason: + raise ValueError( + f"use GITHUB_TOKEN env to bypass rate limit ({err})" + ) from err + raise RuntimeError(f"cannot retrieve data from {req_url}") from err releases = json.loads(result.decode()) if not quiet: @@ -161,7 +247,9 @@ def get_releases(owner=None, repo=None, quiet=False, per_page=None) -> List[str] return avail_releases -def get_release(owner=None, repo=None, tag="latest", quiet=False) -> dict: +def get_release( + owner=None, repo=None, tag="latest", quiet=False, tries=None, delay=None +) -> dict: """Get info about a particular release.""" owner = default_owner if owner is None else owner repo = default_repo if repo is None else repo @@ -172,40 +260,30 @@ def get_release(owner=None, repo=None, tag="latest", quiet=False) -> dict: else f"{api_url}/releases/tags/{tag}" ) request = get_request(req_url) - releases = None - num_tries = 0 - while True: - num_tries += 1 - try: - with urlopen(request, timeout=10, quiet=quiet) as resp: - result = resp.read() - remaining = resp.headers.get("x-ratelimit-remaining", None) - if remaining and int(remaining) <= 10: - warnings.warn( - f"Only {remaining} GitHub API requests remaining " - "before rate-limiting" - ) - break - except urllib.error.HTTPError as err: - if err.code == 401 and os.environ.get("GITHUB_TOKEN"): - raise ValueError("GITHUB_TOKEN env is invalid") from err - elif err.code == 403 and "rate limit exceeded" in err.reason: + try: + result, headers = fetch( + request, timeout=10, quiet=quiet, tries=tries, delay=delay + ) + remaining = headers.get("x-ratelimit-remaining", None) + if remaining and int(remaining) <= 10: + warnings.warn( + f"Only {remaining} GitHub API requests remaining before rate-limiting" + ) + except urllib.error.HTTPError as err: + if err.code == 401 and os.environ.get("GITHUB_TOKEN"): + raise ValueError("GITHUB_TOKEN env is invalid") from err + elif err.code == 403 and "rate limit exceeded" in err.reason: + raise ValueError( + f"use GITHUB_TOKEN env to bypass rate limit ({err})" + ) from err + elif err.code == 404: + releases = get_releases(owner, repo, quiet, tries=tries, delay=delay) + if tag not in releases: raise ValueError( - f"use GITHUB_TOKEN env to bypass rate limit ({err})" + f"Release {tag} not found (choose from {', '.join(releases)})" ) from err - elif err.code == 404: - if releases is None: - releases = get_releases(owner, repo, quiet) - if tag not in releases: - raise ValueError( - f"Release {tag} not found (choose from {', '.join(releases)})" - ) - elif err.code == 503 and num_tries < max_http_tries: - # GitHub sometimes returns this error for valid URLs, so retry - warnings.warn(f"URL request {num_tries} did not work ({err})") - continue - raise RuntimeError(f"cannot retrieve data from {req_url}") from err + raise RuntimeError(f"cannot retrieve data from {req_url}") from err release = json.loads(result.decode()) tag_name = release["tag_name"] @@ -314,6 +392,8 @@ def run_main( downloads_dir=None, force=False, quiet=False, + retries=None, + retry_delay=None, _is_cli=False, ): """Run main method to get MODFLOW and related programs. @@ -344,6 +424,12 @@ def run_main( previously downloaded in ``downloads_dir``. quiet : bool, default False If True, show fewer messages. + retries : int, optional + Attempts per network request. Defaults to the ``GET_MODFLOW_RETRIES`` + environment variable, or 3. + retry_delay : float, optional + Base seconds between retries, growing exponentially. Defaults to the + ``GET_MODFLOW_RETRY_DELAY`` environment variable, or 2.0. _is_cli : bool, default False Control behavior of method if this is run as a command-line interface or as a Python function. @@ -423,7 +509,9 @@ def run_main( raise KeyError(f"repo {repo!r} not supported; choose one of {available_repos}") # get the selected release - release = get_release(owner, repo, release_id, quiet) + release = get_release( + owner, repo, release_id, quiet, tries=retries, delay=retry_delay + ) assets = release.get("assets", []) for asset in assets: asset_name = asset["name"] @@ -464,11 +552,14 @@ def run_main( else: if not quiet: print(f"downloading '{download_url}' to '{download_pth}'") - with urlopen( - urllib.request.Request(download_url), timeout=120, quiet=quiet - ) as resp: - with open(download_pth, "wb") as f: - shutil.copyfileobj(resp, f) + download( + urllib.request.Request(download_url), + download_pth, + timeout=120, + quiet=quiet, + tries=retries, + delay=retry_delay, + ) if subset: if isinstance(subset, str): @@ -739,6 +830,20 @@ def cli_main(): "previously downloaded in downloads-dir.", ) parser.add_argument("--quiet", action="store_true", help="Show fewer messages.") + parser.add_argument( + "--retries", + type=int, + default=None, + help="Number of attempts per network request; default is " + f"{max_http_tries} (env: GET_MODFLOW_RETRIES). Use 1 to disable retries.", + ) + parser.add_argument( + "--retry-delay", + type=float, + default=None, + help="Base seconds between retries, grows exponentially; default is " + f"{http_retry_delay} (env: GET_MODFLOW_RETRY_DELAY).", + ) args = vars(parser.parse_args()) try: run_main(**args, _is_cli=True)