Use amd-smi Python API instead of CLI (#1334)

* Use amd-smi Python API instead of CLI

Formatting fix

python path

* Update CHANGELOG

* Create amdsmi interface

* Added amdsmi tests

* Removed run

* Prioritize rocm's amdsmi python API

* address review comments

* update changelog

* fix ruff formatting

---------

Co-authored-by: Vignesh Edithal <Vignesh.Edithal@amd.com>

projects/rocprofiler-compute/src/utils/amdsmi_interface.py

@@ -0,0 +1,136 @@
+##############################################################################
+# MIT License
+#
+# Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+##############################################################################
+
+import os
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+from utils.logger import (
+    console_debug,
+    console_error,
+    console_warning,
+)
+
+sys.path.insert(0, os.getenv("ROCM_PATH", "/opt/rocm") + "/share/amd_smi")
+
+try:
+    import amdsmi
+except ImportError as e:
+    console_warning(f"Unhandled import error: {e}")
+    console_error("Failed to import the amdsmi Python library.")
+
+
+@contextmanager
+def amdsmi_ctx() -> Iterator[None]:
+    """Context manager to initialize and shutdown amdsmi."""
+    try:
+        amdsmi.amdsmi_init()
+        yield
+    except Exception as e:
+        console_warning(f"amd-smi init failed: {e}")
+    finally:
+        try:
+            amdsmi.amdsmi_shut_down()
+        except Exception as e:
+            console_warning(f"amd-smi shutdown failed: {e}")
+
+
+def get_device_handle() -> "amdsmi.ProcessorHandle | None":
+    """Get the first AMD device handle."""
+    try:
+        devices = amdsmi.amdsmi_get_processor_handles()
+        if len(devices) == 0:
+            console_warning("No AMD GPU detected!")
+            return None
+        console_debug(f"Found {len(devices)} AMD device(s).")
+        return devices[0]
+    except Exception as e:
+        console_warning(f"Error getting device handle: {e}")
+        return None
+
+
+def get_mem_max_clock() -> float:
+    """Get the maximum memory clock of the device."""
+    try:
+        return amdsmi.amdsmi_get_clock_info(
+            get_device_handle(), amdsmi.AmdSmiClkType.GFX
+        )["max_clk"]
+    except Exception as e:
+        console_warning(f"Error getting memory clocks: {e}")
+        return 0.0
+
+
+def get_gpu_model() -> str:
+    """Get the GPU model name."""
+    try:
+        gpu_model_info = (
+            # board -> product_name
+            amdsmi.amdsmi_get_gpu_board_info(get_device_handle())["product_name"],
+            # asic -> market_name
+            amdsmi.amdsmi_get_gpu_asic_info(get_device_handle())["market_name"],
+            # vbios -> name
+            amdsmi.amdsmi_get_gpu_vbios_info(get_device_handle())["name"],
+        )
+        console_debug(f"gpu model info: {str(gpu_model_info)}")
+        return gpu_model_info
+    except Exception as e:
+        console_warning(f"Error getting gpu model info: {e}")
+        return "N/A"
+
+
+def get_gpu_vbios_part_number() -> str:
+    """Get the GPU VBIOS part number."""
+    try:
+        vbios_part_number = amdsmi.amdsmi_get_gpu_vbios_info(get_device_handle())[
+            "part_number"
+        ]
+        console_debug(f"GPU VBIOS Part Number: {vbios_part_number}")
+        return vbios_part_number
+    except Exception as e:
+        console_warning(f"Error getting GPU VBIOS part number: {e}")
+        return "N/A"
+
+
+def get_gpu_compute_partition() -> str:
+    """Get the GPU compute partition."""
+    try:
+        compute_partition = amdsmi.amdsmi_get_gpu_compute_partition(get_device_handle())
+        console_debug(f"GPU Compute Partition: {compute_partition}")
+        return compute_partition
+    except Exception as e:
+        console_warning(f"Error getting GPU compute partition: {e}")
+        return "N/A"
+
+
+def get_gpu_memory_partition() -> str:
+    """Get the GPU memory partition."""
+    try:
+        memory_partition = amdsmi.amdsmi_get_gpu_memory_partition(get_device_handle())
+        console_debug(f"GPU Memory Partition: {memory_partition}")
+        return memory_partition
+    except Exception as e:
+        console_warning(f"Error getting GPU memory partition: {e}")
+        return "N/A"

projects/rocprofiler-compute/src/utils/specs.py

@@ -28,7 +28,6 @@ from __future__ import annotations
 
 import argparse
 import importlib
-import json
 import os
 import re
 import socket
@@ -42,6 +41,12 @@ from typing import Any, Optional, TypeVar
 import pandas as pd
 
 import config
+from utils.amdsmi_interface import (
+    amdsmi_ctx,
+    get_gpu_compute_partition,
+    get_gpu_memory_partition,
+    get_gpu_vbios_part_number,
+)
 from utils.logger import (
     console_debug,
     console_error,
@@ -252,57 +257,18 @@ def extract_gpu_info() -> dict[str, Any]:
         "memory_partition": None,
     }
 
-    # Load amd-smi static data for GPU 0
-    static_output = run(["amd-smi", "static", "--gpu=0", "--json"], exit_on_error=True)
-    if static_output is None:
-        return result
-
-    try:
-        static_data = json.loads(static_output)
-    except json.JSONDecodeError as e:
-        console_warning(f"Failed to parse amd-smi static output: {e}")
-        return result
-
-    # Extract GPU data
-    gpu_list = (
-        static_data
-        if isinstance(static_data, list)
-        else static_data.get("gpu_data", [])
-    )
-    gpu_data = gpu_list[0] if gpu_list else {}
-    result["vbios"] = gpu_data.get("vbios", {}).get("part_number")
-
-    # Load amd-smi partition data for GPU 0 (amd-smi >= 26.0.0)
-    partition_output = run(
-        ["amd-smi", "partition", "--gpu=0", "--json"], exit_on_error=False
-    )
-    partition_data = {}
-
-    if partition_output:
-        try:
-            partition_data = json.loads(partition_output)
-        except json.JSONDecodeError:
-            partition_data = {}
-
-    current_partition = partition_data.get("current_partition", [{}])[0]
-
-    # Extract partition values with gpu_data fallback (amd-smi < 26.0.0)
-    result["compute_partition"] = (
-        current_partition.get("accelerator_type")
-        or gpu_data.get("partition", {}).get("accelerator_partition")
-        or gpu_data.get("partition", {}).get("compute_partition")
-    )
-    result["memory_partition"] = current_partition.get("memory") or gpu_data.get(
-        "partition", {}
-    ).get("memory_partition")
+    with amdsmi_ctx():
+        result["vbios"] = get_gpu_vbios_part_number()
+        result["compute_partition"] = get_gpu_compute_partition()
+        result["memory_partition"] = get_gpu_memory_partition()
 
     # Apply defaults and warnings
-    if not result["compute_partition"]:
+    if result["compute_partition"] == "N/A" or not result["compute_partition"]:
         console_warning("Cannot detect accelerator partition from amd-smi.")
         console_warning("Applying default accelerator partition: SPX")
         result["compute_partition"] = "SPX"
 
-    if not result["memory_partition"]:
+    if result["memory_partition"] == "N/A" or not result["memory_partition"]:
         console_warning("Cannot detect memory partition from amd-smi.")
 
     console_debug(
@@ -891,12 +857,8 @@ def run(cmd: list[str], exit_on_error: bool = False) -> str:
             'Try passing a path to an existing workload results in "analyze" mode.'
         )
 
-    if exit_on_error:
-        if cmd[0] == "amd-smi":
-            if p.returncode != 2 and p.returncode != 0:  # type: ignore
-                console_error("No GPU detected. Unable to load amd-smi")
-        elif p.returncode != 0:  # type: ignore
-            console_error(f"Command {cmd} failed with non-zero exit code")
+    if exit_on_error and p.returncode != 0:  # type: ignore
+        console_error(f"Command {cmd} failed with non-zero exit code")
     return p.stdout.decode("utf-8")  # type: ignore