Use GitHub authentication and retries for metadata requests

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)

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()