MLFlow Metric Storage

MLFlow Metric Storage#

class rizemind.logging.mlflow.metric_storage.MLFLowMetricStorage(experiment_name: str, run_name: str, mlflow_uri: str)[source]#

Bases: BaseMetricStorage

A concrete implementation of BaseMetricStorage that logs metrics and models to an MLflow tracking server.

This class integrates Flower federated learning with MLflow, enabling centralized tracking of experiments, metrics, and model artifacts. Upon initialization, it connects to a specified MLflow tracking URI, sets up an experiment, and creates a new run to store all subsequent data.

experiment_name#

The name of the MLflow experiment.

run_name#

The name of the MLflow run.

mlflow_uri#

The URI for the MLflow tracking server.

mlflow_client#

The MLflow client for interacting with the API.

run_id#

The unique ID of the MLflow run created for this session.

update_best_model(server_round: int, loss: float)[source]#

Saves the current model as an MLflow artifact if its loss is the lowest seen so far.

It compares the provided loss with its internally tracked best loss. If the new loss is lower, it updates the best loss and serializes the in-memory model parameters to a temporary .npz file. This file is then uploaded as an artifact to the MLflow run. It also logs the best round and loss as metrics.

Parameters:

  • server_round – The server round that produced this model.

  • loss – The loss value of the current model, used to determine if it is the new best model.

update_current_round_model(parameters: Parameters)[source]#

Temporarily stores the model parameters for the current round in memory.

This method holds the latest model parameters so they can be saved as an MLflow artifact later by update_best_model if this model proves to be the best one based on its loss.

Parameters:

parameters – The model parameters from the current round.

write_metrics(server_round: int, metrics: dict[str, bool | bytes | float | int | str])[source]#

Logs a dictionary of metrics to the MLflow run for a specific server round.

This method iterates through the provided metrics and logs each one to the active MLflow run, using the server round as the step.

Parameters:

  • server_round – The current round of federated learning, used as the ‘step’ in MLflow.

  • metrics – A dictionary mapping metric names (e.g., “accuracy”) to their scalar values.

Contents