From 42deccd41087eb45abcc67464d0c9101f5204bd7 Mon Sep 17 00:00:00 2001 From: RapidPoseidon Date: Thu, 9 Jul 2026 09:00:12 +0000 Subject: [PATCH 1/2] feat(auth): add set_token and document push/pull token distribution Co-Authored-By: lino --- docs/distributed_training.md | 61 ++++++++----------- openapi/templates/rest.mustache | 8 +++ src/rapidata/api_client/rest.py | 8 +++ .../rapidata_client/rapidata_client.py | 15 +++++ 4 files changed, 58 insertions(+), 34 deletions(-) diff --git a/docs/distributed_training.md b/docs/distributed_training.md index 90e97b535..4582e48d2 100644 --- a/docs/distributed_training.md +++ b/docs/distributed_training.md @@ -49,50 +49,43 @@ coordinator.maintain_token_file("/shared/rapidata_token.json").join() # (2)! that your training job alone can access, and restrict the file permissions accordingly. -### Writing the file yourself +### Building your own token distribution -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. +The shared file is just one transport. The SDK exposes the underlying +primitives directly, so you can move the token around however fits your +infrastructure — a key-value store, an RPC from the launcher, a secret +manager, a message queue: ```python -import json -import os -import time - from rapidata import RapidataClient -TOKEN_FILE = "/shared/rapidata_token.json" - -coordinator = RapidataClient(leeway=300) +# 1. Export — the credential-holding client hands out its current token, +# refreshing it first if it is about to expire. +token = coordinator.get_token() -os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True) +# ... distribute it through any transport you like ... +# 2. Bootstrap — create a worker client directly from a token object. +worker = RapidataClient(token=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: - write_token(coordinator.get_token()) - time.sleep(60) +# 3. Renew — inject a newer token into a running client at any time; +# it is used from the next request on. +worker.set_token(fresh_token) ``` -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. +With these three calls you can build a **push** system (the coordinator +distributes a fresh token to every worker before the old one expires, each +worker applies it with `set_token`) or a **pull** system (each worker +periodically fetches the current token from your own endpoint). `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. + +Whatever the transport, pass the **complete token object** around and keep +its absolute `expires_at` field — it's how a client decides when a token has +expired. If you produce tokens 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. And if your transport is a shared file, write it +atomically (temp file + `os.replace`) so readers never see a partial token. ## Workers diff --git a/openapi/templates/rest.mustache b/openapi/templates/rest.mustache index 4a3a98c3c..de5452f01 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -132,6 +132,14 @@ class RESTClientObject: raise ApiValueError("The current token is expired and cannot be refreshed.") return dict(self.session.token) + def set_token(self, token: dict): + """Replace the token the session authenticates with.""" + if not self.session: + raise ApiValueError( + "OAuth2 session is not initialized. Please initialize it before making requests." + ) + self.session.token = token + def request( self, method, diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 8d4782859..1917526b6 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -142,6 +142,14 @@ def get_token(self) -> dict: raise ApiValueError("The current token is expired and cannot be refreshed.") return dict(self.session.token) + def set_token(self, token: dict): + """Replace the token the session authenticates with.""" + if not self.session: + raise ApiValueError( + "OAuth2 session is not initialized. Please initialize it before making requests." + ) + self.session.token = token + def request( self, method, diff --git a/src/rapidata/rapidata_client/rapidata_client.py b/src/rapidata/rapidata_client/rapidata_client.py index c57df1472..66a6738a9 100644 --- a/src/rapidata/rapidata_client/rapidata_client.py +++ b/src/rapidata/rapidata_client/rapidata_client.py @@ -215,6 +215,21 @@ def get_token(self) -> dict[str, Any]: """ return self._openapi_service.api_client.rest_client.get_token() + def set_token(self, token: dict[str, Any]) -> None: + """Replace the token this client authenticates with, effective from + the next request. + + The counterpart to ``get_token``: inject a fresh token — obtained from + another client and distributed through a transport of your choice — + into a running client, without reconstructing it. Expects the complete + token object (``access_token``, ``token_type`` and an absolute + ``expires_at`` timestamp) — see the Distributed Training guide. + + Args: + token (dict): The complete token object to authenticate with. + """ + self._openapi_service.api_client.rest_client.set_token(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 From b6c7e9f94e3212fef22f91a6b3d540ce98583ba3 Mon Sep 17 00:00:00 2001 From: Lino Giger <68745352+LinoGiger@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:52:50 +0200 Subject: [PATCH 2/2] added warning log if token is not complete --- openapi/templates/rest.mustache | 8 ++++++++ src/rapidata/api_client/rest.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/openapi/templates/rest.mustache b/openapi/templates/rest.mustache index de5452f01..a9fbb1b24 100644 --- a/openapi/templates/rest.mustache +++ b/openapi/templates/rest.mustache @@ -138,6 +138,14 @@ class RESTClientObject: raise ApiValueError( "OAuth2 session is not initialized. Please initialize it before making requests." ) + missing = [k for k in ("access_token", "token_type", "expires_at") if k not in token] + if missing: + _logger.warning( + "set_token received an incomplete token object (missing %s). " + "Expect the complete token from get_token() — access_token, " + "token_type, and an absolute expires_at timestamp.", + ", ".join(missing), + ) self.session.token = token def request( diff --git a/src/rapidata/api_client/rest.py b/src/rapidata/api_client/rest.py index 1917526b6..26906d053 100644 --- a/src/rapidata/api_client/rest.py +++ b/src/rapidata/api_client/rest.py @@ -148,6 +148,14 @@ def set_token(self, token: dict): raise ApiValueError( "OAuth2 session is not initialized. Please initialize it before making requests." ) + missing = [k for k in ("access_token", "token_type", "expires_at") if k not in token] + if missing: + _logger.warning( + "set_token received an incomplete token object (missing %s). " + "Expect the complete token from get_token() — access_token, " + "token_type, and an absolute expires_at timestamp.", + ", ".join(missing), + ) self.session.token = token def request(