From b69273ae55c90583198cbc1b97f93c9782eaf8e7 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Sun, 9 Aug 2026 12:22:49 +0545
Subject: [PATCH 1/6] 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 <uditmahato29271@gmail.com>
---
 monai/handlers/mlflow_handler.py      | 95 +++++++++++++++++++++++++++
 tests/handlers/test_handler_mlflow.py | 81 +++++++++++++++++++++++
 2 files changed, 176 insertions(+)

diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py
index 1cd26d5287..33c0f92a16 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
@@ -43,6 +44,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
@@ -136,6 +140,18 @@ class MLFlowHandler:
             or the ``MLFLOW_TRACKING_URI`` environment variable), it defaults to an ``mlruns`` directory
             next to the database file; for other backends ``None`` lets MLflow decide based on the
             ``tracking_uri``. Has no effect if the experiment already exists.
+        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.
 
@@ -144,6 +160,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,
@@ -165,6 +185,9 @@ def __init__(
         optimizer_param_names: str | Sequence[str] = "lr",
         close_on_complete: bool = False,
         artifact_location: str | None = None,
+        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
@@ -210,11 +233,19 @@ def __init__(
                 f"tracking_uri={effective_tracking_uri!r}. Use a SQLite URI "
                 "(sqlite:///<path>/mlruns.db) or a remote tracking URI instead."
             )
+        # The system metrics monitor reads the run through the global tracking uri,
+        # so remember the uri that was actually settled on rather than the argument.
+        self.tracking_uri = effective_tracking_uri
         # Only the argument is passed to the client; when `MLFLOW_TRACKING_URI` took priority it
         # is left None so MLflow resolves the environment variable itself.
         self.client = mlflow.MlflowClient(tracking_uri=None if env_tracking_uri else tracking_uri)
         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
@@ -295,6 +326,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:
@@ -394,6 +485,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:
@@ -430,6 +523,8 @@ def close(self) -> None:
         Stop current running logger of MLFlow and release local SQLite resources.
 
         """
+        self._stop_system_metrics_monitor()
+
         try:
             if self.cur_run:
                 self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status)
diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py
index a396227eb9..fd5a1a123e 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -335,6 +335,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 tempfile.TemporaryDirectory() as tempdir:

From 71e5fed86657a4e4afdd3e261de88dc10c4633c7 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Mon, 10 Aug 2026 13:30:12 +0545
Subject: [PATCH 2/6] 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 <uditmahato29271@gmail.com>
---
 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 33c0f92a16..1c61efca0b 100644
--- a/monai/handlers/mlflow_handler.py
+++ b/monai/handlers/mlflow_handler.py
@@ -242,6 +242,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
@@ -348,10 +354,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
@@ -359,6 +361,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 fd5a1a123e..3c6ed65dc0 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -336,6 +336,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):
@@ -352,6 +355,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):
@@ -382,6 +389,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):
@@ -416,6 +427,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 tempfile.TemporaryDirectory() as tempdir:

From d3f4ef1bbc638caa701c7d3bbe60f16930c78745 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Tue, 18 Aug 2026 11:13:30 +0545
Subject: [PATCH 3/6] 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 <uditmahato29271@gmail.com>
---
 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 3c6ed65dc0..13379fc954 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -375,7 +375,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)
 
@@ -408,7 +411,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)
 
@@ -427,6 +433,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 4cf57b7efd4d8166d2248f15aeb22085c30d9a30 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Tue, 18 Aug 2026 11:22:28 +0545
Subject: [PATCH 4/6] 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 <uditmahato29271@gmail.com>
---
 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 13379fc954..8d8a427c30 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -335,24 +335,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):
         """
@@ -360,11 +363,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,
@@ -372,7 +371,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 (
@@ -384,12 +383,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):
         """
@@ -397,11 +398,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 = [
@@ -418,8 +415,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:]:
@@ -433,17 +435,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,
@@ -452,11 +450,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 d7ca6b89c3f8e4c06cd510b8174613849335e0b2 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Tue, 18 Aug 2026 11:27:57 +0545
Subject: [PATCH 5/6] Document the training step helper used by the tests

Signed-off-by: uditmahato <uditmahato29271@gmail.com>
---
 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 8d8a427c30..1179437fab 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -337,6 +337,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):

From 489e9fc5b220bb59863f2bcbf4f66cfac4ba3929 Mon Sep 17 00:00:00 2001
From: uditmahato <uditmahato29271@gmail.com>
Date: Mon, 31 Aug 2026 11:10:06 +0545
Subject: [PATCH 6/6] Restore the tracking uri after sampling, and follow the
 sqlite backend

Rebasing on dev turned up a real problem with this feature. Pointing the
global tracking uri at ours so mlflow can read the run also writes
MLFLOW_TRACKING_URI, and dev now reads that in preference to the
tracking_uri argument. Enabling system metrics therefore changed where
every handler built afterwards logged, for the rest of the process.

The previous value is now put back when sampling stops, including the
case where it was unset, so the override lasts only as long as the run
it is for.

The uri remembered for the monitor is the one the handler settled on
rather than the argument, so the default sqlite backend is sampled
correctly, and the tests move to sqlite tracking uris with the rest of
the suite.

Signed-off-by: uditmahato <uditmahato29271@gmail.com>
---
 monai/handlers/mlflow_handler.py      | 28 ++++++++++++++++++++++-----
 tests/handlers/test_handler_mlflow.py | 17 ++++++++++------
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py
index 1c61efca0b..9a024869fc 100644
--- a/monai/handlers/mlflow_handler.py
+++ b/monai/handlers/mlflow_handler.py
@@ -144,9 +144,10 @@ class MLFlowHandler:
             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.
+            run through the global tracking URI to sample it, so this points the global tracking URI
+            (and with it `MLFLOW_TRACKING_URI`, which later handlers read in preference to their own
+            argument) at `tracking_uri` while the run is sampled, and puts the previous value back
+            when sampling stops.
         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,
@@ -252,6 +253,8 @@ def __init__(
         self.system_metrics_samples_before_logging = system_metrics_samples_before_logging
         self.system_metrics_monitor = None
         self._monitored_run_id: str | None = None
+        self._previous_tracking_uri: str | None = None
+        self._tracking_uri_overridden = False
         self.experiment = None
         self.cur_run = None
         self.dataset_dict = dataset_dict
@@ -360,17 +363,28 @@ def _start_system_metrics_monitor(self) -> None:
             if self.system_metrics_samples_before_logging is not None:
                 kwargs["samples_before_logging"] = self.system_metrics_samples_before_logging
 
+            # mlflow reads the run to sample through the global tracking uri, not through
+            # the client, so it has to be pointed at ours for as long as we sample. Setting
+            # it also writes MLFLOW_TRACKING_URI, which every later handler reads in
+            # preference to its own argument, so the previous value is put back when
+            # sampling stops. `None` is a meaningful previous value, meaning it was unset,
+            # and passing it back to mlflow restores exactly that.
+            previous_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI")
+            overridden = False
             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)
+                    overridden = True
                 monitor = SystemMetricsMonitor(run_id, **kwargs)
                 monitor.start()
             except Exception as e:
                 # a workflow should not fail because its resource usage cannot be recorded
+                if overridden:
+                    mlflow.set_tracking_uri(previous_tracking_uri)
                 warnings.warn(f"Failed to record the system metrics: {e}")
                 return
+            self._previous_tracking_uri = previous_tracking_uri
+            self._tracking_uri_overridden = overridden
 
             MLFlowHandler._monitored_run_ids.add(run_id)
             self.system_metrics_monitor = monitor
@@ -388,6 +402,10 @@ def _stop_system_metrics_monitor(self) -> None:
                 self.system_metrics_monitor.finish()
             except Exception as e:
                 warnings.warn(f"Failed to stop recording the system metrics: {e}")
+            if self._tracking_uri_overridden:
+                mlflow.set_tracking_uri(self._previous_tracking_uri)
+                self._tracking_uri_overridden = False
+            self._previous_tracking_uri = None
             MLFlowHandler._monitored_run_ids.discard(self._monitored_run_id)
             self.system_metrics_monitor = None
             self._monitored_run_id = None
diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py
index 1179437fab..ad3a70f339 100644
--- a/tests/handlers/test_handler_mlflow.py
+++ b/tests/handlers/test_handler_mlflow.py
@@ -357,7 +357,9 @@ def test_system_metrics_disabled_by_default(self):
         with tempfile.TemporaryDirectory() as tempdir:
             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 = MLFlowHandler(
+                iteration_log=False, tracking_uri=path_to_sqlite_uri(test_path), close_on_complete=True
+            )
             with (
                 patch("monai.handlers.mlflow_handler.SystemMetricsMonitor") as monitor_class,
                 patch("monai.handlers.mlflow_handler.has_system_metrics", True),
@@ -377,7 +379,7 @@ def test_system_metrics_monitor_life_cycle(self):
             test_path = os.path.join(tempdir, "mlflow_system_metrics")
             handler = MLFlowHandler(
                 iteration_log=False,
-                tracking_uri=path_to_uri(test_path),
+                tracking_uri=path_to_sqlite_uri(test_path),
                 log_system_metrics=True,
                 system_metrics_sampling_interval=1,
                 system_metrics_samples_before_logging=1,
@@ -413,7 +415,10 @@ def test_system_metrics_monitor_shared_by_handlers(self):
             # 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
+                    iteration_log=False,
+                    tracking_uri=path_to_sqlite_uri(test_path),
+                    run_name="shared",
+                    log_system_metrics=True,
                 )
                 for _ in range(3)
             ]
@@ -455,7 +460,7 @@ def test_system_metrics_warns_when_unavailable(self):
             test_path = os.path.join(tempdir, "mlflow_system_metrics_unavailable")
             handler = MLFlowHandler(
                 iteration_log=False,
-                tracking_uri=path_to_uri(test_path),
+                tracking_uri=path_to_sqlite_uri(test_path),
                 log_system_metrics=True,
                 close_on_complete=True,
             )
@@ -475,7 +480,7 @@ def test_system_metrics_start_failure_does_not_stop_the_workflow(self):
             test_path = os.path.join(tempdir, "mlflow_system_metrics_start_failure")
             handler = MLFlowHandler(
                 iteration_log=False,
-                tracking_uri=path_to_uri(test_path),
+                tracking_uri=path_to_sqlite_uri(test_path),
                 log_system_metrics=True,
                 close_on_complete=True,
             )
@@ -503,7 +508,7 @@ def test_system_metrics_stop_failure_is_reported(self):
             test_path = os.path.join(tempdir, "mlflow_system_metrics_stop_failure")
             handler = MLFlowHandler(
                 iteration_log=False,
-                tracking_uri=path_to_uri(test_path),
+                tracking_uri=path_to_sqlite_uri(test_path),
                 log_system_metrics=True,
                 close_on_complete=True,
             )