genesis-flow 1.0.0__py3-none-any.whl

mlflow/system_metrics/__init__.py

@@ -0,0 +1,61 @@
+"""System metrics logging module."""
+
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING,
+    MLFLOW_SYSTEM_METRICS_NODE_ID,
+    MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING,
+    MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL,
+)
+
+
+def disable_system_metrics_logging():
+    """Disable system metrics logging globally.
+
+    Calling this function will disable system metrics logging globally, but users can still opt in
+    system metrics logging for individual runs by `mlflow.start_run(log_system_metrics=True)`.
+    """
+    MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.set(False)
+
+
+def enable_system_metrics_logging():
+    """Enable system metrics logging globally.
+
+    Calling this function will enable system metrics logging globally, but users can still opt out
+    system metrics logging for individual runs by `mlflow.start_run(log_system_metrics=False)`.
+    """
+    MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.set(True)
+
+
+def set_system_metrics_sampling_interval(interval):
+    """Set the system metrics sampling interval.
+
+    Every `interval` seconds, the system metrics will be collected. By default `interval=10`.
+    """
+    if interval is None:
+        MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL.unset()
+    else:
+        MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL.set(interval)
+
+
+def set_system_metrics_samples_before_logging(samples):
+    """Set the number of samples before logging system metrics.
+
+    Every time `samples` samples have been collected, the system metrics will be logged to mlflow.
+    By default `samples=1`.
+    """
+    if samples is None:
+        MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING.unset()
+    else:
+        MLFLOW_SYSTEM_METRICS_SAMPLES_BEFORE_LOGGING.set(samples)
+
+
+def set_system_metrics_node_id(node_id):
+    """Set the system metrics node id.
+
+    node_id is the identifier of the machine where the metrics are collected. This is useful in
+    multi-node (distributed training) setup.
+    """
+    if node_id is None:
+        MLFLOW_SYSTEM_METRICS_NODE_ID.unset()
+    else:
+        MLFLOW_SYSTEM_METRICS_NODE_ID.set(node_id)

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()

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()}

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()}

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()}

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)

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()

mlflow/system_metrics/system_metrics_monitor.py

@@ -0,0 +1,198 @@
+"""Class for monitoring system stats."""
+
+import logging
+import threading
+from typing import Optional
+
+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".
+    """
+
+    def __init__(
+        self,
+        run_id,
+        sampling_interval=10,
+        samples_before_logging=1,
+        resume_logging=False,
+        node_id=None,
+    ):
+        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.mlflow_logger = BatchMetricsLogger(self._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._logging_step = self._get_next_logging_step(run_id) if resume_logging else 0
+
+    def _get_next_logging_step(self, run_id):
+        from mlflow.tracking.client import MlflowClient
+
+        client = MlflowClient()
+        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.fluent import get_run
+
+        while not self._shutdown_event.is_set():
+            for _ in range(self.samples_before_logging):
+                self.collect_metrics()
+                self._shutdown_event.wait(self.sampling_interval)
+                try:
+                    # Get the MLflow run to check if the run is not RUNNING.
+                    run = 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.
+                    return
+            metrics = self.aggregate_metrics()
+            try:
+                self.publish_metrics(metrics)
+            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()
+            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) -> Optional[BaseMetricsMonitor]:
+        # 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

mlflow/tracing/__init__.py

@@ -0,0 +1,16 @@
+from mlflow.tracing.config import configure
+from mlflow.tracing.display import disable_notebook_display, enable_notebook_display
+from mlflow.tracing.provider import disable, enable, reset, set_destination
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+
+__all__ = [
+    "configure",
+    "disable",
+    "enable",
+    "disable_notebook_display",
+    "enable_notebook_display",
+    "set_span_chat_messages",
+    "set_span_chat_tools",
+    "set_destination",
+    "reset",
+]