rocm-systems/tools/amdsmi_quick_start.py

# 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.

# This script is used to quickly debug and check amdsmi Python APIs.
# This is not meant to serve best practices for development.
# Run this post install with python3 -i quick_start.py

import atexit
import logging
import signal
import sys

from amdsmi import *
from pathlib import Path

# Make exit & quit work without parens because it's annoying
type(exit).__repr__ = sys.exit

def signal_handler(sig, frame):
    logging.debug(f"Handling signal: {sig}")
    sys.exit(0)

amdsmi_init()
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
atexit.register(amdsmi_shut_down)

gpus = amdsmi_get_processor_handles()
cpus = amdsmi_get_cpusocket_handles()

print(f"gpus variable populated with:{gpus}")
print(f"cpus variable populated with:{cpus}")