Determine cliqueID from NVML not node label

Signed-off-by: Kevin Klues <[email protected]>

Author: klueska

cmd/compute-domain-kubelet-plugin/device_state.go

@@ -23,7 +23,6 @@ import (
 	"sync"
 
 	resourceapi "k8s.io/api/resource/v1beta1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/klog/v2"
 	drapbv1 "k8s.io/kubelet/pkg/apis/dra/v1beta1"
@@ -89,12 +88,11 @@ func NewDeviceState(ctx context.Context, config *Config) (*DeviceState, error) {
 		return nil, fmt.Errorf("unable to create CDI handler: %w", err)
 	}
 
-	node, err := config.clientsets.Core.CoreV1().Nodes().Get(ctx, config.flags.nodeName, metav1.GetOptions{})
+	cliqueID, err := nvdevlib.getCliqueID()
 	if err != nil {
-		return nil, fmt.Errorf("error getting Node: %w", err)
+		return nil, fmt.Errorf("error getting cliqueID: %w", err)
 	}
 
-	cliqueID := node.Labels[CliqueIDLabelKey]
 	computeDomainManager := NewComputeDomainManager(config, ComputeDomainDaemonSettingsRoot, cliqueID)
 
 	if err := cdi.CreateStandardDeviceSpecFile(allocatable); err != nil {

cmd/compute-domain-kubelet-plugin/nvlib.go

@@ -27,8 +27,11 @@ import (
 
 	"golang.org/x/sys/unix"
 
+	"k8s.io/klog/v2"
 	"k8s.io/mount-utils"
 
+	"github.com/google/uuid"
+
 	nvdev "github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 )
@@ -88,6 +91,21 @@ func newDeviceLib(driverRoot root) (*deviceLib, error) {
 	return &d, nil
 }
 
+func (l deviceLib) init() error {
+	ret := l.nvmllib.Init()
+	if ret != nvml.SUCCESS {
+		return fmt.Errorf("error initializing NVML: %v", ret)
+	}
+	return nil
+}
+
+func (l deviceLib) alwaysShutdown() {
+	ret := l.nvmllib.Shutdown()
+	if ret != nvml.SUCCESS {
+		klog.Warningf("error shutting down NVML: %v", ret)
+	}
+}
+
 func (l deviceLib) enumerateAllPossibleDevices(config *Config) (AllocatableDevices, error) {
 	alldevices := make(AllocatableDevices)
 
@@ -142,6 +160,66 @@ func (l deviceLib) enumerateComputeDomainDaemons(config *Config) (AllocatableDev
 	return devices, nil
 }
 
+func (l deviceLib) getCliqueID() (string, error) {
+	if err := l.init(); err != nil {
+		return "", fmt.Errorf("error initializing deviceLib: %w", err)
+	}
+	defer l.alwaysShutdown()
+
+	uniqueClusterUUIDs := make(map[string]struct{})
+	uniqueCliqueIDs := make(map[string]struct{})
+
+	err := l.VisitDevices(func(i int, d nvdev.Device) error {
+		isFabricAttached, err := d.IsFabricAttached()
+		if err != nil {
+			return fmt.Errorf("error checking if device is fabric attached: %w", err)
+		}
+		if !isFabricAttached {
+			return nil
+		}
+
+		info, ret := d.GetGpuFabricInfo()
+		if ret != nvml.SUCCESS {
+			return fmt.Errorf("failed to get GPU fabric info: %w", ret)
+		}
+
+		clusterUUID, err := uuid.FromBytes(info.ClusterUuid[:])
+		if err != nil {
+			return fmt.Errorf("invalid cluster UUID: %w", err)
+		}
+
+		cliqueID := fmt.Sprintf("%d", info.CliqueId)
+
+		uniqueClusterUUIDs[clusterUUID.String()] = struct{}{}
+		uniqueCliqueIDs[cliqueID] = struct{}{}
+
+		return nil
+	})
+	if err != nil {
+		return "", fmt.Errorf("error getting fabric information from one or more devices: %w", err)
+	}
+
+	if len(uniqueClusterUUIDs) == 0 && len(uniqueCliqueIDs) == 0 {
+		return "", nil
+	}
+
+	if len(uniqueClusterUUIDs) != 1 {
+		return "", fmt.Errorf("unexpected number of unique ClusterUUIDs found on devices")
+	}
+
+	if len(uniqueCliqueIDs) != 1 {
+		return "", fmt.Errorf("unexpected number of unique CliqueIDs found on devices")
+	}
+
+	for clusterUUID := range uniqueClusterUUIDs {
+		for cliqueID := range uniqueCliqueIDs {
+			return fmt.Sprintf("%s.%s", clusterUUID, cliqueID), nil
+		}
+	}
+
+	return "", fmt.Errorf("unexpected return")
+}
+
 func (l deviceLib) getImexChannelCount() (int, error) {
 	// TODO: Pull this value from /proc/driver/nvidia/params
 	return 2048, nil