1  """Class for monitoring disk stats."""
 2  
 3  import os
 4  
 5  import psutil
 6  
 7  from mlflow.system_metrics.metrics.base_metrics_monitor import BaseMetricsMonitor
 8  
 9  
10  class DiskMonitor(BaseMetricsMonitor):
11      """Class for monitoring disk stats."""
12  
13      def collect_metrics(self):
14          # Get disk usage metrics.
15          disk_usage = psutil.disk_usage(os.sep)
16          self._metrics["disk_usage_percentage"].append(disk_usage.percent)
17          self._metrics["disk_usage_megabytes"].append(disk_usage.used / 1e6)
18          self._metrics["disk_available_megabytes"].append(disk_usage.free / 1e6)
19  
20      def aggregate_metrics(self):
21          return {k: round(sum(v) / len(v), 1) for k, v in self._metrics.items()}