diff --git a/projects/amdsmi/amdsmi_cli/amdsmi_helpers.py b/projects/amdsmi/amdsmi_cli/amdsmi_helpers.py index e3647eac31..b1e6308724 100644 --- a/projects/amdsmi/amdsmi_cli/amdsmi_helpers.py +++ b/projects/amdsmi/amdsmi_cli/amdsmi_helpers.py @@ -26,6 +26,7 @@ import os import platform import sys import time +import re from subprocess import run from subprocess import PIPE, STDOUT @@ -305,7 +306,21 @@ class AMDSMIHelpers(): return (gpu_choices, gpu_choices_str) - def get_device_handles_from_gpu_selections(self, gpu_selections: List[str], gpu_choices=None): + @staticmethod + def is_UUID(uuid_question: str) -> bool: + """Determine if given string is of valid UUID format + Args: + uuid_question (str): the given string to be evaluated. + Returns: + True or False: wether the UUID given matches the UUID format. + """ + UUID_pattern = re.compile("^[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}$", flags=re.IGNORECASE) + if re.match(UUID_pattern, uuid_question) is None: + return False + return True + + + def get_device_handles_from_gpu_selections(self, gpu_selections: List[str], gpu_choices=None) -> tuple: """Convert provided gpu_selections to device_handles Args: @@ -313,9 +328,9 @@ class AMDSMIHelpers(): ex: ID:0 | BDF:0000:23:00.0 | UUID:ffffffff-0000-1000-0000-000000000000 gpu_choices (dict{gpu_choices}): This is a dictionary of the possible gpu_choices Returns: - (True, list[device_handles]): Returns a list of all the gpu_selections converted to + (True, True, list[device_handles]): Returns a list of all the gpu_selections converted to amdsmi device_handles - (False, str): Return False, and the first input that failed to be converted + (False, valid_gpu_format, str): Return False, whether the format of the GPU input is valid, and the first input that failed to be converted """ if 'all' in gpu_selections: return (True, amdsmi_interface.amdsmi_get_processor_handles()) @@ -324,6 +339,7 @@ class AMDSMIHelpers(): gpu_selections = [gpu_selections] if gpu_choices is None: + # obtains dictionary of possible gpu choices gpu_choices = self.get_gpu_choices()[0] selected_device_handles = [] @@ -332,6 +348,7 @@ class AMDSMIHelpers(): for gpu_id, gpu_info in gpu_choices.items(): bdf = gpu_info['BDF'] + is_bdf = True uuid = gpu_info['UUID'] device_handle = gpu_info['Device Handle'] @@ -347,13 +364,16 @@ class AMDSMIHelpers(): valid_gpu_choice = True break except Exception: - # Ignore exception when checking if the gpu_choice is a BDF + is_bdf = False pass if not valid_gpu_choice: logging.debug(f"AMDSMIHelpers.get_device_handles_from_gpu_selections - Unable to convert {gpu_selection}") - return False, gpu_selection - return True, selected_device_handles + valid_gpu_format = True + if not self.is_UUID(gpu_selection) and not gpu_selection.isdigit() and not is_bdf: + valid_gpu_format = False + return False, valid_gpu_format, gpu_selection + return True, True, selected_device_handles def get_device_handles_from_cpu_selections(self, cpu_selections: List[str], cpu_choices=None): @@ -390,8 +410,11 @@ class AMDSMIHelpers(): break if not valid_cpu_choice: logging.debug(f"AMDSMIHelpers.get_device_handles_from_cpu_selections - Unable to convert {cpu_selection}") - return False, cpu_selection - return True, selected_device_handles + valid_cpu_format = True + if not cpu_selection.isdigit(): + valid_cpu_format = False + return False, valid_cpu_format, cpu_selection + return True, True, selected_device_handles def get_device_handles_from_core_selections(self, core_selections: List[str], core_choices=None): @@ -428,8 +451,11 @@ class AMDSMIHelpers(): break if not valid_core_choice: logging.debug(f"AMDSMIHelpers.get_device_handles_from_core_selections - Unable to convert {core_selection}") - return False, core_selection - return True, selected_device_handles + valid_core_format = True + if not core_selection.isdigit(): + valid_core_format = False + return False, valid_core_format, core_selection + return True, True, selected_device_handles def handle_gpus(self, args, logger, subcommand): diff --git a/projects/amdsmi/amdsmi_cli/amdsmi_parser.py b/projects/amdsmi/amdsmi_cli/amdsmi_parser.py index 4b9c55124b..4f97ad8f7c 100644 --- a/projects/amdsmi/amdsmi_cli/amdsmi_parser.py +++ b/projects/amdsmi/amdsmi_cli/amdsmi_parser.py @@ -262,13 +262,16 @@ class AMDSMIParser(argparse.ArgumentParser): def __call__(self, parser, args, values, option_string=None): if "all" in gpu_choices: del gpu_choices["all"] - status, selected_device_handles = amdsmi_helpers.get_device_handles_from_gpu_selections(gpu_selections=values, + status, gpu_format, selected_device_handles = amdsmi_helpers.get_device_handles_from_gpu_selections(gpu_selections=values, gpu_choices=gpu_choices) if status: setattr(args, self.dest, selected_device_handles) else: if selected_device_handles == '': raise amdsmi_cli_exceptions.AmdSmiMissingParameterValueException("--gpu", _GPUSelectAction.ouputformat) + elif not gpu_format: + raise amdsmi_cli_exceptions.AmdSmiInvalidParameterValueException(selected_device_handles, + _GPUSelectAction.ouputformat) else: raise amdsmi_cli_exceptions.AmdSmiDeviceNotFoundException(selected_device_handles, _GPUSelectAction.ouputformat, @@ -289,13 +292,16 @@ class AMDSMIParser(argparse.ArgumentParser): def __call__(self, parser, args, values, option_string=None): if "all" in cpu_choices: del cpu_choices["all"] - status, selected_device_handles = amdsmi_helpers.get_device_handles_from_cpu_selections(cpu_selections=values, + status, cpu_format, selected_device_handles = amdsmi_helpers.get_device_handles_from_cpu_selections(cpu_selections=values, cpu_choices=cpu_choices) if status: setattr(args, self.dest, selected_device_handles) else: if selected_device_handles == '': raise amdsmi_cli_exceptions.AmdSmiMissingParameterValueException("--cpu", _CPUSelectAction.ouputformat) + elif not cpu_format: + raise amdsmi_cli_exceptions.AmdSmiInvalidParameterValueException(selected_device_handles, + _CPUSelectAction.ouputformat) else: raise amdsmi_cli_exceptions.AmdSmiDeviceNotFoundException(selected_device_handles, _CPUSelectAction.ouputformat, @@ -315,13 +321,16 @@ class AMDSMIParser(argparse.ArgumentParser): def __call__(self, parser, args, values, option_string=None): if "all" in core_choices: del core_choices["all"] - status, selected_device_handles = amdsmi_helpers.get_device_handles_from_core_selections(core_selections=values, + status, core_format, selected_device_handles = amdsmi_helpers.get_device_handles_from_core_selections(core_selections=values, core_choices=core_choices) if status: setattr(args, self.dest, selected_device_handles) else: if selected_device_handles == '': raise amdsmi_cli_exceptions.AmdSmiMissingParameterValueException("--core", _CoreSelectAction.ouputformat) + elif not core_format: + raise amdsmi_cli_exceptions.AmdSmiInvalidParameterValueException(selected_device_handles, + _CoreSelectAction.ouputformat) else: raise amdsmi_cli_exceptions.AmdSmiDeviceNotFoundException(selected_device_handles, _CoreSelectAction.ouputformat,