Skip to content

Commit b9f4837

Browse files
committed
[ET-VK][CI] Retry rate-limited HF snapshot downloads
Summary: The Vulkan LLM job failed on PR #23892 when Hugging Face returned HTTP 429 while listing the pinned Phi-4-mini checkpoint. Qwen and five Llama recipes passed. Hub repository pagination bypasses its existing HTTP backoff helper. Retry snapshot downloads up to five times for HTTP 429. Honor Retry-After (seconds or HTTP-date) and exhausted RateLimit quota resets, with a one-second rounding margin. Use bounded exponential waits when reset headers are absent or invalid. Preserve pinned download arguments, checksum validation, and path-only stdout; retry messages go to stderr. Other errors still fail promptly. Read hf_download.py first, then test_export_vulkan_llm.py. Regression coverage uses the real Hub repository-listing client against a local HTTP server that returns 429 followed by success, and checks retry exhaustion/error propagation. Authored with OpenAI Codex. Test Plan: ## Observed Results | Correctness validation | Environment | Result | |---|---|---| | New retry tests against parent helper | Python 3.11, Hub 1.33.0, local HTTP fixture | 9 expected failures, 4 passed | | Retry/header/error/checksum subset | Python 3.11, Hub 1.33.0 (failed CI version) | 16 passed, 0 failed, 13 deselected | | Same subset with older Hub | Python 3.11, Hub 0.36.2 | 16 passed, 0 failed, 13 deselected | | Lintrunner on both changed files | macOS | Passed | | Seven-model NVIDIA workflow with this fix | GitHub CI | Pending at publication | | Performance validation | Result | |---|---| | Runtime/operator performance and end-to-end A/B | Not run; download retry change, no performance claim | ## Reproduction Guide ### Checkout and dependency setup Install Git and Python 3.11. These focused HTTP/download tests need no GPU, model checkpoint, native build, or full ExecuTorch installation. ```bash git clone https://github.com/pytorch/executorch.git cd executorch git checkout ssjia/vulkan-hf-rate-limit python3.11 -m venv .venv-hf-retry .venv-hf-retry/bin/python -m pip install pytest pyyaml httpx==0.28.1 huggingface_hub==1.33.0 ``` ### Fixed implementation Run the retry tests and existing pinned-download/checksum guards. The conftest cutoff isolates HTTP tests from the repository-wide PyTorch fixture. A local HTTP server emits the protocol responses; sleeps are recorded without waiting. ```bash .venv-hf-retry/bin/python -m pytest --confcutdir=.ci/scripts/tests -q .ci/scripts/tests/vulkan/test_export_vulkan_llm.py -k 'rate_limit or other_http_errors or existing_hf_download' ``` ### Parent reproducer Copy the current tests into a temporary directory and retrieve the parent helper. The retry cases fail because the parent aborts on the first 429. This keeps the checked-out implementation intact. ```bash ET_HF_PARENT=$(mktemp -d) mkdir -p "$ET_HF_PARENT/.ci/scripts/tests/vulkan" "$ET_HF_PARENT/examples/models/llama" cp .ci/scripts/tests/vulkan/test_export_vulkan_llm.py "$ET_HF_PARENT/.ci/scripts/tests/vulkan/test_export_vulkan_llm.py" git show HEAD^:examples/models/llama/hf_download.py > "$ET_HF_PARENT/examples/models/llama/hf_download.py" .venv-hf-retry/bin/python -m pytest --confcutdir="$ET_HF_PARENT/.ci/scripts/tests" -q "$ET_HF_PARENT/.ci/scripts/tests/vulkan/test_export_vulkan_llm.py" -k 'rate_limit or other_http_errors' ``` ### Older Hub compatibility ```bash .venv-hf-retry/bin/python -m pip install huggingface_hub==0.36.2 .venv-hf-retry/bin/python -m pytest --confcutdir=.ci/scripts/tests -q .ci/scripts/tests/vulkan/test_export_vulkan_llm.py -k 'rate_limit or other_http_errors or existing_hf_download' ``` ### Lint ```bash python3.11 -m pip install lintrunner lintrunner-adapters lintrunner init lintrunner lint --output oneline examples/models/llama/hf_download.py .ci/scripts/tests/vulkan/test_export_vulkan_llm.py ``` ### GPU workflow The existing Vulkan workflow installs dependencies, builds the runner, then exports/tests all seven pinned recipes on NVIDIA. Check its job and download the vulkan-llm-e2e artifact to inspect results.tsv and per-model download logs. The regression's failing run is internet-accessible: https://github.com/pytorch/executorch/actions/runs/39289788119/job/118187613858 No runtime performance validation was run. Reviewers:
1 parent c8b72ff commit b9f4837

2 files changed

Lines changed: 205 additions & 7 deletions

File tree

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

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,11 @@
1313
import shlex
1414
import subprocess
1515
import sys
16+
import time
17+
from datetime import datetime, timezone
18+
from http.server import BaseHTTPRequestHandler, HTTPServer
1619
from pathlib import Path
20+
from threading import Thread
1721
from types import SimpleNamespace
1822

1923
import pytest
@@ -128,6 +132,158 @@ def test_existing_hf_download_accepts_matching_hash(tmp_path, monkeypatch):
128132
) == str(tmp_path)
129133

130134

135+
@pytest.mark.parametrize(
136+
"headers, expected_wait",
137+
[
138+
(
139+
{
140+
"RateLimit": '"api";r=0;t=36',
141+
"RateLimit-Policy": '"fixed window";"api";q=500;w=300',
142+
},
143+
37,
144+
),
145+
({"Retry-After": "36"}, 37),
146+
({"Retry-After": "Mon, 19 Oct 2026 17:50:36 GMT"}, 37),
147+
({"RateLimit": '"api";r=0;t=36', "Retry-After": "50"}, 51),
148+
({"RateLimit": '"api";r=0;t=36', "Retry-After": "invalid"}, 37),
149+
({"RateLimit": '"api";r=10;t=900', "Retry-After": "36"}, 37),
150+
],
151+
)
152+
def test_hf_snapshot_retries_rate_limited_repo_tree(
153+
tmp_path, monkeypatch, capsys, headers, expected_wait
154+
):
155+
from huggingface_hub import HfApi
156+
157+
module = load_existing_module("examples/models/llama/hf_download.py")
158+
requests = []
159+
calls = []
160+
waits = []
161+
monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost")
162+
monkeypatch.setattr(time, "sleep", waits.append)
163+
monkeypatch.setattr(
164+
time,
165+
"time",
166+
lambda: datetime(2026, 10, 19, 17, 50, tzinfo=timezone.utc).timestamp(),
167+
)
168+
169+
class Handler(BaseHTTPRequestHandler):
170+
def do_GET(self):
171+
requests.append(self.path)
172+
self.send_response(429 if len(requests) == 1 else 200)
173+
if len(requests) == 1:
174+
for name, value in headers.items():
175+
self.send_header(name, value)
176+
self.send_header("Content-Type", "application/json")
177+
self.end_headers()
178+
self.wfile.write(
179+
b'{"error": "Rate limited"}' if len(requests) == 1 else b"[]"
180+
)
181+
182+
def log_message(self, *args):
183+
pass
184+
185+
server = HTTPServer(("127.0.0.1", 0), Handler)
186+
thread = Thread(target=server.serve_forever)
187+
thread.start()
188+
api = HfApi(endpoint=f"http://127.0.0.1:{server.server_port}")
189+
190+
def snapshot(**kwargs):
191+
calls.append(kwargs)
192+
list(
193+
api.list_repo_tree(
194+
kwargs["repo_id"],
195+
revision=kwargs["revision"],
196+
token=kwargs["token"],
197+
recursive=True,
198+
)
199+
)
200+
return str(tmp_path)
201+
202+
monkeypatch.setattr(module, "snapshot_download", snapshot)
203+
try:
204+
assert module.download_hf_snapshot("test/model", "r" * 40) == str(tmp_path)
205+
finally:
206+
server.shutdown()
207+
server.server_close()
208+
thread.join()
209+
assert len(requests) == 2 and requests[0] == requests[1]
210+
assert (
211+
calls
212+
== [
213+
{
214+
"repo_id": "test/model",
215+
"revision": "r" * 40,
216+
"allow_patterns": ["*.safetensors", "*.json"],
217+
"token": False,
218+
}
219+
]
220+
* 2
221+
)
222+
assert waits == [expected_wait]
223+
output = capsys.readouterr()
224+
assert output.out == ""
225+
assert str(expected_wait) in output.err
226+
227+
228+
@pytest.mark.parametrize(
229+
"headers", [{}, {"Retry-After": "invalid"}, {"RateLimit": '"api";r=10;t=900'}]
230+
)
231+
def test_hf_snapshot_rate_limit_retry_exhaustion(monkeypatch, headers):
232+
import httpx
233+
from huggingface_hub.errors import HfHubHTTPError
234+
235+
module = load_existing_module("examples/models/llama/hf_download.py")
236+
response = httpx.Response(
237+
429,
238+
headers=headers,
239+
request=httpx.Request("GET", "https://huggingface.co/api/models/test/model"),
240+
)
241+
error = HfHubHTTPError("Rate limited", response=response)
242+
calls = []
243+
waits = []
244+
245+
def rate_limited(**kwargs):
246+
calls.append(kwargs)
247+
raise error
248+
249+
monkeypatch.setattr(module, "snapshot_download", rate_limited)
250+
monkeypatch.setattr(time, "sleep", waits.append)
251+
with pytest.raises(HfHubHTTPError) as raised:
252+
module.download_hf_snapshot("test/model", "r" * 40)
253+
assert raised.value is error
254+
assert len(calls) == 6
255+
assert waits == [60, 120, 240, 300, 300]
256+
257+
258+
@pytest.mark.parametrize("status", [401, 404, 503, None])
259+
def test_hf_snapshot_does_not_retry_other_http_errors(monkeypatch, status):
260+
import httpx
261+
from huggingface_hub.errors import HfHubHTTPError
262+
263+
module = load_existing_module("examples/models/llama/hf_download.py")
264+
response = httpx.Response(
265+
status or 429,
266+
request=httpx.Request("GET", "https://huggingface.co/api/models/test/model"),
267+
)
268+
error = HfHubHTTPError("Download failed", response=response)
269+
if status is None:
270+
error.response = None
271+
calls = []
272+
waits = []
273+
274+
def failed(**kwargs):
275+
calls.append(kwargs)
276+
raise error
277+
278+
monkeypatch.setattr(module, "snapshot_download", failed)
279+
monkeypatch.setattr(time, "sleep", waits.append)
280+
with pytest.raises(HfHubHTTPError) as raised:
281+
module.download_hf_snapshot("test/model", "r" * 40)
282+
assert raised.value is error
283+
assert len(calls) == 1
284+
assert waits == []
285+
286+
131287
def test_snapshot_cli_exposes_pinned_download_and_checksum(
132288
tmp_path, monkeypatch, capsys
133289
):

‎examples/models/llama/hf_download.py‎

Lines changed: 49 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,39 @@
99
import hashlib
1010
import json
1111
import os
12+
import re
1213
import sys
14+
import time
1315

1416
from contextlib import redirect_stdout
17+
from email.utils import parsedate_to_datetime
1518
from importlib.metadata import EntryPoint
1619
from pathlib import Path
17-
from typing import Callable
20+
from typing import Callable, Mapping
1821

1922
from huggingface_hub import snapshot_download
23+
from huggingface_hub.errors import HfHubHTTPError
24+
25+
_MAX_RATE_LIMIT_RETRIES = 5
26+
27+
28+
def _rate_limit_retry_delay(headers: Mapping[str, str], attempt: int) -> float:
29+
delay = None
30+
retry_after = headers.get("Retry-After", "").strip()
31+
if retry_after.isdigit():
32+
delay = float(retry_after)
33+
elif retry_after:
34+
try:
35+
delay = max(0, parsedate_to_datetime(retry_after).timestamp() - time.time())
36+
except (TypeError, ValueError, OverflowError):
37+
pass
38+
rate_limit = re.search(r"\br=0;\s*t=(\d+)\b", headers.get("RateLimit", ""))
39+
if rate_limit:
40+
delay = max(delay or 0, float(rate_limit[1]))
41+
if delay is not None:
42+
# Quota reset times are rounded to whole seconds.
43+
return delay + 1
44+
return min(60 * 2**attempt, 300)
2045

2146

2247
def download_hf_snapshot(
@@ -26,12 +51,29 @@ def download_hf_snapshot(
2651
metadata_path: Path | None = None,
2752
) -> str:
2853
"""Download a pinned Safetensors checkpoint and optionally verify its hash."""
29-
snapshot = snapshot_download(
30-
repo_id=repo_id,
31-
revision=revision,
32-
allow_patterns=["*.safetensors", "*.json"],
33-
token=False,
34-
)
54+
for attempt in range(_MAX_RATE_LIMIT_RETRIES + 1):
55+
try:
56+
snapshot = snapshot_download(
57+
repo_id=repo_id,
58+
revision=revision,
59+
allow_patterns=["*.safetensors", "*.json"],
60+
token=False,
61+
)
62+
break
63+
except HfHubHTTPError as error:
64+
if (
65+
error.response is None
66+
or error.response.status_code != 429
67+
or attempt == _MAX_RATE_LIMIT_RETRIES
68+
):
69+
raise
70+
delay = _rate_limit_retry_delay(error.response.headers, attempt)
71+
print(
72+
f"Rate limited downloading {repo_id}; retrying in {delay:g}s "
73+
f"({attempt + 1}/{_MAX_RATE_LIMIT_RETRIES})",
74+
file=sys.stderr,
75+
)
76+
time.sleep(delay)
3577
if expected_sha256:
3678
checksum = hashlib.sha256()
3779
with (Path(snapshot) / "model.safetensors").open("rb") as checkpoint:

0 commit comments

Comments
 (0)