/mlflow/system_metrics/system_metrics_monitor.py

"""Class for monitoring system stats."""

import logging
import threading

from mlflow.environment_variables import (
    MLFLOW_SYSTEM_METRICS_NODE_ID,
    MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING,
    MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL,
)
from mlflow.exceptions import MlflowException
from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
from mlflow.system_metrics.metrics.cpu_monitor import CPUMonitor
from mlflow.system_metrics.metrics.disk_monitor import DiskMonitor
from mlflow.system_metrics.metrics.gpu_monitor import GPUMonitor
from mlflow.system_metrics.metrics.network_monitor import NetworkMonitor
from mlflow.system_metrics.metrics.rocm_monitor import ROCMMonitor

_logger = logging.getLogger(__name__)


class SystemMetricsMonitor:
    """Class for monitoring system stats.

    This class is used for pulling system metrics and logging them to MLflow. Calling `start()` will
    spawn a thread that logs system metrics periodically. Calling `finish()` will stop the thread.
    Logging is done on a different frequency from pulling metrics, so that the metrics are
    aggregated over the period. Users can change the logging frequency by setting
    `MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL` and `MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING`
    environment variables, e.g., run `export MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL=10` in terminal
    will set the sampling interval to 10 seconds.

    System metrics are logged with a prefix "system/", e.g., "system/cpu_utilization_percentage".

    Args:
        run_id: string, the MLflow run ID.
        sampling_interval: float, default to 10. The interval (in seconds) at which to pull system
            metrics. Will be overridden by `MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL` environment
            variable.
        samples_before_logging: int, default to 1. The number of samples to aggregate before
            logging. Will be overridden by `MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING`
            evnironment variable.
        resume_logging: bool, default to False. If True, we will resume the system metrics logging
            from the `run_id`, and the first step to log will be the last step of `run_id` + 1, if
            False, system metrics logging will start from step 0.
        node_id: string, default to None. The node ID of the machine where the metrics are
            collected. Will be overridden by `MLFLOW_SYSTEM_METRICS_NODE_ID`
            evnironment variable. This is useful in multi-node training to distinguish the metrics
            from different nodes. For example, if you set node_id to "node_0", the system metrics
            getting logged will be of format "system/node_0/cpu_utilization_percentage".
        tracking_uri: string, default to None. The tracking URI of the MLflow server, or `None` to
            use whatever is set via `mlflow.set_tracking_uri()`.
    """

    def __init__(
        self,
        run_id,
        sampling_interval=10,
        samples_before_logging=1,
        resume_logging=False,
        node_id=None,
        tracking_uri=None,
    ):
        from mlflow.tracking import get_tracking_uri
        from mlflow.utils.autologging_utils import BatchMetricsLogger

        # Instantiate default monitors.
        self.monitors = [CPUMonitor(), DiskMonitor(), NetworkMonitor()]

        if gpu_monitor := self._initialize_gpu_monitor():
            self.monitors.append(gpu_monitor)

        self.sampling_interval = MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL.get() or sampling_interval
        self.samples_before_logging = (
            MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING.get() or samples_before_logging
        )

        self._run_id = run_id
        self._shutdown_event = threading.Event()
        self._process = None
        self._metrics_prefix = "system/"
        self.node_id = MLFLOW_SYSTEM_METRICS_NODE_ID.get() or node_id
        self._tracking_uri = tracking_uri or get_tracking_uri()
        self._logging_step = self._get_next_logging_step(run_id) if resume_logging else 0
        self.mlflow_logger = BatchMetricsLogger(self._run_id, tracking_uri=self._tracking_uri)

    def _get_next_logging_step(self, run_id):
        from mlflow.tracking.client import MlflowClient

        client = MlflowClient(self._tracking_uri)
        try:
            run = client.get_run(run_id)
        except MlflowException:
            return 0
        system_metric_name = None
        for metric_name in run.data.metrics.keys():
            if metric_name.startswith(self._metrics_prefix):
                system_metric_name = metric_name
                break
        if system_metric_name is None:
            return 0
        metric_history = client.get_metric_history(run_id, system_metric_name)
        return metric_history[-1].step + 1

    def start(self):
        """Start monitoring system metrics."""
        try:
            self._process = threading.Thread(
                target=self.monitor,
                daemon=True,
                name="SystemMetricsMonitor",
            )
            self._process.start()
            _logger.info("Started monitoring system metrics.")
        except Exception as e:
            _logger.warning(f"Failed to start monitoring system metrics: {e}")
            self._process = None

    def monitor(self):
        """Main monitoring loop, which consistently collect and log system metrics."""
        from mlflow.tracking.client import MlflowClient

        _logger.debug(
            f"Monitor loop started (interval={self.sampling_interval}s, "
            f"samples={self.samples_before_logging})"
        )
        while not self._shutdown_event.is_set():
            for sample_idx in range(self.samples_before_logging):
                self.collect_metrics()
                _logger.debug(f"Collected sample {sample_idx + 1}/{self.samples_before_logging}")
                self._shutdown_event.wait(self.sampling_interval)
                try:
                    # Get the MLflow run to check if the run is not RUNNING.
                    run = MlflowClient(self._tracking_uri).get_run(self._run_id)
                except Exception as e:
                    _logger.warning(f"Failed to get mlflow run: {e}.")
                    return
                if run.info.status != "RUNNING" or self._shutdown_event.is_set():
                    # If the mlflow run is terminated or receives the shutdown signal, stop
                    # monitoring.
                    _logger.debug(f"Monitor loop exiting (status={run.info.status})")
                    return
            metrics = self.aggregate_metrics()
            _logger.debug(f"Aggregated {len(metrics)} metrics")
            try:
                self.publish_metrics(metrics)
                _logger.debug(
                    f"Queued metrics at step {self._logging_step - 1} "
                    f"(training_time={self.mlflow_logger.total_training_time:.3f}s, "
                    f"log_batch_time={self.mlflow_logger.total_log_batch_time:.3f}s, "
                    f"batch_size={len(self.mlflow_logger.data)})"
                )
            except Exception as e:
                _logger.warning(
                    f"Failed to log system metrics: {e}, this is expected if the experiment/run is "
                    "already terminated."
                )
                return

    def collect_metrics(self):
        """Collect system metrics."""
        metrics = {}
        for monitor in self.monitors:
            monitor.collect_metrics()
            metrics.update(monitor._metrics)
        return metrics

    def aggregate_metrics(self):
        """Aggregate collected metrics."""
        metrics = {}
        for monitor in self.monitors:
            metrics.update(monitor.aggregate_metrics())
        return metrics

    def publish_metrics(self, metrics):
        """Log collected metrics to MLflow."""
        # Add prefix "system/" to the metrics name for grouping. If `self.node_id` is not None, also
        # add it to the metrics name.
        prefix = self._metrics_prefix + (self.node_id + "/" if self.node_id else "")
        metrics = {prefix + k: v for k, v in metrics.items()}

        self.mlflow_logger.record_metrics(metrics, self._logging_step)
        self._logging_step += 1
        for monitor in self.monitors:
            monitor.clear_metrics()

    def finish(self):
        """Stop monitoring system metrics."""
        if self._process is None:
            return
        _logger.info("Stopping system metrics monitoring...")
        self._shutdown_event.set()
        try:
            self._process.join()
            _logger.debug(f"Flushing {len(self.mlflow_logger.data)} metrics")
            self.mlflow_logger.flush()
            _logger.info("Successfully terminated system metrics monitoring!")
        except Exception as e:
            _logger.error(f"Error terminating system metrics monitoring process: {e}.")
        self._process = None

    def _initialize_gpu_monitor(self) -> BaseMetricsMonitor | None:
        # NVIDIA GPU
        try:
            return GPUMonitor()
        except Exception:
            _logger.debug("Failed to initialize GPU monitor for NVIDIA GPU.", exc_info=True)

        # Falling back to pyrocml (AMD/HIP GPU)
        try:
            return ROCMMonitor()
        except Exception:
            _logger.debug("Failed to initialize GPU monitor for AMD/HIP GPU.", exc_info=True)

        _logger.info("Skip logging GPU metrics. Set logger level to DEBUG for more details.")
        return None