Fix GetComputeRunningProcesses on CUDA 10.x

Due to additional fields added to the nvmlProcessInfo_st in CUDA 11.x there is a mismatch in the size of the elements of the slice returned by the NVML call to get the running Compute / Graphics processes. This change attempts to detect whether the returned slice needs to be reinterpreted as a slice of smaller elements. Signed-off-by: Evan Lezar <[email protected]>
NVIDIA · Aug 12, 2021 · d566199 · d566199
1 parent 10a3a25
commit d566199
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 8 deletions.
diff --git a/gen/nvml/cgo_helpers.go b/gen/nvml/cgo_helpers.go
@@ -15,6 +15,9 @@
 package nvml
 
 import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
 	"unsafe"
 )
 
@@ -62,3 +65,39 @@ func unpackPCharString(str string) (*C.char, *struct{}) {
 	h := (*stringHeader)(unsafe.Pointer(&str))
 	return (*C.char)(h.Data), cgoAllocsUnknown
 }
+
+// adjustProcessInfoSlice can be used to adjust a ProcessInfo slice to account for
+// differences in the structure across multiple NVML versions. This handles fields that
+// were added across versions, for example.
+func adjustProcessInfoSlice(in []ProcessInfo) ([]ProcessInfo, error) {
+	type v1ProcessInfo struct {
+		pid           uint32
+		usedGpuMemory uint64
+	}
+
+	b := &bytes.Buffer{}
+	err := binary.Write(b, binary.LittleEndian, in)
+	if err != nil {
+		return nil, fmt.Errorf("error creating temporary buffer: %v", err)
+	}
+
+	intermediate := make([]v1ProcessInfo, len(in)*2)
+	err = binary.Read(b, binary.LittleEndian, intermediate)
+	if err != nil {
+		return nil, fmt.Errorf("error reading intermediate values: %v", err)
+	}
+
+	var out []ProcessInfo
+	for i := range in {
+		pin := intermediate[i]
+
+		pout := ProcessInfo{
+			Pid:           pin.pid,
+			UsedGpuMemory: pin.usedGpuMemory,
+		}
+
+		out = append(out, pout)
+	}
+
+	return out, nil
+}
diff --git a/gen/nvml/device.go b/gen/nvml/device.go
@@ -933,18 +933,32 @@ func (Device Device) GetBridgeChipInfo() (BridgeChipHierarchy, Return) {
 
 // nvml.DeviceGetComputeRunningProcesses()
 func DeviceGetComputeRunningProcesses(Device Device) ([]ProcessInfo, Return) {
+	var Infos []ProcessInfo
 	var InfoCount uint32 = 1 // Will be reduced upon returning
 	for {
-		Infos := make([]ProcessInfo, InfoCount)
+		Infos = make([]ProcessInfo, InfoCount)
 		ret := nvmlDeviceGetComputeRunningProcesses(Device, &InfoCount, &Infos[0])
 		if ret == SUCCESS {
-			return Infos[:InfoCount], ret
+			break
 		}
 		if ret != ERROR_INSUFFICIENT_SIZE {
 			return nil, ret
 		}
 		InfoCount *= 2
 	}
+
+	if InfoCount == 0 {
+		return []ProcessInfo{}, SUCCESS
+	}
+	if usesNvmlDeviceGetComputeRunningProcesses_v1 {
+		// in the case of the _v1 API we need to adjust the size of the process info data structure
+		adjusted, err := adjustProcessInfoSlice(Infos[:InfoCount])
+		if err != nil {
+			return nil, ERROR_UNKNOWN
+		}
+		return adjusted, SUCCESS
+	}
+	return Infos[:InfoCount], SUCCESS
 }
 
 func (Device Device) GetComputeRunningProcesses() ([]ProcessInfo, Return) {
@@ -953,18 +967,32 @@ func (Device Device) GetComputeRunningProcesses() ([]ProcessInfo, Return) {
 
 // nvml.DeviceGetGraphicsRunningProcesses()
 func DeviceGetGraphicsRunningProcesses(Device Device) ([]ProcessInfo, Return) {
+	var Infos []ProcessInfo
 	var InfoCount uint32 = 1 // Will be reduced upon returning
 	for {
-		Infos := make([]ProcessInfo, InfoCount)
+		Infos = make([]ProcessInfo, InfoCount)
 		ret := nvmlDeviceGetGraphicsRunningProcesses(Device, &InfoCount, &Infos[0])
 		if ret == SUCCESS {
-			return Infos[:InfoCount], ret
+			break
 		}
 		if ret != ERROR_INSUFFICIENT_SIZE {
 			return nil, ret
 		}
 		InfoCount *= 2
 	}
+
+	if InfoCount == 0 {
+		return []ProcessInfo{}, SUCCESS
+	}
+	if usesNvmlDeviceGetGraphicsRunningProcesses_v1 {
+		// in the case of the _v1 API we need to adjust the size of the process info data structure
+		adjusted, err := adjustProcessInfoSlice(Infos[:InfoCount])
+		if err != nil {
+			return nil, ERROR_UNKNOWN
+		}
+		return adjusted, SUCCESS
+	}
+	return Infos[:InfoCount], SUCCESS
 }
 
 func (Device Device) GetGraphicsRunningProcesses() ([]ProcessInfo, Return) {

diff --git a/gen/nvml/init.go b/gen/nvml/init.go
@@ -94,6 +94,9 @@ var nvmlComputeInstanceGetInfo = nvmlComputeInstanceGetInfo_v1
 var nvmlDeviceGetComputeRunningProcesses = nvmlDeviceGetComputeRunningProcesses_v1
 var nvmlDeviceGetGraphicsRunningProcesses = nvmlDeviceGetGraphicsRunningProcesses_v1
 
+var usesNvmlDeviceGetComputeRunningProcesses_v1 = true
+var usesNvmlDeviceGetGraphicsRunningProcesses_v1 = true
+
 // updateVersionedSymbols()
 func updateVersionedSymbols() {
 	err := nvml.Lookup("nvmlInit_v2")
@@ -153,10 +156,12 @@ func updateVersionedSymbols() {
 	err = nvml.Lookup("nvmlDeviceGetComputeRunningProcesses_v2")
 	if err == nil {
 		nvmlDeviceGetComputeRunningProcesses = nvmlDeviceGetComputeRunningProcesses_v2
+		usesNvmlDeviceGetComputeRunningProcesses_v1 = false
 	}
 	err = nvml.Lookup("nvmlDeviceGetGraphicsRunningProcesses_v2")
 	if err == nil {
 		nvmlDeviceGetGraphicsRunningProcesses = nvmlDeviceGetGraphicsRunningProcesses_v2
+		usesNvmlDeviceGetGraphicsRunningProcesses_v1 = false
 	}
 
 }
diff --git a/pkg/nvml/cgo_helpers.go b/pkg/nvml/cgo_helpers.go
@@ -15,6 +15,9 @@
 package nvml
 
 import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
 	"unsafe"
 )
 
@@ -62,3 +65,39 @@ func unpackPCharString(str string) (*C.char, *struct{}) {
 	h := (*stringHeader)(unsafe.Pointer(&str))
 	return (*C.char)(h.Data), cgoAllocsUnknown
 }
+
+// adjustProcessInfoSlice can be used to adjust a ProcessInfo slice to account for
+// differences in the structure across multiple NVML versions. This handles fields that
+// were added across versions, for example.
+func adjustProcessInfoSlice(in []ProcessInfo) ([]ProcessInfo, error) {
+	type v1ProcessInfo struct {
+		pid           uint32
+		usedGpuMemory uint64
+	}
+
+	b := &bytes.Buffer{}
+	err := binary.Write(b, binary.LittleEndian, in)
+	if err != nil {
+		return nil, fmt.Errorf("error creating temporary buffer: %v", err)
+	}
+
+	intermediate := make([]v1ProcessInfo, len(in)*2)
+	err = binary.Read(b, binary.LittleEndian, intermediate)
+	if err != nil {
+		return nil, fmt.Errorf("error reading intermediate values: %v", err)
+	}
+
+	var out []ProcessInfo
+	for i := range in {
+		pin := intermediate[i]
+
+		pout := ProcessInfo{
+			Pid:           pin.pid,
+			UsedGpuMemory: pin.usedGpuMemory,
+		}
+
+		out = append(out, pout)
+	}
+
+	return out, nil
+}
diff --git a/pkg/nvml/device.go b/pkg/nvml/device.go
@@ -933,18 +933,32 @@ func (Device Device) GetBridgeChipInfo() (BridgeChipHierarchy, Return) {
 
 // nvml.DeviceGetComputeRunningProcesses()
 func DeviceGetComputeRunningProcesses(Device Device) ([]ProcessInfo, Return) {
+	var Infos []ProcessInfo
 	var InfoCount uint32 = 1 // Will be reduced upon returning
 	for {
-		Infos := make([]ProcessInfo, InfoCount)
+		Infos = make([]ProcessInfo, InfoCount)
 		ret := nvmlDeviceGetComputeRunningProcesses(Device, &InfoCount, &Infos[0])
 		if ret == SUCCESS {
-			return Infos[:InfoCount], ret
+			break
 		}
 		if ret != ERROR_INSUFFICIENT_SIZE {
 			return nil, ret
 		}
 		InfoCount *= 2
 	}
+
+	if InfoCount == 0 {
+		return []ProcessInfo{}, SUCCESS
+	}
+	if usesNvmlDeviceGetComputeRunningProcesses_v1 {
+		// in the case of the _v1 API we need to adjust the size of the process info data structure
+		adjusted, err := adjustProcessInfoSlice(Infos[:InfoCount])
+		if err != nil {
+			return nil, ERROR_UNKNOWN
+		}
+		return adjusted, SUCCESS
+	}
+	return Infos[:InfoCount], SUCCESS
 }
 
 func (Device Device) GetComputeRunningProcesses() ([]ProcessInfo, Return) {
@@ -953,18 +967,32 @@ func (Device Device) GetComputeRunningProcesses() ([]ProcessInfo, Return) {
 
 // nvml.DeviceGetGraphicsRunningProcesses()
 func DeviceGetGraphicsRunningProcesses(Device Device) ([]ProcessInfo, Return) {
+	var Infos []ProcessInfo
 	var InfoCount uint32 = 1 // Will be reduced upon returning
 	for {
-		Infos := make([]ProcessInfo, InfoCount)
+		Infos = make([]ProcessInfo, InfoCount)
 		ret := nvmlDeviceGetGraphicsRunningProcesses(Device, &InfoCount, &Infos[0])
 		if ret == SUCCESS {
-			return Infos[:InfoCount], ret
+			break
 		}
 		if ret != ERROR_INSUFFICIENT_SIZE {
 			return nil, ret
 		}
 		InfoCount *= 2
 	}
+
+	if InfoCount == 0 {
+		return []ProcessInfo{}, SUCCESS
+	}
+	if usesNvmlDeviceGetGraphicsRunningProcesses_v1 {
+		// in the case of the _v1 API we need to adjust the size of the process info data structure
+		adjusted, err := adjustProcessInfoSlice(Infos[:InfoCount])
+		if err != nil {
+			return nil, ERROR_UNKNOWN
+		}
+		return adjusted, SUCCESS
+	}
+	return Infos[:InfoCount], SUCCESS
 }
 
 func (Device Device) GetGraphicsRunningProcesses() ([]ProcessInfo, Return) {

diff --git a/pkg/nvml/init.go b/pkg/nvml/init.go
@@ -94,6 +94,9 @@ var nvmlComputeInstanceGetInfo = nvmlComputeInstanceGetInfo_v1
 var nvmlDeviceGetComputeRunningProcesses = nvmlDeviceGetComputeRunningProcesses_v1
 var nvmlDeviceGetGraphicsRunningProcesses = nvmlDeviceGetGraphicsRunningProcesses_v1
 
+var usesNvmlDeviceGetComputeRunningProcesses_v1 = true
+var usesNvmlDeviceGetGraphicsRunningProcesses_v1 = true
+
 // updateVersionedSymbols()
 func updateVersionedSymbols() {
 	err := nvml.Lookup("nvmlInit_v2")
@@ -153,10 +156,12 @@ func updateVersionedSymbols() {
 	err = nvml.Lookup("nvmlDeviceGetComputeRunningProcesses_v2")
 	if err == nil {
 		nvmlDeviceGetComputeRunningProcesses = nvmlDeviceGetComputeRunningProcesses_v2
+		usesNvmlDeviceGetComputeRunningProcesses_v1 = false
 	}
 	err = nvml.Lookup("nvmlDeviceGetGraphicsRunningProcesses_v2")
 	if err == nil {
 		nvmlDeviceGetGraphicsRunningProcesses = nvmlDeviceGetGraphicsRunningProcesses_v2
+		usesNvmlDeviceGetGraphicsRunningProcesses_v1 = false
 	}
 
 }