fix(scheduler): count a shared DRA GPU device once per node

When one physical GPU is shared by several pods through a single
ResourceClaim (status.reservedFor lists more than one consumer, as with
MPS or time-slicing), every consuming pod carries the same allocated
device. The node accounting added that device to the used vector once
per pod, so a 1-GPU node shared by two pods reported used: 2, capacity:
1. IdleVector then went negative and MaxNodeResourcesPredicate rejected
the node for every task, even ones requesting no GPU at all, which
showed up as a confusing "didn't have enough resources: GPUs" on a node
that clearly had one.

Keep a per-device reference count on the node, keyed by
driver/pool/device from the allocation result, and skip a device that
another pod on the node already contributed. Removal mirrors this: the
device stays counted until the last consumer leaves. Exclusive claims
and non-DRA GPUs are untouched.

Reservation pods are a special case worth calling out: addTaskResources
zeroes their GPU index before the dedup runs, so counting their shared
devices would subtract from zero and drive the used count negative.
Both the dedup and release paths now leave a zero-GPU task's accounting
alone.

Signed-off-by: Raya Solano <raya@mbinf.de>

Author: Raya Solano

.changes/unreleased/Fixed-20260717-233428.yaml

@@ -0,0 +1,6 @@
+kind: Fixed
+body: Count a GPU shared by multiple pods through one DRA ResourceClaim once per node, preventing negative idle GPUs.
+time: 2026-07-17T23:34:28.999387137Z
+custom:
+    Author: TensorRaya
+    Issue: "1930"

pkg/scheduler/api/node_info/dra_shared_device_info.go

@@ -0,0 +1,120 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package node_info
+
+import (
+	resourceapi "k8s.io/api/resource/v1"
+
+	"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
+)
+
+// draDeviceKey uniquely identifies a physical DRA device on the node.
+func draDeviceKey(result resourceapi.DeviceRequestAllocationResult) string {
+	return result.Driver + "/" + result.Pool + "/" + result.Device
+}
+
+// allocatedGPUDeviceKeys returns the keys of all GPU devices allocated to the
+// task via DRA ResourceClaims. Non-GPU devices are ignored: they are not part
+// of the GPU accounting that this dedup protects.
+func (ni *NodeInfo) allocatedGPUDeviceKeys(task *pod_info.PodInfo) []string {
+	var keys []string
+	for _, claimAllocation := range task.ResourceClaimInfo {
+		if claimAllocation == nil || claimAllocation.Allocation == nil {
+			continue
+		}
+		for _, result := range claimAllocation.Allocation.Devices.Results {
+			if !resources.IsGPUDeviceClass(result.Driver) {
+				continue
+			}
+			keys = append(keys, draDeviceKey(result))
+		}
+	}
+	return keys
+}
+
+// sharedDRAGpuDiscount returns the number of GPU devices the task requests via
+// DRA claims that are already counted on this node for other pods. A task that
+// shares an allocated device with a running pod does not need additional GPU
+// capacity for that device.
+func (ni *NodeInfo) sharedDRAGpuDiscount(task *pod_info.PodInfo) float64 {
+	discount := 0.0
+	for _, key := range ni.allocatedGPUDeviceKeys(task) {
+		if ni.DRASharedDeviceRefCount[key] > 0 {
+			discount++
+		}
+	}
+	return discount
+}
+
+// dedupSharedDRAGpus removes from resourcesToTrack the GPU count that would
+// double-count physical DRA devices already referenced by other pods on the
+// node. It also updates the node's per-device reference count. It must be
+// called once per addTaskResources, before the vector is added to UsedVector.
+func (ni *NodeInfo) dedupSharedDRAGpus(task *pod_info.PodInfo, resourcesToTrack resource_info.ResourceVector) {
+	current := resourcesToTrack.Get(resource_info.GPUIndex)
+	if current <= 0 {
+		// The task contributes no GPUs to the used vector (e.g. a resource
+		// reservation task whose GPU index was zeroed). Tracking its devices
+		// would both risk a negative deduction below and mask the reference
+		// count of the real consuming pods, so leave the accounting untouched.
+		return
+	}
+
+	alreadyCounted := 0.0
+	for _, key := range ni.allocatedGPUDeviceKeys(task) {
+		if ni.DRASharedDeviceRefCount[key] > 0 {
+			// Another pod on this node already contributed this physical
+			// device to the used vector: do not count it again.
+			alreadyCounted++
+		}
+		ni.DRASharedDeviceRefCount[key]++
+	}
+
+	if alreadyCounted > current {
+		// Never deduct more than the task's own GPU contribution.
+		alreadyCounted = current
+	}
+	if alreadyCounted > 0 {
+		resourcesToTrack.Set(resource_info.GPUIndex, current-alreadyCounted)
+	}
+}
+
+// releaseSharedDRAGpus is the inverse of dedupSharedDRAGpus: it decrements the
+// per-device reference count and adds back the GPU count for devices that
+// remain referenced by other pods (and were therefore never subtracted on this
+// task's removal path). It must be called once per removeTaskResources.
+func (ni *NodeInfo) releaseSharedDRAGpus(task *pod_info.PodInfo, resourcesToTrack resource_info.ResourceVector) {
+	current := resourcesToTrack.Get(resource_info.GPUIndex)
+	if current <= 0 {
+		// Mirror of dedupSharedDRAGpus: a task that contributed no GPUs never
+		// incremented the reference count, so it must not decrement it here.
+		return
+	}
+
+	stillShared := 0.0
+	for _, key := range ni.allocatedGPUDeviceKeys(task) {
+		if ni.DRASharedDeviceRefCount[key] > 1 {
+			// The device stays referenced by another pod after this removal:
+			// it must remain in the used vector, so this task's removal must
+			// not subtract it.
+			stillShared++
+		}
+		if ni.DRASharedDeviceRefCount[key] > 0 {
+			ni.DRASharedDeviceRefCount[key]--
+		}
+		if ni.DRASharedDeviceRefCount[key] == 0 {
+			delete(ni.DRASharedDeviceRefCount, key)
+		}
+	}
+
+	if stillShared > current {
+		// Never add back more than the task's own GPU contribution.
+		stillShared = current
+	}
+	if stillShared > 0 {
+		resourcesToTrack.Set(resource_info.GPUIndex, current-stillShared)
+	}
+}

pkg/scheduler/api/node_info/dra_shared_device_info_test.go

@@ -0,0 +1,226 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package node_info
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+	v1 "k8s.io/api/core/v1"
+	resourceapi "k8s.io/api/resource/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
+
+	commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_affinity"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
+)
+
+const gpuDeviceClass = "gpu.nvidia.com"
+
+// sharedGPUClaim builds a ResourceClaim that requests one GPU device and is
+// allocated to the given physical device. The same claim object is referenced
+// by every pod that shares the device (status.reservedFor with multiple
+// entries), matching a DRA time-slicing / MPS setup.
+func sharedGPUClaim(name, namespace, driver, pool, device string) *resourceapi.ResourceClaim {
+	return &resourceapi.ResourceClaim{
+		ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: namespace},
+		Spec: resourceapi.ResourceClaimSpec{
+			Devices: resourceapi.DeviceClaim{
+				Requests: []resourceapi.DeviceRequest{
+					{
+						Name: "gpu",
+						Exactly: &resourceapi.ExactDeviceRequest{
+							DeviceClassName: gpuDeviceClass,
+							AllocationMode:  resourceapi.DeviceAllocationModeExactCount,
+							Count:           1,
+						},
+					},
+				},
+			},
+		},
+		Status: resourceapi.ResourceClaimStatus{
+			Allocation: &resourceapi.AllocationResult{
+				Devices: resourceapi.DeviceAllocationResult{
+					Results: []resourceapi.DeviceRequestAllocationResult{
+						{Request: "gpu", Driver: driver, Pool: pool, Device: device},
+					},
+				},
+			},
+		},
+	}
+}
+
+// draConsumerPod builds a running pod that consumes the given claim by name.
+func draConsumerPod(name, namespace, nodeName, claimName string) *v1.Pod {
+	pod := common_info.BuildPod(namespace, name, nodeName, v1.PodRunning,
+		common_info.BuildResourceList("1000m", "1G"), []metav1.OwnerReference{},
+		map[string]string{}, map[string]string{
+			pod_info.ReceivedResourceTypeAnnotationName: string(pod_info.ReceivedTypeRegular),
+			commonconstants.PodGroupAnnotationForPod:    common_info.FakePogGroupId,
+		})
+	pod.Spec.ResourceClaims = []v1.PodResourceClaim{
+		{Name: claimName, ResourceClaimName: ptr.To(claimName)},
+	}
+	return pod
+}
+
+// newGPUNodeInfo builds a NodeInfo for a node with the given whole-GPU count,
+// wired with a mock pod-affinity that expects addPods AddPod and rmPods
+// RemovePod calls, and a vectorMap that knows the DRA GPU device class.
+func newGPUNodeInfo(t *testing.T, name, gpuCount string, addPods, rmPods int) (*NodeInfo, *resource_info.ResourceVectorMap) {
+	node := common_info.BuildNode(name, common_info.BuildResourceListWithGPU("8000m", "16G", gpuCount))
+
+	ctrl := gomock.NewController(t)
+	affinity := pod_affinity.NewMockNodePodAffinityInfo(ctrl)
+	affinity.EXPECT().AddPod(gomock.Any()).Times(addPods)
+	affinity.EXPECT().RemovePod(gomock.Any()).Times(rmPods)
+
+	vectorMap := resource_info.NewResourceVectorMap()
+	for resourceName := range node.Status.Allocatable {
+		vectorMap.AddResource(resourceName)
+	}
+	// DRA GPU counts are tracked under the device-class resource name.
+	vectorMap.AddResource(v1.ResourceName(gpuDeviceClass))
+
+	return NewNodeInfo(node, affinity, vectorMap), vectorMap
+}
+
+// TestAddTask_SharedDRAClaimCountedOnce is the regression test for the
+// shared-ResourceClaim double count. Two pods share one physical GPU through a
+// single claim (reservedFor has both). Naive per-task accounting adds the GPU
+// once per pod, driving the node's used GPU count to 2 on a 1-GPU node and
+// IdleVector negative, which makes the node unschedulable for every task. The
+// fix must keep the used count at exactly 1.
+//
+// This exercises the real accounting path (AddTask -> addTaskResources ->
+// UsedVector), not the dedup helper in isolation, so it fails if the dedup is
+// not wired into AddTask.
+func TestAddTask_SharedDRAClaimCountedOnce(t *testing.T) {
+	ni, vectorMap := newGPUNodeInfo(t, "atlas", "1", 2, 0)
+
+	claim := sharedGPUClaim("voice-atlas-shared", "voice-pipeline", gpuDeviceClass, "atlas", "gpu-0")
+	pod1 := draConsumerPod("voice-tts", "voice-pipeline", "atlas", "voice-atlas-shared")
+	pod2 := draConsumerPod("voice-worker", "voice-pipeline", "atlas", "voice-atlas-shared")
+
+	task1 := pod_info.NewTaskInfo(pod1, vectorMap, pod_info.TaskInfoOptions{
+		DraPodClaims: []*resourceapi.ResourceClaim{claim},
+	})
+	task2 := pod_info.NewTaskInfo(pod2, vectorMap, pod_info.TaskInfoOptions{
+		DraPodClaims: []*resourceapi.ResourceClaim{claim},
+	})
+
+	assert.NoError(t, ni.AddTask(task1))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"first consumer of the shared device must be counted")
+
+	assert.NoError(t, ni.AddTask(task2))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"second consumer of the same shared device must not be double-counted")
+
+	idleGPUs, _ := ni.GetSumOfIdleGPUs()
+	assert.Equal(t, 0.0, idleGPUs, "idle GPUs must be 0, never negative, on a fully-shared 1-GPU node")
+}
+
+// TestAddRemoveTask_SharedDRAClaimSymmetry verifies the used count follows the
+// number of distinct physical devices as consumers are added and removed: it
+// stays at 1 while any consumer of the shared device remains, and returns to 0
+// only when the last one leaves.
+func TestAddRemoveTask_SharedDRAClaimSymmetry(t *testing.T) {
+	ni, vectorMap := newGPUNodeInfo(t, "atlas", "1", 2, 2)
+
+	claim := sharedGPUClaim("voice-atlas-shared", "voice-pipeline", gpuDeviceClass, "atlas", "gpu-0")
+	pod1 := draConsumerPod("voice-tts", "voice-pipeline", "atlas", "voice-atlas-shared")
+	pod2 := draConsumerPod("voice-worker", "voice-pipeline", "atlas", "voice-atlas-shared")
+	task1 := pod_info.NewTaskInfo(pod1, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claim}})
+	task2 := pod_info.NewTaskInfo(pod2, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claim}})
+
+	assert.NoError(t, ni.AddTask(task1))
+	assert.NoError(t, ni.AddTask(task2))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex))
+
+	assert.NoError(t, ni.RemoveTask(task2))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"removing one of two consumers must keep the shared device counted")
+
+	assert.NoError(t, ni.RemoveTask(task1))
+	assert.Equal(t, 0.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"removing the last consumer must release the shared device")
+}
+
+// draReservationPod builds a resource-reservation pod that references the given
+// claim. addTaskResources zeroes such a pod's GPU index, so it must not affect
+// the shared-device GPU accounting at all.
+func draReservationPod(name, namespace, nodeName, claimName string) *v1.Pod {
+	pod := common_info.BuildPod(namespace, name, nodeName, v1.PodRunning,
+		common_info.BuildResourceList("1000m", "1G"), []metav1.OwnerReference{},
+		map[string]string{
+			commonconstants.AppLabelName: "kai-resource-reservation",
+		}, map[string]string{
+			pod_info.ReceivedResourceTypeAnnotationName: string(pod_info.ReceivedTypeRegular),
+			commonconstants.PodGroupAnnotationForPod:    common_info.FakePogGroupId,
+		})
+	pod.Spec.ResourceClaims = []v1.PodResourceClaim{
+		{Name: claimName, ResourceClaimName: ptr.To(claimName)},
+	}
+	return pod
+}
+
+// TestAddTask_DistinctDRADevicesEachCounted guards against over-dedup: two pods
+// on two different physical GPUs (own claim each) must both be counted.
+func TestAddTask_DistinctDRADevicesEachCounted(t *testing.T) {
+	ni, vectorMap := newGPUNodeInfo(t, "nyx", "2", 2, 0)
+
+	claimA := sharedGPUClaim("coder-claim", "vllm", gpuDeviceClass, "nyx", "gpu-4")
+	claimB := sharedGPUClaim("gemma-claim", "vllm", gpuDeviceClass, "nyx", "gpu-6")
+	podA := draConsumerPod("coder", "vllm", "nyx", "coder-claim")
+	podB := draConsumerPod("gemma", "vllm", "nyx", "gemma-claim")
+	taskA := pod_info.NewTaskInfo(podA, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claimA}})
+	taskB := pod_info.NewTaskInfo(podB, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claimB}})
+
+	assert.NoError(t, ni.AddTask(taskA))
+	assert.NoError(t, ni.AddTask(taskB))
+	assert.Equal(t, 2.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"two distinct physical devices must both be counted")
+}
+
+// TestAddRemoveTask_ReservationPodDoesNotCorruptSharedDRAAccounting guards the
+// dedup against resource-reservation tasks. addTaskResources zeroes a
+// reservation pod's GPU index, so it contributes 0 GPUs; if the dedup still
+// tracked that pod's shared device it would subtract a device from a 0 vector
+// (driving UsedVector negative and corrupting node capacity) and inflate the
+// reference count, masking the real consumers. A reservation pod referencing
+// the same shared claim as a real consumer must leave the GPU used count at 1
+// on add and on remove, in any order.
+func TestAddRemoveTask_ReservationPodDoesNotCorruptSharedDRAAccounting(t *testing.T) {
+	ni, vectorMap := newGPUNodeInfo(t, "atlas", "1", 2, 2)
+
+	claim := sharedGPUClaim("voice-atlas-shared", "voice-pipeline", gpuDeviceClass, "atlas", "gpu-0")
+	consumer := draConsumerPod("voice-tts", "voice-pipeline", "atlas", "voice-atlas-shared")
+	reservation := draReservationPod("kai-resource-reservation-abc", "voice-pipeline", "atlas", "voice-atlas-shared")
+	consumerTask := pod_info.NewTaskInfo(consumer, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claim}})
+	reservationTask := pod_info.NewTaskInfo(reservation, vectorMap, pod_info.TaskInfoOptions{DraPodClaims: []*resourceapi.ResourceClaim{claim}})
+
+	assert.NoError(t, ni.AddTask(consumerTask))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"real consumer of the shared device must be counted once")
+
+	assert.NoError(t, ni.AddTask(reservationTask))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"reservation pod contributes no GPU and must not change the used count")
+
+	idleGPUs, _ := ni.GetSumOfIdleGPUs()
+	assert.Equal(t, 0.0, idleGPUs, "idle GPUs must stay 0, never negative")
+
+	assert.NoError(t, ni.RemoveTask(reservationTask))
+	assert.Equal(t, 1.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"removing the reservation pod must leave the real consumer's device counted")
+
+	assert.NoError(t, ni.RemoveTask(consumerTask))
+	assert.Equal(t, 0.0, ni.UsedVector.Get(resource_info.GPUIndex),
+		"removing the last real consumer must release the shared device")
+}

pkg/scheduler/api/node_info/gpu_sharing_node_info.go

@@ -21,6 +21,12 @@ type GpuSharingNodeInfo struct {
 	UsedSharedGPUsMemory      map[string]int64
 	ReleasingSharedGPUsMemory map[string]int64
 	AllocatedSharedGPUsMemory map[string]int64
+
+	// DRASharedDeviceRefCount counts how many pods on the node reference each
+	// physical DRA GPU device (keyed by driver/pool/device). A device shared
+	// by several pods through one ResourceClaim (status.reservedFor with more
+	// than one entry) must contribute to the node's used GPU count only once.
+	DRASharedDeviceRefCount map[string]int
 }
 
 func newGpuSharingNodeInfo() *GpuSharingNodeInfo {
@@ -30,6 +36,8 @@ func newGpuSharingNodeInfo() *GpuSharingNodeInfo {
 		UsedSharedGPUsMemory:      make(map[string]int64),
 		ReleasingSharedGPUsMemory: make(map[string]int64),
 		AllocatedSharedGPUsMemory: make(map[string]int64),
+
+		DRASharedDeviceRefCount: make(map[string]int),
 	}
 }
 
@@ -48,6 +56,9 @@ func (g *GpuSharingNodeInfo) Clone() *GpuSharingNodeInfo {
 	for k, v := range g.AllocatedSharedGPUsMemory {
 		gpuSharingNodeInfo.AllocatedSharedGPUsMemory[k] = v
 	}
+	for k, v := range g.DRASharedDeviceRefCount {
+		gpuSharingNodeInfo.DRASharedDeviceRefCount[k] = v
+	}
 
 	return gpuSharingNodeInfo
 }