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..90e97b535 --- /dev/null +++ b/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)) +``` 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..4a3a98c3c 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -54,9 +54,16 @@ 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( @@ -64,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() @@ -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): diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 01cb2625e..8d4782859 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -64,9 +64,16 @@ 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( @@ -74,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() @@ -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): diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index fb5aed018..c57df1472 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,7 +110,8 @@ 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. - leeway (int, optional): An optional leeway to use to determine if a token is expired. Defaults to 60 seconds. + 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): 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. @@ -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, ) @@ -195,6 +203,65 @@ 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 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") diff --git a/src/rapidata/service/openapi_service.py b/src/rapidata/service/openapi_service.py index 615408eff..9cdec414d 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 @@ -107,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": @@ -128,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") @@ -249,6 +257,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: