nvidia

zeus.device.gpu.nvidia

NVIDIA GPUs.

NVIDIAGPU

Bases: GPU

Implementation of GPU for NVIDIA GPUs.

Source code in zeus/device/gpu/nvidia.py

class NVIDIAGPU(gpu_common.GPU):
    """Implementation of `GPU` for NVIDIA GPUs."""

    def __init__(self, gpu_index: int) -> None:
        """Initialize the GPU object."""
        super().__init__(gpu_index)
        self._get_handle()
        self._supportsGetTotalEnergyConsumption = None

        # NVML power scopes for instant and average power field value queries,
        # resolved lazily on first use by `_resolve_power_field_scope`. The two
        # fields may resolve to different scopes, so they're probed independently.
        self._instant_power_scope: int | None = None
        self._average_power_scope: int | None = None

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

    @_handle_nvml_errors
    def _get_handle(self):
        self.handle = pynvml.nvmlDeviceGetHandleByIndex(self.gpu_index)

    @_handle_nvml_errors
    def get_name(self) -> str:
        """Return the name of the GPU model."""
        return pynvml.nvmlDeviceGetName(self.handle)

    @property
    def supports_nonblocking_setters(self) -> bool:
        """Return True if the GPU object supports non-blocking configuration setters."""
        return False

    @_handle_nvml_errors
    def get_power_management_limit_constraints(self) -> tuple[int, int]:
        """Return the minimum and maximum power management limits. Units: mW."""
        min_, max_ = pynvml.nvmlDeviceGetPowerManagementLimitConstraints(self.handle)
        return (min_, max_)

    @_handle_nvml_errors
    def get_power_management_limit(self) -> int:
        """Return the current power management limit. Units: mW."""
        return pynvml.nvmlDeviceGetPowerManagementLimit(self.handle)

    @_handle_nvml_errors
    def set_power_management_limit(self, power_limit_mw: int, block: bool = True) -> None:
        """Set the GPU's power management limit. Unit: mW."""
        current_limit = self.get_power_management_limit()
        if current_limit != power_limit_mw:
            self._warn_sys_admin()
            pynvml.nvmlDeviceSetPowerManagementLimit(self.handle, power_limit_mw)

    @_handle_nvml_errors
    def reset_power_management_limit(self, block: bool = True) -> None:
        """Reset the GPU's power management limit to the default value."""
        default_limit = pynvml.nvmlDeviceGetPowerManagementDefaultLimit(self.handle)
        current_limit = self.get_power_management_limit()
        if current_limit != default_limit:
            self._warn_sys_admin()
            pynvml.nvmlDeviceSetPowerManagementLimit(self.handle, default_limit)

    @_handle_nvml_errors
    def set_persistence_mode(self, enabled: bool, block: bool = True) -> None:
        """Set persistence mode."""
        current_mode = pynvml.nvmlDeviceGetPersistenceMode(self.handle)
        desired_mode = pynvml.NVML_FEATURE_ENABLED if enabled else pynvml.NVML_FEATURE_DISABLED
        if current_mode != desired_mode:
            self._warn_sys_admin()
            pynvml.nvmlDeviceSetPersistenceMode(self.handle, desired_mode)

    @_handle_nvml_errors
    def get_persistence_mode(self) -> bool:
        """Return whether persistence mode is currently enabled."""
        return pynvml.nvmlDeviceGetPersistenceMode(self.handle) == pynvml.NVML_FEATURE_ENABLED

    @_handle_nvml_errors
    def get_supported_memory_clocks(self) -> list[int]:
        """Return a list of supported memory clock frequencies. Units: MHz."""
        return pynvml.nvmlDeviceGetSupportedMemoryClocks(self.handle)

    @_handle_nvml_errors
    def set_memory_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
        """Lock the memory clock to a specified range. Units: MHz."""
        self._warn_sys_admin()
        pynvml.nvmlDeviceSetMemoryLockedClocks(self.handle, min_clock_mhz, max_clock_mhz)

    @_handle_nvml_errors
    def reset_memory_locked_clocks(self, block: bool = True) -> None:
        """Reset the locked memory clocks to the default."""
        self._warn_sys_admin()
        pynvml.nvmlDeviceResetMemoryLockedClocks(self.handle)

    @_handle_nvml_errors
    def get_supported_graphics_clocks(self, memory_clock_mhz: int | None = None) -> list[int]:
        """Return a list of supported graphics clock frequencies. Units: MHz.

        Args:
            memory_clock_mhz: Memory clock frequency to use. Some GPUs have
                different supported graphics clocks depending on the memory clock.
        """
        pass
        return pynvml.nvmlDeviceGetSupportedGraphicsClocks(self.handle, memory_clock_mhz)

    @_handle_nvml_errors
    def set_gpu_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
        """Lock the GPU clock to a specified range. Units: MHz."""
        self._warn_sys_admin()
        pynvml.nvmlDeviceSetGpuLockedClocks(self.handle, min_clock_mhz, max_clock_mhz)

    @_handle_nvml_errors
    def reset_gpu_locked_clocks(self, block: bool = True) -> None:
        """Reset the locked GPU clocks to the default."""
        self._warn_sys_admin()
        pynvml.nvmlDeviceResetGpuLockedClocks(self.handle)

    def _resolve_power_field_scope(self, field: int) -> int:
        """Determine the NVML power scope to use for a power field value query.

        This is a GPU device, so GPU-scope power is the expected outcome. We
        therefore probe `NVML_POWER_SCOPE_GPU` first and use it if it works.

        Some C2C superchips (e.g., Grace Hopper) do not support GPU-scope power
        queries and instead report power for the entire module (CPU + GPU). On
        those, the GPU-scope query fails, so we warn and fall back to
        `NVML_POWER_SCOPE_MODULE`. If neither scope works, something is wrong
        and we raise.

        The scope does not change over the lifetime of the GPU, so callers cache
        the result and only probe on the first power query.

        Args:
            field: The NVML power field to resolve the scope for, e.g.,
                `NVML_FI_DEV_POWER_INSTANT` or `NVML_FI_DEV_POWER_AVERAGE`.
        """
        metric = pynvml.nvmlDeviceGetFieldValues(self.handle, [(field, pynvml.NVML_POWER_SCOPE_GPU)])[0]
        if metric.nvmlReturn == pynvml.NVML_SUCCESS:
            return pynvml.NVML_POWER_SCOPE_GPU

        warnings.warn(
            f"GPU {self.gpu_index} does not support GPU-scope power queries for NVML field "
            f"{field}. Falling back to module scope, so reported power and energy will be for "
            "the entire module (CPU + GPU). This is expected on C2C superchips like Grace Hopper.",
            stacklevel=2,
        )
        metric = pynvml.nvmlDeviceGetFieldValues(self.handle, [(field, pynvml.NVML_POWER_SCOPE_MODULE)])[0]
        if metric.nvmlReturn == pynvml.NVML_SUCCESS:
            return pynvml.NVML_POWER_SCOPE_MODULE

        raise pynvml.NVMLError(metric.nvmlReturn)

    @_handle_nvml_errors
    def get_average_power_usage(self) -> int:
        """Return the average power draw of the GPU. Units: mW."""
        if self._average_power_scope is None:
            self._average_power_scope = self._resolve_power_field_scope(pynvml.NVML_FI_DEV_POWER_AVERAGE)
        fields = [(pynvml.NVML_FI_DEV_POWER_AVERAGE, self._average_power_scope)]

        metric = pynvml.nvmlDeviceGetFieldValues(self.handle, fields)[0]
        if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
            raise pynvml.NVMLError(ret)
        return metric.value.uiVal

    @_handle_nvml_errors
    def get_instant_power_usage(self) -> int:
        """Return the current power draw of the GPU. Units: mW."""
        if self._instant_power_scope is None:
            self._instant_power_scope = self._resolve_power_field_scope(pynvml.NVML_FI_DEV_POWER_INSTANT)
        fields = [(pynvml.NVML_FI_DEV_POWER_INSTANT, self._instant_power_scope)]

        metric = pynvml.nvmlDeviceGetFieldValues(self.handle, fields)[0]
        if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
            raise pynvml.NVMLError(ret)
        return metric.value.uiVal

    @_handle_nvml_errors
    def get_average_memory_power_usage(self) -> int:
        """Return the average power draw of the GPU's memory. Units: mW.

        !!! Warning
            This isn't exactly documented in NVML at the time of writing, but `nvidia-smi`
            makes use of this API.

            Confirmed working on H100 80GB HBM3. Confirmed not working on A40.
        """
        metric = pynvml.nvmlDeviceGetFieldValues(
            self.handle,
            [(pynvml.NVML_FI_DEV_POWER_AVERAGE, pynvml.NVML_POWER_SCOPE_MEMORY)],
        )[0]
        if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
            raise pynvml.NVMLError(ret)
        power = metric.value.uiVal
        if power == 0:
            warnings.warn(
                "Average memory power returned 0. The current GPU may not be supported.",
                stacklevel=1,
            )
        return power

    @_handle_nvml_errors
    def supports_get_total_energy_consumption(self) -> bool:
        """Check if the GPU supports retrieving total energy consumption."""
        # Supported on Volta or newer microarchitectures
        if self._supportsGetTotalEnergyConsumption is None:
            self._supportsGetTotalEnergyConsumption = (
                pynvml.nvmlDeviceGetArchitecture(self.handle) >= pynvml.NVML_DEVICE_ARCH_VOLTA
            )

        return self._supportsGetTotalEnergyConsumption

    @_handle_nvml_errors
    def get_total_energy_consumption(self) -> int:
        """Return the total energy consumption of the specified GPU. Units: mJ."""
        return pynvml.nvmlDeviceGetTotalEnergyConsumption(self.handle)

    @_handle_nvml_errors
    def get_gpu_temperature(self) -> int:
        """Return the current GPU temperature. Units: Celsius."""
        temperature = pynvml.nvmlDeviceGetTemperatureV(self.handle, pynvml.NVML_TEMPERATURE_GPU)
        return temperature

supports_nonblocking_setters property

supports_nonblocking_setters

Return True if the GPU object supports non-blocking configuration setters.

__init__

__init__(gpu_index)

Source code in zeus/device/gpu/nvidia.py

def __init__(self, gpu_index: int) -> None:
    """Initialize the GPU object."""
    super().__init__(gpu_index)
    self._get_handle()
    self._supportsGetTotalEnergyConsumption = None

    # NVML power scopes for instant and average power field value queries,
    # resolved lazily on first use by `_resolve_power_field_scope`. The two
    # fields may resolve to different scopes, so they're probed independently.
    self._instant_power_scope: int | None = None
    self._average_power_scope: int | None = None

get_name

get_name()

Return the name of the GPU model.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_name(self) -> str:
    """Return the name of the GPU model."""
    return pynvml.nvmlDeviceGetName(self.handle)

get_power_management_limit_constraints

get_power_management_limit_constraints()

Return the minimum and maximum power management limits. Units: mW.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_power_management_limit_constraints(self) -> tuple[int, int]:
    """Return the minimum and maximum power management limits. Units: mW."""
    min_, max_ = pynvml.nvmlDeviceGetPowerManagementLimitConstraints(self.handle)
    return (min_, max_)

get_power_management_limit

get_power_management_limit()

Return the current power management limit. Units: mW.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_power_management_limit(self) -> int:
    """Return the current power management limit. Units: mW."""
    return pynvml.nvmlDeviceGetPowerManagementLimit(self.handle)

set_power_management_limit

set_power_management_limit(power_limit_mw, block=True)

Set the GPU's power management limit. Unit: mW.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def set_power_management_limit(self, power_limit_mw: int, block: bool = True) -> None:
    """Set the GPU's power management limit. Unit: mW."""
    current_limit = self.get_power_management_limit()
    if current_limit != power_limit_mw:
        self._warn_sys_admin()
        pynvml.nvmlDeviceSetPowerManagementLimit(self.handle, power_limit_mw)

reset_power_management_limit

reset_power_management_limit(block=True)

Reset the GPU's power management limit to the default value.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def reset_power_management_limit(self, block: bool = True) -> None:
    """Reset the GPU's power management limit to the default value."""
    default_limit = pynvml.nvmlDeviceGetPowerManagementDefaultLimit(self.handle)
    current_limit = self.get_power_management_limit()
    if current_limit != default_limit:
        self._warn_sys_admin()
        pynvml.nvmlDeviceSetPowerManagementLimit(self.handle, default_limit)

set_persistence_mode

set_persistence_mode(enabled, block=True)

Set persistence mode.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def set_persistence_mode(self, enabled: bool, block: bool = True) -> None:
    """Set persistence mode."""
    current_mode = pynvml.nvmlDeviceGetPersistenceMode(self.handle)
    desired_mode = pynvml.NVML_FEATURE_ENABLED if enabled else pynvml.NVML_FEATURE_DISABLED
    if current_mode != desired_mode:
        self._warn_sys_admin()
        pynvml.nvmlDeviceSetPersistenceMode(self.handle, desired_mode)

get_persistence_mode

get_persistence_mode()

Return whether persistence mode is currently enabled.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_persistence_mode(self) -> bool:
    """Return whether persistence mode is currently enabled."""
    return pynvml.nvmlDeviceGetPersistenceMode(self.handle) == pynvml.NVML_FEATURE_ENABLED

get_supported_memory_clocks

get_supported_memory_clocks()

Return a list of supported memory clock frequencies. Units: MHz.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_supported_memory_clocks(self) -> list[int]:
    """Return a list of supported memory clock frequencies. Units: MHz."""
    return pynvml.nvmlDeviceGetSupportedMemoryClocks(self.handle)

set_memory_locked_clocks

set_memory_locked_clocks(
    min_clock_mhz, max_clock_mhz, block=True
)

Lock the memory clock to a specified range. Units: MHz.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def set_memory_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
    """Lock the memory clock to a specified range. Units: MHz."""
    self._warn_sys_admin()
    pynvml.nvmlDeviceSetMemoryLockedClocks(self.handle, min_clock_mhz, max_clock_mhz)

reset_memory_locked_clocks

reset_memory_locked_clocks(block=True)

Reset the locked memory clocks to the default.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def reset_memory_locked_clocks(self, block: bool = True) -> None:
    """Reset the locked memory clocks to the default."""
    self._warn_sys_admin()
    pynvml.nvmlDeviceResetMemoryLockedClocks(self.handle)

get_supported_graphics_clocks

get_supported_graphics_clocks(memory_clock_mhz=None)

Return a list of supported graphics clock frequencies. Units: MHz.

Parameters:

Name	Type	Description	Default
memory_clock_mhz	int | None	Memory clock frequency to use. Some GPUs have different supported graphics clocks depending on the memory clock.	None

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_supported_graphics_clocks(self, memory_clock_mhz: int | None = None) -> list[int]:
    """Return a list of supported graphics clock frequencies. Units: MHz.

    Args:
        memory_clock_mhz: Memory clock frequency to use. Some GPUs have
            different supported graphics clocks depending on the memory clock.
    """
    pass
    return pynvml.nvmlDeviceGetSupportedGraphicsClocks(self.handle, memory_clock_mhz)

set_gpu_locked_clocks

set_gpu_locked_clocks(
    min_clock_mhz, max_clock_mhz, block=True
)

Lock the GPU clock to a specified range. Units: MHz.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def set_gpu_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
    """Lock the GPU clock to a specified range. Units: MHz."""
    self._warn_sys_admin()
    pynvml.nvmlDeviceSetGpuLockedClocks(self.handle, min_clock_mhz, max_clock_mhz)

reset_gpu_locked_clocks

reset_gpu_locked_clocks(block=True)

Reset the locked GPU clocks to the default.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def reset_gpu_locked_clocks(self, block: bool = True) -> None:
    """Reset the locked GPU clocks to the default."""
    self._warn_sys_admin()
    pynvml.nvmlDeviceResetGpuLockedClocks(self.handle)

_resolve_power_field_scope

_resolve_power_field_scope(field)

Determine the NVML power scope to use for a power field value query.

This is a GPU device, so GPU-scope power is the expected outcome. We therefore probe NVML_POWER_SCOPE_GPU first and use it if it works.

Some C2C superchips (e.g., Grace Hopper) do not support GPU-scope power queries and instead report power for the entire module (CPU + GPU). On those, the GPU-scope query fails, so we warn and fall back to NVML_POWER_SCOPE_MODULE. If neither scope works, something is wrong and we raise.

The scope does not change over the lifetime of the GPU, so callers cache the result and only probe on the first power query.

Parameters:

Name	Type	Description	Default
field	int	The NVML power field to resolve the scope for, e.g., NVML_FI_DEV_POWER_INSTANT or NVML_FI_DEV_POWER_AVERAGE.	required

Source code in zeus/device/gpu/nvidia.py

def _resolve_power_field_scope(self, field: int) -> int:
    """Determine the NVML power scope to use for a power field value query.

    This is a GPU device, so GPU-scope power is the expected outcome. We
    therefore probe `NVML_POWER_SCOPE_GPU` first and use it if it works.

    Some C2C superchips (e.g., Grace Hopper) do not support GPU-scope power
    queries and instead report power for the entire module (CPU + GPU). On
    those, the GPU-scope query fails, so we warn and fall back to
    `NVML_POWER_SCOPE_MODULE`. If neither scope works, something is wrong
    and we raise.

    The scope does not change over the lifetime of the GPU, so callers cache
    the result and only probe on the first power query.

    Args:
        field: The NVML power field to resolve the scope for, e.g.,
            `NVML_FI_DEV_POWER_INSTANT` or `NVML_FI_DEV_POWER_AVERAGE`.
    """
    metric = pynvml.nvmlDeviceGetFieldValues(self.handle, [(field, pynvml.NVML_POWER_SCOPE_GPU)])[0]
    if metric.nvmlReturn == pynvml.NVML_SUCCESS:
        return pynvml.NVML_POWER_SCOPE_GPU

    warnings.warn(
        f"GPU {self.gpu_index} does not support GPU-scope power queries for NVML field "
        f"{field}. Falling back to module scope, so reported power and energy will be for "
        "the entire module (CPU + GPU). This is expected on C2C superchips like Grace Hopper.",
        stacklevel=2,
    )
    metric = pynvml.nvmlDeviceGetFieldValues(self.handle, [(field, pynvml.NVML_POWER_SCOPE_MODULE)])[0]
    if metric.nvmlReturn == pynvml.NVML_SUCCESS:
        return pynvml.NVML_POWER_SCOPE_MODULE

    raise pynvml.NVMLError(metric.nvmlReturn)

get_average_power_usage

get_average_power_usage()

Return the average power draw of the GPU. Units: mW.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_average_power_usage(self) -> int:
    """Return the average power draw of the GPU. Units: mW."""
    if self._average_power_scope is None:
        self._average_power_scope = self._resolve_power_field_scope(pynvml.NVML_FI_DEV_POWER_AVERAGE)
    fields = [(pynvml.NVML_FI_DEV_POWER_AVERAGE, self._average_power_scope)]

    metric = pynvml.nvmlDeviceGetFieldValues(self.handle, fields)[0]
    if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
        raise pynvml.NVMLError(ret)
    return metric.value.uiVal

get_instant_power_usage

get_instant_power_usage()

Return the current power draw of the GPU. Units: mW.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_instant_power_usage(self) -> int:
    """Return the current power draw of the GPU. Units: mW."""
    if self._instant_power_scope is None:
        self._instant_power_scope = self._resolve_power_field_scope(pynvml.NVML_FI_DEV_POWER_INSTANT)
    fields = [(pynvml.NVML_FI_DEV_POWER_INSTANT, self._instant_power_scope)]

    metric = pynvml.nvmlDeviceGetFieldValues(self.handle, fields)[0]
    if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
        raise pynvml.NVMLError(ret)
    return metric.value.uiVal

get_average_memory_power_usage

get_average_memory_power_usage()

Return the average power draw of the GPU's memory. Units: mW.

Warning

This isn't exactly documented in NVML at the time of writing, but nvidia-smi makes use of this API.

Confirmed working on H100 80GB HBM3. Confirmed not working on A40.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_average_memory_power_usage(self) -> int:
    """Return the average power draw of the GPU's memory. Units: mW.

    !!! Warning
        This isn't exactly documented in NVML at the time of writing, but `nvidia-smi`
        makes use of this API.

        Confirmed working on H100 80GB HBM3. Confirmed not working on A40.
    """
    metric = pynvml.nvmlDeviceGetFieldValues(
        self.handle,
        [(pynvml.NVML_FI_DEV_POWER_AVERAGE, pynvml.NVML_POWER_SCOPE_MEMORY)],
    )[0]
    if (ret := metric.nvmlReturn) != pynvml.NVML_SUCCESS:
        raise pynvml.NVMLError(ret)
    power = metric.value.uiVal
    if power == 0:
        warnings.warn(
            "Average memory power returned 0. The current GPU may not be supported.",
            stacklevel=1,
        )
    return power

supports_get_total_energy_consumption

supports_get_total_energy_consumption()

Check if the GPU supports retrieving total energy consumption.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def supports_get_total_energy_consumption(self) -> bool:
    """Check if the GPU supports retrieving total energy consumption."""
    # Supported on Volta or newer microarchitectures
    if self._supportsGetTotalEnergyConsumption is None:
        self._supportsGetTotalEnergyConsumption = (
            pynvml.nvmlDeviceGetArchitecture(self.handle) >= pynvml.NVML_DEVICE_ARCH_VOLTA
        )

    return self._supportsGetTotalEnergyConsumption

get_total_energy_consumption

get_total_energy_consumption()

Return the total energy consumption of the specified GPU. Units: mJ.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_total_energy_consumption(self) -> int:
    """Return the total energy consumption of the specified GPU. Units: mJ."""
    return pynvml.nvmlDeviceGetTotalEnergyConsumption(self.handle)

get_gpu_temperature

get_gpu_temperature()

Return the current GPU temperature. Units: Celsius.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def get_gpu_temperature(self) -> int:
    """Return the current GPU temperature. Units: Celsius."""
    temperature = pynvml.nvmlDeviceGetTemperatureV(self.handle, pynvml.NVML_TEMPERATURE_GPU)
    return temperature

ZeusdNVIDIAGPU

Bases: NVIDIAGPU

An NVIDIAGPU that sets GPU knobs that require SYS_ADMIN via zeusd.

Some NVML APIs (e.g., setting persistence mode, power limit, frequency) requires the Linux security capability SYS_ADMIN, which is virtually sudo. This class overrides those methods so that they send a request to the Zeus daemon.

See here for details on system privileges required.

Source code in zeus/device/gpu/nvidia.py

class ZeusdNVIDIAGPU(NVIDIAGPU):
    """An NVIDIAGPU that sets GPU knobs that require `SYS_ADMIN` via zeusd.

    Some NVML APIs (e.g., setting persistence mode, power limit, frequency)
    requires the Linux security capability `SYS_ADMIN`, which is virtually `sudo`.
    This class overrides those methods so that they send a request to the
    Zeus daemon.

    See [here](https://ml.energy/zeus/getting_started/#system-privileges)
    for details on system privileges required.
    """

    def __init__(self, gpu_index: int, client: ZeusdClient) -> None:
        """Initialize the GPU object backed by a Zeusd daemon.

        Args:
            gpu_index: Index of the GPU.
            client: ZeusdClient connected to the daemon.
        """
        super().__init__(gpu_index)
        self._client = client
        self._gpu_index = gpu_index
        # Reads stay local to NVML (they are unprivileged), so only the
        # gpu-control API group is required from the daemon.
        require_capabilities(client, control_gpu=True)

    @property
    def supports_nonblocking_setters(self) -> bool:
        """Return True if the GPU object supports non-blocking configuration setters."""
        return True

    def set_power_management_limit(self, power_limit_mw: int, block: bool = True) -> None:
        """Set the GPU's power management limit. Unit: mW."""
        current_limit = self.get_power_management_limit()
        if current_limit == power_limit_mw:
            return
        self._client.set_power_limit([self._gpu_index], power_limit_mw, block)

    @_handle_nvml_errors
    def reset_power_management_limit(self, block: bool = True) -> None:
        """Reset the GPU's power management limit to the default value."""
        self.set_power_management_limit(
            pynvml.nvmlDeviceGetPowerManagementDefaultLimit(self.handle),
            block,
        )

    def set_persistence_mode(self, enabled: bool, block: bool = True) -> None:
        """Set persistence mode."""
        self._client.set_persistence_mode([self._gpu_index], enabled, block)

    def set_memory_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
        """Lock the memory clock to a specified range. Units: MHz."""
        self._client.set_mem_locked_clocks([self._gpu_index], min_clock_mhz, max_clock_mhz, block)

    def reset_memory_locked_clocks(self, block: bool = True) -> None:
        """Reset the locked memory clocks to the default."""
        self._client.reset_mem_locked_clocks([self._gpu_index], block)

    def set_gpu_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
        """Lock the GPU clock to a specified range. Units: MHz."""
        self._client.set_gpu_locked_clocks([self._gpu_index], min_clock_mhz, max_clock_mhz, block)

    def reset_gpu_locked_clocks(self, block: bool = True) -> None:
        """Reset the locked GPU clocks to the default."""
        self._client.reset_gpu_locked_clocks([self._gpu_index], block)

supports_nonblocking_setters property

supports_nonblocking_setters

Return True if the GPU object supports non-blocking configuration setters.

__init__

__init__(gpu_index, client)

Parameters:

Name	Type	Description	Default
gpu_index	int	Index of the GPU.	required
client	ZeusdClient	ZeusdClient connected to the daemon.	required

Source code in zeus/device/gpu/nvidia.py

def __init__(self, gpu_index: int, client: ZeusdClient) -> None:
    """Initialize the GPU object backed by a Zeusd daemon.

    Args:
        gpu_index: Index of the GPU.
        client: ZeusdClient connected to the daemon.
    """
    super().__init__(gpu_index)
    self._client = client
    self._gpu_index = gpu_index
    # Reads stay local to NVML (they are unprivileged), so only the
    # gpu-control API group is required from the daemon.
    require_capabilities(client, control_gpu=True)

set_power_management_limit

set_power_management_limit(power_limit_mw, block=True)

Set the GPU's power management limit. Unit: mW.

Source code in zeus/device/gpu/nvidia.py

def set_power_management_limit(self, power_limit_mw: int, block: bool = True) -> None:
    """Set the GPU's power management limit. Unit: mW."""
    current_limit = self.get_power_management_limit()
    if current_limit == power_limit_mw:
        return
    self._client.set_power_limit([self._gpu_index], power_limit_mw, block)

reset_power_management_limit

reset_power_management_limit(block=True)

Reset the GPU's power management limit to the default value.

Source code in zeus/device/gpu/nvidia.py

@_handle_nvml_errors
def reset_power_management_limit(self, block: bool = True) -> None:
    """Reset the GPU's power management limit to the default value."""
    self.set_power_management_limit(
        pynvml.nvmlDeviceGetPowerManagementDefaultLimit(self.handle),
        block,
    )

set_persistence_mode

set_persistence_mode(enabled, block=True)

Set persistence mode.

Source code in zeus/device/gpu/nvidia.py

def set_persistence_mode(self, enabled: bool, block: bool = True) -> None:
    """Set persistence mode."""
    self._client.set_persistence_mode([self._gpu_index], enabled, block)

set_memory_locked_clocks

set_memory_locked_clocks(
    min_clock_mhz, max_clock_mhz, block=True
)

Lock the memory clock to a specified range. Units: MHz.

Source code in zeus/device/gpu/nvidia.py

def set_memory_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
    """Lock the memory clock to a specified range. Units: MHz."""
    self._client.set_mem_locked_clocks([self._gpu_index], min_clock_mhz, max_clock_mhz, block)

reset_memory_locked_clocks

reset_memory_locked_clocks(block=True)

Reset the locked memory clocks to the default.

Source code in zeus/device/gpu/nvidia.py

def reset_memory_locked_clocks(self, block: bool = True) -> None:
    """Reset the locked memory clocks to the default."""
    self._client.reset_mem_locked_clocks([self._gpu_index], block)

set_gpu_locked_clocks

set_gpu_locked_clocks(
    min_clock_mhz, max_clock_mhz, block=True
)

Lock the GPU clock to a specified range. Units: MHz.

Source code in zeus/device/gpu/nvidia.py

def set_gpu_locked_clocks(self, min_clock_mhz: int, max_clock_mhz: int, block: bool = True) -> None:
    """Lock the GPU clock to a specified range. Units: MHz."""
    self._client.set_gpu_locked_clocks([self._gpu_index], min_clock_mhz, max_clock_mhz, block)

reset_gpu_locked_clocks

reset_gpu_locked_clocks(block=True)

Reset the locked GPU clocks to the default.

Source code in zeus/device/gpu/nvidia.py

def reset_gpu_locked_clocks(self, block: bool = True) -> None:
    """Reset the locked GPU clocks to the default."""
    self._client.reset_gpu_locked_clocks([self._gpu_index], block)

NVIDIAGPUs

Bases: GPUs

Implementation of GPUs for NVIDIA GPUs.

CUDA_VISIBLE_DEVICES environment variable is respected if set. For example, if there are 4 GPUs on the node and CUDA_VISIBLE_DEVICES=0,2, only GPUs 0 and 2 are instantiated. In this case, to access GPU of CUDA index 0, use the index 0, and for CUDA index 2, use the index 1.

If you have the Zeus daemon deployed, make sure you have set the ZEUSD_SOCK_PATH environment variable to the path of the Zeus daemon socket. This class will automatically use ZeusdNVIDIAGPU if ZEUSD_SOCK_PATH is set.

Note

Some C2C superchips (e.g., Grace Hopper) do not support GPU-scope power queries. In that case, a warning is issued and the whole module's power is reported.

Source code in zeus/device/gpu/nvidia.py

class NVIDIAGPUs(gpu_common.GPUs):
    """Implementation of `GPUs` for NVIDIA GPUs.

    `CUDA_VISIBLE_DEVICES` environment variable is respected if set.
    For example, if there are 4 GPUs on the node and `CUDA_VISIBLE_DEVICES=0,2`,
    only GPUs 0 and 2 are instantiated. In this case, to access
    GPU of CUDA index 0, use the index 0, and for CUDA index 2, use the index 1.

    If you have the Zeus daemon deployed, make sure you have set the `ZEUSD_SOCK_PATH`
    environment variable to the path of the Zeus daemon socket. This class will
    automatically use [`ZeusdNVIDIAGPU`][zeus.device.gpu.nvidia.ZeusdNVIDIAGPU]
    if `ZEUSD_SOCK_PATH` is set.

    !!! Note
        Some C2C superchips (e.g., Grace Hopper) do not support GPU-scope power queries.
        In that case, a warning is issued and the whole module's power is reported.
    """

    def __init__(self, ensure_homogeneous: bool = False) -> None:
        """Initialize NVML and sets up the GPUs.

        Args:
            ensure_homogeneous (bool): If True, ensures that all tracked GPUs have the same name.
        """
        try:
            pynvml.nvmlInit()
            self._init_gpus()
            if ensure_homogeneous:
                self._ensure_homogeneous()
        except pynvml.NVMLError as e:
            exception_class = NVIDIAGPU._exception_map.get(
                e.value,  # type: ignore
                gpu_common.ZeusBaseGPUError,
            )
            raise exception_class(str(e)) from e

    @property
    def gpus(self) -> Sequence[NVIDIAGPU]:
        """Return a list of NVIDIAGPU objects being tracked."""
        return self._gpus

    def _init_gpus(self) -> None:
        # Must respect `CUDA_VISIBLE_DEVICES` if set
        if (visible_device := os.environ.get("CUDA_VISIBLE_DEVICES")) is not None:
            if not visible_device:
                raise gpu_common.ZeusGPUInitError(
                    "CUDA_VISIBLE_DEVICES is set to an empty string. "
                    "It should either be unset or a comma-separated list of GPU indices."
                )
            if visible_device.startswith("MIG"):
                raise gpu_common.ZeusGPUInitError(
                    "CUDA_VISIBLE_DEVICES contains MIG devices. NVML (the library used by Zeus) "
                    "currently does not support measuring the power or energy consumption of MIG "
                    "slices. You can still measure the whole GPU by temporarily setting "
                    "CUDA_VISIBLE_DEVICES to integer GPU indices and restoring it afterwards."
                )
            visible_indices = [int(idx) for idx in visible_device.split(",")]
        else:
            visible_indices = list(range(pynvml.nvmlDeviceGetCount()))

        # If Zeusd env vars are set, use ZeusdNVIDIAGPU backed by a shared client.
        config = ZeusdConfig.from_env()
        if config is not None:
            try:
                client = ZeusdClient(config)
                self._gpus = [ZeusdNVIDIAGPU(gpu_num, client) for gpu_num in visible_indices]
            except ZeusBaseError as e:
                raise gpu_common.ZeusGPUInitError(str(e)) from e
            for gpu in self._gpus:
                gpu._disable_sys_admin_warning = True
        else:
            self._gpus = [NVIDIAGPU(gpu_num) for gpu_num in visible_indices]

    def __del__(self) -> None:
        """Shut down NVML."""
        with contextlib.suppress(pynvml.NVMLError):
            pynvml.nvmlShutdown()

gpus property

gpus

Return a list of NVIDIAGPU objects being tracked.

__init__

__init__(ensure_homogeneous=False)

Parameters:

Name	Type	Description	Default
ensure_homogeneous	bool	If True, ensures that all tracked GPUs have the same name.	False

Source code in zeus/device/gpu/nvidia.py

def __init__(self, ensure_homogeneous: bool = False) -> None:
    """Initialize NVML and sets up the GPUs.

    Args:
        ensure_homogeneous (bool): If True, ensures that all tracked GPUs have the same name.
    """
    try:
        pynvml.nvmlInit()
        self._init_gpus()
        if ensure_homogeneous:
            self._ensure_homogeneous()
    except pynvml.NVMLError as e:
        exception_class = NVIDIAGPU._exception_map.get(
            e.value,  # type: ignore
            gpu_common.ZeusBaseGPUError,
        )
        raise exception_class(str(e)) from e

__del__

__del__()

Shut down NVML.

Source code in zeus/device/gpu/nvidia.py

def __del__(self) -> None:
    """Shut down NVML."""
    with contextlib.suppress(pynvml.NVMLError):
        pynvml.nvmlShutdown()

nvml_is_available cached

nvml_is_available()

Check if NVML is available.

Source code in zeus/device/gpu/nvidia.py

@lru_cache(maxsize=1)
def nvml_is_available() -> bool:
    """Check if NVML is available."""
    try:
        import pynvml
    except ImportError:
        logger.info("Failed to import `pynvml`. Make sure you have `nvidia-ml-py` installed.")
        return False

    # Detect unofficial pynvml packages.
    # If detected, this should be a critical error.
    if not hasattr(pynvml, "_nvmlGetFunctionPointer"):
        logger.error("Unoffical pynvml package detected!")
        raise ImportError(
            "Unofficial pynvml package detected! "
            "This causes conflicts with the official NVIDIA bindings. "
            "Please remove with `pip uninstall pynvml` and instead use the official "
            "bindings from NVIDIA: `nvidia-ml-py`. "
        )

    try:
        pynvml.nvmlInit()
        logger.info("pynvml is available and initialized.")
        return True
    except pynvml.NVMLError as e:
        logger.info("pynvml is available but could not initialize NVML: %s.", e)
        return False