TheRock/rocm-systems/projects/amdsmi/example/amd_smi_partition_example.py

#!/usr/bin/env python3
# Copyright (C) Advanced Micro Devices. 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.

# NOTE: Partition changes alter device topology -- AMD SMI must re-initialize
#
# Changing either memory partition (NPS mode) or accelerator (compute)
# partition changes the number of logical devices visible to the OS and to
# AMD SMI. AMD SMI builds its internal device table at amdsmi_init(); it does
# not update live. After any partition change, callers must call
# amdsmi_shut_down() followed by amdsmi_init() to pull in the updated
# topology; any processor handle obtained before re-initialization is
# invalid and must not be used.
#
#   Memory partition (NPS mode):
#     Requires driver reload to take effect. The driver reload
#     tears down and rebuilds all kernel device objects, which also destroys
#     any existing handles. amdsmi_shut_down() + amdsmi_init() is therefore
#     mandatory before querying or changing anything further.
#
#   Accelerator (compute) partition:
#     Takes effect immediately without a driver reload, but still changes the
#     logical device count (e.g. SPX: 1 handle per physical GPU vs. DPX: 2).
#     amdsmi_shut_down() + amdsmi_init() is required to get the correct
#     post-change handle list before issuing any further queries or sets.

import amdsmi
from typing import Any, Dict, List, Tuple

# ---------------------------------------------------------------------------
# Partition name converters
# ---------------------------------------------------------------------------


# Returns a comma-separated list of supported NPS modes (e.g. "NPS1,NPS2,NPS4").
def nps_caps_str(caps: List[str]) -> str:
    filtered = [c for c in caps if c != "N/A"]
    return ",".join(filtered) if filtered else "none"


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

_FILL = "-" * 70
_NPS_ORDER = {"NPS1": 1, "NPS2": 2, "NPS4": 4, "NPS8": 8}


def print_separator(title: str = "") -> None:
    if not title:
        print(_FILL)
        return
    prefix_len = 5 + len(title)
    remaining = _FILL[prefix_len:] if prefix_len < len(_FILL) else ""
    print(f"\n--- {title} {remaining}")


# ---------------------------------------------------------------------------
# AMD SMI lifecycle
# ---------------------------------------------------------------------------


# Scoped amdsmi session. Calls amdsmi_init on __enter__ and amdsmi_shut_down
# on __exit__. Re-enumerate (new instance) after any partition change that
# alters device topology:
#   - Memory partition change requires a driver reload, which destroys all
#     kernel device objects and invalidates every processor handle.
#   - Accelerator partition change changes the number of visible logical
#     devices (e.g. SPX: 1 per GPU -> DPX: 2 per GPU), so the handle list
#     built at amdsmi_init is stale even without a driver reload.
# Treat every processor handle as valid only within the session that
# obtained it.
class AmdsmiSession:
    def __init__(self) -> None:
        self._ok = False

    def is_ok(self) -> bool:
        return self._ok

    def __enter__(self) -> "AmdsmiSession":
        try:
            amdsmi.amdsmi_init()
            self._ok = True
            print("AMDSMI session started (amdsmi_init called)")
        except amdsmi.AmdSmiException as e:
            print(f"amdsmi_init failed: {e}")
        return self

    def __exit__(self, *_: Any) -> None:
        if not self._ok:
            return
        amdsmi.amdsmi_shut_down()
        print("AMDSMI session ended (amdsmi_shut_down called)")


# ---------------------------------------------------------------------------
# GPU enumeration
# ---------------------------------------------------------------------------


# Enumerate all AMD GPU handles into a flat list.
def get_all_gpu_handles() -> List:
    try:
        return amdsmi.amdsmi_get_processor_handles()
    except amdsmi.AmdSmiException as e:
        print(f"amdsmi_get_processor_handles failed: {e}")
        return []


# Returns only the "primary" handles from a full handle list -- i.e., those
# that successfully respond to amdsmi_get_gpu_accelerator_partition_profile_config.
# In a partitioned layout (e.g. CPX/NPS4) amdsmi enumerates one handle per
# *logical* partition. Only the root handle of each physical GPU owns the
# partition config; sub-partition handles return NOT_SUPPORTED. Calling set on
# a sub-partition handle is a no-op (NOT_SUPPORTED), so filter them out first.
# Each entry is (amd-smi GPU index, handle).
def get_primary_gpu_handles(all_gpus: List) -> List[Tuple[int, Any]]:
    primary = []
    for i, gpu in enumerate(all_gpus):
        try:
            amdsmi.amdsmi_get_gpu_accelerator_partition_profile_config(gpu)
            primary.append((i, gpu))
        except amdsmi.AmdSmiException:
            pass  # sub-partition handles are expected to raise here
    return primary


# ---------------------------------------------------------------------------
# Per-step display helpers
# ---------------------------------------------------------------------------


def print_current_partition(idx: int, gpu: Any) -> None:
    print(f"  GPU {idx}:")
    try:
        cur_profile = amdsmi.amdsmi_get_gpu_accelerator_partition_profile(gpu)
        pp = cur_profile.get("partition_profile", {})
        print(f"    Accelerator profile type : {pp.get('profile_type', 'N/A')}")
        print(f"    Profile index            : {pp.get('profile_index', 'N/A')}")
        print(f"    Num partitions           : {pp.get('num_partitions', 'N/A')}")
        print(f"    Partition ID             : {cur_profile.get('partition_id', ['N/A'])[0]}")
        print(f"    Compatible NPS           : {nps_caps_str(pp.get('memory_caps', []))}")
    except amdsmi.AmdSmiException as e:
        print(f"    amdsmi_get_gpu_accelerator_partition_profile: {e}")
    try:
        mem_config = amdsmi.amdsmi_get_gpu_memory_partition_config(gpu)
        print(f"    Memory partition : {mem_config.get('mp_mode', 'N/A')}")
    except amdsmi.AmdSmiException as e:
        print(f"    amdsmi_get_gpu_memory_partition_config: {e}")


def print_available_modes(idx: int, gpu: Any) -> None:
    print(f"  GPU {idx}:")
    try:
        mem_config = amdsmi.amdsmi_get_gpu_memory_partition_config(gpu)
        caps = mem_config.get("partition_caps", [])
        print(f"    Supported NPS modes     : {nps_caps_str(caps)}")
    except amdsmi.AmdSmiException as e:
        print(f"    amdsmi_get_gpu_memory_partition_config: {e}")
    try:
        acc_config = amdsmi.amdsmi_get_gpu_accelerator_partition_profile_config(gpu)
        profiles = acc_config.get("profiles", [])
        print(f"    Available accelerator profiles ({len(profiles)}):")
        for profile in profiles:
            print(
                f"      Index {profile.get('profile_index')} "
                f"({profile.get('profile_type')}, "
                f"{profile.get('num_partitions')} partition(s), "
                f"compatible NPS: {nps_caps_str(profile.get('memory_caps', []))})"
            )
            for r, res in enumerate(profile.get("resources", [])):
                print(
                    f"        Resource {r}: {res.get('resource_type')}"
                    f", per_partition={res.get('partition_resource')}"
                    f", shared_by={res.get('num_partitions_share_resource')}"
                )
    except amdsmi.AmdSmiException as e:
        print(f"    amdsmi_get_gpu_accelerator_partition_profile_config: {e}")


# ---------------------------------------------------------------------------
# NPS -> accelerator-profile map
# ---------------------------------------------------------------------------


# Maps each NPS mode to the accelerator profiles that are compatible with it.
# Built by inverting the memory_caps list on every profile returned by
# amdsmi_get_gpu_accelerator_partition_profile_config.
def build_nps_to_profiles_map(acc_config: Dict) -> Dict[str, List[Dict]]:
    nps_map = {}
    for prof in acc_config.get("profiles", []):
        for nps in prof.get("memory_caps", []):
            if nps != "N/A":
                nps_map.setdefault(nps, []).append(prof)
    return nps_map


# ---------------------------------------------------------------------------
# Workflow helpers
# ---------------------------------------------------------------------------


def print_current_partition_info(gpus: List) -> None:
    print_separator("Current partition settings")
    for idx, gpu in enumerate(gpus):
        print_current_partition(idx, gpu)


def print_available_partition_modes(gpus: List) -> None:
    print_separator("Available partition modes")
    for idx, gpu in enumerate(gpus):
        print_available_modes(idx, gpu)


# Returns True if the memory partition was set successfully.
def set_memory_partition(gpu0: Any, target: amdsmi.AmdSmiMemoryPartitionType) -> bool:
    print_separator("Set memory partition")
    # Memory partition is hive-wide -- setting it on one device affects all.
    # Change the target mode to another supported mode as needed.
    target_name = target.name
    print(f"  Attempting: {target_name}")
    try:
        amdsmi.amdsmi_set_gpu_memory_partition_mode(gpu0, target)
        print(f"  amdsmi_set_gpu_memory_partition_mode(GPU 0, {target_name}): success")
        return True
    except amdsmi.AmdSmiException as e:
        print(f"  amdsmi_set_gpu_memory_partition_mode(GPU 0, {target_name}): {e}")
        return False


# For each GPU: build the NPS->profiles map, print every entry, then demonstrate
# iterating over only the profiles compatible with the currently active NPS mode.
def print_profiles_by_nps(gpus: List) -> None:
    print_separator("Accelerator profiles grouped by NPS mode")
    for gpu_idx, gpu in enumerate(gpus):
        print(f"  GPU {gpu_idx}:")

        # --- query available accelerator profiles ---
        try:
            acc_config = amdsmi.amdsmi_get_gpu_accelerator_partition_profile_config(gpu)
        except amdsmi.AmdSmiException as e:
            print(f"    amdsmi_get_gpu_accelerator_partition_profile_config: {e}")
            continue

        # --- query current NPS mode ---
        current_nps = None
        try:
            mem_config = amdsmi.amdsmi_get_gpu_memory_partition_config(gpu)
            current_nps = mem_config.get("mp_mode")
        except amdsmi.AmdSmiException:
            pass

        # --- build map ---
        nps_map = build_nps_to_profiles_map(acc_config)

        # Print every NPS mode and its compatible profiles
        print("    All profiles by NPS mode:")
        for nps, profiles in sorted(nps_map.items(), key=lambda x: _NPS_ORDER.get(x[0], 99)):
            is_current = " [current]" if nps == current_nps else ""
            print(f"      {nps}{is_current} ({len(profiles)} profile(s)):")
            for prof in profiles:
                print(
                    f"        index={prof.get('profile_index')}"
                    f"  type={prof.get('profile_type')}"
                    f"  partitions={prof.get('num_partitions')}"
                )

        # Iterate over profiles for the *current* NPS mode
        if current_nps:
            print(f"    Profiles available for current NPS ({current_nps}):")
            current_profiles = nps_map.get(current_nps, [])
            if current_profiles:
                for prof in current_profiles:
                    print(
                        f"      -> index={prof.get('profile_index')}"
                        f"  type={prof.get('profile_type')}"
                        f"  partitions={prof.get('num_partitions')}"
                    )
            else:
                print("      (none)")


# Set the accelerator partition on every *primary* handle.
#
# Why primary-only:
#   In a partitioned layout (e.g. CPX+NPS4) amdsmi enumerates one handle per
#   logical partition (e.g. 64 handles for 8 physical GPUs in CPX). Only the
#   root handle of each physical GPU owns the partition config; the rest return
#   NOT_SUPPORTED. Calling set on them is noise, so we filter first.
def set_accelerator_partition_all_devices(
    all_gpus: List, profile_index: int, profile_type_name: str
) -> None:
    print_separator(f"Set accelerator partition -> index={profile_index} ({profile_type_name})")
    primary = get_primary_gpu_handles(all_gpus)
    print(f"  Primary handles: {len(primary)} of {len(all_gpus)} total")
    for phy_idx, (gpu_num, gpu) in enumerate(primary):
        try:
            amdsmi.amdsmi_set_gpu_accelerator_partition_profile(gpu, profile_index)
            result = "success"
        except amdsmi.AmdSmiException as e:
            result = str(e)
        print(
            f"  Physical GPU {phy_idx} (amd-smi GPU {gpu_num})"
            f": amdsmi_set_gpu_accelerator_partition_profile -> {result}"
        )


def main() -> None:
    mem_changed = False

    # -----------------------------------------------------------------------
    # Phase 1: Query state, set memory partition, reload driver if changed.
    #
    # If memory partition is successfully changed, a driver reload follows and
    # the session is closed (amdsmi_shut_down()) before Phase 2 re-initializes
    # with fresh handles. If memory partition cannot be changed, the driver
    # reload is skipped and accelerator partition changes are attempted
    # immediately within this same session.
    # -----------------------------------------------------------------------
    with AmdsmiSession() as session:
        if not session.is_ok():
            return
        gpus = get_all_gpu_handles()
        if not gpus:
            print("No GPUs found.")
            return
        print(f"Found {len(gpus)} GPU(s).")

        print_current_partition_info(gpus)
        print_available_partition_modes(gpus)

        # In real-world usage, the target mode may be dictated by workload requirements.
        # For this demo, we pick the highest supported NPS mode to show the most
        # significant partition change.

        # Pick the highest supported NPS mode, preferring NPS4 > NPS2 > NPS1.
        # Query the caps from GPU 0 (memory partition is hive-wide, so any GPU works).
        target_memory_partition = amdsmi.AmdSmiMemoryPartitionType.NPS1
        current_nps = None
        try:
            mem_config = amdsmi.amdsmi_get_gpu_memory_partition_config(gpus[0])
            caps = mem_config.get("partition_caps", [])
            current_nps = mem_config.get("mp_mode")
            if "NPS4" in caps:
                target_memory_partition = amdsmi.AmdSmiMemoryPartitionType.NPS4
            elif "NPS2" in caps:
                target_memory_partition = amdsmi.AmdSmiMemoryPartitionType.NPS2
        except amdsmi.AmdSmiException as e:
            print(f"  [warn] Could not query NPS capabilities: {e}; defaulting to NPS1")
        print(f"  Selected NPS target: {target_memory_partition.name}")
        if current_nps == target_memory_partition.name:
            print(
                f"\n[info] Current NPS mode is already {current_nps}; "
                f"skipping set and driver reload."
            )
        else:
            mem_changed = set_memory_partition(gpus[0], target_memory_partition)
            if mem_changed:
                print(
                    "  Memory partition staged. Run to apply:\n"
                    "    sudo modprobe -r amdgpu && sudo modprobe amdgpu"
                )
            else:
                print("\n[info] Memory partition unchanged; skipping driver reload.")

    # -----------------------------------------------------------------------
    # Phase 2: Re-enumerate to get fresh handles and accurate device count after
    # the driver reload (triggered by the memory partition change). Display the
    # current state and capture the accelerator profiles compatible with the
    # active NPS mode.
    #
    # Accelerator-partition profile dicts are plain data with no embedded
    # handles, so they remain valid across amdsmi_shut_down / amdsmi_init
    # boundaries and can be saved here for use in Phase 3.
    # -----------------------------------------------------------------------
    profiles_to_test = []
    with AmdsmiSession() as session:
        if not session.is_ok():
            return
        gpus = get_all_gpu_handles()
        if not gpus:
            print("No GPUs found after driver reload.")
            return
        print(f"Found {len(gpus)} GPU(s) after driver reload.")

        print_current_partition_info(gpus)
        print_available_partition_modes(gpus)
        print_profiles_by_nps(gpus)

        # Build NPS->profiles map and save the profiles for the active NPS mode.
        # These are plain dicts and remain valid after the session is torn down.
        try:
            mem_config = amdsmi.amdsmi_get_gpu_memory_partition_config(gpus[0])
            acc_config = amdsmi.amdsmi_get_gpu_accelerator_partition_profile_config(gpus[0])
            nps_map = build_nps_to_profiles_map(acc_config)
            current_nps = mem_config.get("mp_mode")
            if current_nps in nps_map:
                profiles_to_test = nps_map[current_nps]
                print(
                    f"\n[info] {len(profiles_to_test)} accelerator profile(s) to iterate "
                    f"for current NPS mode ({current_nps})."
                )
        except amdsmi.AmdSmiException as e:
            print(f"[warn] Could not build accelerator profile map: {e}")

    # -----------------------------------------------------------------------
    # Phase 3: For each accelerator profile compatible with the current NPS mode:
    #   a) Open a session, set the profile on every device, close the session.
    #   b) Open a NEW session to re-enumerate before reading anything.
    #      Accelerator partition changes alter the number of logical devices
    #      visible to the OS (e.g. SPX: N handles -> DPX: 2*N handles), so the
    #      handle list from step (a) is immediately stale. A new amdsmi_init
    #      is required to get the correct post-change device count and handles
    #      before displaying topology or issuing further queries.
    # -----------------------------------------------------------------------
    if not profiles_to_test:
        print("[info] No accelerator profiles available for current NPS mode; skipping Phase 3.")
        return

    for target in profiles_to_test:
        profile_index = target.get("profile_index")
        profile_type = target.get("profile_type")

        # a) Set accelerator partition on all devices.
        with AmdsmiSession() as session:
            if not session.is_ok():
                continue
            gpus = get_all_gpu_handles()
            if not gpus:
                print(
                    f"\n[warn] No GPUs found; skipping profile "
                    f"index={profile_index} ({profile_type})."
                )
                continue
            print(
                f"Found {len(gpus)} GPU(s) BEFORE re-initialization for accelerator profile "
                f"change to index={profile_index} ({profile_type})."
            )
            set_accelerator_partition_all_devices(gpus, profile_index, profile_type)

        # b) Re-enumerate and show resulting device topology.
        with AmdsmiSession() as session:
            if not session.is_ok():
                continue
            gpus = get_all_gpu_handles()
            if not gpus:
                print(
                    f"\n[warn] No GPUs found; skipping profile "
                    f"index={profile_index} ({profile_type})."
                )
                continue
            print(
                f"Found {len(gpus)} GPU(s) AFTER re-initialization for accelerator profile "
                f"change to index={profile_index} ({profile_type})."
            )
            print_current_partition_info(gpus)
            print_available_partition_modes(gpus)
            print_profiles_by_nps(gpus)


if __name__ == "__main__":
    main()