From 3aae92c5d6a13d0b5ee911925c237b63a59c7af0 Mon Sep 17 00:00:00 2001 From: uditmahato Date: Sun, 9 Aug 2026 12:22:49 +0545 Subject: [PATCH 1/5] 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/5] 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: From 03f57c5911e936fb35a12cc5a85244d502ce453f Mon Sep 17 00:00:00 2001 From: uditmahato Date: Tue, 18 Aug 2026 11:13:30 +0545 Subject: [PATCH 3/5] Make the system metrics tests independent of the installed mlflow The monitor is only started when the mlflow system metrics module is importable, so the tests that assert it starts were relying on that being true in the environment they run in. Pin it for those tests and cover the case where it is missing, where the workflow should carry on with a warning. Signed-off-by: uditmahato --- tests/handlers/test_handler_mlflow.py | 35 +++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 090fd0343f3..afbdc4b2ea8 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -270,7 +270,10 @@ def _train_func(engine, batch): close_on_complete=True, ) monitor = MagicMock() - with patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class: + with ( + patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class, + patch("monai.handlers.mlflow_handler.has_system_metrics", True), + ): handler.attach(engine) engine.run(range(3), max_epochs=1) @@ -303,7 +306,10 @@ def _train_func(engine, batch): for _ in range(3) ] monitor = MagicMock() - with patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class: + with ( + patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor) as monitor_class, + patch("monai.handlers.mlflow_handler.has_system_metrics", True), + ): for handler in handlers: handler.start(engine) @@ -322,6 +328,31 @@ def _train_func(engine, batch): for handler in handlers: handler.close() + def test_system_metrics_warns_when_mlflow_is_too_old(self): + """ + Test that a workflow still runs, with a warning, when the installed mlflow cannot + record the system metrics. + """ + 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_unavailable") + handler = MLFlowHandler( + iteration_log=False, + tracking_uri=path_to_uri(test_path), + log_system_metrics=True, + close_on_complete=True, + ) + with patch("monai.handlers.mlflow_handler.has_system_metrics", False): + with self.assertWarns(Warning): + handler.attach(engine) + engine.run(range(3), max_epochs=1) + + self.assertIsNone(handler.system_metrics_monitor) + def test_system_metrics_settings_are_validated(self): """ Test that a sampling setting that mlflow does not define a behaviour for is rejected. From 939efd2ec493c3649af3beddbd51c7607814d0e4 Mon Sep 17 00:00:00 2001 From: uditmahato Date: Tue, 18 Aug 2026 11:22:28 +0545 Subject: [PATCH 4/5] Strengthen the system metrics tests The test for the default settings was asserting on state that close() clears anyway, so it would have passed even if the monitor had run. Assert that the monitor is never constructed instead. Also assert the monitor is given the run of the handler, that handlers sharing a run all resolve the same one, and match the expected warning rather than any warning. Cover the two failure paths: a monitor that cannot start, and one that cannot stop, neither of which should stop the workflow. Signed-off-by: uditmahato --- tests/handlers/test_handler_mlflow.py | 114 +++++++++++++++++++------- 1 file changed, 83 insertions(+), 31 deletions(-) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index afbdc4b2ea8..a82eba86ca0 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -230,24 +230,27 @@ def _update_metric(engine): else: self.assertEqual(handler._default_iteration_log.call_count, 2) # 2 = len([1, 3]) from event_filter + @staticmethod + def _train_func(engine, batch): + return [batch + 1.0] + def test_system_metrics_disabled_by_default(self): """ - Test that a handler left at its default settings does not sample the system metrics. + Test that a handler left at its default settings does not sample the system metrics, + even where mlflow is able to. """ with tempfile.TemporaryDirectory() as tempdir: - - def _train_func(engine, batch): - return [batch + 1.0] - - engine = Engine(_train_func) + engine = Engine(self._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) + with ( + patch("monai.handlers.mlflow_handler.SystemMetricsMonitor") as monitor_class, + patch("monai.handlers.mlflow_handler.has_system_metrics", 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) + monitor_class.assert_not_called() def test_system_metrics_monitor_life_cycle(self): """ @@ -255,11 +258,7 @@ def test_system_metrics_monitor_life_cycle(self): and stops when the workflow completes. """ with tempfile.TemporaryDirectory() as tempdir: - - def _train_func(engine, batch): - return [batch + 1.0] - - engine = Engine(_train_func) + engine = Engine(self._train_func) test_path = os.path.join(tempdir, "mlflow_system_metrics") handler = MLFlowHandler( iteration_log=False, @@ -267,7 +266,7 @@ def _train_func(engine, batch): log_system_metrics=True, system_metrics_sampling_interval=1, system_metrics_samples_before_logging=1, - close_on_complete=True, + close_on_complete=False, ) monitor = MagicMock() with ( @@ -279,12 +278,14 @@ def _train_func(engine, batch): # the monitor samples the run of the handler, with the requested sampling settings monitor_class.assert_called_once() + self.assertEqual(monitor_class.call_args.args[0], handler.cur_run.info.run_id) 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) + handler.close() def test_system_metrics_monitor_shared_by_handlers(self): """ @@ -292,11 +293,7 @@ def test_system_metrics_monitor_shared_by_handlers(self): until the handler that started the sampling completes. """ with tempfile.TemporaryDirectory() as tempdir: - - def _train_func(engine, batch): - return [batch + 1.0] - - engine = Engine(_train_func) + engine = Engine(self._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 = [ @@ -313,8 +310,13 @@ def _train_func(engine, batch): for handler in handlers: handler.start(engine) + run_ids = {handler.cur_run.info.run_id for handler in handlers} + self.assertEqual(len(run_ids), 1) + # the run is sampled by the first handler only monitor_class.assert_called_once() + self.assertEqual(monitor_class.call_args.args[0], run_ids.pop()) + monitor.start.assert_called_once() # the handlers that do not sample the run leave it running when they complete for handler in handlers[1:]: @@ -328,17 +330,13 @@ def _train_func(engine, batch): for handler in handlers: handler.close() - def test_system_metrics_warns_when_mlflow_is_too_old(self): + def test_system_metrics_warns_when_unavailable(self): """ - Test that a workflow still runs, with a warning, when the installed mlflow cannot - record the system metrics. + Test that a workflow still runs, with a warning, when the installed mlflow does not + support recording the system metrics. """ with tempfile.TemporaryDirectory() as tempdir: - - def _train_func(engine, batch): - return [batch + 1.0] - - engine = Engine(_train_func) + engine = Engine(self._train_func) test_path = os.path.join(tempdir, "mlflow_system_metrics_unavailable") handler = MLFlowHandler( iteration_log=False, @@ -347,11 +345,65 @@ def _train_func(engine, batch): close_on_complete=True, ) with patch("monai.handlers.mlflow_handler.has_system_metrics", False): - with self.assertWarns(Warning): + with self.assertWarnsRegex(Warning, "Please install mlflow>=2.8.0 to record the system metrics."): + handler.attach(engine) + engine.run(range(3), max_epochs=1) + + self.assertIsNone(handler.system_metrics_monitor) + + def test_system_metrics_start_failure_does_not_stop_the_workflow(self): + """ + Test that a workflow still runs, with a warning, when the monitor cannot be started. + """ + with tempfile.TemporaryDirectory() as tempdir: + engine = Engine(self._train_func) + test_path = os.path.join(tempdir, "mlflow_system_metrics_start_failure") + handler = MLFlowHandler( + iteration_log=False, + tracking_uri=path_to_uri(test_path), + log_system_metrics=True, + close_on_complete=True, + ) + with ( + patch( + "monai.handlers.mlflow_handler.SystemMetricsMonitor", side_effect=RuntimeError("no monitor for you") + ), + patch("monai.handlers.mlflow_handler.has_system_metrics", True), + ): + with self.assertWarnsRegex(Warning, "Failed to record the system metrics"): + handler.attach(engine) + engine.run(range(3), max_epochs=1) + + self.assertEqual(engine.state.epoch, 1) + self.assertIsNone(handler.system_metrics_monitor) + self.assertEqual(len(MLFlowHandler._monitored_run_ids), 0) + + def test_system_metrics_stop_failure_is_reported(self): + """ + Test that a monitor which fails to stop is reported and released, so that the run can + be sampled again. + """ + with tempfile.TemporaryDirectory() as tempdir: + engine = Engine(self._train_func) + test_path = os.path.join(tempdir, "mlflow_system_metrics_stop_failure") + handler = MLFlowHandler( + iteration_log=False, + tracking_uri=path_to_uri(test_path), + log_system_metrics=True, + close_on_complete=True, + ) + monitor = MagicMock() + monitor.finish.side_effect = RuntimeError("monitor will not stop") + with ( + patch("monai.handlers.mlflow_handler.SystemMetricsMonitor", return_value=monitor), + patch("monai.handlers.mlflow_handler.has_system_metrics", True), + ): + with self.assertWarnsRegex(Warning, "Failed to stop recording the system metrics"): handler.attach(engine) engine.run(range(3), max_epochs=1) self.assertIsNone(handler.system_metrics_monitor) + self.assertEqual(len(MLFlowHandler._monitored_run_ids), 0) def test_system_metrics_settings_are_validated(self): """ From d2a13a3ea078bef06112867f1a3c4b8db440a0f6 Mon Sep 17 00:00:00 2001 From: uditmahato Date: Tue, 18 Aug 2026 11:27:57 +0545 Subject: [PATCH 5/5] Document the training step helper used by the tests Signed-off-by: uditmahato --- tests/handlers/test_handler_mlflow.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index a82eba86ca0..3807e325d9f 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -232,6 +232,16 @@ def _update_metric(engine): @staticmethod def _train_func(engine, batch): + """ + Produce the output of one training step, for an engine that does no real work. + + Args: + engine: the ignite engine running the step, unused. + batch: the batch of the current step. + + Returns: + The batch shifted by one, as the single output of the step. + """ return [batch + 1.0] def test_system_metrics_disabled_by_default(self):