Skip to content

Commit

Permalink
moved exception_map out of function
Browse files Browse the repository at this point in the history
  • Loading branch information
@parthraut
parthraut committed Apr 30, 2024
1 parent 13d4ee8 commit 5767ae9
Showing 1 changed file with 48 additions and 46 deletions.
94 changes: 48 additions & 46 deletions zeus/device/gpu.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,35 +273,13 @@ def getTotalEnergyConsumption(self) -> int:


def _handle_nvml_errors(func):
_exception_map = {
pynvml.NVML_ERROR_UNINITIALIZED: ZeusGPUInitError,
pynvml.NVML_ERROR_INVALID_ARGUMENT: ZeusGPUInvalidArgError,
pynvml.NVML_ERROR_NOT_SUPPORTED: ZeusGPUNotSupportedError,
pynvml.NVML_ERROR_NO_PERMISSION: ZeusGPUNoPermissionError,
pynvml.NVML_ERROR_ALREADY_INITIALIZED: ZeusGPUAlreadyInitializedError,
pynvml.NVML_ERROR_NOT_FOUND: ZeusGPUNotFoundError,
pynvml.NVML_ERROR_INSUFFICIENT_SIZE: ZeusGPUInsufficientSizeError,
pynvml.NVML_ERROR_INSUFFICIENT_POWER: ZeusGPUInsufficientPowerError,
pynvml.NVML_ERROR_DRIVER_NOT_LOADED: ZeusGPUDriverNotLoadedError,
pynvml.NVML_ERROR_TIMEOUT: ZeusGPUTimeoutError,
pynvml.NVML_ERROR_IRQ_ISSUE: ZeusGPUIRQError,
pynvml.NVML_ERROR_LIBRARY_NOT_FOUND: ZeusGPULibraryNotFoundError,
pynvml.NVML_ERROR_FUNCTION_NOT_FOUND: ZeusGPUFunctionNotFoundError,
pynvml.NVML_ERROR_CORRUPTED_INFOROM: ZeusGPUCorruptedInfoROMError,
pynvml.NVML_ERROR_GPU_IS_LOST: ZeusGPULostError,
pynvml.NVML_ERROR_RESET_REQUIRED: ZeusGPUResetRequiredError,
pynvml.NVML_ERROR_OPERATING_SYSTEM: ZeusGPUOperatingSystemError,
pynvml.NVML_ERROR_LIB_RM_VERSION_MISMATCH: ZeusGPULibRMVersionMismatchError,
pynvml.NVML_ERROR_MEMORY: ZeusGPUMemoryError,
pynvml.NVML_ERROR_UNKNOWN: ZeusGPUUnknownError,
}

@functools.wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except pynvml.NVMLError as e:
exception_class = _exception_map.get(e.value, ZeusGPUUnknownError)
exception_class = NVIDIAGPU._exception_map.get(e.value, ZeusGPUUnknownError)
raise exception_class(str(e)) from e

return wrapper
Expand All @@ -320,6 +298,29 @@ def __init__(self, gpu_index: int) -> None:
super().__init__(gpu_index)
self._get_handle()
self._supportsGetTotalEnergyConsumption = None

_exception_map = {
pynvml.NVML_ERROR_UNINITIALIZED: ZeusGPUInitError,
pynvml.NVML_ERROR_INVALID_ARGUMENT: ZeusGPUInvalidArgError,
pynvml.NVML_ERROR_NOT_SUPPORTED: ZeusGPUNotSupportedError,
pynvml.NVML_ERROR_NO_PERMISSION: ZeusGPUNoPermissionError,
pynvml.NVML_ERROR_ALREADY_INITIALIZED: ZeusGPUAlreadyInitializedError,
pynvml.NVML_ERROR_NOT_FOUND: ZeusGPUNotFoundError,
pynvml.NVML_ERROR_INSUFFICIENT_SIZE: ZeusGPUInsufficientSizeError,
pynvml.NVML_ERROR_INSUFFICIENT_POWER: ZeusGPUInsufficientPowerError,
pynvml.NVML_ERROR_DRIVER_NOT_LOADED: ZeusGPUDriverNotLoadedError,
pynvml.NVML_ERROR_TIMEOUT: ZeusGPUTimeoutError,
pynvml.NVML_ERROR_IRQ_ISSUE: ZeusGPUIRQError,
pynvml.NVML_ERROR_LIBRARY_NOT_FOUND: ZeusGPULibraryNotFoundError,
pynvml.NVML_ERROR_FUNCTION_NOT_FOUND: ZeusGPUFunctionNotFoundError,
pynvml.NVML_ERROR_CORRUPTED_INFOROM: ZeusGPUCorruptedInfoROMError,
pynvml.NVML_ERROR_GPU_IS_LOST: ZeusGPULostError,
pynvml.NVML_ERROR_RESET_REQUIRED: ZeusGPUResetRequiredError,
pynvml.NVML_ERROR_OPERATING_SYSTEM: ZeusGPUOperatingSystemError,
pynvml.NVML_ERROR_LIB_RM_VERSION_MISMATCH: ZeusGPULibRMVersionMismatchError,
pynvml.NVML_ERROR_MEMORY: ZeusGPUMemoryError,
pynvml.NVML_ERROR_UNKNOWN: ZeusGPUUnknownError,
}

@_handle_nvml_errors
def _get_handle(self):
Expand Down Expand Up @@ -427,35 +428,13 @@ class UnprivilegedNVIDIAGPU(NVIDIAGPU):


def _handle_amdsmi_errors(func):
_exception_map = {
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INVAL: ZeusGPUInvalidArgError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_SUPPORTED: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_TIMEOUT: ZeusGPUTimeoutError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NO_PERM: ZeusGPUNoPermissionError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_OUT_OF_RESOURCES: ZeusGPUMemoryError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INIT_ERROR: ZeusGPUInitError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_FOUND: ZeusGPUNotFoundError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_INIT: ZeusGPUInitError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_DRIVER_NOT_LOADED: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INSUFFICIENT_SIZE: ZeusGPUInsufficientSizeError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_ENERGY_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_MSR_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_SUP: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_MSG_SUP: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_HSMP_TIMEOUT: ZeusGPUTimeoutError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_FILE_NOT_FOUND: ZeusGPULibraryNotFoundError,
amdsmi.amdsmi_wrapper.AMDSMI_ARG_PTR_NULL: ZeusGPUInvalidArgError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_UNKNOWN_ERROR: ZeusGPUUnknownError,
}

@functools.wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except amdsmi.AmdSmiLibraryException as e:
exception_class = _exception_map.get(
exception_class = AMDGPU._exception_map.get(
e.get_error_code(), ZeusGPUUnknownError
)
raise exception_class(e.get_error_info()) from e
Expand All @@ -477,6 +456,29 @@ def __init__(self, gpu_index: int) -> None:
super().__init__(gpu_index)
self._get_handle()
self._supportsGetTotalEnergyConsumption = None

_exception_map = {
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INVAL: ZeusGPUInvalidArgError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_SUPPORTED: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_TIMEOUT: ZeusGPUTimeoutError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NO_PERM: ZeusGPUNoPermissionError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_OUT_OF_RESOURCES: ZeusGPUMemoryError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INIT_ERROR: ZeusGPUInitError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_FOUND: ZeusGPUNotFoundError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_NOT_INIT: ZeusGPUInitError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_DRIVER_NOT_LOADED: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_INSUFFICIENT_SIZE: ZeusGPUInsufficientSizeError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_ENERGY_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_MSR_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_SUP: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_HSMP_MSG_SUP: ZeusGPUNotSupportedError,
amdsmi.amdsmi_wrapper.AMDSMI_HSMP_TIMEOUT: ZeusGPUTimeoutError,
amdsmi.amdsmi_wrapper.AMDSMI_NO_DRV: ZeusGPUDriverNotLoadedError,
amdsmi.amdsmi_wrapper.AMDSMI_FILE_NOT_FOUND: ZeusGPULibraryNotFoundError,
amdsmi.amdsmi_wrapper.AMDSMI_ARG_PTR_NULL: ZeusGPUInvalidArgError,
amdsmi.amdsmi_wrapper.AMDSMI_STATUS_UNKNOWN_ERROR: ZeusGPUUnknownError,
}

@_handle_amdsmi_errors
def _get_handle(self):
Expand Down

0 comments on commit 5767ae9

Please sign in to comment.