support mig devices

.go-version

@@ -1 +1,2 @@
-1.22.4
+1.22.6
+

README.md

@@ -14,6 +14,9 @@ RPC. GPUs can be excluded from fingerprinting by setting the `ignored_gpu_ids`
 field (see below). Plugin sends statistics for fingerprinted devices every
 `stats_period` period.
 
+The plugin will be able to distinguish whether the GPU has [`Multi-Instance GPU (MIG)`](https://www.nvidia.com/en-us/technologies/multi-instance-gpu/) enabled.
+When enabled all instances will be fingerprinted as individual GPUs that can be addressed accordingly.
+
 ## Config
 
 The plugin is configured in the Nomad client's

device_test.go

@@ -6,7 +6,7 @@ package nvidia
 import (
 	"testing"
 
-	hclog "github.com/hashicorp/go-hclog"
+	"github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/nomad-device-nvidia/nvml"
 	"github.com/hashicorp/nomad/plugins/device"
 	"github.com/shoenig/test/must"

nvml/client.go

@@ -4,7 +4,9 @@
 package nvml
 
 import (
+	"cmp"
 	"fmt"
+	"slices"
 )
 
 // DeviceData represents common fields for Nvidia device
@@ -95,28 +97,32 @@ func (c *nvmlClient) GetFingerprintData() (*FingerprintData, error) {
 	*/
 
 	// Assumed that this method is called with receiver retrieved from
-	// NewNvmlClient
-	// because this method handles initialization of NVML library
+	// NewNvmlClient because this method handles initialization of NVML library
 
 	driverVersion, err := c.driver.SystemDriverVersion()
 	if err != nil {
 		return nil, fmt.Errorf("nvidia nvml SystemDriverVersion() error: %v\n", err)
 	}
 
-	numDevices, err := c.driver.DeviceCount()
+	deviceUUIDs, err := c.driver.ListDeviceUUIDs()
 	if err != nil {
-		return nil, fmt.Errorf("nvidia nvml DeviceCount() error: %v\n", err)
+		return nil, fmt.Errorf("nvidia nvml ListDeviceUUIDs() error: %v\n", err)
 	}
 
-	allNvidiaGPUResources := make([]*FingerprintDeviceData, numDevices)
+	allNvidiaGPUResources := make([]*FingerprintDeviceData, 0, len(deviceUUIDs))
 
-	for i := 0; i < int(numDevices); i++ {
-		deviceInfo, err := c.driver.DeviceInfoByIndex(uint(i))
+	for uuid, mode := range deviceUUIDs {
+		// do not care about phsyical parents of MIGs
+		if mode == parent {
+			continue
+		}
+
+		deviceInfo, err := c.driver.DeviceInfoByUUID(uuid)
 		if err != nil {
-			return nil, fmt.Errorf("nvidia nvml DeviceInfoByIndex() error: %v\n", err)
+			return nil, fmt.Errorf("nvidia nvml DeviceInfoByUUID() error: %v\n", err)
 		}
 
-		allNvidiaGPUResources[i] = &FingerprintDeviceData{
+		allNvidiaGPUResources = append(allNvidiaGPUResources, &FingerprintDeviceData{
 			DeviceData: &DeviceData{
 				DeviceName: deviceInfo.Name,
 				UUID:       deviceInfo.UUID,
@@ -130,8 +136,13 @@ func (c *nvmlClient) GetFingerprintData() (*FingerprintData, error) {
 			DisplayState:       deviceInfo.DisplayState,
 			PersistenceMode:    deviceInfo.PersistenceMode,
 			PCIBusID:           deviceInfo.PCIBusID,
-		}
+		})
+
+		slices.SortFunc(allNvidiaGPUResources, func(a, b *FingerprintDeviceData) int {
+			return cmp.Compare(a.DeviceData.UUID, b.DeviceData.UUID)
+		})
 	}
+
 	return &FingerprintData{
 		Devices:       allNvidiaGPUResources,
 		DriverVersion: driverVersion,
@@ -156,23 +167,32 @@ func (c *nvmlClient) GetStatsData() ([]*StatsData, error) {
 	*/
 
 	// Assumed that this method is called with receiver retrieved from
-	// NewNvmlClient
-	// because this method handles initialization of NVML library
+	// NewNvmlClient because this method handles initialization of NVML library
 
-	numDevices, err := c.driver.DeviceCount()
+	deviceUUIDs, err := c.driver.ListDeviceUUIDs()
 	if err != nil {
-		return nil, fmt.Errorf("nvidia nvml DeviceCount() error: %v\n", err)
+		return nil, fmt.Errorf("nvidia nvml ListDeviceUUIDs() error: %v\n", err)
 	}
 
-	allNvidiaGPUStats := make([]*StatsData, numDevices)
+	allNvidiaGPUStats := make([]*StatsData, 0, len(deviceUUIDs))
 
-	for i := 0; i < int(numDevices); i++ {
-		deviceInfo, deviceStatus, err := c.driver.DeviceInfoAndStatusByIndex(uint(i))
+	for uuid, mode := range deviceUUIDs {
+
+		// A30/A100 MIG devices have no stats.
+		//
+		// https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#telemetry
+		//
+		// Is this fixed on H100 or later? Maybe?
+		if mode == mig || mode == parent {
+			continue
+		}
+
+		deviceInfo, deviceStatus, err := c.driver.DeviceInfoAndStatusByUUID(uuid)
 		if err != nil {
-			return nil, fmt.Errorf("nvidia nvml DeviceInfoAndStatusByIndex() error: %v\n", err)
+			return nil, fmt.Errorf("nvidia nvml DeviceInfoAndStatusByUUID() error: %v\n", err)
 		}
 
-		allNvidiaGPUStats[i] = &StatsData{
+		allNvidiaGPUStats = append(allNvidiaGPUStats, &StatsData{
 			DeviceData: &DeviceData{
 				DeviceName: deviceInfo.Name,
 				UUID:       deviceInfo.UUID,
@@ -191,7 +211,11 @@ func (c *nvmlClient) GetStatsData() ([]*StatsData, error) {
 			ECCErrorsL1Cache:   deviceStatus.ECCErrorsL1Cache,
 			ECCErrorsL2Cache:   deviceStatus.ECCErrorsL2Cache,
 			ECCErrorsDevice:    deviceStatus.ECCErrorsDevice,
-		}
+		})
+
+		slices.SortFunc(allNvidiaGPUStats, func(a, b *StatsData) int {
+			return cmp.Compare(a.DeviceData.UUID, b.DeviceData.UUID)
+		})
 	}
 	return allNvidiaGPUStats, nil
 }