anemoi.utils.mlflow package

Submodules

anemoi.utils.mlflow.auth module

class anemoi.utils.mlflow.auth.UserInfo(*, name: str | None = None, email: str | None = None, preferred_username: str | None = None)

Bases: BaseModel

User info associated with a token.

name: str | None

email: str | None

username: str | None

classmethod from_jwt(token: str) → UserInfo: Create a UserInfo object from a JWT. The token payload must contain the user info fields.

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.utils.mlflow.auth.ServerConfig(*, refresh_token: str | None = None, refresh_expires: int = 0)

Bases: BaseModel

refresh_token: str | None

refresh_expires: int

classmethod to_int(value: float | int) → int

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.utils.mlflow.auth.ServerStore(root: RootModelRootType = PydanticUndefined)

Bases: RootModel

root: dict[str, ServerConfig]

get(url: str) → ServerConfig | None

items()

update(url, config: ServerConfig) → None: Update the server configuration for a given URL.

property servers: list[tuple[str, int]]: List of servers in the store, as a tuple (url, refresh_expires). Ordered most recently used first.

classmethod load_legacy_format(data: dict) → dict: Convert legacy single-server config format to multi-server.

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.utils.mlflow.auth.AuthBase(*args, **kwargs)

Bases: ABC

Base class for authentication implementations.

abstractmethod save(**kwargs)

abstractmethod login(force_credentials: bool = False, **kwargs)

abstractmethod authenticate(**kwargs)

abstractmethod user_info() → UserInfo

class anemoi.utils.mlflow.auth.NoAuth(*args, **kwargs)

Bases: AuthBase

No-op authentication class.

save(**kwargs)

login(force_credentials: bool = False, **kwargs)

authenticate(**kwargs)

user_info() → UserInfo

class anemoi.utils.mlflow.auth.TokenAuth(url: str, enabled: bool = True, target_env_var: str = 'MLFLOW_TRACKING_TOKEN')

Bases: AuthBase

Manage authentication with a keycloak token server.

access_token: str | None

property refresh_token: str

static get_servers() → list[tuple[str, int]]: List of all saved servers, as a tuple (url, refresh_expires). Ordered most recently used first.

static load_config() → dict

Load the last used server configuration

Returns:: config – Dictionary with the following keys: url, refresh_token, refresh_expires. If no configuration is found, an empty dictionary is returned.
Return type:: dict

enabled() → Callable: Decorator to call or ignore a function based on the enabled flag.

login(force_credentials: bool = False, **kwargs: dict) → None

Acquire a new refresh token and save it to disk.

If an existing valid refresh token is already on disk it will be used. If not, or the token has expired, the user will be asked to obtain one from the API.

Refresh token expiry time is set in the REFRESH_EXPIRE_DAYS constant (default 29 days).

This function should be called once, interactively, right before starting a training run.

Parameters:

force_credentials (bool, optional) – Force a credential login even if a refreh token is available, by default False.
kwargs (dict) – Additional keyword arguments.

Raises:

RuntimeError – A new refresh token could not be acquired.

authenticate(**kwargs: dict) → None

Check the access token and refresh it if necessary. A new refresh token will also be acquired upon refresh.

This requires a valid refresh token to be available, obtained from the login method.

The access token is stored in memory and in an environment variable. If the access token is still valid, this function does nothing.

This function should be called before every MLflow API request.

Raises:: RuntimeError – No refresh token is available or the token request failed.

save(**kwargs: dict) → None: Save the latest refresh token to disk.

user_info() → UserInfo: Get user info embedded in the access token.

anemoi.utils.mlflow.client module

anemoi.utils.mlflow.utils module

anemoi.utils.mlflow.utils.health_check(tracking_uri: str) → None

Query the health endpoint of an MLflow server.

If the server is not reachable, raise an error and remind the user that authentication may be required.

Raises:: ConnectionError – If the server is not reachable.