Add AMDSMIReader for monitoring AMD GPUs

README.md

@@ -37,6 +37,10 @@ pip install wattameter[benchmark]
 pip install wattameter[postprocessing,mqtt]
 ```
 
+### AMD SMI support
+
+For AMD GPU monitoring, WattAMeter relies on the AMD SMI Python package. Follow the instructions at https://rocm.docs.amd.com/projects/amdsmi/en/latest/install/install.html to install the AMD SMI Python package from your ROCm installation. If you have a ROCm installation, you can try to use [setup-amdsmi](src/wattameter/utils/setup_amdsmi.sh) to automatically set up the AMD SMI environment, which includes installing the AMD SMI Python package to your Python environment.
+
 ## Usage
 
 ### As a Python module

src/wattameter/cli/utils.py

@@ -9,6 +9,11 @@
 from ..readers import RAPLReader, NVMLReader
 from ..readers import Energy, DataThroughput, Utilization, Power, Temperature
 
+try:
+    from ..readers import AMDSMIReader
+except ImportError:
+    AMDSMIReader = None
+
 signal_handled = threading.Event()
 
 
@@ -41,11 +46,24 @@ def parse_tracker_spec(spec_string):
         "nvml-util": (NVMLReader, Utilization),
         "nvml-nvlink": (NVMLReader, DataThroughput),
     }
+    if AMDSMIReader is not None:
+        metric_map.update(
+            {
+                "amdsmi-power": (AMDSMIReader, Power),
+                "amdsmi-temp": (AMDSMIReader, Temperature),
+                "amdsmi-util": (AMDSMIReader, Utilization),
+            }
+        )
 
     # Group metrics by reader type
     _metrics = {}
     for metric_name in parts[1:]:
         metric_name_lower = metric_name.strip().lower()
+        if metric_name_lower.startswith("amdsmi-") and AMDSMIReader is None:
+            raise argparse.ArgumentTypeError(
+                "AMD SMI metrics require the ROCm-provided amdsmi package. "
+                "Run: bash src/wattameter/utils/setup_amdsmi.sh"
+            )
         if metric_name_lower not in metric_map:
             raise argparse.ArgumentTypeError(
                 f"Unknown metric: {metric_name}. Valid metrics: {', '.join(metric_map.keys())}"

src/wattameter/readers/__init__.py

@@ -13,6 +13,11 @@
     Utilization,
 )
 
+try:
+    from .amdsmi import AMDSMIReader
+except ModuleNotFoundError:
+    AMDSMIReader = None
+
 __all__ = [
     "NVMLReader",
     "RAPLReader",
@@ -28,3 +33,6 @@
     "Utilization",
     "DataThroughput",
 ]
+
+if AMDSMIReader is not None:
+    __all__.append("AMDSMIReader")

src/wattameter/readers/amdsmi.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# SPDX-FileCopyrightText: 2025, Alliance for Energy Innovation, LLC
+
+import amdsmi
+import logging
+
+from .base import BaseReader
+from .utils import (
+    Power,
+    Temperature,
+    Quantity,
+    Watt,
+    Celsius,
+    Unit,
+    Utilization,
+)
+
+# Module-level logger
+logger = logging.getLogger(__name__)
+
+
+class AMDSMIReader(BaseReader):
+    """Reader for AMD System Management Interface (AMDSMI) to monitor AMD GPUs.
+
+    .. attribute:: devices
+
+        List of AMDSMI device handles for available GPUs.
+
+    """
+
+    UNITS = {
+        Temperature: Celsius(),
+        Power: Watt("m"),
+    }  #: Dictionary of measurement units for physical quantities.
+
+    def __init__(self, quantities=(Power,)) -> None:
+        super().__init__(quantities)
+
+        self.devices = []
+
+        # Initialize AMDSMI
+        try:
+            ret = amdsmi.amdsmi_init()
+            if ret != 0:
+                logger.warning(
+                    f"AMDSMI initialization returned non-zero status: {ret}. "
+                    f"Continuing without AMDSMI support."
+                )
+            else:
+                logger.info("AMDSMI initialized successfully.")
+        except amdsmi.AmdSmiException as e:
+            logger.warning(
+                f"Failed to initialize AMDSMI: {e}. Continuing without AMDSMI support."
+            )
+            return
+
+        # Get the handles for all available devices
+        try:
+            self.devices = amdsmi.amdsmi_get_processor_handles()
+            logger.info(f"Found {len(self.devices)} AMD GPU(s) on the system.")
+        except amdsmi.AmdSmiException as e:
+            logger.error(f"Failed to get device handles: {e}")
+            self.devices = []
+
+        # Set the quantities to read
+        invalid_quantities = [
+            q for q in quantities if q not in self.UNITS and q != Utilization
+        ]
+        if invalid_quantities:
+            raise ValueError(
+                f"Unsupported quantities: {invalid_quantities}. "
+                f"Supported quantities are: {list(self.UNITS.keys())}."
+            )
+
+    def __del__(self):
+        """Shutdown AMDSMI on deletion."""
+        try:
+            amdsmi.amdsmi_shut_down()
+            logger.info("AMDSMI shutdown successfully.")
+        except amdsmi.AmdSmiException as e:
+            logger.warning(f"Failed to shutdown AMDSMI: {e}")
+
+    @property
+    def tags(self) -> list[str]:
+        _tags = []
+        for q in self.quantities:
+            if q == Utilization:
+                _tags.extend([f"gpu-{i}[%gpu]" for i in range(len(self.devices))])
+                _tags.extend([f"gpu-{i}[%mem]" for i in range(len(self.devices))])
+            else:
+                unit = self.get_unit(q)
+                _tags.extend([f"gpu-{i}[{unit}]" for i in range(len(self.devices))])
+        return _tags
+
+    def get_unit(self, quantity: type[Quantity]) -> Unit:
+        if quantity in self.UNITS:
+            return self.UNITS[quantity]
+        else:
+            logger.warning(
+                f"The quantity: {quantity} is either unsupported or has no associated unit. "
+                f"Supported quantities with units are: {list(self.UNITS.keys())}."
+            )
+            return Unit()  # Return a default Unit instance
+
+    def read_temperature_on_device(self, i: int) -> int:
+        """Read the temperature for the i-th device."""
+        try:
+            return amdsmi.amdsmi_get_temp_metric(
+                self.devices[i],
+                amdsmi.AmdSmiTemperatureType.EDGE,
+                amdsmi.AmdSmiTemperatureMetric.CURRENT,
+            )
+        except amdsmi.AmdSmiException as e:
+            logger.error(f"Failed to get temperature for device {i}: {e}")
+            return 0
+        except IndexError:
+            logger.error(f"Device index {i} out of range.")
+            return 0
+
+    def read_power_on_device(self, i: int) -> int:
+        """Read the current power usage for the i-th device."""
+        try:
+            return int(
+                amdsmi.amdsmi_get_power_info(self.devices[i])["average_socket_power"]
+            )
+        except amdsmi.AmdSmiException as e:
+            logger.error(f"Failed to get power usage for device {i}: {e}")
+            return 0
+        except IndexError:
+            logger.error(f"Device index {i} out of range.")
+            return 0
+
+    def read_utilization_on_device(self, i: int) -> tuple[int, int]:
+        """Read the current utilization for the i-th device."""
+        try:
+            utilization = amdsmi.amdsmi_get_utilization_count(
+                self.devices[i],
+                [
+                    amdsmi.AmdSmiUtilizationCounterType.COARSE_GRAIN_GFX_ACTIVITY,
+                    amdsmi.AmdSmiUtilizationCounterType.COARSE_GRAIN_MEM_ACTIVITY,
+                ],
+            )
+            return utilization[0]["value"], utilization[1][
+                "value"
+            ]  # GPU utilization, Memory utilization
+        except amdsmi.AmdSmiException as e:
+            logger.error(f"Failed to get utilization for device {i}: {e}")
+            return 0, 0
+        except IndexError:
+            logger.error(f"Device index {i} out of range.")
+            return 0, 0
+
+    def read_temperature(self) -> list[int]:
+        """Read the current temperature for all devices."""
+        return [self.read_temperature_on_device(i) for i in range(len(self.devices))]
+
+    def read_power(self) -> list[int]:
+        """Read the current power usage for all devices."""
+        return [self.read_power_on_device(i) for i in range(len(self.devices))]
+
+    def read_utilization(self) -> list[tuple[int, int]]:
+        """Read the current utilization for all devices."""
+        return [self.read_utilization_on_device(i) for i in range(len(self.devices))]
+
+    def read(self) -> list[int]:
+        """Read the specified quantities for all devices."""
+        res = []
+        for q in self.quantities:
+            if q == Temperature:
+                res = res + self.read_temperature()
+            elif q == Power:
+                res = res + self.read_power()
+            elif q == Utilization:
+                util = self.read_utilization()
+                res = res + [u[0] for u in util]  # GPU utilization
+                res = res + [u[1] for u in util]  # Memory utilization
+            else:
+                logger.warning(f"Unsupported quantity requested: {q}. Skipping.")
+        return res

src/wattameter/utils/setup_amdsmi.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+# SPDX-License-Identifier: BSD-3-Clause
+# SPDX-FileCopyrightText: 2025, Alliance for Energy Innovation, LLC
+#
+# This script sets up the AMD SMI Python environment by installing the amdsmi package
+
+if [ -z "$ROCMPATH" ]; then
+    echo "ROCMPATH is not set. Trying to find it through rocm-smi..."
+    ROCMPATH=$(which rocm-smi | xargs readlink -f | xargs dirname | xargs dirname | xargs dirname)
+    if [ -z "$ROCMPATH" ]; then
+        echo "Could not find ROCMPATH. Please set it manually."
+        exit 1
+    fi
+fi
+
+# Install AMD SMI from the target ROCm instance.
+# Copy to a writable temp dir first because pip's egg_info might not write
+# into a read-only system path where amdsmi.egg-info already exists.
+AMDSMI_TMP=$(mktemp -d)
+cp -r "$ROCMPATH/share/amd_smi/." "$AMDSMI_TMP/"
+python -m pip install "$AMDSMI_TMP"
+rm -rf "$AMDSMI_TMP"