fix: Improve reliability of S3 credential refresh

[dhedey]

What does this PR do?

S3Client / R2Client refetch temporary bucket credentials from the control plane on a timer. Four things on that path turned one blip into a dead run — reported by a user streaming from /teamspace/s3_connections/* across 4 ranks × 8 workers, dying 8–9 hours in on a burst of boto3 failures:

• Retries the /v1/auth/login POST, which urllib3 was never retrying.
• Keeps serving the current credentials when a refresh fails, instead of killing the worker.
• Retries a failed refresh on a timer rather than on every read.
• Waits out a control-plane outage when creating the first client, where there is nothing to fall back on.
• Retries only failures that can clear, so a bad storage_options key still raises at once.
• Jitters the refresh interval per process, so forked workers stop refreshing in lockstep.
• Moves the interval from 55 to 45 minutes, for a usable margin inside the 1 hour credential TTL.

Root cause

_login_and_get_temp_bucket_credentials makes two calls, and mounts a Retry on the session for both. urllib3 defaults allowed_methods to a set that excludes POST, so only the credentials GET was ever retried — a single 429 or 502 on the login POST failed the whole refresh. A non-200 login was also read straight as JSON, so it surfaced as a JSONDecodeError, or as a missing-token error that named neither the status nor the login call.

That failure then propagated out of the client property into the caller. Nothing between there and the DataLoader catches it, so one failed refresh killed the worker and the run — even though the credentials in hand were still valid for another five minutes. And because _last_time only advanced on success while _client kept the old client, every subsequent read re-attempted the refresh: one outage became a login storm from every worker, which is the burst the user saw.

Two things made that window small and likely to be hit by everyone at once. The credentials carry no expiry for us to read, and 3300s was set against a TTL of 3600s, leaving five minutes to recover. Workers are forked together, so they mint together and expire together.

After: the login POST is retried; a failed refresh logs, keeps the current client, and backs off for 60s; only once it is _REFRESH_GRACE_PERIOD (900s) past the interval are the credentials assumed dead and the error raised. The interval is jittered downwards and re-rolled per process, in __setstate__ and on a pid change, so forked workers spread out.

Creating the first client has no credentials to fall back on, so it gets the same 60s timer and 900s bound rather than the caller's error — otherwise a control plane that is down at job start kills the job in the time it takes to lose one retry budget. Each attempt logs how long it has been waiting and that data loading is blocked meanwhile.

Retrying only helps failures that can clear, so credential fetches raise a _CredentialsError — _CredentialsUnavailableError when the control plane or IMDS was unreachable or answered 5xx/429/408, _CredentialsConfigurationError for missing config or a rejected 4xx — and only those are retried. Anything else is a local mistake: a bad storage_options key raises on the first attempt, as it did before. Initial creation waits only for the unavailable kind; a refresh gives both the grace period, because it holds working credentials and a mid-deploy 403 from a proxy is not worth dying on. A real revocation is still caught by the deadline.

The two budgets are separate because they answer to different constraints. _REFRESH_GRACE_PERIOD is 600s so that 2700 + 600 stays ~5 minutes under the 1 hour TTL rather than running exactly to it; waiting for the first client holds nothing, so _INITIAL_RETRY_BUDGET keeps 900s.

Scope

The per-request retry budget drops from 2880 to 4, which shortens a single call from ~96h of backoff to ~7s. That budget was never the thing protecting a refresh — login is the first call and POST was never retried, so one 429 there was already instantly fatal — and a refresh holds the client lock for its whole duration, so blocking every thread in the process for days is not a fallback, it is a hang. Waiting is now done by the surrounding timer, where it is bounded, logged, and lets the job keep reading. broadcast.py has its own copy of that constant and is untouched.

Not addressed here: workers still refresh independently, so a Studio makes N calls per interval rather than one. Sharing across a process tree is a bigger change and I'd rather do it separately.

Fixes a pre-existing bug this change would otherwise have been built on: _CustomRetryAdapter.send used kwargs.get("timeout", self.timeout), but requests always passes timeout explicitly — as None when the caller gave none — so the default never applied and the login POST had no timeout at all. The retry budget bounded attempts but not wall time.

Declares urllib3 >=1.26 in requirements.txt. client.py has always imported urllib3.util.retry.Retry directly while urllib3 was only ever a transitive dependency, and the effective floor from requests and botocore allows 1.25.x, which has neither DEFAULT_ALLOWED_METHODS nor the allowed_methods kwarg — both arrived in 1.26.0. Latent before; retrying the login POST makes it load-bearing. The bound sits inside both existing constraints, so it excludes nothing anything else needs.

Also drops a print and a comment in get_r2_bucket_credentials that promised a fallback to hardcoded credentials that does not exist.

Tests

Twenty new tests in tests/streaming/test_client.py cover POST being retried; a failed refresh keeping the current client, not retrying on every access, riding out a 403, and raising past the grace period; the jitter being bounded and re-rolled on unpickle; initial creation retrying until the control plane returns and giving up when it does not; local and permanent failures not being retried; and the adapter applying its default timeout. They fail without the change, the behavioural ones with the RuntimeError the fix is meant to absorb. The retry interval is a module constant so the tests set it to zero — the file runs in 7.5s.

Five existing tests needed updating: one for the new default interval, and four whose mocks never set a login status_code. The old "login failure" test asserted a missing-token error while mocking something that was never a login failure; it is split into a rejected-login case and a 200-without-token case.

tests/streaming/test_client.py, test_downloader.py and test_fs_provider.py are green locally (81 passed), serially and under CI's -n 2 --dist=loadgroup, as are ruff, ruff format and mypy.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 91.11111% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 82%. Comparing base (8805bae) to head (66c3ca7).
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@          Coverage Diff          @@
##            main    #891   +/-   ##
=====================================
  Coverage     82%     82%           
=====================================
  Files         65      65           
  Lines      13377   13439   +62     
=====================================
+ Hits       10949   11013   +64     
+ Misses      2428    2426    -2     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.