Add HuggingFace auth and 429 retry handling to HTTP downloads - #815
warmbowski wants to merge 5 commits into
Conversation
warmbowski
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ecb4e74ef4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| save_folder=opts.save_folder, | ||
| save_overwrite=True, | ||
| load_path="https://storage.googleapis.com/ai2-llm/checkpoints/OLMo25/step596047/", | ||
| load_path="https://huggingface.co/buckets/allenai/ai2-llm/resolve/checkpoints/OLMo25/step596047/", |
There was a problem hiding this comment.
For these new https://huggingface.co/... checkpoint paths, the first checkpoint reads still go through cached_path() before the new HTTP range helper is used: trainer state is loaded with cached_path(path) in Checkpoint.load(), and DCP .metadata is loaded via resource_path(...)/cached_path() in RemoteFileSystemReader.read_metadata(). Since _http_auth_headers() is only wired into _http_file_exists() and range reads, a cold-cache load from a private or rate-limited HF bucket still sends no HF_TOKEN for those metadata/trainer-state requests and can fail with 401/429 before tensor range reads begin; please pass the same headers into cached_path/resource_path or use an HF-aware scheme before switching these official load paths.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
I will look into this and see about making it work for cached/resource path.
There was a problem hiding this comment.
okay, looks like cached_path has a headers arg that only applies to http requests. This fits perfectly
headers :
Custom headers to add to HTTP requests.
Example: ``{"Authorization": "Bearer YOUR_TOKEN"}`` for private resources.
Only used for HTTP/HTTPS resources.
1b602e7 to
60c16ae
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 60c16ae07e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| else: | ||
| return cached_path(f"{folder}/{fname}", quiet=True) | ||
| url = f"{folder}/{fname}" | ||
| return cached_path(url, quiet=True, headers=_http_auth_headers(url)) |
There was a problem hiding this comment.
After this patch, the fresh evidence is that this new cached_path(...headers=...) path still bypasses the new _http_retry_condition; cached_path does the cold-cache HEAD/download for resource_path() before RemoteFileSystemReader reaches the range-read helper, and the trainer-state cached_path call has the same shape. For rate-limited Hugging Face checkpoint URLs, a transient 429 on .metadata or train/rank*.pt can still abort checkpoint loading without the retry behavior this commit adds for direct HTTP range/head reads.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
looking into this now.
There was a problem hiding this comment.
code added to support 429 retry in the cached_path
warmbowski
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 92a9f4c3c9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| and exc.response.status_code == 502 | ||
| ), | ||
| ) | ||
| @retriable(retry_condition=_http_retry_condition) |
There was a problem hiding this comment.
When a checkpoint is loaded without pre_download, RemoteFileSystemReader._get_bytes() calls get_bytes_range() for each tensor shard, so a Hugging Face 429 on this _http_get_bytes_range() path is retried by the generic retriable wrapper. That wrapper only sleeps _wait_before_retry() (capped at 3s) and never reads exc.response.headers['Retry-After'], unlike the new cached-path session; if HF returns a longer reset window, all retries are spent before the window opens and the load still fails. Please either use the retry-aware session here or teach retriable to honor the HTTP header for this condition.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Trying to avoid scope creep on this. I can take care of this if you really want me to, but this should suffice for now.
| if response.status_code == 404: | ||
| raise FileNotFoundError(url) |
There was a problem hiding this comment.
this is a bit of an exiting bug fix that brings it in line with the other http* methods. I think we want to raise this before checking the content-length
|
I am unsure that this PR should go in with the http auth code added. It turns out that it isn't solving the issue of allowing external users to get better rate limits on bucket access when using an hf token. I am going to put this in draft until we figure out if what we want is even possible. |
|
It looks like connecting to our allenai bucket with an hf token is possible with the correct token permissions. This will get higher rate limits for end users and just about eliminate 429 errors in the run, even on the free tier. Anonymous access still face a lot of 429 errors. |
2 participants
Add HuggingFace auth and 429 retry handling to HTTP downloads
Why
Checkpoints now load from public HuggingFace bucket URLs. A distributed checkpoint load issues many concurrent ranged GETs, which trips HF's per-window rate limit: nothing sent an
HF_TOKEN(anonymous requests get the lowest rate tier) and 429 responses weren't retried.Changes
_http_auth_headers(), which attachesAuthorization: Bearer $HF_TOKENfor hosts in ahost → env varregistry (huggingface.co,hf.co). Returns no header when the env var is unset, so anonymous access is unchanged._http_file_size(),_http_get_bytes_range(), and_http_file_exists().cached_path()calls that read trainer state and checkpoint metadata (resource_path()inio.py,Checkpoint.load()and the ephemeral-check read intrain/checkpoint.py), which previously requested anonymously.HTTP_RETRY_STATUS_CODES = (429, 502, 503, 504)andHTTP_MAX_RETRIESas the single retry policy for HTTP._build_http_session(), which builds a pooled session and optionally attaches a urllib3 retry policy honoringRetry-After._get_http_session()now delegates to it._RetryableHttpClient, a cached-path scheme client forhttp(s)://that retriesHTTP_RETRY_STATUS_CODES; registered inadd_cached_path_clients(). Subclasses cached-path'sHttpClientbecauseget_scheme_client()only forwards custom headers to clients derived from it, and treats exhausted retries as recoverable so a cached copy can be used._http_retry_condition()to readHTTP_RETRY_STATUS_CODESinstead of a hardcoded(429, 502), and apply it to all three_http_*helpers.response.raise_for_status()to_http_file_size()so error statuses surface as HTTP errors rather than a missingcontent-lengtherror, and raiseFileNotFoundErroron 404 to match the GCS and S3 equivalents.io.pycached-path section header to cover both clients.Notes
HF_TOKEN(read scope) raises the rate limit tier; unset leaves behavior as before.Retry-Afterbut not HF'sRateLimit: t=<reset>header.@retriable, so retries don't stack.🤖 Generated with Claude Code