Enable labels for ClusterUUID and CliqueId

Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>

internal/lm/nvml.go

@@ -80,13 +80,19 @@ func NewDeviceLabeler(manager resource.Manager, config *spec.Config) (Labeler, e
 		return nil, fmt.Errorf("error creating resource labeler: %v", err)
 	}
 
+	imexLabeler, err := newImexDomainLabeler(devices)
+	if err != nil {
+		return nil, fmt.Errorf("error creating imex domain labeler: %v", err)
+	}
+
 	l := Merge(
 		machineTypeLabeler,
 		versionLabeler,
 		migCapabilityLabeler,
 		sharingLabeler,
 		resourceLabeler,
 		gpuModeLabeler,
+		imexLabeler,
 	)
 
 	return l, nil
@@ -218,6 +224,45 @@ func newGPUModeLabeler(devices []resource.Device) (Labeler, error) {
 	return labels, nil
 }
 
+func newImexDomainLabeler(devices []resource.Device) (Labeler, error) {
+	var commonClusterUUID, commonCliqueID string
+
+	for _, d := range devices {
+		clusterUuid, err := d.GetClusterUUID()
+		if err != nil {
+			return nil, err
+		}
+		if commonClusterUUID == "" {
+			commonClusterUUID = clusterUuid
+		}
+		if commonClusterUUID != clusterUuid {
+			klog.Warningf("cluster UUID mismatch: %s != %s", commonClusterUUID, clusterUuid)
+			return nil, fmt.Errorf("cluster UUID mismatch: %s != %s", commonClusterUUID, clusterUuid)
+		}
+
+		cliqueId, err := d.GetCliqueID()
+		if err != nil {
+			return nil, err
+		}
+		if commonCliqueID == "" {
+			commonCliqueID = cliqueId
+		}
+		if commonCliqueID != cliqueId {
+			klog.Warningf("clique ID mismatch: %s != %s", commonCliqueID, cliqueId)
+			return nil, fmt.Errorf("clique ID mismatch: %s != %s", commonCliqueID, cliqueId)
+		}
+	}
+
+	domain := fmt.Sprintf("%s-%s", commonClusterUUID, commonCliqueID)
+	labels := Labels{
+		"nvidia.com/gpu.clusteruuid": commonClusterUUID,
+		"nvidia.com/gpu.cliqueid":    commonCliqueID,
+		"nvidia.com/gpu.imex-domain": domain,
+	}
+
+	return labels, nil
+}
+
 func getModeForClasses(classes []uint32) string {
 	if len(classes) == 0 {
 		return "unknown"

internal/resource/cuda-device.go

@@ -100,3 +100,11 @@ func (d *cudaDevice) IsMigEnabled() (bool, error) {
 func (d *cudaDevice) GetPCIClass() (uint32, error) {
 	return 0, nil
 }
+
+func (d *cudaDevice) GetClusterUUID() (string, error) {
+	return "", nil
+}
+
+func (d *cudaDevice) GetCliqueID() (string, error) {
+	return "", nil
+}

internal/resource/nvml-device.go

@@ -17,7 +17,9 @@
 package resource
 
 import (
+	"encoding/hex"
 	"fmt"
+	"strconv"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
@@ -99,3 +101,35 @@ func (d nvmlDevice) GetPCIClass() (uint32, error) {
 	}
 	return nvDevice.Class, nil
 }
+
+func (d nvmlDevice) GetClusterUUID() (string, error) {
+	gfInfo, ret := d.GetGpuFabricInfo()
+	if ret != nvml.SUCCESS {
+		return "", ret
+	}
+
+	// Convert the array to a byte slice
+	byteSlice := gfInfo.ClusterUuid[:]
+
+	// Encode the byte slice as a hex string
+	hexStr := hex.EncodeToString(byteSlice)
+
+	// Format the hex string with dashes to match UUID format
+	uuid := fmt.Sprintf("%s-%s-%s-%s-%s",
+		hexStr[0:8],   // First 8 characters
+		hexStr[8:12],  // Next 4 characters
+		hexStr[12:16], // Next 4 characters
+		hexStr[16:20], // Next 4 characters
+		hexStr[20:32]) // Last 12 characters
+
+	return uuid, nil
+}
+
+func (d nvmlDevice) GetCliqueID() (string, error) {
+	gfInfo, ret := d.GetGpuFabricInfo()
+	if ret != nvml.SUCCESS {
+		return "", ret
+	}
+
+	return strconv.FormatUint(uint64(gfInfo.CliqueId), 10), nil
+}

internal/resource/nvml-mig-device.go

@@ -138,3 +138,11 @@ func (d nvmlMigDevice) GetPCIClass() (uint32, error) {
 	// GPU devices that support MIG do not support switching mode between graphics and compute, so they are always in compute mode.
 	return nvpci.PCI3dControllerClass, nil
 }
+
+func (d nvmlMigDevice) GetClusterUUID() (string, error) {
+	return "", fmt.Errorf("GetClusterUUID is not supported for MIG devices")
+}
+
+func (d nvmlMigDevice) GetCliqueID() (string, error) {
+	return "", fmt.Errorf("GetCliqueID is not supported for MIG devices")
+}

internal/resource/sysfs-device.go

@@ -68,3 +68,11 @@ func (d vfioDevice) IsMigCapable() (bool, error) {
 func (d vfioDevice) GetPCIClass() (uint32, error) {
 	return d.nvidiaPCIDevice.Class, nil
 }
+
+func (d vfioDevice) GetClusterUUID() (string, error) {
+	return "", nil
+}
+
+func (d vfioDevice) GetCliqueID() (string, error) {
+	return "", nil
+}

internal/resource/types.go

@@ -40,4 +40,6 @@ type Device interface {
 	GetDeviceHandleFromMigDeviceHandle() (Device, error)
 	GetCudaComputeCapability() (int, int, error)
 	GetPCIClass() (uint32, error)
+	GetClusterUUID() (string, error)
+	GetCliqueID() (string, error)
 }