diff --git a/src/rapids_metadata/remote.py b/src/rapids_metadata/remote.py index b867f1a..da947ea 100644 --- a/src/rapids_metadata/remote.py +++ b/src/rapids_metadata/remote.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import time +import urllib.error import urllib.request +from typing import Optional import pydantic @@ -21,12 +25,45 @@ __all__ = ["fetch_latest"] _GITHUB_METADATA_URL = "https://raw.githubusercontent.com/rapidsai/rapids-metadata/main/rapids-metadata.json" +_GITHUB_API_METADATA_URL = "https://api.github.com/repos/rapidsai/rapids-metadata/contents/rapids-metadata.json" +_GITHUB_API_VERSION = "2026-03-10" +_MAX_ATTEMPTS = 3 +_RETRYABLE_HTTP_STATUS_CODES = {408, 429, 500, 502, 503, 504} -def _fetch_from_url(url: str) -> RAPIDSMetadata: - with urllib.request.urlopen(url) as f: - return pydantic.TypeAdapter(RAPIDSMetadata).validate_json(f.read()) +def _fetch_from_url( + url: str, *, headers: Optional[dict[str, str]] = None +) -> RAPIDSMetadata: + request = urllib.request.Request(url, headers=headers or {}) + for attempt in range(_MAX_ATTEMPTS): + try: + with urllib.request.urlopen(request) as f: + return pydantic.TypeAdapter(RAPIDSMetadata).validate_json(f.read()) + except urllib.error.HTTPError as error: + if ( + attempt == _MAX_ATTEMPTS - 1 + or error.code not in _RETRYABLE_HTTP_STATUS_CODES + ): + raise + default_delay = (60 if error.code == 429 else 1) * 2**attempt + delay = max(0, int(error.headers.get("Retry-After", default_delay))) + except urllib.error.URLError: + if attempt == _MAX_ATTEMPTS - 1: + raise + delay = 2**attempt + time.sleep(delay) + + raise AssertionError("unreachable") def fetch_latest() -> RAPIDSMetadata: + if token := os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN"): + return _fetch_from_url( + _GITHUB_API_METADATA_URL, + headers={ + "Accept": "application/vnd.github.raw+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": _GITHUB_API_VERSION, + }, + ) return _fetch_from_url(_GITHUB_METADATA_URL) diff --git a/tests/test_remote.py b/tests/test_remote.py index d77f488..e67e945 100644 --- a/tests/test_remote.py +++ b/tests/test_remote.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,8 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import patch +import io +import urllib.error +from unittest.mock import call, patch +import pytest import rapids_metadata.remote as rapids_remote from pydantic import TypeAdapter from pytest_httpserver import HTTPServer @@ -31,8 +34,109 @@ def test_fetch_from_url(httpserver: HTTPServer): ) -def test_fetch_latest(): +def test_fetch_latest_without_token(monkeypatch): + monkeypatch.delenv("GH_TOKEN", raising=False) + monkeypatch.delenv("GITHUB_TOKEN", raising=False) with patch("rapids_metadata.remote._fetch_from_url") as patch_fetch_from_url: - return_value = rapids_remote.fetch_latest() + rapids_remote.fetch_latest() patch_fetch_from_url.assert_called_once_with(rapids_remote._GITHUB_METADATA_URL) - assert return_value == patch_fetch_from_url() + + +@pytest.mark.parametrize( + ("environment", "expected_token"), + [ + ({"GITHUB_TOKEN": "github-token"}, "github-token"), + ( + {"GH_TOKEN": "gh-token", "GITHUB_TOKEN": "github-token"}, + "gh-token", + ), + ], +) +def test_fetch_latest_with_token(monkeypatch, environment, expected_token): + monkeypatch.delenv("GH_TOKEN", raising=False) + monkeypatch.delenv("GITHUB_TOKEN", raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + with patch("rapids_metadata.remote._fetch_from_url") as patch_fetch_from_url: + rapids_remote.fetch_latest() + + patch_fetch_from_url.assert_called_once_with( + rapids_remote._GITHUB_API_METADATA_URL, + headers={ + "Accept": "application/vnd.github.raw+json", + "Authorization": f"Bearer {expected_token}", + "X-GitHub-Api-Version": rapids_remote._GITHUB_API_VERSION, + }, + ) + + +def _metadata_response(): + return io.BytesIO(TypeAdapter(RAPIDSMetadata).dump_json(all_metadata)) + + +def _http_error(status, headers=None): + return urllib.error.HTTPError( + rapids_remote._GITHUB_METADATA_URL, + status, + "request failed", + headers or {}, + None, + ) + + +def test_fetch_retries_429(): + with patch("rapids_metadata.remote.urllib.request.urlopen") as patch_urlopen: + with patch("rapids_metadata.remote.time.sleep") as patch_sleep: + patch_urlopen.side_effect = [ + _http_error(429, {"Retry-After": "7"}), + _metadata_response(), + ] + + assert rapids_remote._fetch_from_url("https://example.com") == all_metadata + assert patch_urlopen.call_count == 2 + patch_sleep.assert_called_once_with(7) + + +@pytest.mark.parametrize( + "error", + [_http_error(503), urllib.error.URLError("connection reset")], +) +def test_fetch_retries_transient_error(error): + with patch("rapids_metadata.remote.urllib.request.urlopen") as patch_urlopen: + with patch("rapids_metadata.remote.time.sleep") as patch_sleep: + patch_urlopen.side_effect = [error, _metadata_response()] + + assert rapids_remote._fetch_from_url("https://example.com") == all_metadata + patch_sleep.assert_called_once_with(1) + + +def test_fetch_stops_after_max_attempts(): + errors = [ + _http_error(429, {"Retry-After": "0"}) + for _ in range(rapids_remote._MAX_ATTEMPTS) + ] + with patch("rapids_metadata.remote.urllib.request.urlopen") as patch_urlopen: + with patch("rapids_metadata.remote.time.sleep") as patch_sleep: + patch_urlopen.side_effect = errors + + with pytest.raises(urllib.error.HTTPError) as raised: + rapids_remote._fetch_from_url("https://example.com") + + assert raised.value is errors[-1] + assert patch_urlopen.call_count == rapids_remote._MAX_ATTEMPTS + assert patch_sleep.call_args_list == [call(0), call(0)] + + +def test_fetch_does_not_retry_nonretryable_error(): + error = _http_error(404) + with patch("rapids_metadata.remote.urllib.request.urlopen") as patch_urlopen: + with patch("rapids_metadata.remote.time.sleep") as patch_sleep: + patch_urlopen.side_effect = error + + with pytest.raises(urllib.error.HTTPError) as raised: + rapids_remote._fetch_from_url("https://example.com") + + assert raised.value is error + patch_urlopen.assert_called_once() + patch_sleep.assert_not_called()