"""Reusable CloudFerro authentication for MLflow tracking with automatic token refresh. MLflow reads MLFLOW_TRACKING_TOKEN once, when the tracking store is created, so rotating that variable in a running process does not work. Instead, this module registers a custom MLflow request auth provider (named "cloudferro") that MLflow calls before *every* request. The provider returns a requests-compatible auth object which keeps the CloudFerro SSO session alive via the refresh token grant. This covers all MLflow traffic, including framework autologging (TensorFlow, Keras, sklearn, ...), which calls MLflow internally and cannot be wrapped in your own retry logic. Refresh strategy: each refresh is scheduled at half of whichever deadline comes first — the access token lifetime or the REFRESH TOKEN'S OWN expiry, which encodes the SSO session idle deadline. Refreshing right before the access token expires is too late whenever the session idle timeout is shorter than the access token lifetime: refresh tokens die with the idle timeout (`invalid_grant: Token is not active`). Each refresh resets the idle timer, so staying ahead of it keeps the session alive indefinitely (up to the SSO session max). If a refresh still fails, the session falls back to a full re-login (password grant), prompting for a fresh TOTP code instead of killing the run. Usage: import cloudferro_auth cloudferro_auth.setup_mlflow("https://mlflow-.mlops.cloudferro.com") mlflow.set_experiment("my-experiment") ... Required environment variables: MLFLOW_USERNAME - platform username (same as the web UI login) MLFLOW_PASSWORD - platform password Optional: MLFLOW_TOTP - 6-digit 2FA code; prompted interactively when unset. CLOUDFERRO_TOKEN_URL - defaults to the CloudFerro Cloud production realm CLOUDFERRO_CLIENT_ID - defaults to "mlflow-cli-authservice" """ import base64 import json import os import sys import threading import time import requests from mlflow.tracking.request_auth.abstract_request_auth_provider import RequestAuthProvider # Defaults point at the CloudFerro Cloud production realm. Override with # CLOUDFERRO_TOKEN_URL / CLOUDFERRO_CLIENT_ID for other environments. DEFAULT_TOKEN_URL = ( "https://identity.cloudferro.com/auth/realms/CloudFerro-Cloud/protocol/openid-connect/token" ) DEFAULT_CLIENT_ID = "mlflow-cli-authservice" PROVIDER_NAME = "cloudferro" # Schedule each refresh at half of the relevant lifetime (access token or # refresh token — see module docstring), never more often than the floor. REFRESH_LIFETIME_FRACTION = 0.5 MIN_REFRESH_DELAY_SECONDS = 30 def _jwt_exp(token) -> float | None: """Read the exp claim from a JWT without verifying the signature.""" try: payload = token.split(".")[1] payload += "=" * (-len(payload) % 4) return float(json.loads(base64.urlsafe_b64decode(payload))["exp"]) except Exception: return None class CloudFerroSession(requests.auth.AuthBase): """Holds the SSO session and injects a fresh Bearer token into every request.""" def __init__(self): self.token_url = os.environ.get("CLOUDFERRO_TOKEN_URL", DEFAULT_TOKEN_URL) self.client_id = os.environ.get("CLOUDFERRO_CLIENT_ID", DEFAULT_CLIENT_ID) self.username = os.environ["MLFLOW_USERNAME"] self.password = os.environ["MLFLOW_PASSWORD"] self._lock = threading.Lock() self._access_token = None self._refresh_token = None self._expires_at = 0.0 self._refresh_after = 0.0 self._refresh_exp = None self._last_expires_in = 0 def _request(self, payload): response = requests.post(self.token_url, data=payload, timeout=30) if response.status_code != 200: try: error = response.json() detail = f"{error.get('error')}: {error.get('error_description')}" except ValueError: detail = response.text raise RuntimeError( f"Token request failed (HTTP {response.status_code}): {detail}" ) return response.json() def _store(self, data): self._access_token = data["access_token"] # The server rotates the refresh token on every refresh — always keep the latest one. self._refresh_token = data.get("refresh_token", self._refresh_token) self._last_expires_in = data.get("expires_in", 300) now = time.time() self._expires_at = now + self._last_expires_in # The refresh token's exp encodes the SSO session idle deadline. # Refresh at half of whichever lifetime runs out first, so the # session idle timer never trips before we use the token. self._refresh_exp = _jwt_exp(self._refresh_token) access_half = self._last_expires_in * REFRESH_LIFETIME_FRACTION if self._refresh_exp is not None: refresh_half = max( (self._refresh_exp - now) * REFRESH_LIFETIME_FRACTION, MIN_REFRESH_DELAY_SECONDS, ) self._refresh_after = now + min(access_half, refresh_half) else: self._refresh_after = now + access_half def _password_grant(self, reprompt: bool = False): totp = os.environ.get("MLFLOW_TOTP", "").strip() if sys.stdin.isatty(): prompt = ( "Fresh TOTP code for re-login (6 digits, Enter to reuse MLFLOW_TOTP or skip): " if reprompt else "TOTP code (6 digits, Enter to skip if your account has no 2FA): " ) entered = input(prompt).strip() totp = entered or totp payload = { "client_id": self.client_id, "grant_type": "password", "username": self.username, "password": self.password, } if totp: payload["totp"] = totp self._store(self._request(payload)) refresh_ttl = ( int(self._refresh_exp - time.time()) if self._refresh_exp is not None else "unknown" ) print( f"[cloudferro_auth] logged in: access token valid {self._last_expires_in}s, " f"refresh token valid ~{refresh_ttl}s", file=sys.stderr, ) def _refresh_grant(self): self._store( self._request( { "client_id": self.client_id, "grant_type": "refresh_token", "refresh_token": self._refresh_token, } ) ) print( f"[cloudferro_auth] access token refreshed " f"(next refresh in {int(self._refresh_after - time.time())}s)", file=sys.stderr, ) def token(self) -> str: """Return a valid access token, logging in, refreshing, or re-logging in as needed.""" with self._lock: if self._access_token is None: self._password_grant() elif time.time() >= self._refresh_after: try: self._refresh_grant() except RuntimeError as exc: # The SSO session is gone (e.g. idle timeout). Do not kill # the run — fall back to a full re-login with a fresh TOTP. print( f"[cloudferro_auth] refresh failed ({exc}); falling back to re-login", file=sys.stderr, ) self._password_grant(reprompt=True) return self._access_token def __call__(self, request): request.headers["Authorization"] = f"Bearer {self.token()}" return request _session = None def get_session() -> CloudFerroSession: global _session if _session is None: _session = CloudFerroSession() return _session class CloudFerroRequestAuthProvider(RequestAuthProvider): """MLflow request auth provider. get_auth() is called before every MLflow request.""" def get_name(self): return PROVIDER_NAME def get_auth(self): return get_session() def _register_provider(): # No public registration API — the registry is a module-level singleton # that MLflow consults via fetch_auth() on every request. from mlflow.tracking.request_auth.registry import _request_auth_provider_registry if not any(p.get_name() == PROVIDER_NAME for p in _request_auth_provider_registry): _request_auth_provider_registry.register(CloudFerroRequestAuthProvider) def setup_mlflow(tracking_uri: str) -> None: """Register the auth provider, perform the initial login, and set the tracking URI. Fails fast on bad credentials or TOTP, before any training work starts. """ import mlflow _register_provider() # A static token would take precedence over the auth provider — make sure it is gone. os.environ.pop("MLFLOW_TRACKING_TOKEN", None) os.environ["MLFLOW_TRACKING_AUTH"] = PROVIDER_NAME get_session().token() mlflow.set_tracking_uri(tracking_uri)