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

docs/authentication.md

@@ -39,6 +39,10 @@ with no arguments. On a workstation, calling `RapidataClient()` with no
 credentials instead opens a browser login once and caches the token in
 `~/.config/rapidata/credentials.json`.
 
+Running many workers in parallel (e.g. a distributed training job)? Don't
+give each worker its own credentials — authenticate once and share the token
+via a file, as described in [Distributed Training](distributed_training.md).
+
 ## Direct token request
 
 To call the API without the SDK, request a token from the token endpoint and

docs/distributed_training.md

@@ -0,0 +1,139 @@
+# Distributed Training
+
+When you use Rapidata from a large distributed job — for example a ranking
+flow queried from hundreds or thousands of GPU workers — don't let every
+worker authenticate on its own. Each `RapidataClient()` exchanges the client
+credentials for an access token at startup, and because all those tokens
+expire at the same time (~1 hour later), every worker hits the auth server
+again in the same instant. At scale this looks like a coordinated burst and
+leads to rate limiting and failed requests.
+
+The fix is to authenticate **once** and share the token:
+
+- A single **coordinator** process (your launcher, or rank 0) holds the
+  client credentials. It fetches the access token, refreshes it *before* it
+  expires, and writes it to a file on storage all workers can read (a shared
+  filesystem like NFS, or a mounted volume).
+- Every **worker** points the SDK at that file. Workers never see the client
+  secret, make no auth-server calls of their own, and automatically re-read
+  the file when the current token expires.
+
+## Coordinator
+
+Start the coordinator first — it authenticates with the client credentials
+as usual (see [Authentication](authentication.md)) and keeps the shared file
+fresh with `maintain_token_file()`:
+
+```python
+from rapidata import RapidataClient
+
+coordinator = RapidataClient(leeway=300)  # (1)!
+coordinator.maintain_token_file("/shared/rapidata_token.json").join()  # (2)!
+```
+
+1. The coordinator is the only process that holds the client credentials —
+   via `RAPIDATA_CLIENT_ID` / `RAPIDATA_CLIENT_SECRET` or the cached
+   credentials file. `leeway=300` renews the token once it is within
+   5 minutes of expiry, well before workers need a new one.
+2. Writes the file immediately, then keeps rewriting it atomically from a
+   background thread (every 60 seconds by default), creating the directory
+   if needed. `.join()` blocks the process forever — drop it if your
+   coordinator also does other work, e.g. when rank 0 both trains and
+   refreshes the token. The root-level `/shared` is just for the example —
+   use whatever storage all your workers actually mount (an NFS path, a
+   shared volume, …).
+
+!!! warning "Keep the token file secure"
+    The file contains a bearer token — anyone who can read it can act on
+    your Rapidata account until the token expires. Write it only to storage
+    that your training job alone can access, and restrict the file
+    permissions accordingly.
+
+### Writing the file yourself
+
+If the built-in helper doesn't fit (different storage backend, custom
+rotation), export the token with `get_token()` and write it however you
+like. `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.
+
+```python
+import json
+import os
+import time
+
+from rapidata import RapidataClient
+
+TOKEN_FILE = "/shared/rapidata_token.json"
+
+coordinator = RapidataClient(leeway=300)
+
+os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True)
+
+
+def write_token(token: dict) -> None:
+    tmp_path = TOKEN_FILE + ".tmp"
+    with open(tmp_path, "w") as f:
+        json.dump(token, f)
+    # Atomic replace: a worker reading concurrently sees either the old
+    # token or the new one, never a partially written file.
+    os.replace(tmp_path, TOKEN_FILE)
+
+
+while True:
+    write_token(coordinator.get_token())
+    time.sleep(60)
+```
+
+Two things to keep when rolling your own:
+
+- **Write atomically** (temp file + `os.replace`, as above) so workers never
+  read a half-written file.
+- **Keep the absolute `expires_at` field** in the JSON — it's how workers
+  decide when to re-read the file. If you produce the token without the SDK
+  (e.g. a `curl` call to the token endpoint), the response only carries the
+  relative `expires_in`, so add `expires_at` yourself.
+
+## Workers
+
+Once the coordinator has written the file, workers pass its path to the
+client — or set the `RAPIDATA_TOKEN_FILE` environment variable and construct
+the client with no arguments:
+
+```python
+from rapidata import RapidataClient
+
+client = RapidataClient(token_file="/shared/rapidata_token.json")
+```
+
+The SDK reads the token from the file at startup and re-reads it whenever the
+in-memory token is within 60 seconds of expiry (configurable via `leeway`).
+As long as the coordinator keeps the file fresh, workers never need to be
+restarted or re-authenticated.
+
+## Practical tips
+
+- **Stagger worker startup.** Even with a shared token, thousands of workers
+  constructing their first client and firing their first API call in the
+  same second is an unnecessary burst. A few seconds of random jitter before
+  the first call smooths it out.
+- **Create one client per process and keep it.** The client refreshes its
+  token from the file transparently; there is no need to reconstruct it.
+- **Keep the coordinator alive for the whole run.** If it stops refreshing,
+  workers fail with an invalid-token error once the token in the file
+  expires.
+
+## Older SDK versions
+
+Before `token_file` was available, the same pattern worked by passing the
+token dict directly — but the SDK never re-reads the file, so each worker
+must construct a new client when the token expires:
+
+```python
+import json
+
+from rapidata import RapidataClient
+
+with open("/shared/rapidata_token.json") as f:
+    client = RapidataClient(token=json.load(f))
+```

mkdocs.yml

@@ -80,6 +80,7 @@ plugins:
           - mri_advanced.md
         Flows:
           - flows.md
+          - distributed_training.md
         AI Agents:
           - ai_agents.md
       enable_markdown_urls: true
@@ -138,7 +139,9 @@ nav:
       - Select Words: examples/select_words_job.md
       - Free Text: examples/free_text_job.md
       - Ranking: examples/ranking_job.md
-    - Ranking Flows: flows.md
+    - Ranking Flows:
+      - Getting Started: flows.md
+      - Distributed Training: distributed_training.md
     - Model Ranking:
       - Getting Started: mri.md
       - Advanced: mri_advanced.md

openapi/templates/rest.mustache

@@ -54,18 +54,27 @@ class RESTClientObject:
         self.configuration = configuration
 
         self.session: Optional[OAuth2Client] = None
+        self._token_file: Optional[str] = None
+        self._token_leeway: int = 60
 
     def setup_oauth_client_credentials(
-        self, client_id: str, client_secret: str, token_endpoint: str, scope: str
+        self,
+        client_id: str,
+        client_secret: str,
+        token_endpoint: str,
+        scope: str,
+        leeway: int = 60,
     ):
         client_args = self._get_session_defaults()
         self.session = OAuth2Client(
             client_id=client_id,
             client_secret=client_secret,
             token_endpoint=token_endpoint,
             scope=scope,
+            leeway=leeway,
             **client_args,
         )
+        self._token_leeway = leeway
 
         try:
             self.session.fetch_token()
@@ -85,6 +94,7 @@ class RESTClientObject:
         token: dict,
         token_endpoint: str,
         leeway: int = 60,
+        token_file: str | None = None,
     ):
         client_args = self._get_session_defaults()
         self.session = OAuth2Client(
@@ -95,6 +105,32 @@ class RESTClientObject:
             leeway=leeway,
             **client_args,
         )
+        self._token_file = token_file
+        self._token_leeway = leeway
+
+    def _reload_token_from_file_if_expired(self, token_file: str):
+        # A pre-supplied token has no refresh path of its own; an external
+        # process keeps the token file fresh, so re-read it once ours expires.
+        assert self.session is not None
+        token = self.session.token
+        if token is not None and not token.is_expired(leeway=self._token_leeway):
+            return
+        with open(token_file) as f:
+            self.session.token = json.load(f)
+
+    def get_token(self) -> dict:
+        """Return the complete token object, refreshing it first if expired."""
+        if not self.session:
+            raise ApiValueError(
+                "OAuth2 session is not initialized. Please initialize it before making requests."
+            )
+        if self._token_file is not None:
+            self._reload_token_from_file_if_expired(self._token_file)
+        elif self.session.token is None or not self.session.ensure_active_token(
+            self.session.token
+        ):
+            raise ApiValueError("The current token is expired and cannot be refreshed.")
+        return dict(self.session.token)
 
     def request(
         self,
@@ -135,6 +171,9 @@ class RESTClientObject:
                 "OAuth2 session is not initialized. Please initialize it before making requests."
             )
 
+        if self._token_file is not None:
+            self._reload_token_from_file_if_expired(self._token_file)
+
         timeout = self._build_timeout(_request_timeout)
 
         for attempt in range(_TRANSIENT_RETRY_MAX_ATTEMPTS + 1):

src/rapidata/api_client/rest.py

@@ -64,18 +64,27 @@ def __init__(self, configuration) -> None:
         self.configuration = configuration
 
         self.session: Optional[OAuth2Client] = None
+        self._token_file: Optional[str] = None
+        self._token_leeway: int = 60
 
     def setup_oauth_client_credentials(
-        self, client_id: str, client_secret: str, token_endpoint: str, scope: str
+        self,
+        client_id: str,
+        client_secret: str,
+        token_endpoint: str,
+        scope: str,
+        leeway: int = 60,
     ):
         client_args = self._get_session_defaults()
         self.session = OAuth2Client(
             client_id=client_id,
             client_secret=client_secret,
             token_endpoint=token_endpoint,
             scope=scope,
+            leeway=leeway,
             **client_args,
         )
+        self._token_leeway = leeway
 
         try:
             self.session.fetch_token()
@@ -95,6 +104,7 @@ def setup_oauth_with_token(
         token: dict,
         token_endpoint: str,
         leeway: int = 60,
+        token_file: str | None = None,
     ):
         client_args = self._get_session_defaults()
         self.session = OAuth2Client(
@@ -105,6 +115,32 @@ def setup_oauth_with_token(
             leeway=leeway,
             **client_args,
         )
+        self._token_file = token_file
+        self._token_leeway = leeway
+
+    def _reload_token_from_file_if_expired(self, token_file: str):
+        # A pre-supplied token has no refresh path of its own; an external
+        # process keeps the token file fresh, so re-read it once ours expires.
+        assert self.session is not None
+        token = self.session.token
+        if token is not None and not token.is_expired(leeway=self._token_leeway):
+            return
+        with open(token_file) as f:
+            self.session.token = json.load(f)
+
+    def get_token(self) -> dict:
+        """Return the complete token object, refreshing it first if expired."""
+        if not self.session:
+            raise ApiValueError(
+                "OAuth2 session is not initialized. Please initialize it before making requests."
+            )
+        if self._token_file is not None:
+            self._reload_token_from_file_if_expired(self._token_file)
+        elif self.session.token is None or not self.session.ensure_active_token(
+            self.session.token
+        ):
+            raise ApiValueError("The current token is expired and cannot be refreshed.")
+        return dict(self.session.token)
 
     def request(
         self,
@@ -145,6 +181,9 @@ def request(
                 "OAuth2 session is not initialized. Please initialize it before making requests."
             )
 
+        if self._token_file is not None:
+            self._reload_token_from_file_if_expired(self._token_file)
+
         timeout = self._build_timeout(_request_timeout)
 
         for attempt in range(_TRANSIENT_RETRY_MAX_ATTEMPTS + 1):