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

A single physical GPU shared by several pods through one ResourceClaim
(status.reservedFor with more than one entry) was counted once per
consuming pod in the node used vector. On a node whose shared device
pushed the used GPU count above physical capacity, IdleVector went
negative and the node became unschedulable for every task, including
GPU-less ones, surfacing as a misleading "didn't have enough resources:
GPUs" message.

Track a per-device reference count on the node (keyed by
driver/pool/device) so a shared physical device contributes to the used
GPU count exactly once, with symmetric add/remove accounting.

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

Author: Raya Solano

.changes/unreleased/0029-fixed.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.
+custom:
+  Issue: "1930"
+  Author: TensorRaya

pkg/common/constants/constants.go

@@ -9,6 +9,7 @@ const (
 	AppLabelName              = "app"
 	NvidiaGpuResource         = "nvidia.com/gpu"
 	NvidiaGpuMemory           = "nvidia.com/gpu.memory"
+	NvidiaGpuDraDriver        = "gpu.nvidia.com"
 	NvidiaMigResourcePrefix   = "nvidia.com/mig-"
 	GpuResource               = "gpu"
 	UnlimitedResourceQuantity = float64(-1)

pkg/scheduler/api/node_info/dra_shared_device_info.go

@@ -0,0 +1,122 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package node_info
+
+import (
+	resourceapi "k8s.io/api/resource/v1"
+
+	commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
+)
+
+// A single physical DRA device may be shared by several pods through one
+// ResourceClaim with more than one entry in status.reservedFor. Each such pod
+// carries the same allocated device in its ResourceClaimInfo, so counting the
+// device once per pod (the naive per-task accounting) inflates the node's used
+// GPU count above physical capacity and drives IdleVector negative, making the
+// whole node unschedulable. draSharedDeviceRefCount tracks how many pods on the
+// node currently reference each allocated device (keyed by driver/pool/device)
+// so a shared device contributes to UsedVector exactly once.
+
+// 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 !isGPUDRADriver(result.Driver) {
+				continue
+			}
+			keys = append(keys, draDeviceKey(result))
+		}
+	}
+	return keys
+}
+
+// isGPUDRADriver reports whether the DRA driver name belongs to an NVIDIA GPU
+// device. The device-class name is not present on the allocation result, so the
+// driver name is used instead.
+func isGPUDRADriver(driver string) bool {
+	return driver == commonconstants.NvidiaGpuDraDriver
+}
+
+// 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
 }