Re-wrap most of the enums in cuda.bindings.nvml for cuda.core.system.

Author: mdboom

cuda_core/cuda/core/system/_clock.pxi

@@ -3,9 +3,67 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-ClockId = nvml.ClockId
-ClocksEventReasons = nvml.ClocksEventReasons
-ClockType = nvml.ClockType
+class ClockId(StrEnum):
+    """
+    Clock Ids. These are used in combination with :class:`ClockType` to specify a single clock value.
+    """
+    CURRENT = "current"
+    CUSTOMER_BOOST_MAX = "customer_boost_max"
+    # APP_CLOCK_TARGET and APP_CLOCK_DEFAULT are deprecated so not included here
+ClockId.CURRENT.__doc__ = """
+Current actual clock value.
+"""
+ClockId.CUSTOMER_BOOST_MAX.__doc__ = """
+OEM-defined maximum clock rate
+"""
+cdef dict _CLOCK_ID_MAPPING = {
+    ClockId.CURRENT: nvml.ClockId.CURRENT,
+    ClockId.CUSTOMER_BOOST_MAX: nvml.ClockId.CUSTOMER_BOOST_MAX,
+}
+
+
+class ClocksEventReasons(StrEnum):
+    """
+    Reasons for a clocks event.  These are used in combination with :class:`ClockType` to specify the reason for a clocks event.
+    """
+    NONE = "none"
+    GPU_IDLE = "gpu_idle"
+    APPLICATIONS_CLOCKS_SETTING = "applications_clocks_setting"
+    SW_POWER_CAP = "sw_power_cap"
+    HW_SLOWDOWN = "hw_slowdown"
+    SYNC_BOOST = "sync_boost"
+    SW_THERMAL_SLOWDOWN = "sw_thermal_slowdown"
+    HW_THERMAL_SLOWDOWN = "hw_thermal_slowdown"
+    HW_POWER_BRAKE_SLOWDOWN = "hw_power_brake_slowdown"
+    DISPLAY_CLOCK_SETTING = "display_clock_setting"
+cdef dict _CLOCKS_EVENT_REASONS_MAPPING = {
+    nvml.ClocksEventReasons.EVENT_REASON_NONE: ClocksEventReasons.NONE,
+    nvml.ClocksEventReasons.EVENT_REASON_GPU_IDLE: ClocksEventReasons.GPU_IDLE,
+    nvml.ClocksEventReasons.EVENT_REASON_APPLICATIONS_CLOCKS_SETTING: ClocksEventReasons.APPLICATIONS_CLOCKS_SETTING,
+    nvml.ClocksEventReasons.EVENT_REASON_SW_POWER_CAP: ClocksEventReasons.SW_POWER_CAP,
+    nvml.ClocksEventReasons.THROTTLE_REASON_HW_SLOWDOWN: ClocksEventReasons.HW_SLOWDOWN,
+    nvml.ClocksEventReasons.EVENT_REASON_SYNC_BOOST: ClocksEventReasons.SYNC_BOOST,
+    nvml.ClocksEventReasons.EVENT_REASON_SW_THERMAL_SLOWDOWN: ClocksEventReasons.SW_THERMAL_SLOWDOWN,
+    nvml.ClocksEventReasons.THROTTLE_REASON_HW_THERMAL_SLOWDOWN: ClocksEventReasons.HW_THERMAL_SLOWDOWN,
+    nvml.ClocksEventReasons.THROTTLE_REASON_HW_POWER_BRAKE_SLOWDOWN: ClocksEventReasons.HW_POWER_BRAKE_SLOWDOWN,
+    nvml.ClocksEventReasons.EVENT_REASON_DISPLAY_CLOCK_SETTING: ClocksEventReasons.DISPLAY_CLOCK_SETTING,
+}
+
+
+class ClockType(StrEnum):
+    """
+    Clock types. All speeds are in Mhz.
+    """
+    GRAPHICS = "graphics"
+    SM = "sm"
+    MEMORY = "memory"
+    VIDEO = "video"
+cdef dict _CLOCK_TYPE_MAPPING = {
+    ClockType.GRAPHICS: nvml.ClockType.CLOCK_GRAPHICS,
+    ClockType.SM: nvml.ClockType.CLOCK_SM,
+    ClockType.MEMORY: nvml.ClockType.CLOCK_MEM,
+    ClockType.VIDEO: nvml.ClockType.CLOCK_VIDEO,
+}
 
 
 cdef class ClockOffsets:
@@ -48,26 +106,40 @@ cdef class ClockInfo:
     cdef intptr_t _handle
     cdef int _clock_type
 
-    def __init__(self, handle, clock_type: ClockType):
+    def __init__(self, handle, clock_type: ClockType | str):
         self._handle = handle
+        try:
+            clock_type = _CLOCK_TYPE_MAPPING[clock_type]
+        except KeyError:
+            raise ValueError(
+                f"Invalid clock type: {clock_type}. "
+                f"Must be one of {list(ClockType.__members__.values())}"
+            ) from None
         self._clock_type = int(clock_type)
 
-    def get_current_mhz(self, clock_id: ClockId = ClockId.CURRENT) -> int:
+    def get_current_mhz(self, clock_id: ClockId | str = ClockId.CURRENT) -> int:
         """
         Get the current clock speed of a specific clock domain, in MHz.
 
         For Kepler™ or newer fully supported devices.
 
         Parameters
         ----------
-        clock_id: :class:`ClockId`
-            The clock ID to query.
+        clock_id: :class:`ClockId` | str
+            The clock ID to query.  Defaults to the current clock value.
 
         Returns
         -------
         int
             The clock speed in MHz.
         """
+        try:
+            clock_id = _CLOCK_ID_MAPPING[clock_id]
+        except KeyError:
+            raise ValueError(
+                f"Invalid clock ID: {clock_id}. "
+                f"Must be one of {list(ClockId.__members__.values())}"
+            ) from None
         return nvml.device_get_clock(self._handle, self._clock_type, clock_id)
 
     def get_max_mhz(self) -> int:
@@ -99,37 +171,41 @@ cdef class ClockInfo:
         """
         return nvml.device_get_max_customer_boost_clock(self._handle, self._clock_type)
 
-    def get_min_max_clock_of_pstate_mhz(self, pstate: Pstates) -> tuple[int, int]:
+    def get_min_max_clock_of_pstate_mhz(self, pstate: int) -> tuple[int, int]:
         """
         Get the minimum and maximum clock speeds for this clock domain
         at a given performance state (Pstate), in MHz.
 
         Parameters
         ----------
-        pstate: :class:`Pstates`
-            The performance state to query.
+        pstate: int
+            The performance state to query.  Must be an int between 0 and 15,
+            where 0 is the highest performance state (P0) and 15 is the lowest
+            (P15).
 
         Returns
         -------
         tuple[int, int]
             A tuple containing the minimum and maximum clock speeds in MHz.
         """
-        return nvml.device_get_min_max_clock_of_p_state(self._handle, self._clock_type, pstate)
+        return nvml.device_get_min_max_clock_of_p_state(self._handle, self._clock_type, _pstate_to_enum(pstate))
 
-    def get_offsets(self, pstate: Pstates) -> ClockOffsets:
+    def get_offsets(self, pstate: int) -> ClockOffsets:
         """
         Retrieve min, max and current clock offset of some clock domain for a given Pstate.
 
         For Maxwell™ or newer fully supported devices.
 
         Parameters
         ----------
-        pstate: :class:`Pstates`
-            The performance state to query.
+        pstate: int
+            The performance state to query.  Must be an int between 0 and 15,
+            where 0 is the highest performance state (P0) and 15 is the lowest
+            (P15).
 
         Returns
         -------
         :obj:`~_device.ClockOffsets`
             An object with the min, max and current clock offset.
         """
-        return ClockOffsets(nvml.device_get_clock_offsets(self._handle, self._clock_type, pstate))
+        return ClockOffsets(nvml.device_get_clock_offsets(self._handle, self._clock_type, _pstate_to_enum(pstate)))

cuda_core/cuda/core/system/_cooler.pxi

@@ -3,8 +3,50 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-CoolerControl = nvml.CoolerControl
-CoolerTarget = nvml.CoolerTarget
+class CoolerControl(StrEnum):
+    """
+    Cooler control type.
+    """
+    TOGGLE = "toggle"
+    VARIABLE = "variable"
+CoolerControl.TOGGLE.__doc__ = """
+This cooler can only be toggled either ON or OFF (e.g. a switch).
+"""
+CoolerControl.VARIABLE.__doc__ = """
+This cooler's level can be adjusted from some minimum to some maximum (e.g. a knob).
+"""
+cdef dict _COOLER_CONTROL_MAPPING = {
+    nvml.CoolerControl.THERMAL_COOLER_SIGNAL_TOGGLE: CoolerControl.TOGGLE,
+    nvml.CoolerControl.THERMAL_COOLER_SIGNAL_VARIABLE: CoolerControl.VARIABLE,
+}
+
+
+class CoolerTarget(StrEnum):
+    """
+    Cooler target.
+    """
+    NONE = "none"
+    GPU = "gpu"
+    MEMORY = "memory"
+    POWER_SUPPLY = "power_supply"
+CoolerTarget.NONE.__doc__ = """
+This cooler controls nothing.
+"""
+CoolerTarget.GPU.__doc__ = """
+This cooler can cool the GPU.
+"""
+CoolerTarget.MEMORY.__doc__ = """
+This cooler can cool the memory.
+"""
+CoolerTarget.POWER_SUPPLY.__doc__ = """
+This cooler can cool the power supply.
+"""
+cdef dict _COOLER_TARGET_MAPPING = {
+    nvml.CoolerTarget.THERMAL_NONE: CoolerTarget.NONE,
+    nvml.CoolerTarget.THERMAL_GPU: CoolerTarget.GPU,
+    nvml.CoolerTarget.THERMAL_MEMORY: CoolerTarget.MEMORY,
+    nvml.CoolerTarget.THERMAL_POWER_SUPPLY: CoolerTarget.POWER_SUPPLY,
+}
 
 
 cdef class CoolerInfo:
@@ -14,14 +56,13 @@ cdef class CoolerInfo:
         self._cooler_info = cooler_info
 
     @property
-    def signal_type(self) -> CoolerControl:
+    def signal_type(self) -> CoolerControl | None:
         """
         The cooler's control signal characteristics.
 
-        The possible types are restricted, variable and toggle.  See
-        :class:`CoolerControl` for details.
+        The possible types are variable and toggle.
         """
-        return CoolerControl(self._cooler_info.signal_type)
+        return _COOLER_CONTROL_MAPPING.get(self._cooler_info.signal_type, None)
 
     @property
     def target(self) -> list[CoolerTarget]:
@@ -32,4 +73,11 @@ cdef class CoolerInfo:
         :class:`CoolerTarget` for details.
         """
         cdef uint64_t[1] targets = [self._cooler_info.target]
-        return [CoolerTarget(1 << ev) for ev in _unpack_bitmask(targets)]
+        output_targets = []
+        for target in _unpack_bitmask(targets):
+            try:
+                output_target = _COOLER_TARGET_MAPPING[1 << target]
+            except KeyError:
+                raise ValueError(f"Unknown cooler target bit: {1 << target}")
+            output_targets.append(output_target)
+        return output_targets