Skip to content

feat(mlflow): add StaticTokenAuth for HTTP bearer-token servers - #338

Open
HCookie wants to merge 1 commit into
mainfrom
feat/custom-mlflow
Open

feat(mlflow): add StaticTokenAuth for HTTP bearer-token servers#338
HCookie wants to merge 1 commit into
mainfrom
feat/custom-mlflow

Conversation

@HCookie

@HCookie HCookie commented Aug 4, 2026

Copy link
Copy Markdown
Member

Summary

Adds StaticTokenAuth so anemoi.utils.mlflow can authenticate to MLflow servers that accept a static, pre-issued HTTP bearer token, instead of only the interactive Keycloak refresh-token flow (TokenAuth).

What's included

  • StaticTokenAuth (auth.py): an AuthBase that stores the token in MLFLOW_TRACKING_TOKEN (so MLflow sends Authorization: Bearer <token>). It reuses the existing on-disk store (~/.anemoi/mlflow-token.json), keyed by server URL, so a token supplied once via login() is remembered across sessions. No token server / refresh flow.
  • AnemoiMlflowClient(static_token=...) (client.py): pass a token string to use it directly, True to load a previously saved static token, or leave unset to fall back to TokenAuth (unchanged default behaviour).
  • Shared enabled_guard decorator (factored out of TokenAuth); ServerConfig gains a static_token field. The deprecated load_config() contract is preserved (the new field is excluded from its output).
  • Tests: StaticTokenAuth unit tests + client selection-logic tests (34 pass locally; pre-commit clean).

Usage

from anemoi.utils.mlflow.client import AnemoiMlflowClient

client = AnemoiMlflowClient(
    "https://my-mlflow.example.com",
    authentication=True,
    static_token="<http bearer token>",
)
client.search_experiments()

Notes / caveats

  • If the server sits behind a proxy/gateway that enforces per-method permissions, the token must allow the HTTP methods the MLflow client uses. The client issues several read operations as POST, so a read-only token may be insufficient.

Backwards compatibility

No behaviour change unless static_token is passed; TokenAuth remains the default. Draft pending review of the API surface (static_token: str | bool).

Add StaticTokenAuth, an AuthBase that authenticates with a static, pre-issued
bearer token instead of the interactive Keycloak refresh-token flow. It reuses
the existing on-disk token store (~/.anemoi/mlflow-token.json), keyed by server
URL, so a token supplied once via login() is remembered across sessions.

AnemoiMlflowClient gains a `static_token` parameter: pass a token string to use
it directly, or True to load a previously saved static token; otherwise it
falls back to TokenAuth as before.

This enables using anemoi-utils against MLflow servers that accept a long-lived
HTTP access/bearer token.

- auth.py: StaticTokenAuth, shared enabled_guard decorator, ServerConfig gains
  a static_token field (legacy load_config() contract preserved)
- client.py: static_token parameter wiring
- tests: StaticTokenAuth unit tests + client selection-logic tests
@HCookie
HCookie force-pushed the feat/custom-mlflow branch from c316291 to b31c03e Compare August 4, 2026 19:49
@HCookie
HCookie requested a review from gmertes August 4, 2026 19:54
@HCookie HCookie self-assigned this Aug 4, 2026
@HCookie
HCookie marked this pull request as ready for review August 4, 2026 19:54

@gmertes gmertes left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any benefit to sharing the same token file? Given that the new class is already standalone and not inheriting from the other one, would it not make more sense to also keep the token file separate?

I'm also not sure if the login logic will even work with a shared file, because the login command in anemoi-training will iterate over all urls in the store and attempt to login to them. Let's discuss when you're back from leave?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: To be triaged

Development

Successfully merging this pull request may close these issues.

2 participants