fix: enhance DevicePluginDetector with normalization of time-sliced GPU

pkg/hypervisor/backend/kubernetes/external_dp/detector_test.go

@@ -281,6 +281,130 @@ var _ = Describe("DevicePluginDetector", func() {
 			Expect(err).NotTo(HaveOccurred())
 			Expect(mockAPI.AssertExpectations(GinkgoT())).To(BeTrue())
 		})
+
+		It("should de-duplicate time-sliced GPU IDs when device is added", func() {
+			mockAPI := new(MockAPIServer)
+
+			checkpointData := `{
+  "Data": {
+    "PodDeviceEntries": [
+      {
+        "PodUID": "a7461dc1-023a-4bd5-a403-c738bb1d7db4",
+        "ContainerName": "web",
+        "ResourceName": "nvidia.com/gpu",
+        "DeviceIDs": {
+          "-1": [
+            "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3::0",
+            "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3::1",
+            "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3::2",
+            "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3::3"
+          ]
+        }
+      }
+    ],
+    "RegisteredDevices": {
+      "nvidia.com/gpu": [
+        "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3"
+      ]
+    }
+  }
+}`
+
+			tmpFile, err := os.CreateTemp("", "checkpoint-*.json")
+			Expect(err).NotTo(HaveOccurred())
+			defer func() { _ = os.Remove(tmpFile.Name()) }()
+
+			_, err = tmpFile.WriteString(checkpointData)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(tmpFile.Close()).To(Succeed())
+
+			gpu := &tfv1.GPU{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3",
+				},
+				Status: tfv1.GPUStatus{
+					UsedBy: tfv1.UsedByTensorFusion,
+				},
+			}
+
+			mockAPI.On("GetGPU", "gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3").Return(gpu, nil).Once()
+			mockAPI.On("UpdateGPUStatus", mock.MatchedBy(func(gpu *tfv1.GPU) bool {
+				return gpu.Status.UsedBy == UsedByNvidiaDevicePlugin
+			})).Return(nil).Once()
+
+			detector := &DevicePluginDetector{
+				ctx:               context.Background(),
+				checkpointPath:    tmpFile.Name(),
+				apiClient:         mockAPI,
+				vendorDetectors:   make(map[string]VendorDetector),
+				previousDeviceIDs: make(map[string]string),
+			}
+			nvdpDetector := NewNvidiaDevicePluginDetector()
+			for _, prefix := range nvdpDetector.GetResourceNamePrefixes() {
+				detector.vendorDetectors[prefix] = nvdpDetector
+			}
+
+			err = detector.processDeviceState(false)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(mockAPI.AssertExpectations(GinkgoT())).To(BeTrue())
+		})
+
+		It("should normalize time-sliced GPU IDs when device is removed", func() {
+			mockAPI := new(MockAPIServer)
+
+			checkpointData := `{
+  "Data": {
+    "PodDeviceEntries": [],
+    "RegisteredDevices": {
+      "nvidia.com/gpu": [
+        "GPU-355c151e-7dd4-5a89-f270-cbda6915d7a3"
+      ]
+    }
+  }
+}`
+
+			tmpFile, err := os.CreateTemp("", "checkpoint-*.json")
+			Expect(err).NotTo(HaveOccurred())
+			defer func() { _ = os.Remove(tmpFile.Name()) }()
+
+			_, err = tmpFile.WriteString(checkpointData)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(tmpFile.Close()).To(Succeed())
+
+			gpu := &tfv1.GPU{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3",
+				},
+				Status: tfv1.GPUStatus{
+					UsedBy: UsedByNvidiaDevicePlugin,
+				},
+			}
+
+			mockAPI.On("GetGPU", "gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3").Return(gpu, nil).Once()
+			mockAPI.On("UpdateGPUStatus", mock.MatchedBy(func(gpu *tfv1.GPU) bool {
+				return gpu.Status.UsedBy == tfv1.UsedByTensorFusion
+			})).Return(nil).Once()
+
+			detector := &DevicePluginDetector{
+				ctx:            context.Background(),
+				checkpointPath: tmpFile.Name(),
+				apiClient:      mockAPI,
+				vendorDetectors: map[string]VendorDetector{
+					"nvidia.com/gpu": NewNvidiaDevicePluginDetector(),
+					"nvidia.com/mig": NewNvidiaDevicePluginDetector(),
+				},
+				previousDeviceIDs: map[string]string{
+					"gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3::0": "nvidia.com/gpu",
+					"gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3::1": "nvidia.com/gpu",
+					"gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3::2": "nvidia.com/gpu",
+					"gpu-355c151e-7dd4-5a89-f270-cbda6915d7a3::3": "nvidia.com/gpu",
+				},
+			}
+
+			err = detector.processDeviceState(false)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(mockAPI.AssertExpectations(GinkgoT())).To(BeTrue())
+		})
 	})
 
 	Describe("findEntryForDevice", func() {

pkg/hypervisor/backend/kubernetes/external_dp/kubelet_checkpoint.go

@@ -9,7 +9,6 @@ import (
 	"net"
 	"os"
 	"path/filepath"
-	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -84,7 +83,7 @@ type DevicePluginDetector struct {
 	ctx               context.Context
 	checkpointPath    string
 	apiClient         APIClientInterface
-	vendorDetectors   map[string]VendorDetector // key: resource name
+	vendorDetectors   map[string]VendorDetector // key: resource name prefix
 	previousDeviceIDs map[string]string
 	mu                sync.RWMutex
 	watcher           *fsnotify.Watcher
@@ -277,10 +276,13 @@ func (d *DevicePluginDetector) processDeviceState(patchAllDevices bool) error {
 			allocated = allocatedDevices
 		}
 	}
+	allocated, allocatedUsedBySystems := d.normalizeDeviceIDs(allocated)
+	registeredDeviceIDs, _ = d.normalizeDeviceIDs(registeredDeviceIDs)
 
 	// Determine added and removed devices
 	previousDeviceIDs := make(map[string]string, len(d.previousDeviceIDs))
 	maps.Copy(previousDeviceIDs, d.previousDeviceIDs)
+	previousDeviceIDs, _ = d.normalizeDeviceIDs(previousDeviceIDs)
 
 	var addedDevices, removedDevices map[string]string
 
@@ -314,40 +316,51 @@ func (d *DevicePluginDetector) processDeviceState(patchAllDevices bool) error {
 	// Process added devices using vendor-specific detectors
 	hasError := false
 	for deviceID, resName := range addedDevices {
-		for _, detector := range d.vendorDetectors {
-			resourceNamePrefixes := detector.GetResourceNamePrefixes()
-			if slices.Contains(resourceNamePrefixes, resName) {
-				usedBySystem, realDeviceID := detector.GetUsedBySystemAndRealDeviceID(deviceID, resName)
-				klog.V(4).Infof(
-					"Device added: %s, resource: %s, patching with usedBy: %s, realDeviceID: %s",
+		detector := d.getVendorDetector(resName)
+		if detector == nil {
+			continue
+		}
+
+		usedBySystem := allocatedUsedBySystems[deviceID]
+		if usedBySystem == "" {
+			var normalizedDeviceID string
+			usedBySystem, normalizedDeviceID = detector.GetUsedBySystemAndRealDeviceID(deviceID, resName)
+			if normalizedDeviceID != deviceID {
+				klog.Warningf(
+					"Device %s for resource %s should have been normalized as %s before patching",
 					deviceID,
 					resName,
-					usedBySystem,
-					realDeviceID,
+					normalizedDeviceID,
 				)
-				if err := d.patchGPUResource(realDeviceID, usedBySystem); err != nil {
-					klog.Errorf("Failed to patch GPU resource for added device %s: %v", deviceID, err)
-					hasError = true
-				}
 			}
 		}
+
+		klog.V(4).Infof(
+			"Device added: %s, resource: %s, patching with usedBy: %s",
+			deviceID,
+			resName,
+			usedBySystem,
+		)
+		if err := d.patchGPUResource(deviceID, usedBySystem); err != nil {
+			klog.Errorf("Failed to patch GPU resource for added device %s: %v", deviceID, err)
+			hasError = true
+		}
 	}
 
 	// Process removed devices
 	for deviceID, resName := range removedDevices {
-		for _, detector := range d.vendorDetectors {
-			resourceNamePrefixes := detector.GetResourceNamePrefixes()
-			if slices.Contains(resourceNamePrefixes, resName) {
-				klog.V(4).Infof(
-					"Device plugin allocated container removed: %s, resource: %s, patching usedBy field to tensor fusion",
-					deviceID,
-					resName,
-				)
-				if err := d.patchGPUResource(deviceID, string(tfv1.UsedByTensorFusion)); err != nil {
-					klog.Errorf("Failed to patch GPU resource usedBy field to tensor fusion for removed device %s: %v", deviceID, err)
-					hasError = true
-				}
-			}
+		if d.getVendorDetector(resName) == nil {
+			continue
+		}
+
+		klog.V(4).Infof(
+			"Device plugin allocated container removed: %s, resource: %s, patching usedBy field to tensor fusion",
+			deviceID,
+			resName,
+		)
+		if err := d.patchGPUResource(deviceID, string(tfv1.UsedByTensorFusion)); err != nil {
+			klog.Errorf("Failed to patch GPU resource usedBy field to tensor fusion for removed device %s: %v", deviceID, err)
+			hasError = true
 		}
 	}
 
@@ -448,6 +461,60 @@ func (d *DevicePluginDetector) extractDeviceIDs(
 	return allocated, registered
 }
 
+func (d *DevicePluginDetector) normalizeDeviceIDs(
+	deviceIDs map[string]string,
+) (normalized map[string]string, usedBySystems map[string]string) {
+	normalized = make(map[string]string, len(deviceIDs))
+	usedBySystems = make(map[string]string, len(deviceIDs))
+
+	for deviceID, resourceName := range deviceIDs {
+		normalizedDeviceID := deviceID
+		usedBySystem := ""
+
+		detector := d.getVendorDetector(resourceName)
+		if detector != nil {
+			usedBySystem, normalizedDeviceID = detector.GetUsedBySystemAndRealDeviceID(deviceID, resourceName)
+		}
+
+		if existingResourceName, exists := normalized[normalizedDeviceID]; exists && existingResourceName != resourceName {
+			klog.Warningf(
+				"Conflicting resource names for normalized device %s: %s vs %s",
+				normalizedDeviceID,
+				existingResourceName,
+				resourceName,
+			)
+		}
+		normalized[normalizedDeviceID] = resourceName
+
+		if usedBySystem == "" {
+			continue
+		}
+		if existingSystem, exists := usedBySystems[normalizedDeviceID]; exists && existingSystem != usedBySystem {
+			klog.Warningf(
+				"Conflicting usedBy systems for normalized device %s: %s vs %s",
+				normalizedDeviceID,
+				existingSystem,
+				usedBySystem,
+			)
+		}
+		usedBySystems[normalizedDeviceID] = usedBySystem
+	}
+
+	return normalized, usedBySystems
+}
+
+func (d *DevicePluginDetector) getVendorDetector(resourceName string) VendorDetector {
+	if detector, exists := d.vendorDetectors[resourceName]; exists {
+		return detector
+	}
+	for resourceNamePrefix, detector := range d.vendorDetectors {
+		if strings.HasPrefix(resourceName, resourceNamePrefix) {
+			return detector
+		}
+	}
+	return nil
+}
+
 // findEntryForDevice finds the pod device entry for a given device ID
 func (d *DevicePluginDetector) findEntryForDevice(checkpoint *KubeletCheckpoint, deviceID string) PodDeviceEntry {
 	deviceIDLower := strings.ToLower(deviceID)

pkg/hypervisor/backend/kubernetes/external_dp/nvdp_detector.go

@@ -1,6 +1,8 @@
 package external_dp
 
 import (
+	"strings"
+
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 )
 
@@ -32,6 +34,12 @@ func (n *NvidiaDevicePluginDetector) GetUsedBySystemAndRealDeviceID(
 	resourceName string,
 ) (system string, realDeviceID string) {
 	if resourceName == resourceNvidiaGPU {
+		// NVIDIA time-slicing appends "::replicaIndex" to the same physical GPU UUID.
+		// Keep usedBy as nvidia-device-plugin and strip the replica suffix.
+		if baseID, _, hasReplicaIndex := strings.Cut(deviceID, "::"); hasReplicaIndex {
+			return string(UsedByNvidiaDevicePlugin), baseID
+		}
+
 		// Some external device plugin's device ID is GPU-(UUID)-0, 1, 2, 3 (e.g. HAMI)
 		// Need to recover to real device ID
 		if len(deviceID) > realDeviceIDLength {