From 3c93dae75b80b91565ed908a31422634aae82717 Mon Sep 17 00:00:00 2001
From: deorth-kku <deorth@outlook.com>
Date: Tue, 15 Oct 2024 12:13:02 +0800
Subject: [PATCH] Add gpu.thermal_target property

---
 src/pynvraw/gpu.py       | 18 ++++++++++++++++++
 src/pynvraw/nvapi_api.py | 19 +++++++++++++++++++
 test.py                  | 11 ++++++++++-
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/pynvraw/gpu.py b/src/pynvraw/gpu.py
index 802a558..ae36d8e 100644
--- a/src/pynvraw/gpu.py
+++ b/src/pynvraw/gpu.py
@@ -210,6 +210,24 @@ def power_limit(self, value: float):
         status.entries[0].power = int(value * 1000)
         self.api.NvAPI_GPU_ClientPowerPoliciesSetStatus(self.handle, ctypes.pointer(status))
 
+    @property
+    def thermal_target(self) ->float:
+        '''Reads thermal target limit in °C.'''
+        status = self.api.get_thermal_status(self.handle)
+        if status.count == 0:
+            return None
+        return max(e.value for e in status.entries[:status.count]) / 256
+    
+    @thermal_target.setter
+    def thermal_target(self, value: float):
+        '''Sets thermal target limit in °C.'''
+        status = self.api.get_thermal_status(self.handle)
+        for e in status.entries[:status.count]:
+            if e.controller != 0:
+                e.value = int(value * 256)
+        self.api.NvAPI_GPU_SetThermalPoliciesStatus(self.handle, ctypes.pointer(status))
+
+
     @property
     def power(self) -> float:
         '''Reads current power consumption in %.'''
diff --git a/src/pynvraw/nvapi_api.py b/src/pynvraw/nvapi_api.py
index 2a75f50..0aa9204 100644
--- a/src/pynvraw/nvapi_api.py
+++ b/src/pynvraw/nvapi_api.py
@@ -140,6 +140,16 @@ class NV_GPU_THERMAL_EX(NvVersioned):
     @property
     def sensors(self):
         return tuple(x / 256.0 for x in self._sensors)
+    
+class NV_GPU_THERMAL_POLICIES_STATUS(NvVersioned):
+    class NV_GPU_THERMAL_POLICIES_STATUS_ENTRY(StrStructure):
+        _fields_ = [('controller',ctypes.c_uint32),
+                    ('value',ctypes.c_uint32),
+                    ('flags',ctypes.c_uint32)]
+    _nv_version_ = 2
+    _fields_ = [('version', ctypes.c_uint32),
+                ('count', ctypes.c_uint32),
+                ('entries', NV_GPU_THERMAL_POLICIES_STATUS_ENTRY * 4)]
 
 class NV_COOLER_TARGET(enum.IntFlag):
     NONE = 0
@@ -945,6 +955,10 @@ class NvAPI:
     NvAPI_GPU_RestoreCoolerSettings = NvMethod(0x8F6ED0FB, 'NvAPI_GPU_RestoreCoolerSettings', NvPhysicalGpu, ctypes.POINTER(ctypes.c_uint32), ctypes.c_uint32)
     NvAPI_GPU_GetPstates20 = NvMethod(0x6FF81213, 'NvAPI_GPU_GetPstates20', NvPhysicalGpu, ctypes.POINTER(NV_GPU_PERF_PSTATES20_INFO))
     NvAPI_GPU_SetPstates20 = NvMethod(0x0F4DAE6B, 'NvAPI_GPU_SetPstates20', NvPhysicalGpu, ctypes.POINTER(NV_GPU_PERF_PSTATES20_INFO))
+
+    NvAPI_GPU_GetThermalPoliciesStatus = NvMethod(0x0E9C425A1, 'NvAPI_GPU_GetThermalPoliciesStatus', NvPhysicalGpu, ctypes.POINTER(NV_GPU_THERMAL_POLICIES_STATUS))
+    NvAPI_GPU_SetThermalPoliciesStatus = NvMethod(0x034C0B13D, 'NvAPI_GPU_SetThermalPoliciesStatus', NvPhysicalGpu, ctypes.POINTER(NV_GPU_THERMAL_POLICIES_STATUS))
+
     NvAPI_GPU_ClientPowerPoliciesGetInfo = NvMethod(0x34206D86, 'NvAPI_GPU_ClientPowerPoliciesGetInfo', NvPhysicalGpu, ctypes.POINTER(NV_GPU_POWER_INFO))
     NvAPI_GPU_ClientPowerPoliciesGetStatus = NvMethod(0x70916171, 'NvAPI_GPU_ClientPowerPoliciesGetStatus', NvPhysicalGpu, ctypes.POINTER(NV_GPU_POWER_STATUS))
     NvAPI_GPU_ClientPowerPoliciesSetStatus = NvMethod(0xAD95F5ED, 'NvAPI_GPU_ClientPowerPoliciesSetStatus', NvPhysicalGpu, ctypes.POINTER(NV_GPU_POWER_STATUS))
@@ -1069,6 +1083,11 @@ def get_power_status(self, dev: NvPhysicalGpu) -> NV_GPU_POWER_STATUS:
         value = NV_GPU_POWER_STATUS()
         self.NvAPI_GPU_ClientPowerPoliciesGetStatus(dev, ctypes.pointer(value))
         return value
+    
+    def get_thermal_status(self,dev:NvPhysicalGpu) -> NV_GPU_THERMAL_POLICIES_STATUS:
+        value = NV_GPU_THERMAL_POLICIES_STATUS()
+        self.NvAPI_GPU_GetThermalPoliciesStatus(dev, ctypes.pointer(value))
+        return value
 
     def get_topology_status(self, dev: NvPhysicalGpu) -> NV_GPU_TOPOLOGY_STATUS:
         value = NV_GPU_TOPOLOGY_STATUS()
diff --git a/test.py b/test.py
index 60398aa..d70249d 100644
--- a/test.py
+++ b/test.py
@@ -1,4 +1,4 @@
-from pynvraw import api, NvError, get_phys_gpu, nvapi_api
+from src.pynvraw import api, NvError, get_phys_gpu, nvapi_api
 
 def main():
     print(api.get_driver_version())
@@ -68,6 +68,15 @@ def main():
         print(f'current power: {gpu.power}%')
         print(f'perf limit: {gpu.perf_limit!s}')
 
+        thermal_target=gpu.thermal_target
+        print(f'thermal target: {gpu.thermal_target!s}')
+        try:
+            gpu.thermal_target=thermal_target-1
+            gpu.thermal_target=thermal_target
+        except Exception as e:
+            print('set thermal target failed',e)
+        else:
+            print('set thermal target ok')
         '''
         mask = api.get_boost_mask(gpu.handle)
         vfp = api.get_vfp_curve(gpu.handle, mask)