From 38fbc7d1a233802a43db27ce5d9287d181cc2b31 Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Wed, 8 Jul 2026 08:24:02 +0000 Subject: [PATCH 1/4] feat(auth): support shared token file for distributed workers Co-Authored-By: lino --- docs/authentication.md | 4 + docs/distributed_training.md | 120 ++++++++++++++++++ mkdocs.yml | 5 +- openapi/templates/rest.mustache | 18 +++ src/rapidata/api_client/rest.py | 18 +++ .../rapidata_client/rapidata_client.py | 16 ++- src/rapidata/service/openapi_service.py | 16 ++- 7 files changed, 191 insertions(+), 6 deletions(-) create mode 100644 docs/distributed_training.md diff --git a/docs/authentication.md b/docs/authentication.md index 0e8bb4642..2b9507425 100644 --- a/docs/authentication.md +++ b/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 diff --git a/docs/distributed_training.md b/docs/distributed_training.md new file mode 100644 index 000000000..fbb72ad7b --- /dev/null +++ b/docs/distributed_training.md @@ -0,0 +1,120 @@ +# 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. + +## Workers + +Pass the file 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. + +## Coordinator + +The coordinator exchanges the client credentials for a token (see +[Authentication](authentication.md)) and rewrites the shared file ahead of +expiry: + +```python +import json +import os +import time + +import requests + +TOKEN_URL = "https://auth.rapidata.ai/connect/token" +TOKEN_FILE = "/shared/rapidata_token.json" + + +def fetch_token() -> dict: + response = requests.post( + TOKEN_URL, + data={ + "grant_type": "client_credentials", + "client_id": os.environ["RAPIDATA_CLIENT_ID"], + "client_secret": os.environ["RAPIDATA_CLIENT_SECRET"], + "scope": "openid roles email api", + }, + ) + response.raise_for_status() + token = response.json() + # Workers decide when to re-read the file based on this absolute + # timestamp — the relative expires_in is meaningless to a later reader. + token["expires_at"] = int(time.time()) + token["expires_in"] + return token + + +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: + token = fetch_token() + write_token(token) + # Refresh 5 minutes before expiry so workers never read a stale token. + time.sleep(max(token["expires_in"] - 300, 60)) +``` + +Two details matter: + +1. **Write an absolute `expires_at`** (Unix timestamp, seconds). Without it + the SDK cannot tell that a token read from the file is about to expire. +2. **Write atomically** (temp file + `os.replace`), so workers never parse a + half-written file. + +## 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)) +``` diff --git a/mkdocs.yml b/mkdocs.yml index 92b3ef387..81079e358 100644 --- a/mkdocs.yml +++ b/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 diff --git a/openapi/templates/rest.mustache b/openapi/templates/rest.mustache index d64626059..f624ae53e 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -54,6 +54,8 @@ 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 @@ -85,6 +87,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 +98,18 @@ 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 request( self, @@ -135,6 +150,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): diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 01cb2625e..4d3a6bbcc 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -64,6 +64,8 @@ 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 @@ -95,6 +97,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 +108,18 @@ 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 request( self, @@ -145,6 +160,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): diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index fb5aed018..0fc5e9f43 100644 --- a/src/rapidata/rapidata_client/rapidata_client.py +++ b/src/rapidata/rapidata_client/rapidata_client.py @@ -75,19 +75,23 @@ def __init__( oauth_scope: str = "openid roles email api", cert_path: str | None = None, token: dict | None = None, + token_file: str | None = None, leeway: int = 60, ): """Initialize the RapidataClient. Credentials are resolved in the following order: - 1. ``client_id`` / ``client_secret`` passed explicitly to this + 1. A ``token`` or ``token_file`` passed explicitly to this + constructor (or the ``RAPIDATA_TOKEN_FILE`` environment + variable), which bypasses the token exchange entirely. + 2. ``client_id`` / ``client_secret`` passed explicitly to this constructor. - 2. The ``RAPIDATA_CLIENT_ID`` / ``RAPIDATA_CLIENT_SECRET`` + 3. The ``RAPIDATA_CLIENT_ID`` / ``RAPIDATA_CLIENT_SECRET`` environment variables (useful for headless / container deployments). - 3. Credentials stored under ``~/.config/rapidata/credentials.json``. - 4. Interactive browser login, which then saves credentials to the + 4. Credentials stored under ``~/.config/rapidata/credentials.json``. + 5. Interactive browser login, which then saves credentials to the file above so you don't have to log in again. The ``environment`` argument follows the same pattern: when omitted @@ -106,6 +110,7 @@ def __init__( oauth_scope (str, optional): The scopes to use for authentication. In general this does not need to be changed. cert_path (str, optional): An optional path to a certificate file useful for development. token (dict, optional): If you already have a token that the client should use for authentication. Important, if set, this needs to be the complete token object containing the access token, token type and expiration time. + token_file (str, optional): Path to a JSON file containing the token object described above (with an absolute ``expires_at`` timestamp). The file is re-read whenever the current token expires, so an external process can keep it fresh — useful to share one token across many workers (e.g. distributed training). Falls back to the ``RAPIDATA_TOKEN_FILE`` environment variable when omitted. leeway (int, optional): An optional leeway to use to determine if a token is expired. Defaults to 60 seconds. Attributes: @@ -133,6 +138,8 @@ def __init__( client_id = os.environ.get("RAPIDATA_CLIENT_ID") or None if client_secret is None: client_secret = os.environ.get("RAPIDATA_CLIENT_SECRET") or None + if token is None and token_file is None: + token_file = os.environ.get("RAPIDATA_TOKEN_FILE") or None if environment is None: environment = os.environ.get("RAPIDATA_ENVIRONMENT") or "rapidata.ai" @@ -155,6 +162,7 @@ def __init__( oauth_scope=oauth_scope, cert_path=cert_path, token=token, + token_file=token_file, leeway=leeway, ) diff --git a/src/rapidata/service/openapi_service.py b/src/rapidata/service/openapi_service.py index 615408eff..871e89c5a 100644 --- a/src/rapidata/service/openapi_service.py +++ b/src/rapidata/service/openapi_service.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import subprocess from importlib.metadata import version, PackageNotFoundError from typing import TYPE_CHECKING @@ -35,6 +36,7 @@ def __init__( oauth_scope: str, cert_path: str | None = None, token: dict | None = None, + token_file: str | None = None, leeway: int = 60, ): self.environment = environment @@ -80,7 +82,10 @@ def __init__( self._signal: SignalService | None = None self._translation: TranslationService | None = None - if token: + if token or token_file: + if token is None: + assert token_file is not None + token = _read_token_file(token_file) logger.debug("Using token for authentication") self.api_client.rest_client.setup_oauth_with_token( token=token, @@ -88,6 +93,7 @@ def __init__( client_id=client_id, client_secret=client_secret, leeway=leeway, + token_file=token_file, ) logger.debug("Token authentication setup complete") return @@ -249,6 +255,14 @@ def __str__(self) -> str: return f"OpenAPIService(environment={self.environment})" +def _read_token_file(token_file: str) -> dict: + try: + with open(token_file) as f: + return json.load(f) + except (OSError, ValueError) as e: + raise ValueError(f"Failed to read token file '{token_file}': {e}") from e + + def _get_local_certificate() -> str | None: result = subprocess.run(["mkcert", "-CAROOT"], capture_output=True) if result.returncode != 0: From 2efaa7708495206b5d20c2cf738c41ac6e097af0 Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Wed, 8 Jul 2026 08:43:35 +0000 Subject: [PATCH 2/4] feat(auth): expose get_token for coordinator token export Co-Authored-By: lino --- docs/distributed_training.md | 51 ++++++++----------- openapi/templates/rest.mustache | 14 +++++ src/rapidata/api_client/rest.py | 14 +++++ .../rapidata_client/rapidata_client.py | 12 +++++ 4 files changed, 60 insertions(+), 31 deletions(-) diff --git a/docs/distributed_training.md b/docs/distributed_training.md index fbb72ad7b..d6e3db818 100644 --- a/docs/distributed_training.md +++ b/docs/distributed_training.md @@ -36,37 +36,20 @@ restarted or re-authenticated. ## Coordinator -The coordinator exchanges the client credentials for a token (see -[Authentication](authentication.md)) and rewrites the shared file ahead of -expiry: +The coordinator authenticates with the client credentials as usual (see +[Authentication](authentication.md)), exports its token with `get_token()`, +and rewrites the shared file ahead of expiry: ```python import json import os import time -import requests +from rapidata import RapidataClient -TOKEN_URL = "https://auth.rapidata.ai/connect/token" TOKEN_FILE = "/shared/rapidata_token.json" - -def fetch_token() -> dict: - response = requests.post( - TOKEN_URL, - data={ - "grant_type": "client_credentials", - "client_id": os.environ["RAPIDATA_CLIENT_ID"], - "client_secret": os.environ["RAPIDATA_CLIENT_SECRET"], - "scope": "openid roles email api", - }, - ) - response.raise_for_status() - token = response.json() - # Workers decide when to re-read the file based on this absolute - # timestamp — the relative expires_in is meaningless to a later reader. - token["expires_at"] = int(time.time()) + token["expires_in"] - return token +coordinator = RapidataClient() # (1)! def write_token(token: dict) -> None: @@ -79,18 +62,24 @@ def write_token(token: dict) -> None: while True: - token = fetch_token() + token = coordinator.get_token() # (2)! write_token(token) - # Refresh 5 minutes before expiry so workers never read a stale token. - time.sleep(max(token["expires_in"] - 300, 60)) + # Rewrite 5 minutes before expiry so workers never read a stale token. + time.sleep(max(token["expires_at"] - time.time() - 300, 60)) ``` -Two details matter: - -1. **Write an absolute `expires_at`** (Unix timestamp, seconds). Without it - the SDK cannot tell that a token read from the file is about to expire. -2. **Write atomically** (temp file + `os.replace`), so workers never parse a - half-written file. +1. The coordinator is the only process that holds the client credentials — + via `RAPIDATA_CLIENT_ID` / `RAPIDATA_CLIENT_SECRET` or the cached + credentials file. +2. Returns the complete token object, refreshing it first if it has + expired. It includes the absolute `expires_at` timestamp workers use to + decide when to re-read the file. + +**Write the file atomically** (temp file + `os.replace`, as above) so +workers never parse a half-written file. If you produce the token without +the SDK (e.g. a `curl` call to the token endpoint), you must also add the +absolute `expires_at` yourself — the raw response only carries the relative +`expires_in`, which is meaningless to a later reader. ## Practical tips diff --git a/openapi/templates/rest.mustache b/openapi/templates/rest.mustache index f624ae53e..f6868ccf1 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -111,6 +111,20 @@ class RESTClientObject: 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, method, diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 4d3a6bbcc..5d12c09a3 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -121,6 +121,20 @@ def _reload_token_from_file_if_expired(self, token_file: str): 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, method, diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index 0fc5e9f43..98e0e7285 100644 --- a/src/rapidata/rapidata_client/rapidata_client.py +++ b/src/rapidata/rapidata_client/rapidata_client.py @@ -203,6 +203,18 @@ def __init__( self._check_beta_features() # can't be in the trace for some reason + def get_token(self) -> dict[str, Any]: + """Return the complete token object this client authenticates with, + refreshing it first if it has expired. + + The returned dict has the shape expected by the ``token`` and + ``token_file`` constructor arguments (``access_token``, ``token_type`` + and an absolute ``expires_at`` timestamp), so it can be written to a + shared token file that other workers consume — see the Distributed + Training guide. + """ + return self._openapi_service.api_client.rest_client.get_token() + def reset_credentials(self): """Reset the credentials saved in the configuration file for the current environment.""" logger.info("Resetting credentials") From 6a637e42179418a9386d95c49ef830ae1524c69b Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Wed, 8 Jul 2026 09:17:46 +0000 Subject: [PATCH 3/4] fix(auth): honor leeway for client credentials and reorder distributed docs Co-Authored-By: lino --- docs/distributed_training.md | 67 ++++++++++--------- openapi/templates/rest.mustache | 9 ++- src/rapidata/api_client/rest.py | 9 ++- .../rapidata_client/rapidata_client.py | 2 +- src/rapidata/service/openapi_service.py | 2 + 5 files changed, 56 insertions(+), 33 deletions(-) diff --git a/docs/distributed_training.md b/docs/distributed_training.md index d6e3db818..771edde19 100644 --- a/docs/distributed_training.md +++ b/docs/distributed_training.md @@ -18,27 +18,11 @@ The fix is to authenticate **once** and share the token: secret, make no auth-server calls of their own, and automatically re-read the file when the current token expires. -## Workers - -Pass the file 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. - ## Coordinator -The coordinator authenticates with the client credentials as usual (see -[Authentication](authentication.md)), exports its token with `get_token()`, -and rewrites the shared file ahead of expiry: +Start the coordinator first — it authenticates with the client credentials +as usual (see [Authentication](authentication.md)), exports its token with +`get_token()`, and keeps the shared file fresh: ```python import json @@ -47,9 +31,11 @@ import time from rapidata import RapidataClient -TOKEN_FILE = "/shared/rapidata_token.json" +TOKEN_FILE = "/shared/rapidata_token.json" # (1)! + +coordinator = RapidataClient(leeway=300) # (2)! -coordinator = RapidataClient() # (1)! +os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True) def write_token(token: dict) -> None: @@ -62,18 +48,22 @@ def write_token(token: dict) -> None: while True: - token = coordinator.get_token() # (2)! - write_token(token) - # Rewrite 5 minutes before expiry so workers never read a stale token. - time.sleep(max(token["expires_at"] - time.time() - 300, 60)) + write_token(coordinator.get_token()) # (3)! + time.sleep(60) ``` -1. The coordinator is the only process that holds the client credentials — +1. For simplicity this example writes to a root-level `/shared` directory — + point it at whatever storage all your workers actually mount (an NFS + path, a shared volume, …). +2. The coordinator is the only process that holds the client credentials — via `RAPIDATA_CLIENT_ID` / `RAPIDATA_CLIENT_SECRET` or the cached - credentials file. -2. Returns the complete token object, refreshing it first if it has - expired. It includes the absolute `expires_at` timestamp workers use to - decide when to re-read the file. + credentials file. `leeway=300` treats the token as expired once it is + within 5 minutes of its actual expiry, so the file is renewed well before + workers need it. +3. `get_token()` returns the current token — including the absolute + `expires_at` timestamp workers rely on — and only fetches a new one when + it is within `leeway` of expiry. Polling it every minute therefore does + **not** spam the auth server; most iterations just rewrite the same token. **Write the file atomically** (temp file + `os.replace`, as above) so workers never parse a half-written file. If you produce the token without @@ -81,6 +71,23 @@ the SDK (e.g. a `curl` call to the token endpoint), you must also add the absolute `expires_at` yourself — the raw response only carries the relative `expires_in`, which is meaningless to a later reader. +## 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 diff --git a/openapi/templates/rest.mustache b/openapi/templates/rest.mustache index f6868ccf1..4a3a98c3c 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -58,7 +58,12 @@ class RESTClientObject: 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( @@ -66,8 +71,10 @@ class RESTClientObject: client_secret=client_secret, token_endpoint=token_endpoint, scope=scope, + leeway=leeway, **client_args, ) + self._token_leeway = leeway try: self.session.fetch_token() diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 5d12c09a3..8d4782859 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -68,7 +68,12 @@ def __init__(self, configuration) -> 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( @@ -76,8 +81,10 @@ def setup_oauth_client_credentials( client_secret=client_secret, token_endpoint=token_endpoint, scope=scope, + leeway=leeway, **client_args, ) + self._token_leeway = leeway try: self.session.fetch_token() diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index 98e0e7285..8fe7cd0de 100644 --- a/src/rapidata/rapidata_client/rapidata_client.py +++ b/src/rapidata/rapidata_client/rapidata_client.py @@ -111,7 +111,7 @@ def __init__( cert_path (str, optional): An optional path to a certificate file useful for development. token (dict, optional): If you already have a token that the client should use for authentication. Important, if set, this needs to be the complete token object containing the access token, token type and expiration time. token_file (str, optional): Path to a JSON file containing the token object described above (with an absolute ``expires_at`` timestamp). The file is re-read whenever the current token expires, so an external process can keep it fresh — useful to share one token across many workers (e.g. distributed training). Falls back to the ``RAPIDATA_TOKEN_FILE`` environment variable when omitted. - leeway (int, optional): An optional leeway to use to determine if a token is expired. Defaults to 60 seconds. + leeway (int, optional): How many seconds before its actual expiry a token is treated as expired — i.e. how early it is refreshed (or re-read from ``token_file``). Defaults to 60 seconds. Attributes: order (RapidataOrderManager): The RapidataOrderManager instance. diff --git a/src/rapidata/service/openapi_service.py b/src/rapidata/service/openapi_service.py index 871e89c5a..9cdec414d 100644 --- a/src/rapidata/service/openapi_service.py +++ b/src/rapidata/service/openapi_service.py @@ -113,6 +113,7 @@ def __init__( client_secret=client_secret, token_endpoint=f"{auth_endpoint}/connect/token", scope=oauth_scope, + leeway=leeway, ) except OAuthError as e: if e.error != "invalid_client": @@ -134,6 +135,7 @@ def __init__( client_secret=credentials.client_secret, token_endpoint=f"{auth_endpoint}/connect/token", scope=oauth_scope, + leeway=leeway, ) managed_print("Credentials were reset and re-authenticated successfully") From 36c887ef0691ea8e9fdc768bf35dd6d3fa10c6d6 Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Wed, 8 Jul 2026 12:51:53 +0000 Subject: [PATCH 4/4] feat(auth): add maintain_token_file helper for the coordinator loop Co-Authored-By: lino --- docs/distributed_training.md | 69 ++++++++++++------- .../rapidata_client/rapidata_client.py | 47 +++++++++++++ 2 files changed, 93 insertions(+), 23 deletions(-) diff --git a/docs/distributed_training.md b/docs/distributed_training.md index 771edde19..90e97b535 100644 --- a/docs/distributed_training.md +++ b/docs/distributed_training.md @@ -21,8 +21,41 @@ The fix is to authenticate **once** and share the token: ## Coordinator Start the coordinator first — it authenticates with the client credentials -as usual (see [Authentication](authentication.md)), exports its token with -`get_token()`, and keeps the shared file fresh: +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 @@ -31,9 +64,9 @@ import time from rapidata import RapidataClient -TOKEN_FILE = "/shared/rapidata_token.json" # (1)! +TOKEN_FILE = "/shared/rapidata_token.json" -coordinator = RapidataClient(leeway=300) # (2)! +coordinator = RapidataClient(leeway=300) os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True) @@ -48,28 +81,18 @@ def write_token(token: dict) -> None: while True: - write_token(coordinator.get_token()) # (3)! + write_token(coordinator.get_token()) time.sleep(60) ``` -1. For simplicity this example writes to a root-level `/shared` directory — - point it at whatever storage all your workers actually mount (an NFS - path, a shared volume, …). -2. 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` treats the token as expired once it is - within 5 minutes of its actual expiry, so the file is renewed well before - workers need it. -3. `get_token()` returns the current token — including the absolute - `expires_at` timestamp workers rely on — and only fetches a new one when - it is within `leeway` of expiry. Polling it every minute therefore does - **not** spam the auth server; most iterations just rewrite the same token. - -**Write the file atomically** (temp file + `os.replace`, as above) so -workers never parse a half-written file. If you produce the token without -the SDK (e.g. a `curl` call to the token endpoint), you must also add the -absolute `expires_at` yourself — the raw response only carries the relative -`expires_in`, which is meaningless to a later reader. +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 diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index 8fe7cd0de..c57df1472 100644 --- a/src/rapidata/rapidata_client/rapidata_client.py +++ b/src/rapidata/rapidata_client/rapidata_client.py @@ -215,6 +215,53 @@ def get_token(self) -> dict[str, Any]: """ return self._openapi_service.api_client.rest_client.get_token() + def maintain_token_file(self, path: str, interval: float = 60) -> threading.Thread: + """Continuously write this client's token to ``path`` so that other + workers can authenticate from it via ``token_file`` — see the + Distributed Training guide. + + Writes the file once immediately (creating the parent directory if + needed), then keeps rewriting it from a background daemon thread every + ``interval`` seconds. Writes are atomic, and a new token is only + fetched from the auth server shortly before the current one expires + (controlled by the client's ``leeway``). + + The file contains a bearer token — write it only to storage that the + consuming workers alone can read. + + Args: + path (str): Where to write the token file. + interval (float, optional): Seconds between rewrites. Defaults to 60. + + Returns: + threading.Thread: The daemon thread keeping the file fresh. Call + ``.join()`` on it to block the calling process forever. + """ + directory = os.path.dirname(os.path.abspath(path)) + os.makedirs(directory, exist_ok=True) + + def write_once(): + tmp_path = f"{path}.tmp.{os.getpid()}" + with open(tmp_path, "w") as f: + json.dump(self.get_token(), f) + os.replace(tmp_path, path) + + # First write happens synchronously so the file is guaranteed to + # exist (or a failure is raised here) before workers are started. + write_once() + + def loop(): + while True: + time.sleep(interval) + try: + write_once() + except Exception as e: + logger.warning("Failed to refresh token file %s: %s", path, e) + + thread = threading.Thread(target=loop, daemon=True, name="rapidata-token-file") + thread.start() + return thread + def reset_credentials(self): """Reset the credentials saved in the configuration file for the current environment.""" logger.info("Resetting credentials")