Record system resource usage with MLFlowHandler

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,19 @@ 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 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,
+            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 +161,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 +186,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 +234,27 @@ 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
+        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
+        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
@@ -295,6 +335,81 @@ 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
+
+            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
+
+            # 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:
+                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
+            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}")
+            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
+
     def _set_experiment(self):
         experiment = self.experiment
         if not experiment:
@@ -394,6 +509,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 +547,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)

tests/handlers/test_handler_mlflow.py

@@ -335,6 +335,209 @@ 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):
+        """
+        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):
+        """
+        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:
+            engine = Engine(self._train_func)
+            test_path = os.path.join(tempdir, "mlflow_system_metrics_off")
+            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),
+            ):
+                handler.attach(engine)
+                engine.run(range(3), max_epochs=1)
+
+            monitor_class.assert_not_called()
+
+    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:
+            engine = Engine(self._train_func)
+            test_path = os.path.join(tempdir, "mlflow_system_metrics")
+            handler = MLFlowHandler(
+                iteration_log=False,
+                tracking_uri=path_to_sqlite_uri(test_path),
+                log_system_metrics=True,
+                system_metrics_sampling_interval=1,
+                system_metrics_samples_before_logging=1,
+                close_on_complete=False,
+            )
+            monitor = MagicMock()
+            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)
+
+            # 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):
+        """
+        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:
+            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 = [
+                MLFlowHandler(
+                    iteration_log=False,
+                    tracking_uri=path_to_sqlite_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,
+                patch("monai.handlers.mlflow_handler.has_system_metrics", True),
+            ):
+                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:]:
+                    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_system_metrics_warns_when_unavailable(self):
+        """
+        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:
+            engine = Engine(self._train_func)
+            test_path = os.path.join(tempdir, "mlflow_system_metrics_unavailable")
+            handler = MLFlowHandler(
+                iteration_log=False,
+                tracking_uri=path_to_sqlite_uri(test_path),
+                log_system_metrics=True,
+                close_on_complete=True,
+            )
+            with patch("monai.handlers.mlflow_handler.has_system_metrics", False):
+                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_sqlite_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_sqlite_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):
+        """
+        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: