From 3aae92c5d6a13d0b5ee911925c237b63a59c7af0 Mon Sep 17 00:00:00 2001 From: uditmahato Date: Sun, 9 Aug 2026 12:22:49 +0545 Subject: [PATCH 1/2] Record system resource usage with MLFlowHandler Add a log_system_metrics option to MLFlowHandler that samples CPU, memory, disk, network and GPU usage while a workflow runs, so that the resource usage is recorded through the handler rather than next to it. The sampling is done by mlflow itself and lands in the run of the workflow, under the system/ prefix. The handlers of a workflow share a run, so the run is sampled by the first handler that starts it and left alone by the others. Fixes #7405 Signed-off-by: uditmahato --- monai/handlers/mlflow_handler.py | 93 +++++++++++++++++++++++++++ tests/handlers/test_handler_mlflow.py | 83 +++++++++++++++++++++++- 2 files changed, 175 insertions(+), 1 deletion(-) diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 3078d89f97c..6c6db0d03df 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -12,6 +12,7 @@ from __future__ import annotations import os +import threading import time import warnings from collections.abc import Callable, Mapping, Sequence @@ -34,6 +35,9 @@ ) pandas, _ = optional_import("pandas", descriptor="Please install pandas for recording the dataset.") tqdm, _ = optional_import("tqdm", "4.47.0", min_version, "tqdm") +SystemMetricsMonitor, has_system_metrics = optional_import( + "mlflow.system_metrics.system_metrics_monitor", name="SystemMetricsMonitor" +) if TYPE_CHECKING: from ignite.engine import Engine @@ -113,6 +117,18 @@ class MLFlowHandler: optimizer_param_names: parameter names in the optimizer that need to be recorded during running the workflow, default to `'lr'`. close_on_complete: whether to close the mlflow run in `complete` phase in workflow, default to False. + log_system_metrics: whether to record system resource usage (CPU, memory, disk, network and GPU) + while the workflow runs, default to False. The metrics are sampled in a background thread by + MLflow itself and stored in the same run as the workflow metrics, under the `system/` prefix. + Requires `psutil`, and `pynvml` in addition for the GPU metrics. Note that MLflow reads the + run through the global tracking URI to sample it, so enabling this sets the global tracking + URI to `tracking_uri`; a process that tracks to several URIs at the same time should keep + this disabled. + system_metrics_sampling_interval: seconds between two samples of the system metrics, default to + `None`, which keeps the MLflow default (10 seconds). Only used if `log_system_metrics` is True. + system_metrics_samples_before_logging: number of samples to aggregate before they are logged, + default to `None`, which keeps the MLflow default (1 sample). Only used if `log_system_metrics` + is True. For more details of MLFlow usage, please refer to: https://mlflow.org/docs/latest/index.html. @@ -121,6 +137,10 @@ class MLFlowHandler: # parameters that are logged at the start of training default_tracking_params = ["max_epochs", "epoch_length"] + # runs whose system metrics are being sampled, so that handlers sharing a run sample it once + _monitored_run_ids: set[str] = set() + _system_metrics_lock = threading.Lock() + def __init__( self, tracking_uri: str | None = None, @@ -141,6 +161,9 @@ def __init__( artifacts: str | Sequence[Path] | None = None, optimizer_param_names: str | Sequence[str] = "lr", close_on_complete: bool = False, + log_system_metrics: bool = False, + system_metrics_sampling_interval: int | None = None, + system_metrics_samples_before_logging: int | None = None, ) -> None: self.iteration_log = iteration_log self.epoch_log = epoch_log @@ -156,9 +179,15 @@ def __init__( self.experiment_param = experiment_param self.artifacts = ensure_tuple(artifacts) self.optimizer_param_names = ensure_tuple(optimizer_param_names) + self.tracking_uri = tracking_uri self.client = mlflow.MlflowClient(tracking_uri=tracking_uri if tracking_uri else None) self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete + self.log_system_metrics = log_system_metrics + self.system_metrics_sampling_interval = system_metrics_sampling_interval + self.system_metrics_samples_before_logging = system_metrics_samples_before_logging + self.system_metrics_monitor = None + self._monitored_run_id: str | None = None self.experiment = None self.cur_run = None self.dataset_dict = dataset_dict @@ -238,6 +267,66 @@ def start(self, engine: Engine) -> None: else: self._default_dataset_log(self.dataset_dict) + if self.log_system_metrics: + self._start_system_metrics_monitor() + + def _start_system_metrics_monitor(self) -> None: + """ + Start sampling the system resource usage of the current run, if it is not sampled yet. + + A workflow attaches one handler per engine, and those handlers share a run, so the run is + sampled by the first handler that starts and left alone by the other ones. + """ + if self.system_metrics_monitor is not None or self.cur_run is None: + return + + if not has_system_metrics: + warnings.warn("Please install mlflow>=2.8.0 to record the system metrics.") + return + + run_id = self.cur_run.info.run_id + with MLFlowHandler._system_metrics_lock: + if run_id in MLFlowHandler._monitored_run_ids: + return + + # mlflow reads the run to sample through the global tracking URI, not through the client + if self.tracking_uri: + mlflow.set_tracking_uri(self.tracking_uri) + + kwargs = {} + if self.system_metrics_sampling_interval is not None: + kwargs["sampling_interval"] = self.system_metrics_sampling_interval + if self.system_metrics_samples_before_logging is not None: + kwargs["samples_before_logging"] = self.system_metrics_samples_before_logging + + try: + monitor = SystemMetricsMonitor(run_id, **kwargs) + monitor.start() + except Exception as e: + # a workflow should not fail because its resource usage cannot be recorded + warnings.warn(f"Failed to record the system metrics: {e}") + return + + MLFlowHandler._monitored_run_ids.add(run_id) + self.system_metrics_monitor = monitor + self._monitored_run_id = run_id + + def _stop_system_metrics_monitor(self) -> None: + """ + Stop sampling the system resource usage, if this handler is the one sampling it. + """ + if self.system_metrics_monitor is None: + return + + with MLFlowHandler._system_metrics_lock: + try: + self.system_metrics_monitor.finish() + except Exception as e: + warnings.warn(f"Failed to stop recording the system metrics: {e}") + MLFlowHandler._monitored_run_ids.discard(self._monitored_run_id) + self.system_metrics_monitor = None + self._monitored_run_id = None + def _set_experiment(self): experiment = self.experiment if not experiment: @@ -331,6 +420,8 @@ def complete(self) -> None: """ Handler for train or validation/evaluation completed Event. """ + self._stop_system_metrics_monitor() + if self.artifacts and self.cur_run: artifact_list = self._parse_artifacts() for artifact in artifact_list: @@ -341,6 +432,8 @@ def close(self) -> None: Stop current running logger of MLFlow. """ + self._stop_system_metrics_monitor() + if self.cur_run: self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status) self.cur_run = None diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 80630e6f5a2..f14eb8ff0e0 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -17,7 +17,7 @@ import tempfile import unittest from concurrent.futures import ThreadPoolExecutor -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np from ignite.engine import Engine, Events @@ -230,6 +230,87 @@ def _update_metric(engine): else: self.assertEqual(handler._default_iteration_log.call_count, 2) # 2 = len([1, 3]) from event_filter + def test_system_metrics_disabled_by_default(self): + with tempfile.TemporaryDirectory() as tempdir: + + def _train_func(engine, batch): + return [batch + 1.0] + + engine = Engine(_train_func) + test_path = os.path.join(tempdir, "mlflow_system_metrics_off") + handler = MLFlowHandler(iteration_log=False, tracking_uri=path_to_uri(test_path), close_on_complete=True) + handler.attach(engine) + engine.run(range(3), max_epochs=1) + + self.assertIsNone(handler.system_metrics_monitor) + run = handler.client.get_run(handler.cur_run.info.run_id) if handler.cur_run else None + self.assertIsNone(run) + + def test_system_metrics_monitor_life_cycle(self): + with tempfile.TemporaryDirectory() as tempdir: + + def _train_func(engine, batch): + return [batch + 1.0] + + engine = Engine(_train_func) + test_path = os.path.join(tempdir, "mlflow_system_metrics") + handler = MLFlowHandler( + iteration_log=False, + tracking_uri=path_to_uri(test_path), + log_system_metrics=True, + system_metrics_sampling_interval=1, + system_metrics_samples_before_logging=1, + close_on_complete=True, + ) + monitor = MagicMock() + with patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class: + handler.attach(engine) + engine.run(range(3), max_epochs=1) + + # the monitor samples the run of the handler, with the requested sampling settings + monitor_class.assert_called_once() + self.assertEqual(monitor_class.call_args.kwargs["sampling_interval"], 1) + self.assertEqual(monitor_class.call_args.kwargs["samples_before_logging"], 1) + monitor.start.assert_called_once() + # the sampling is stopped when the workflow completes + monitor.finish.assert_called_once() + self.assertIsNone(handler.system_metrics_monitor) + + def test_system_metrics_monitor_shared_by_handlers(self): + with tempfile.TemporaryDirectory() as tempdir: + + def _train_func(engine, batch): + return [batch + 1.0] + + engine = Engine(_train_func) + test_path = os.path.join(tempdir, "mlflow_system_metrics_shared") + # a workflow attaches one handler per engine, all of them sharing a run + handlers = [ + MLFlowHandler( + iteration_log=False, tracking_uri=path_to_uri(test_path), run_name="shared", log_system_metrics=True + ) + for _ in range(3) + ] + monitor = MagicMock() + with patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class: + for handler in handlers: + handler.start(engine) + + # the run is sampled by the first handler only + monitor_class.assert_called_once() + + # the handlers that do not sample the run leave it running when they complete + for handler in handlers[1:]: + handler.complete() + monitor.finish.assert_not_called() + + # the sampling stops when the handler that started it completes + handlers[0].complete() + monitor.finish.assert_called_once() + + for handler in handlers: + handler.close() + def test_multi_thread(self): test_uri_list = ["monai_mlflow_test1", "monai_mlflow_test2"] with ThreadPoolExecutor(2, "Training") as executor: From ad5bc5d983209e3b0365bd1b1c909c7013792e8b Mon Sep 17 00:00:00 2001 From: uditmahato Date: Mon, 10 Aug 2026 13:30:12 +0545 Subject: [PATCH 2/2] Address review: guard the tracking uri call and validate the settings Move the tracking uri call inside the block that catches failures, so that a workflow cannot die because the uri could not be set, which was the intent of that block already. Reject a sampling interval or a sample count that is not positive, as mlflow does not define a behaviour for those, and document the new tests. Signed-off-by: uditmahato --- monai/handlers/mlflow_handler.py | 14 ++++++++++---- tests/handlers/test_handler_mlflow.py | 24 ++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 6c6db0d03df..351346b3a15 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -184,6 +184,12 @@ def __init__( self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete self.log_system_metrics = log_system_metrics + for name, value in ( + ("system_metrics_sampling_interval", system_metrics_sampling_interval), + ("system_metrics_samples_before_logging", system_metrics_samples_before_logging), + ): + if value is not None and value <= 0: + raise ValueError(f"`{name}` must be a positive number, got {value}.") self.system_metrics_sampling_interval = system_metrics_sampling_interval self.system_metrics_samples_before_logging = system_metrics_samples_before_logging self.system_metrics_monitor = None @@ -289,10 +295,6 @@ def _start_system_metrics_monitor(self) -> None: if run_id in MLFlowHandler._monitored_run_ids: return - # mlflow reads the run to sample through the global tracking URI, not through the client - if self.tracking_uri: - mlflow.set_tracking_uri(self.tracking_uri) - kwargs = {} if self.system_metrics_sampling_interval is not None: kwargs["sampling_interval"] = self.system_metrics_sampling_interval @@ -300,6 +302,10 @@ def _start_system_metrics_monitor(self) -> None: kwargs["samples_before_logging"] = self.system_metrics_samples_before_logging try: + # mlflow reads the run to sample through the global tracking URI, + # not through the client + if self.tracking_uri: + mlflow.set_tracking_uri(self.tracking_uri) monitor = SystemMetricsMonitor(run_id, **kwargs) monitor.start() except Exception as e: diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index f14eb8ff0e0..090fd0343f3 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -231,6 +231,9 @@ def _update_metric(engine): self.assertEqual(handler._default_iteration_log.call_count, 2) # 2 = len([1, 3]) from event_filter def test_system_metrics_disabled_by_default(self): + """ + Test that a handler left at its default settings does not sample the system metrics. + """ with tempfile.TemporaryDirectory() as tempdir: def _train_func(engine, batch): @@ -247,6 +250,10 @@ def _train_func(engine, batch): self.assertIsNone(run) def test_system_metrics_monitor_life_cycle(self): + """ + Test that the monitor samples the run of the handler with the requested settings, + and stops when the workflow completes. + """ with tempfile.TemporaryDirectory() as tempdir: def _train_func(engine, batch): @@ -277,6 +284,10 @@ def _train_func(engine, batch): self.assertIsNone(handler.system_metrics_monitor) def test_system_metrics_monitor_shared_by_handlers(self): + """ + Test that handlers sharing a run sample it once, and that the run keeps being sampled + until the handler that started the sampling completes. + """ with tempfile.TemporaryDirectory() as tempdir: def _train_func(engine, batch): @@ -311,6 +322,19 @@ def _train_func(engine, batch): for handler in handlers: handler.close() + def test_system_metrics_settings_are_validated(self): + """ + Test that a sampling setting that mlflow does not define a behaviour for is rejected. + """ + for kwargs in ( + {"system_metrics_sampling_interval": 0}, + {"system_metrics_sampling_interval": -1}, + {"system_metrics_samples_before_logging": 0}, + {"system_metrics_samples_before_logging": -5}, + ): + with self.assertRaises(ValueError): + MLFlowHandler(log_system_metrics=True, **kwargs) + def test_multi_thread(self): test_uri_list = ["monai_mlflow_test1", "monai_mlflow_test2"] with ThreadPoolExecutor(2, "Training") as executor: