adding device::vgpu_scheduler_state/set_vgpu_scheduler_state (#111)

Author: devnexen

nvml-wrapper/src/device.rs

@@ -6403,7 +6403,7 @@ impl<'nvml> Device<'nvml> {
         }
     }
 
-    /// Obtain the n log entries (max 200) of the vGPU scheduler, to be called several time if need
+    /// Obtain the n log entries (max 200) of the vGPU scheduler, to be called several times if need
     /// be.
     pub fn vgpu_scheduler_log(&self) -> Result<VgpuSchedulerLog, NvmlError> {
         let sym = nvml_sym(self.nvml.lib.nvmlDeviceGetVgpuSchedulerLog.as_ref())?;
@@ -6417,6 +6417,29 @@ impl<'nvml> Device<'nvml> {
         }
     }
 
+    /// Obtain the vGPU scheduler state of the device
+    pub fn vgpu_scheduler_state(&self) -> Result<VgpuSchedulerGetState, NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlDeviceGetVgpuSchedulerState.as_ref())?;
+
+        unsafe {
+            let mut scheduler_state: nvmlVgpuSchedulerGetState_t = mem::zeroed();
+
+            nvml_try(sym(self.device, &mut scheduler_state))?;
+
+            Ok(VgpuSchedulerGetState::from(scheduler_state))
+        }
+    }
+
+    /// Set the vGPU scheduler state of the device
+    pub fn set_vgpu_scheduler_state(
+        &self,
+        scheduler_state: VgpuSchedulerSetState,
+    ) -> Result<(), NvmlError> {
+        let sym = nvml_sym(self.nvml.lib.nvmlDeviceSetVgpuSchedulerState.as_ref())?;
+
+        unsafe { nvml_try(sym(self.device, &mut scheduler_state.as_c())) }
+    }
+
     /// Check if the GPU is on vGPU host mode
     pub fn vgpu_host_mode(&self) -> Result<HostVgpuMode, NvmlError> {
         let sym = nvml_sym(self.nvml.lib.nvmlDeviceGetHostVgpuMode.as_ref())?;

nvml-wrapper/src/struct_wrappers/device.rs

@@ -934,6 +934,70 @@ impl From<nvmlVgpuSchedulerLog_t> for VgpuSchedulerLog {
     }
 }
 
+/// Vgpu scheduler state
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct VgpuSchedulerGetState {
+    /// Adaptative Round Robin scheduler mode
+    pub arr_mode: u32,
+    /// Scheduler policy
+    pub scheduler_policy: u32,
+}
+
+impl From<nvmlVgpuSchedulerGetState_t> for VgpuSchedulerGetState {
+    fn from(value: nvmlVgpuSchedulerGetState_t) -> Self {
+        Self {
+            arr_mode: value.arrMode,
+            scheduler_policy: value.schedulerPolicy,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct VgpuSchedulerSetParams {
+    /// Average factor in compensating the timeslice for Adaptive Round Robin mode
+    pub avg_factor: Option<u32>,
+    /// Frequency for Adaptative Mode (when avg_factor is set)/timeslice in ns for each software run list as configured, or the default value otherwise
+    pub frequency_or_timeslice: u32,
+}
+
+impl VgpuSchedulerSetParams {
+    pub fn as_c(&self) -> nvmlVgpuSchedulerSetParams_t {
+        match self.avg_factor {
+            Some(a) => nvmlVgpuSchedulerSetParams_t {
+                vgpuSchedDataWithARR: nvmlVgpuSchedulerSetParams_t__bindgen_ty_1 {
+                    avgFactor: a,
+                    frequency: self.frequency_or_timeslice,
+                },
+            },
+            _ => nvmlVgpuSchedulerSetParams_t {
+                vgpuSchedData: nvmlVgpuSchedulerSetParams_t__bindgen_ty_2 {
+                    timeslice: self.frequency_or_timeslice,
+                },
+            },
+        }
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct VgpuSchedulerSetState {
+    pub scheduler_policy: u32,
+    pub enable_arr_mode: u32,
+    pub scheduler_params: VgpuSchedulerSetParams,
+}
+
+impl VgpuSchedulerSetState {
+    pub fn as_c(&self) -> nvmlVgpuSchedulerSetState_t {
+        nvmlVgpuSchedulerSetState_t {
+            enableARRMode: self.enable_arr_mode,
+            schedulerPolicy: self.scheduler_policy,
+            schedulerParams: self.scheduler_params.as_c(),
+        }
+    }
+}
+
 #[cfg(test)]
 #[allow(unused_variables, unused_imports)]
 mod tests {