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
@@ -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,21 @@ 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
+        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.experiment = None
         self.cur_run = None
         self.dataset_dict = dataset_dict
@@ -239,6 +274,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
+
+            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:
+                # 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:
+                # 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:
@@ -333,6 +428,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:
@@ -343,6 +440,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

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,204 @@ 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_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_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_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_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_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):
+        """
+        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: