Merge pull request #25 from NVIDIA/add-set-compute-mode

Add SetComputeMode method to Device
NVIDIA · Mar 11, 2024 · 6ac4869 · 6ac4869
2 parents 22de91f + 2feaa48
commit 6ac4869
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 0 deletions.
diff --git a/pkg/nvml/device.go b/pkg/nvml/device.go
@@ -207,3 +207,9 @@ func (d nvmlDevice) GetNvLinkRemotePciInfo(link int) (PciInfo, Return) {
 	p, r := nvml.Device(d).GetNvLinkRemotePciInfo(link)
 	return PciInfo(p), Return(r)
 }
+
+// SetComputeMode sets the compute mode for the device.
+func (d nvmlDevice) SetComputeMode(mode ComputeMode) Return {
+	r := nvml.Device(d).SetComputeMode(nvml.ComputeMode(mode))
+	return Return(r)
+}
diff --git a/pkg/nvml/device_mock.go b/pkg/nvml/device_mock.go
diff --git a/pkg/nvml/types.go b/pkg/nvml/types.go
@@ -67,6 +67,7 @@ type Device interface {
 	GetUUID() (string, Return)
 	IsMigDeviceHandle() (bool, Return)
 	RegisterEvents(uint64, EventSet) Return
+	SetComputeMode(ComputeMode) Return
 	SetMigMode(Mode int) (Return, Return)
 	// nvmlDeviceHandle returns a pointer to the underlying NVML device.
 	nvmlDeviceHandle() *nvml.Device
@@ -156,3 +157,6 @@ type GpuTopologyLevel nvml.GpuTopologyLevel
 
 // EnableState represents a generic enable/disable enum
 type EnableState nvml.EnableState
+
+// ComputeMode represents the compute mode for a device
+type ComputeMode nvml.ComputeMode