nvflare.app_opt.tracking.mlflow.mlflow_receiver module

class MLflowReceiver(tracking_uri: str | None = None, kw_args: dict | None = None, artifact_location: str | None = None, events: List[str] | None = None, buffer_flush_time=1)

Bases: AnalyticsReceiver

MLflowReceiver receives log events from clients and deliver them to the MLflow tracking server.

Parameters:

• tracking_uri (Optional[str], optional) – MLflow tracking server URI. When this is not specified, the metrics will be written to the local file system. If the tracking URI is specified, the MLflow tracking server must started before running the job. Defaults to None.

• kw_args (Optional[dict], optional) – keyword arguments: “experiment_name” (str): Specifies the experiment name. If not specified, the default name of “FLARE FL Experiment” will be used. “run_name” (str): Specifies the run name “experiment_tags” (dict): Tags used when creating the MLflow experiment. “mlflow.note.content” is a special MLflow tag. When provided, it displays as experiment description field on the MLflow UI. You can use Markdown syntax for the description. “run_tags” (str): Tags used when creating the MLflow run. “mlflow.note.content” is a special MLflow tag. When provided, it displays as run description field on the MLflow UI. You can use Markdown syntax for the description.

• artifact_location (Optional[str], optional) – Relative location of artifacts. Currently only text is supported at the moment.

• events (optional, List[str]) – A list of event that this receiver will handle.

• buffer_flush_time (int, optional) – The time in seconds between deliveries of event data to the MLflow tracking server. The data is buffered and then delivered to the MLflow tracking server in batches, and the buffer_flush_time controls the frequency of the sending. By default, the buffer flushes every second. You can reduce the time to a fraction of a second if you prefer less delay. Keep in mind that reducing the buffer_flush_time will potentially cause high traffic to the MLflow tracking server, which in some cases can actually cause more latency.

buffer_data(data: AnalyticsData, record_origin: str) → None

Buffer the data to send later.

A buffer for each data_type is in each site_buffer, all of which are in self.buffer

Parameters:

• data (AnalyticsData) – Data.

• record_origin (str) – Origin of the data, or site name.

finalize(fl_ctx: FLContext)

Finalizes the receiver.

Called after EventType.END_RUN.

Parameters:

fl_ctx (FLContext) – fl context.

flush_buffer(log_buffer: List)

flush_buffers(record_origin)

Flush buffers and send all the data to the MLflow tracking server.

Parameters:

record_origin (str) – Origin of the data, or site name.

get_mlflow_client(site_id: str) → MlflowClient

get_run_id(site_id: str) → str | None

get_target_type(data_type: AnalyticsDataType)

initialize(fl_ctx: FLContext)

Initializes MlflowClient for each site.

This method: 1. Sets up the FL context and timing 2. Validates and prepares experiment configuration 3. Determines participating sites 4. Sets up MLflow clients and experiments for each site 5. Initializes data buffers

Parameters:

fl_ctx (FLContext) – The FLContext containing runtime information

Raises:

• ValueError – If experiment name is empty

• RuntimeError – If unable to determine participating sites

save(fl_ctx: FLContext, shareable: Shareable, record_origin: str)

Saves the received data.

Specific implementations of AnalyticsReceiver will implement save in their own way.

Parameters:

• fl_ctx (FLContext) – fl context.

• shareable (Shareable) – the received message.

• record_origin (str) – the sender of this message / record.

class MlflowConstants

Bases: object

EXPERIMENT_NAME = 'experiment_name'

EXPERIMENT_TAG = 'experiment_tag'

RUN_TAG = 'run_tag'

get_current_time_millis()