TheRock/rocm-systems/projects/amdsmi/amdsmi_cli/amdsmi_cli.py

#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK
#
# 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.

import functools
import logging
import sys
import os

try:
    import argcomplete
except ImportError as e:
    logging.debug(f"Unhandled import error: {e}")
    logging.debug("argcomplete module not found. Autocomplete will not work.")

from typing import TYPE_CHECKING

# only used for type checking
# pyright trips up and cannot find amdsmi scripts without it
if TYPE_CHECKING:
    from amdsmi import amdsmi_exception

# Set the environment variable for GPU metrics cache duration
gpu_metrics_cache_ms = os.environ.setdefault("AMDSMI_GPU_METRICS_CACHE_MS", "100")
logging.debug("AMDSMI_GPU_METRICS_CACHE_MS = %sms", gpu_metrics_cache_ms)

# Set the environment variable for ASIC cache duration
asic_info_cache_ms = os.environ.setdefault("AMDSMI_ASIC_INFO_CACHE_MS", "10000")  # 10 seconds
logging.debug("AMDSMI_ASIC_INFO_CACHE_MS = %sms", asic_info_cache_ms)

# Set the environment variable for process info cache duration
process_info_cache_ms = os.environ.setdefault("AMDSMI_PROCESS_INFO_CACHE_MS", "100")
logging.debug("AMDSMI_PROCESS_INFO_CACHE_MS = %sms", process_info_cache_ms)

try:
    from amdsmi_init import *
    from amdsmi_helpers import AMDSMIHelpers
    from amdsmi_commands import AMDSMICommands
    from amdsmi_parser import AMDSMIParser
    from amdsmi_logger import AMDSMILogger
    import amdsmi_cli_exceptions
except ImportError:
    current_path = os.path.dirname(os.path.abspath(__file__))
    cli_files_path = f"{current_path}/../libexec/amdsmi_cli"
    sys.path.append(cli_files_path)
    try:
        from amdsmi_init import *
        from amdsmi_helpers import AMDSMIHelpers
        from amdsmi_commands import AMDSMICommands
        from amdsmi_parser import AMDSMIParser
        from amdsmi_logger import AMDSMILogger
        import amdsmi_cli_exceptions
    except ImportError as e:
        print(f"Unhandled import error: {e}")
        print(f"Unable to import amdsmi_cli files. Check {cli_files_path} if they are present.")
        sys.exit(1)


def _print_error(e, destination):
    if destination in ["stdout", "json", "csv"]:
        print(e)
    else:
        f = open(destination, "w", encoding="utf-8")
        f.write(e)
        f.close()
        print("Error occurred. Result written to " + str(destination) + " file")


def configure_logging_and_execute(args, amd_smi_commands):
    """
    Configures logging based on the provided arguments and executes the subcommand.

    Args:
        args: Parsed command-line arguments.
        amd_smi_commands: Instance of AMDSMICommands.
    """
    # Remove previous log handlers
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    # To enable debug logs in AMD SMI library:
    #   set RSMI_LOGGING = 1 for logging to files
    #   set RSMI_LOGGING = 2 for logging to stdout
    #   set RSMI_LOGGING = 3 for logging to stdout and files
    #   set RSMI_LOGGING = 0 to disable logging
    # Files will be located in /var/log/amd_smi_lib/AMD-SMI-lib.log*

    # log string with the following format:
    # loglevel | YYYY-MM-DD HH:MM:SS.ms | filename:line | message
    logging_dict = {
        "DEBUG": logging.DEBUG,
        "INFO": logging.INFO,
        "WARNING": logging.WARNING,
        "ERROR": logging.ERROR,
        "CRITICAL": logging.CRITICAL,
    }

    time = "%(asctime)s.%(msecs)03d"
    datefmt = "%Y-%m-%d %H:%M:%S"
    logging.basicConfig(
        format="%(levelname)s | " + time + " | %(filename)s:%(lineno)d | %(message)s",
        level=logging_dict[args.loglevel],
        datefmt=datefmt,
    )

    # Disable traceback for non-debug log levels
    if args.loglevel == "DEBUG":
        sys.tracebacklimit = 10
    else:
        sys.tracebacklimit = -1

    logging.debug(args)

    # Check if --rocm-smi flag is set (top-level flag, not a subcommand)
    if hasattr(args, "rocm_smi") and args.rocm_smi:
        amd_smi_commands.rocm_smi(args)
        return

    # Execute subcommands
    args.func(args)


if __name__ == "__main__":
    # Disable traceback before possible init errors in AMDSMICommands and AMDSMIParser
    copy_argv = str(sys.argv.copy()).upper()
    if "DEBUG" in copy_argv:
        sys.tracebacklimit = 10
    else:
        sys.tracebacklimit = -1

    amd_smi_helpers = AMDSMIHelpers()
    amd_smi_commands = AMDSMICommands(helpers=amd_smi_helpers)
    amd_smi_parser = AMDSMIParser(
        amd_smi_commands.version,
        amd_smi_commands.list_devices,
        amd_smi_commands.static,
        amd_smi_commands.firmware,
        amd_smi_commands.bad_pages,
        amd_smi_commands.metric,
        amd_smi_commands.process,
        amd_smi_commands.profile,
        amd_smi_commands.event,
        amd_smi_commands.topology,
        amd_smi_commands.set_value,
        amd_smi_commands.reset,
        amd_smi_commands.monitor,
        amd_smi_commands.xgmi,
        amd_smi_commands.partition,
        amd_smi_commands.ras,
        amd_smi_commands.node,
        amd_smi_commands.fabric,
        amd_smi_commands.rocm_smi,
        amd_smi_commands.default,
        sys_argv=sys.argv,
        helpers=amd_smi_helpers,
    )
    try:
        argcomplete.autocomplete(amd_smi_parser)
    except NameError:
        logging.debug("argcomplete module not found. Autocomplete will not work.")

    # Store possible subcommands & aliases for later errors
    valid_commands = amd_smi_parser.possible_commands
    valid_commands += ["--help", "-h", "--rocm-smi"]

    # Convert arguments to lowercase, but preserve case for folder path values
    processed_argv = []
    # Arguments that should preserve case
    case_sensitive_args = [
        "--folder",
        "--file",
        "--gpu",
        "--cpu",
        "--core",
        "--profile",
        "--cper-file",
    ]
    case_sensitive_prefixes = [
        "--folder=",
        "--file=",
        "--gpu=",
        "--cpu=",
        "--core=",
        "--profile=",
        "--cper-file=",
    ]

    preserve_case_for_next = False
    for i, arg in enumerate(sys.argv):
        if preserve_case_for_next:
            # Preserve case for the next argument value
            processed_argv.append(arg)
            preserve_case_for_next = False
        elif arg in case_sensitive_args:
            # Convert flag to lowercase but preserve next value
            processed_argv.append(arg.lower())
            preserve_case_for_next = True
        elif any(arg.startswith(prefix) for prefix in case_sensitive_prefixes):
            # Handle --arg=value format, preserve case for the value part
            for prefix in case_sensitive_prefixes:
                if arg.startswith(prefix):
                    flag = prefix.rstrip("=")
                    value = arg[len(prefix) :]
                    processed_argv.append(flag.lower() + "=" + value)
                    break
        elif arg.startswith("--") or not arg.startswith("-"):
            # Convert other long options and positional arguments to lowercase
            processed_argv.append(arg.lower())
        else:
            # Preserve case for short options
            processed_argv.append(arg)
    sys.argv = processed_argv

    try:
        if len(sys.argv) == 1:
            args = amd_smi_parser.parse_args(args=["default"])
        elif sys.tracebacklimit == 10 and (sys.argv[1] == "--loglevel"):
            args = amd_smi_parser.parse_args(args=["default", "--loglevel"] + sys.argv[2:])
        elif sys.argv[1] in valid_commands:
            args = amd_smi_parser.parse_args(args=None)
        else:
            raise amdsmi_cli_exceptions.AmdSmiInvalidSubcommandException(
                sys.argv[1], amd_smi_commands.logger.destination
            )

        # Handle command modifiers before subcommand execution
        # human readable is the default output format
        if hasattr(args, "json") and args.json:
            amd_smi_commands.logger.format = amd_smi_commands.logger.LoggerFormat.json.value
        if hasattr(args, "csv") and args.csv:
            amd_smi_commands.logger.format = amd_smi_commands.logger.LoggerFormat.csv.value
        if hasattr(args, "file") and args.file:
            amd_smi_commands.logger.destination = args.file
        configure_logging_and_execute(args, amd_smi_commands)
    except amdsmi_cli_exceptions.AmdSmiException as e:
        _print_error(
            f"{type(e).__module__}.{type(e).__name__}: {str(e)}",
            amd_smi_commands.logger.destination,
        )
        sys.exit(abs(e.value))
    except amdsmi_exception.AmdSmiLibraryException as e:
        exc = amdsmi_cli_exceptions.AmdSmiLibraryErrorException(
            amd_smi_commands.logger.format, e.get_error_code()
        )
        _print_error(
            f"{type(exc).__module__}.{type(exc).__name__}: {str(exc)}",
            amd_smi_commands.logger.destination,
        )
        sys.exit(abs(exc.value))
    except PermissionError as e:
        command = sys.argv[1] if len(sys.argv) > 1 else ""
        outputformat = amd_smi_commands.logger.format
        exc = amdsmi_cli_exceptions.AmdSmiPermissionDeniedException(command, outputformat)
        _print_error(
            f"{type(exc).__module__}.{type(exc).__name__}: {str(exc)}",
            amd_smi_commands.logger.destination,
        )
        sys.exit(abs(exc.value))