feat(auth): support shared token file for distributed workers

[RapidPoseidon]

Why

A customer running the SDK (ranking flows) across thousands of GPU workers hit auth-scale problems: every worker's RapidataClient() exchanged the client credentials for its own token at startup, and since all tokens expire together (~1 h), all workers re-hit the identity service in the same instant. They worked around it with one coordinator process that refreshes the token and writes it to a shared file, with workers reusing it.

The SDK already half-supports this via the token= dict param — but a pre-supplied client-credentials token has no refresh path in authlib (no refresh_token, no grant_type metadata), so once it expires every request raises InvalidTokenError and the client must be rebuilt.

What

• RapidataClient(token_file=...) (or RAPIDATA_TOKEN_FILE env var): the client loads the token JSON from the file at startup and re-reads the file whenever the in-memory token is within leeway (60 s) of expiry. Workers never hold the client secret and never need to be reconstructed — the coordinator just keeps the file fresh.
• Plumbed through OpenAPIService → RESTClientObject.setup_oauth_with_token (both rest.py and its rest.mustache template, per repo convention).
• New docs page Distributed Training under the Ranking Flows section documenting the full pattern: coordinator loop (absolute expires_at, atomic temp-file + os.replace writes), worker usage, startup-jitter tip, and the recreate-the-client fallback for older SDK versions. Cross-linked from the Authentication guide.

Verification

• pyright src/rapidata/rapidata_client src/rapidata/service: 0 errors
• black --check clean on touched files under src/rapidata/rapidata_client
• Functional smoke test: init from an expired token file → first request path reloads the fresh token from the file; valid in-memory token → no reload; RapidataClient() construction via RAPIDATA_TOKEN_FILE alone (no client credentials) works
• uv run --group docs mkdocs build succeeds; new page renders in nav

🔗 Session: https://session-43887903.poseidon.rapidata.internal/

🤖 Generated with Claude Code

[claude]

Review: shared token file for distributed workers

Nice, focused solution to a real scaling problem (thundering-herd on the identity service). The plumbing through RapidataClient → OpenAPIService → RESTClientObject/rest.mustache is consistent, the generated file and template stay in sync per repo convention, precedence ordering (token/token_file bypass the credential exchange entirely) matches the docstring, and the new docs page is clear and appropriately scoped without overselling.

Correctness / robustness

• Unhandled exceptions on reload can leak raw errors past the SDK's error contract (src/rapidata/api_client/rest.py:114-122, mirrored in rest.mustache). _reload_token_from_file_if_expired does a bare open(token_file) / json.load(f) with no error handling, unlike _read_token_file in openapi_service.py:258-263, which wraps OSError/ValueError into a friendly ValueError at startup. If the file is transiently missing/unreadable/malformed at reload time (coordinator mid-restart, NFS hiccup, permission issue), the raw FileNotFoundError/JSONDecodeError/PermissionError propagates out of RESTClientObject.request(). That happens inside call_api(), which is not covered by RapidataApiClient.response_deserialize's except ApiException handler (rapidata_api_client.py:211) — so instead of the documented RapidataError format (per CLAUDE.md), the caller gets a raw stdlib exception. Worth wrapping the same way _read_token_file does, especially since this is exactly the failure mode the distributed-workers use case will hit if the coordinator ever hiccups.

• Reload only happens once per request() call, before the retry loop (rest.py:163-164, loop starts at 168). The retry loop can legitimately run for a while (exponential backoff + Retry-After, up to _RETRY_AFTER_MAX_SECONDS per attempt, up to 3 retries). If the in-memory token crosses expiry during that loop, later attempts within the same request() call won't re-check the file and could fail on a now-expired token even though a fresh one is sitting in the file. Probably rare in practice (retries are for 429/502/503/504, not auth failures), but since this is precisely the class of bug the PR is fixing, it seems worth re-checking expiry inside the retry loop too, not just before it.

Minor / consider

• Once a worker's in-memory token is within leeway of expiry and the coordinator has stopped updating the file (e.g. crashed), every subsequent request from every worker will open()/json.load() the file on every call, indefinitely, until the process is restarted. At "thousands of GPU workers" scale on a shared filesystem this is extra sustained I/O pressure exactly during an incident — not wrong, but worth a mental note (e.g. a log line the first time reload fails would help operators notice quickly).
• Docstring on token_file (rapidata_client.py:113) says the file is "re-read whenever the current token expires" — the actual behavior is re-read when within leeway of expiry (proactive), which is better but slightly different from what's described. Small wording nit.
• Security nit for the new docs' coordinator sample (docs/distributed_training.md): write_token() doesn't set restrictive permissions on the token file. Since that file is a live bearer credential valid for up to an hour, and the whole point is putting it on a shared filesystem readable by many workers, it'd be worth explicitly recommending os.chmod(tmp_path, 0o600) (or documenting that the shared filesystem/directory ACLs must restrict read access to the worker fleet) so the pattern doesn't accidentally widen the credential's blast radius relative to per-worker client-credential auth.

Test coverage

The repo has no automated test suite (no tests/, no pytest dependency), so this isn't a gap specific to this PR, and the description's manual smoke tests are reasonable given that. Just flagging that this is auth-critical, easily-regressed logic (the expiry/reload branch in particular) that would benefit from a unit test whenever test infra exists — e.g. asserting _reload_token_from_file_if_expired is a no-op for a non-expired token and reloads for an expired one.

Everything else

Env var fallback ordering, token or token_file handling in OpenAPIService, mkdocs nav wiring, and docs cross-linking all look correct and match the stated design.

[LinoGiger]

should this not use the sdk client methods?

[RapidPoseidon]

Good call — the SDK had no public way to get the token out, which is why the example went to the raw endpoint. Added RapidataClient.get_token() in 2efaa77: it returns the complete token object (with absolute expires_at), refreshing it first if expired, so the coordinator example is now fully SDK-native. Verified the refresh path: authlib stores grant_type=client_credentials in the session metadata on the initial fetch, so ensure_active_token re-fetches on demand.

[claude]

Overview

This PR adds a token_file / RAPIDATA_TOKEN_FILE mechanism so many workers can share a single OAuth token maintained by one coordinator process, avoiding a thundering-herd auth burst when thousands of tokens expire simultaneously. It plumbs token_file through RapidataClient -> OpenAPIService -> RESTClientObject.setup_oauth_with_token (mirrored correctly in both src/rapidata/api_client/rest.py and openapi/templates/rest.mustache, per repo convention), adds a get_token() accessor for the coordinator to export a refreshed token, and ships a new "Distributed Training" docs page. Overall a well-scoped, clearly motivated change with good documentation.

Potential correctness issues

1. self.session.token = json.load(f) may store a plain dict instead of an OAuth2Token. (rest.py/rest.mustache, _reload_token_from_file_if_expired)
   authlib's OAuth2Client normally holds an OAuth2Token (a dict subclass exposing .is_expired()) as session.token. This code assigns the raw output of json.load() directly to self.session.token. If the underlying setter doesn't itself re-wrap the value into OAuth2Token (I could not execute code in this sandbox to confirm either way for the pinned authlib version), then on the second reload cycle the guard token.is_expired(leeway=...) will raise AttributeError: 'dict' object has no attribute 'is_expired', since a plain dict has no such method. The included smoke test ("expired token file -> reload") only exercises a single reload, so it wouldn't catch a failure that only appears on the second cycle.
   Suggested defensive fix regardless of what the setter does today: self.session.token = OAuth2Token(json.load(f)) (from authlib.oauth2.rfc6749) so the invariant holds explicitly rather than relying on library internals.

2. No synchronization around the shared session.token mutation in a multi-threaded client. The SDK itself uses ThreadPoolExecutor for concurrent work against a single shared client/session (e.g. _batch_asset_uploader.py, _asset_upload_orchestrator.py). Once the token expires, every in-flight thread's next request() call will concurrently hit _reload_token_from_file_if_expired, each doing a racing read-then-assign on self.session.token with no lock. Best case this is just redundant file reads; worst case it's a data race on shared mutable state during concurrent requests. Given this is exactly the "many workers hammering a shared client" scenario the PR is designed for, a small threading.Lock around the check-and-reload critical section would remove the ambiguity.

Minor / polish

• Inconsistent error handling for reload-time file errors. _read_token_file (in openapi_service.py, used at startup) wraps OSError/ValueError into a clear ValueError message. _reload_token_from_file_if_expired (used on every request once expired) has no equivalent wrapping - a transient NFS hiccup or a coordinator that's fallen behind will surface a raw FileNotFoundError/JSONDecodeError from deep inside request() instead of a clear, actionable error. Since this is precisely the operational failure mode distributed workers on shared storage are likely to hit, consider wrapping it similarly.
• get_token() on RESTClientObject/RapidataApiClient duplicates ~15 lines identically between rest.py and rest.mustache - correct per the project's codegen convention, just flagging that any future tweak needs both kept in sync (as already done here).
• The _reload_token_from_file_if_expired comment explains why well; nice touch given the non-obvious authlib refresh-path limitation.

Docs

• docs/distributed_training.md is clear and well-structured (motivation -> workers -> coordinator -> tips -> fallback for older SDKs). The atomic-write guidance (temp file + os.replace) and the callout that raw expires_in responses need to be converted to absolute expires_at are good, non-obvious details worth documenting.
• Code annotations (# (1)!) follow the existing content.code.annotate convention used elsewhere in the docs - consistent.
• mkdocs.yml nav/docs updates look correct and match the new page.

Test coverage

No automated tests are included (the repo doesn't appear to have a test suite at all currently, so this isn't a regression in convention), but this is an auth-critical, security-sensitive code path - a couple of unit tests would be valuable and cheap:

• reload triggers only when the in-memory token is within leeway of expiry, not otherwise;
• the reloaded token type still has .is_expired() after being set from a file (this would directly catch/rule out issue First version #1 above);
• RAPIDATA_TOKEN_FILE env var is picked up when no explicit token/token_file/credentials are given;
• a missing/corrupt token file at startup raises the wrapped ValueError with a useful message.

Security

No credentials are logged; workers never receive the client secret, which is the stated goal. The atomic-write documentation appropriately warns against partially-written files. No other security concerns spotted.

---

Nice, well-motivated feature with solid docs. The main thing I'd want resolved before merge is confirming behavior of reassigning session.token with a raw dict across two consecutive reload cycles, and considering a lock for the multi-threaded case - both are cheap to verify/fix and go to the core reliability guarantee this PR is trying to provide.

[LinoGiger]

can you swap the coordinator and worker around, since the coordinator has to be started first, i think it makes more sense?

additionally add a os.makedir for the /shared and add a note that in this example it's gonna be added as a root file

[RapidPoseidon]

Done in 6a637e4 — coordinator section now comes first, the example creates the directory with os.makedirs(..., exist_ok=True), and an annotation notes that the root-level /shared is just for the example (point it at whatever storage the workers actually mount).

Your sleep(1) test also surfaced something worth fixing: get_token() intentionally returns the cached token while it's valid and only refreshes within leeway of expiry (refreshing every call would be the stampede all over again) — but leeway was silently ignored on the client-credentials path, so the example's '5 minutes early' promise couldn't actually hold. Now plumbed through: the coordinator example uses RapidataClient(leeway=300) + a simple time.sleep(60) loop, which renews the file 5 minutes before expiry.

[claude]

Review: feat(auth): support shared token file for distributed workers

Overview
Adds RapidataClient(token_file=...) / RAPIDATA_TOKEN_FILE so a single coordinator process can refresh an OAuth token and have many workers share it via a file, re-reading it whenever the in-memory token nears expiry (leeway, default 60s). Correctly plumbed through OpenAPIService -> RESTClientObject, mirrored between rest.py and rest.mustache per the repo's codegen convention, and documented in a new Distributed Training doc page. A focused, well-motivated change addressing a real thundering-herd problem against the identity service.

Strengths

• rest.py / rest.mustache kept in sync, per CLAUDE.md convention.
• _read_token_file in openapi_service.py wraps OSError/ValueError into a clear message at startup.
• Precedence rules (token/token_file > explicit credentials > env credentials > cached file > browser login) are clearly documented and match the implementation.
• Docs are concrete and practical (atomic-write coordinator loop, jitter tip, fallback for older SDKs).

Issues / Suggestions

1. Security: the docs example doesn't restrict token file permissions. The coordinator example in distributed_training.md writes the shared token with a plain open(tmp_path, "w") + os.replace. The project already has a stricter, established pattern for exactly this kind of secret in credential_manager.py (temp file opened via os.open(..., 0o600) before the atomic rename, plus os.chmod(config_dir, 0o700)). A live bearer token sitting on a shared NFS mount with default (often 0644) permissions is a broader credential-exposure surface than intended. Worth updating the doc snippet to use 0o600, consistent with the SDK's own convention.

2. Robustness: the runtime reload path has no error handling, unlike the startup path. _read_token_file (used once at construction) wraps read/parse failures into a friendly ValueError. But _reload_token_from_file_if_expired (called on every request once the in-memory token is within leeway of expiry, in both rest.py and rest.mustache) does a bare open(token_file) / json.load(f) with no try/except. A transient hiccup (file briefly missing, NFS latency, a permissions blip) raises a raw FileNotFoundError/JSONDecodeError/OSError straight out of request(), bypassing the existing transient-retry loop (_TRANSIENT_RETRY_MAX_ATTEMPTS/backoff) entirely, with a much less friendly message than the rest of the client gives. Given the target scenario is thousands of workers reading a shared filesystem, this is exactly the kind of transient failure worth tolerating - e.g. wrap it the same way _read_token_file does, and/or let it participate in the retry loop.

3. Scale: no mtime short-circuit, so the thundering herd moves from the identity server to the shared filesystem. Once the in-memory token is within leeway of expiry, _reload_token_from_file_if_expired re-opens and re-parses the file on every single request, not just once - so for the motivating scenario (thousands of GPU workers, high request volume) the last leeway seconds before expiry produce a burst of concurrent reads against the shared file/NFS mount from every worker. It's a smaller problem than the original (file I/O vs. an auth round trip), but it's the same shape of issue this PR sets out to fix. A cheap mitigation: skip the re-read if the file's mtime hasn't changed since the last successful load.

4. Doc/docstring gap: behavior when both token and token_file are supplied isn't described. OpenAPIService.init uses the explicit token for the initial session, but still passes token_file through to setup_oauth_with_token, so the client will silently start reloading from token_file once that initial (explicit) token expires. That's probably reasonable, but it's a bit surprising and isn't called out in RapidataClient.init's docstring, which documents the two args separately.

5. Test coverage: there's no tests/ directory in this repo, so this isn't a new gap this PR introduces, but it's worth flagging that the PR description relies on manual/functional smoke testing rather than an automated regression test. Given the subtlety here - token reload interacting with authlib's OAuth2Client.token/is_expired, atomic-replace races, leeway semantics across two code paths (ensure_active_token vs. file reload) - a small automated test exercising at least two successive reload cycles from a token file (not just one) would catch regressions a one-off manual check might miss, e.g. if a future authlib upgrade changes how session.token assignment behaves.

6. Minor style nit: the new parameters (token_file: str | None = None) in rest.py/rest.mustache use PEP 604 unions, while the rest of that file uses Optional[X] (no 'from future import annotations' there). Harmless given requires-python >=3.10, but a small style inconsistency within an autogenerated file that otherwise sticks to one convention.

Summary
Solid, well-scoped feature with good docs and correct codegen hygiene. The main things worth addressing before merge are #1 (permissions on the example token file) and #2 (error handling parity for the runtime reload path) - both small, low-risk changes. #3-#6 are lower-priority polish.

[jorgeparavicini]

I'm wondering whether we want to implement this as well. I mean it's quite boilerplatey. We could have a method that can take a path and a duration and we would continously write to it. If the customer would want to write the token differently they could still write this part themselves. But i think having a default implementation would make adoption of this pattern easier

[RapidPoseidon]

Agreed — added RapidataClient.maintain_token_file(path, interval=60) in 36c887e. It writes the file once synchronously (so it's guaranteed to exist — or fail loudly — before workers start), creates the directory, then keeps rewriting it atomically from a daemon thread; transient failures inside the loop are logged and retried on the next interval instead of killing the thread. It returns the thread so a dedicated coordinator can just .join() it, while a rank-0-that-also-trains drops the join. The manual loop moved to a 'Writing the file yourself' subsection for custom setups.

[jorgeparavicini]

It's also important to mention that doing this exposes the token and the path where the token is written to should be kept secure from unauthorized access. Yes it's obvious but its important to state the obviousness in these sitatuiosn

[RapidPoseidon]

Added a dedicated warning admonition in 36c887e: the file contains a bearer token, anyone who can read it can act on the account until expiry, so write it only to storage the training job alone can access and restrict permissions. Also noted in the maintain_token_file docstring.

[jorgeparavicini]

This is an implementation detail that only concerns us, not our customers. This documentation should be written with details useful for the customer. The point it makes is valid but it can use some rewording

[RapidPoseidon]

Reworded in 36c887e — the annotation now frames it from the customer's side: get_token() is cheap to call at any frequency; it returns the current token and only contacts the auth server once the token is within leeway of expiry. (With maintain_token_file as the headline API, the manual loop and this note now live in the 'Writing the file yourself' subsection.)

[claude]

Review: feat(auth): support shared token file for distributed workers

Nice feature — clean plumbing through RapidataClient → OpenAPIService → RESTClientObject, mirrored correctly between src/rapidata/api_client/rest.py and openapi/templates/rest.mustache per the repo's codegen convention, and backed by a well-written docs page.

Potential bug: token may not survive a second reload cycle

_reload_token_from_file_if_expired (rest.py / rest.mustache) does:

with open(token_file) as f:
    self.session.token = json.load(f)

This assigns a plain dict to session.token. The very next call reads:

token = self.session.token
if token is not None and not token.is_expired(leeway=self._token_leeway):
    return

token.is_expired(...) requires an authlib OAuth2Token, not a plain dict. If authlib's OAuth2Client.token setter doesn't itself coerce a raw dict into OAuth2Token (I wasn't able to confirm either way from this sandbox), then the second reload cycle raises AttributeError: 'dict' object has no attribute 'is_expired' — i.e. the feature could work once and then break for the remainder of the worker's life. Since the PR's smoke test only describes a single reload ("first request path reloads the fresh token"), this multi-cycle path may not have been exercised.

Suggested fix: be explicit regardless of authlib's internal coercion behavior —

from authlib.oauth2.rfc6749 import OAuth2Token
...
self.session.token = OAuth2Token(json.load(f))

and add a smoke test that drives two reload cycles (expire → reload → expire → reload) to lock in the behavior.

Robustness: hot-path reload has no error handling

openapi_service.py's _read_token_file() wraps open/json.load and raises a clear ValueError on failure. The hot-path equivalent, _reload_token_from_file_if_expired() in rest.py/rest.mustache, does the same open()/json.load() with no error handling. Since this runs inside request() on every call once the in-memory token is within leeway of expiry, a transient hiccup (coordinator mid-restart, NFS blip, torn read despite the atomic-write discipline on the writer side) surfaces as a raw FileNotFoundError/json.JSONDecodeError deep in the request path instead of an actionable error. Worth wrapping with the same context-adding treatment as _read_token_file, and possibly a short retry/backoff given this fires per-request.

Minor / non-blocking

• maintain_token_file's background thread calls self.get_token() (which can itself trigger ensure_active_token()/refresh) on the same OAuth2Client session that the coordinator's own outgoing requests use. If a coordinator process also makes heavy concurrent API calls while running maintain_token_file, there's a theoretical race on session.token between the refresh thread and request-time refresh logic. Likely fine for the documented "dedicated coordinator" use case — maybe worth a one-line callout in the docstring/docs.
• maintain_token_file returns a daemon=True thread with no way to stop the loop short of process exit. Reasonable for the "run forever" coordinator pattern described in the docs; a stop()/threading.Event could be added later if a graceful-shutdown use case shows up.
• setup_oauth_client_credentials now stores self._token_leeway, but that field is only consumed by _reload_token_from_file_if_expired, which is unreachable on the client-credentials path (only setup_oauth_with_token sets _token_file). Harmless dead assignment, not worth blocking on.

Nice catches

• Threading leeway into setup_oauth_client_credentials's OAuth2Client(...) construction looks like a genuine (incidental) fix: previously a caller-supplied leeway was silently ignored on the default (non-token) auth path, since that constructor never forwarded it.
• Atomic writes (temp file + os.replace, PID-suffixed temp names to avoid cross-process collisions), a synchronous first write in maintain_token_file so failures surface immediately, and explicit "keep the token file secure" warnings in the docs are all solid details.

Test coverage

The repo has no existing unit test suite (no tests/ directory anywhere), so relying on the documented manual/functional smoke test is consistent with current project convention. Given the subtlety above, I'd still suggest at minimum a script-based check exercising the two-reload-cycle scenario before merging.

Security

The "authenticate once, share a bearer token via file" model is a reasonable trade-off for the auth-server thundering-herd problem described in the PR, and the risk (anyone with file read access can act as the client until expiry) is called out clearly in the new docs page. No token values are logged in the diff. No other concerns spotted.