[ET-VK][CI] Retry rate-limited HF snapshot downloads

.ci/scripts/tests/vulkan/test_export_vulkan_llm.py

@@ -13,7 +13,11 @@
 import shlex
 import subprocess
 import sys
+import time
+from datetime import datetime, timezone
+from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
+from threading import Thread
 from types import SimpleNamespace
 
 import pytest
@@ -128,6 +132,158 @@ def test_existing_hf_download_accepts_matching_hash(tmp_path, monkeypatch):
     ) == str(tmp_path)
 
 
+@pytest.mark.parametrize(
+    "headers, expected_wait",
+    [
+        (
+            {
+                "RateLimit": '"api";r=0;t=36',
+                "RateLimit-Policy": '"fixed window";"api";q=500;w=300',
+            },
+            37,
+        ),
+        ({"Retry-After": "36"}, 37),
+        ({"Retry-After": "Mon, 19 Oct 2026 17:50:36 GMT"}, 37),
+        ({"RateLimit": '"api";r=0;t=36', "Retry-After": "50"}, 51),
+        ({"RateLimit": '"api";r=0;t=36', "Retry-After": "invalid"}, 37),
+        ({"RateLimit": '"api";r=10;t=900', "Retry-After": "36"}, 37),
+    ],
+)
+def test_hf_snapshot_retries_rate_limited_repo_tree(
+    tmp_path, monkeypatch, capsys, headers, expected_wait
+):
+    from huggingface_hub import HfApi
+
+    module = load_existing_module("examples/models/llama/hf_download.py")
+    requests = []
+    calls = []
+    waits = []
+    monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost")
+    monkeypatch.setattr(time, "sleep", waits.append)
+    monkeypatch.setattr(
+        time,
+        "time",
+        lambda: datetime(2026, 10, 19, 17, 50, tzinfo=timezone.utc).timestamp(),
+    )
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            requests.append(self.path)
+            self.send_response(429 if len(requests) == 1 else 200)
+            if len(requests) == 1:
+                for name, value in headers.items():
+                    self.send_header(name, value)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(
+                b'{"error": "Rate limited"}' if len(requests) == 1 else b"[]"
+            )
+
+        def log_message(self, *args):
+            pass
+
+    server = HTTPServer(("127.0.0.1", 0), Handler)
+    thread = Thread(target=server.serve_forever)
+    thread.start()
+    api = HfApi(endpoint=f"http://127.0.0.1:{server.server_port}")
+
+    def snapshot(**kwargs):
+        calls.append(kwargs)
+        list(
+            api.list_repo_tree(
+                kwargs["repo_id"],
+                revision=kwargs["revision"],
+                token=kwargs["token"],
+                recursive=True,
+            )
+        )
+        return str(tmp_path)
+
+    monkeypatch.setattr(module, "snapshot_download", snapshot)
+    try:
+        assert module.download_hf_snapshot("test/model", "r" * 40) == str(tmp_path)
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()
+    assert len(requests) == 2 and requests[0] == requests[1]
+    assert (
+        calls
+        == [
+            {
+                "repo_id": "test/model",
+                "revision": "r" * 40,
+                "allow_patterns": ["*.safetensors", "*.json"],
+                "token": False,
+            }
+        ]
+        * 2
+    )
+    assert waits == [expected_wait]
+    output = capsys.readouterr()
+    assert output.out == ""
+    assert str(expected_wait) in output.err
+
+
+@pytest.mark.parametrize(
+    "headers", [{}, {"Retry-After": "invalid"}, {"RateLimit": '"api";r=10;t=900'}]
+)
+def test_hf_snapshot_rate_limit_retry_exhaustion(monkeypatch, headers):
+    import httpx
+    from huggingface_hub.errors import HfHubHTTPError
+
+    module = load_existing_module("examples/models/llama/hf_download.py")
+    response = httpx.Response(
+        429,
+        headers=headers,
+        request=httpx.Request("GET", "https://example.com"),
+    )
+    error = HfHubHTTPError("Rate limited", response=response)
+    calls = []
+    waits = []
+
+    def rate_limited(**kwargs):
+        calls.append(kwargs)
+        raise error
+
+    monkeypatch.setattr(module, "snapshot_download", rate_limited)
+    monkeypatch.setattr(time, "sleep", waits.append)
+    with pytest.raises(HfHubHTTPError) as raised:
+        module.download_hf_snapshot("test/model", "r" * 40)
+    assert raised.value is error
+    assert len(calls) == 6
+    assert waits == [60, 120, 240, 300, 300]
+
+
+@pytest.mark.parametrize("status", [401, 404, 503, None])
+def test_hf_snapshot_does_not_retry_other_http_errors(monkeypatch, status):
+    import httpx
+    from huggingface_hub.errors import HfHubHTTPError
+
+    module = load_existing_module("examples/models/llama/hf_download.py")
+    response = httpx.Response(
+        status or 429,
+        request=httpx.Request("GET", "https://example.com"),
+    )
+    error = HfHubHTTPError("Download failed", response=response)
+    if status is None:
+        error.response = None
+    calls = []
+    waits = []
+
+    def failed(**kwargs):
+        calls.append(kwargs)
+        raise error
+
+    monkeypatch.setattr(module, "snapshot_download", failed)
+    monkeypatch.setattr(time, "sleep", waits.append)
+    with pytest.raises(HfHubHTTPError) as raised:
+        module.download_hf_snapshot("test/model", "r" * 40)
+    assert raised.value is error
+    assert len(calls) == 1
+    assert waits == []
+
+
 def test_snapshot_cli_exposes_pinned_download_and_checksum(
     tmp_path, monkeypatch, capsys
 ):

examples/models/llama/hf_download.py

@@ -9,14 +9,39 @@
 import hashlib
 import json
 import os
+import re
 import sys
+import time
 
 from contextlib import redirect_stdout
+from email.utils import parsedate_to_datetime
 from importlib.metadata import EntryPoint
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Mapping
 
 from huggingface_hub import snapshot_download
+from huggingface_hub.errors import HfHubHTTPError
+
+_MAX_RATE_LIMIT_RETRIES = 5
+
+
+def _rate_limit_retry_delay(headers: Mapping[str, str], attempt: int) -> float:
+    delay = None
+    retry_after = headers.get("Retry-After", "").strip()
+    if retry_after.isdigit():
+        delay = float(retry_after)
+    elif retry_after:
+        try:
+            delay = max(0, parsedate_to_datetime(retry_after).timestamp() - time.time())
+        except (TypeError, ValueError, OverflowError):
+            pass
+    rate_limit = re.search(r"\br=0;\s*t=(\d+)\b", headers.get("RateLimit", ""))
+    if rate_limit:
+        delay = max(delay or 0, float(rate_limit[1]))
+    if delay is not None:
+        # Quota reset times are rounded to whole seconds.
+        return delay + 1
+    return min(60 * 2**attempt, 300)
 
 
 def download_hf_snapshot(
@@ -26,12 +51,29 @@ def download_hf_snapshot(
     metadata_path: Path | None = None,
 ) -> str:
     """Download a pinned Safetensors checkpoint and optionally verify its hash."""
-    snapshot = snapshot_download(
-        repo_id=repo_id,
-        revision=revision,
-        allow_patterns=["*.safetensors", "*.json"],
-        token=False,
-    )
+    for attempt in range(_MAX_RATE_LIMIT_RETRIES + 1):
+        try:
+            snapshot = snapshot_download(
+                repo_id=repo_id,
+                revision=revision,
+                allow_patterns=["*.safetensors", "*.json"],
+                token=False,
+            )
+            break
+        except HfHubHTTPError as error:
+            if (
+                error.response is None
+                or error.response.status_code != 429
+                or attempt == _MAX_RATE_LIMIT_RETRIES
+            ):
+                raise
+            delay = _rate_limit_retry_delay(error.response.headers, attempt)
+            print(
+                f"Rate limited downloading {repo_id}; retrying in {delay:g}s "
+                f"({attempt + 1}/{_MAX_RATE_LIMIT_RETRIES})",
+                file=sys.stderr,
+            )
+            time.sleep(delay)
     if expected_sha256:
         checksum = hashlib.sha256()
         with (Path(snapshot) / "model.safetensors").open("rb") as checkpoint: