star

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/base_metrics_monitor.py

@@ -0,0 +1,32 @@
+"""Base class of system metrics monitor."""
+
+import abc
+from collections import defaultdict
+
+
+class BaseMetricsMonitor(abc.ABC):
+    """Base class of system metrics monitor."""
+
+    def __init__(self):
+        self._metrics = defaultdict(list)
+
+    @abc.abstractmethod
+    def collect_metrics(self):
+        """Method to collect metrics.
+
+        Subclass should implement this method to collect metrics and store in `self._metrics`.
+        """
+
+    @abc.abstractmethod
+    def aggregate_metrics(self):
+        """Method to aggregate metrics.
+
+        Subclass should implement this method to aggregate the metrics and return it in a dict.
+        """
+
+    @property
+    def metrics(self):
+        return self._metrics
+
+    def clear_metrics(self):
+        self._metrics.clear()

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/cpu_monitor.py

@@ -0,0 +1,23 @@
+"""Class for monitoring CPU stats."""
+
+import psutil
+
+from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
+
+
+class CPUMonitor(BaseMetricsMonitor):
+    """Class for monitoring CPU stats."""
+
+    def collect_metrics(self):
+        # Get CPU metrics.
+        cpu_percent = psutil.cpu_percent()
+        self._metrics["cpu_utilization_percentage"].append(cpu_percent)
+
+        system_memory = psutil.virtual_memory()
+        self._metrics["system_memory_usage_megabytes"].append(system_memory.used / 1e6)
+        self._metrics["system_memory_usage_percentage"].append(
+            system_memory.used / system_memory.total * 100
+        )
+
+    def aggregate_metrics(self):
+        return {k: round(sum(v) / len(v), 1) for k, v in self._metrics.items()}

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/disk_monitor.py

@@ -0,0 +1,21 @@
+"""Class for monitoring disk stats."""
+
+import os
+
+import psutil
+
+from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
+
+
+class DiskMonitor(BaseMetricsMonitor):
+    """Class for monitoring disk stats."""
+
+    def collect_metrics(self):
+        # Get disk usage metrics.
+        disk_usage = psutil.disk_usage(os.sep)
+        self._metrics["disk_usage_percentage"].append(disk_usage.percent)
+        self._metrics["disk_usage_megabytes"].append(disk_usage.used / 1e6)
+        self._metrics["disk_available_megabytes"].append(disk_usage.free / 1e6)
+
+    def aggregate_metrics(self):
+        return {k: round(sum(v) / len(v), 1) for k, v in self._metrics.items()}

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/gpu_monitor.py

@@ -0,0 +1,71 @@
+"""Class for monitoring GPU stats."""
+
+import logging
+import sys
+
+from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
+
+_logger = logging.getLogger(__name__)
+
+try:
+    import pynvml
+except ImportError:
+    # If `pynvml` is not installed, a warning will be logged at monitor instantiation.
+    # We don't log a warning here to avoid spamming warning at every import.
+    pass
+
+
+class GPUMonitor(BaseMetricsMonitor):
+    """Class for monitoring GPU stats."""
+
+    def __init__(self):
+        if "pynvml" not in sys.modules:
+            # Only instantiate if `pynvml` is installed.
+            raise ImportError(
+                "`pynvml` is not installed, to log GPU metrics please run `pip install pynvml` "
+                "to install it."
+            )
+        try:
+            # `nvmlInit()` will fail if no GPU is found.
+            pynvml.nvmlInit()
+        except pynvml.NVMLError as e:
+            raise RuntimeError(f"Failed to initialize NVML, skip logging GPU metrics: {e}")
+
+        super().__init__()
+        self.num_gpus = pynvml.nvmlDeviceGetCount()
+        self.gpu_handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(self.num_gpus)]
+
+    def collect_metrics(self):
+        # Get GPU metrics.
+        for i, handle in enumerate(self.gpu_handles):
+            try:
+                memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                self._metrics[f"gpu_{i}_memory_usage_percentage"].append(
+                    round(memory.used / memory.total * 100, 1)
+                )
+                self._metrics[f"gpu_{i}_memory_usage_megabytes"].append(memory.used / 1e6)
+            except pynvml.NVMLError as e:
+                _logger.warning(f"Encountered error {e} when trying to collect GPU memory metrics.")
+
+            try:
+                device_utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
+                self._metrics[f"gpu_{i}_utilization_percentage"].append(device_utilization.gpu)
+            except pynvml.NVMLError as e:
+                _logger.warning(
+                    f"Encountered error {e} when trying to collect GPU utilization metrics."
+                )
+
+            try:
+                power_milliwatts = pynvml.nvmlDeviceGetPowerUsage(handle)
+                power_capacity_milliwatts = pynvml.nvmlDeviceGetEnforcedPowerLimit(handle)
+                self._metrics[f"gpu_{i}_power_usage_watts"].append(power_milliwatts / 1000)
+                self._metrics[f"gpu_{i}_power_usage_percentage"].append(
+                    (power_milliwatts / power_capacity_milliwatts) * 100
+                )
+            except pynvml.NVMLError as e:
+                _logger.warning(
+                    f"Encountered error {e} when trying to collect GPU power usage metrics."
+                )
+
+    def aggregate_metrics(self):
+        return {k: round(sum(v) / len(v), 1) for k, v in self._metrics.items()}

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/network_monitor.py

@@ -0,0 +1,34 @@
+"""Class for monitoring network stats."""
+
+import psutil
+
+from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
+
+
+class NetworkMonitor(BaseMetricsMonitor):
+    def __init__(self):
+        super().__init__()
+        self._set_initial_metrics()
+
+    def _set_initial_metrics(self):
+        # Set initial network usage metrics. `psutil.net_io_counters()` counts the stats since the
+        # system boot, so to set network usage metrics as 0 when we start logging, we need to keep
+        # the initial network usage metrics.
+        network_usage = psutil.net_io_counters()
+        self._initial_receive_megabytes = network_usage.bytes_recv / 1e6
+        self._initial_transmit_megabytes = network_usage.bytes_sent / 1e6
+
+    def collect_metrics(self):
+        # Get network usage metrics.
+        network_usage = psutil.net_io_counters()
+        # Usage metrics will be the diff between current and initial metrics.
+        self._metrics["network_receive_megabytes"] = (
+            network_usage.bytes_recv / 1e6 - self._initial_receive_megabytes
+        )
+        self._metrics["network_transmit_megabytes"] = (
+            network_usage.bytes_sent / 1e6 - self._initial_transmit_megabytes
+        )
+
+    def aggregate_metrics(self):
+        # Network metrics don't need to be averaged.
+        return dict(self._metrics)

venv/lib/python3.9/site-packages/mlflow/system_metrics/metrics/rocm_monitor.py

@@ -0,0 +1,123 @@
+"""Class for monitoring GPU stats on HIP devices.
+Inspired by GPUMonitor, but with the pynvml method
+named replaced by pyrsmi method names
+"""
+
+import contextlib
+import io
+import logging
+import sys
+
+from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
+
+_logger = logging.getLogger(__name__)
+
+is_rocml_available = False
+try:
+    from pyrsmi import rocml
+
+    is_rocml_available = True
+except ImportError:
+    # If `pyrsmi` is not installed, a warning will be logged at monitor instantiation.
+    # We don't log a warning here to avoid spamming warning at every import.
+    pass
+
+
+class ROCMMonitor(BaseMetricsMonitor):
+    """
+    Class for monitoring AMD GPU stats. This is
+    class has been modified and has been inspired by
+    the original GPUMonitor class written by MLflow.
+    This class uses the package pyrsmi which is an
+    official ROCM python package which tracks and monitor
+    AMD GPU's, has been tested on AMD MI250x 128GB GPUs
+
+    For more information see:
+    https://github.com/ROCm/pyrsmi
+
+    PyPi package:
+    https://pypi.org/project/pyrsmi/
+
+
+    """
+
+    def __init__(self):
+        if "pyrsmi" not in sys.modules:
+            # Only instantiate if `pyrsmi` is installed.
+            raise ImportError(
+                "`pyrsmi` is not installed, to log GPU metrics please run `pip install pyrsmi` "
+                "to install it."
+            )
+
+        try:
+            rocml.smi_initialize()
+        except RuntimeError:
+            raise RuntimeError("Failed to initialize RSMI, skip logging GPU metrics")
+
+        super().__init__()
+
+        # Check if GPU is virtual. If so, collect power information from physical GPU
+        self.physical_idx = []
+        for i in range(rocml.smi_get_device_count()):
+            try:
+                self.raise_error(rocml.smi_get_device_average_power, i)
+                # physical GPU if no error is raised
+                self.physical_idx.append(i)
+            except SystemError:
+                # virtual if error is raised
+                # all virtual GPUs must share physical GPU with previous virtual/physical GPU
+                assert i >= 1
+                self.physical_idx.append(self.physical_idx[-1])
+
+    @staticmethod
+    def raise_error(func, *args, **kwargs):
+        """Raise error if message containing 'error' is printed out to stdout or stderr."""
+        stdout = io.StringIO()
+        stderr = io.StringIO()
+
+        with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
+            func(*args, **kwargs)
+
+        out = stdout.getvalue()
+        err = stderr.getvalue()
+
+        # Check if there is an error message in either stdout or stderr
+        if "error" in out.lower():
+            raise SystemError(out)
+        if "error" in err.lower():
+            raise SystemError(err)
+
+    def collect_metrics(self):
+        # Get GPU metrics.
+        self.num_gpus = rocml.smi_get_device_count()
+
+        for i in range(self.num_gpus):
+            memory_used = rocml.smi_get_device_memory_used(i)
+            memory_total = rocml.smi_get_device_memory_total(i)
+            self._metrics[f"gpu_{i}_memory_usage_percentage"].append(
+                round(memory_used / memory_total * 100, 1)
+            )
+            self._metrics[f"gpu_{i}_memory_usage_gigabytes"].append(memory_used / 1e9)
+
+            device_utilization = rocml.smi_get_device_utilization(i)
+            self._metrics[f"gpu_{i}_utilization_percentage"].append(device_utilization)
+
+            power_watts = rocml.smi_get_device_average_power(self.physical_idx[i])
+            power_capacity_watts = 500  # hard coded for now, should get this from rocm-smi
+            self._metrics[f"gpu_{i}_power_usage_watts"].append(power_watts)
+            self._metrics[f"gpu_{i}_power_usage_percentage"].append(
+                (power_watts / power_capacity_watts) * 100
+            )
+
+            # TODO:
+            # memory_busy (and other useful metrics) are available in pyrsmi>1.1.0.
+            # We are currently on pyrsmi==1.0.1, so these are not available
+            # memory_busy  = rocml.smi_get_device_memory_busy(i)
+            # self._metrics[f"gpu_{i}_memory_busy_time_percent"].append(memory_busy)
+
+    def aggregate_metrics(self):
+        return {k: round(sum(v) / len(v), 1) for k, v in self._metrics.items()}
+
+    def __del__(self):
+        if is_rocml_available:
+            rocml.smi_shutdown()