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

.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/actions/integration_tests/allocate/allocate_shared_dra_test.go

@@ -0,0 +1,260 @@
+// Copyright 2026 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package allocate
+
+import (
+	"testing"
+	"time"
+
+	resourceapi "k8s.io/api/resource/v1"
+
+	commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+	featuregates "github.com/kai-scheduler/KAI-scheduler/pkg/common/feature_gates"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/dra_fake"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/tasks_fake"
+)
+
+func TestSharedDRADeviceDoesNotBlockCPUOnlyPod(t *testing.T) {
+	featuregates.SetDynamicResourcesEnabledForTest(true)
+	t.Cleanup(func() {
+		featuregates.SetDynamicResourcesEnabledForTest(false)
+	})
+
+	integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
+		{
+			Name: "shared DRA device does not block CPU-only pod",
+			TestTopologyBasic: test_utils.TestTopologyBasic{
+				Name: "shared DRA device does not block CPU-only pod",
+				Jobs: []*jobs_fake.TestJobBasic{
+					{
+						Name:      "shared_dra_job0",
+						Namespace: "test",
+						Priority:  constants.PriorityTrainNumber,
+						QueueName: "queue1",
+						Tasks: []*tasks_fake.TestTaskBasic{
+							{
+								NodeName:           "node0",
+								State:              pod_status.Running,
+								ResourceClaimNames: []string{"shared-claim"},
+							},
+						},
+					},
+					{
+						Name:      "shared_dra_job1",
+						Namespace: "test",
+						Priority:  constants.PriorityTrainNumber,
+						QueueName: "queue1",
+						Tasks: []*tasks_fake.TestTaskBasic{
+							{
+								NodeName:           "node0",
+								State:              pod_status.Running,
+								ResourceClaimNames: []string{"shared-claim"},
+							},
+						},
+					},
+					{
+						Name:      "cpu_only_job",
+						Namespace: "test",
+						Priority:  constants.PriorityTrainNumber,
+						QueueName: "queue1",
+						Tasks: []*tasks_fake.TestTaskBasic{
+							{
+								State:             pod_status.Pending,
+								NodeAffinityNames: []string{"node0"},
+							},
+						},
+					},
+				},
+				TestDRAObjects: dra_fake.TestDRAObjects{
+					DeviceClasses: []string{"nvidia.com/gpu"},
+					ResourceSlices: []*dra_fake.TestResourceSlice{
+						{
+							Name:            "node0-gpu",
+							DeviceClassName: "nvidia.com/gpu",
+							NodeName:        "node0",
+							Count:           1,
+						},
+					},
+					ResourceClaims: []*dra_fake.TestResourceClaim{
+						{
+							Name:            "shared-claim",
+							Namespace:       "test",
+							DeviceClassName: "nvidia.com/gpu",
+							Count:           1,
+							Labels: map[string]string{
+								commonconstants.DefaultQueueLabel: "queue1",
+							},
+							ClaimStatus: &resourceapi.ResourceClaimStatus{
+								Allocation: &resourceapi.AllocationResult{
+									Devices: resourceapi.DeviceAllocationResult{
+										Results: []resourceapi.DeviceRequestAllocationResult{
+											{
+												Request: "request",
+												Driver:  "nvidia.com/gpu",
+												Pool:    "node0",
+												Device:  "0",
+											},
+										},
+									},
+								},
+								ReservedFor: []resourceapi.ResourceClaimConsumerReference{
+									{Resource: "pods", Name: "shared_dra_job0-0", UID: "shared_dra_job0-0"},
+									{Resource: "pods", Name: "shared_dra_job1-0", UID: "shared_dra_job1-0"},
+								},
+							},
+						},
+					},
+				},
+				Nodes: map[string]nodes_fake.TestNodeBasic{
+					"node0": {},
+				},
+				Queues: []test_utils.TestQueueBasic{
+					{
+						Name:         "queue1",
+						DeservedGPUs: 1,
+					},
+				},
+				JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
+					"shared_dra_job0": {
+						NodeName: "node0",
+						Status:   pod_status.Running,
+					},
+					"shared_dra_job1": {
+						NodeName: "node0",
+						Status:   pod_status.Running,
+					},
+					"cpu_only_job": {
+						NodeName: "node0",
+						Status:   pod_status.Running,
+					},
+				},
+				Mocks: &test_utils.TestMock{
+					CacheRequirements: &test_utils.CacheMocking{
+						NumberOfCacheBinds: 1,
+					},
+				},
+			},
+			RoundsUntilMatch:   1,
+			RoundsAfterMatch:   1,
+			SchedulingDuration: time.Millisecond,
+		},
+	})
+}
+
+func TestPendingPodCanUseSharedDRADevice(t *testing.T) {
+	featuregates.SetDynamicResourcesEnabledForTest(true)
+	t.Cleanup(func() {
+		featuregates.SetDynamicResourcesEnabledForTest(false)
+	})
+
+	integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
+		{
+			Name: "pending pod can use shared DRA device",
+			TestTopologyBasic: test_utils.TestTopologyBasic{
+				Name: "pending pod can use shared DRA device",
+				Jobs: []*jobs_fake.TestJobBasic{
+					{
+						Name:      "running_shared_dra_job",
+						Namespace: "test",
+						Priority:  constants.PriorityTrainNumber,
+						QueueName: "queue1",
+						Tasks: []*tasks_fake.TestTaskBasic{
+							{
+								NodeName:           "node0",
+								State:              pod_status.Running,
+								ResourceClaimNames: []string{"shared-claim"},
+							},
+						},
+					},
+					{
+						Name:      "pending_shared_dra_job",
+						Namespace: "test",
+						Priority:  constants.PriorityTrainNumber,
+						QueueName: "queue1",
+						Tasks: []*tasks_fake.TestTaskBasic{
+							{
+								State:              pod_status.Pending,
+								NodeAffinityNames:  []string{"node0"},
+								ResourceClaimNames: []string{"shared-claim"},
+							},
+						},
+					},
+				},
+				TestDRAObjects: dra_fake.TestDRAObjects{
+					DeviceClasses: []string{"nvidia.com/gpu"},
+					ResourceSlices: []*dra_fake.TestResourceSlice{
+						{
+							Name:            "node0-gpu",
+							DeviceClassName: "nvidia.com/gpu",
+							NodeName:        "node0",
+							Count:           1,
+						},
+					},
+					ResourceClaims: []*dra_fake.TestResourceClaim{
+						{
+							Name:            "shared-claim",
+							Namespace:       "test",
+							DeviceClassName: "nvidia.com/gpu",
+							Count:           1,
+							Labels: map[string]string{
+								commonconstants.DefaultQueueLabel: "queue1",
+							},
+							ClaimStatus: &resourceapi.ResourceClaimStatus{
+								Allocation: &resourceapi.AllocationResult{
+									Devices: resourceapi.DeviceAllocationResult{
+										Results: []resourceapi.DeviceRequestAllocationResult{
+											{
+												Request: "request",
+												Driver:  "nvidia.com/gpu",
+												Pool:    "node0",
+												Device:  "0",
+											},
+										},
+									},
+								},
+								ReservedFor: []resourceapi.ResourceClaimConsumerReference{
+									{Resource: "pods", Name: "running_shared_dra_job-0", UID: "running_shared_dra_job-0"},
+									{Resource: "pods", Name: "pending_shared_dra_job-0", UID: "pending_shared_dra_job-0"},
+								},
+							},
+						},
+					},
+				},
+				Nodes: map[string]nodes_fake.TestNodeBasic{
+					"node0": {},
+				},
+				Queues: []test_utils.TestQueueBasic{
+					{
+						Name:         "queue1",
+						DeservedGPUs: 1,
+					},
+				},
+				JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
+					"running_shared_dra_job": {
+						NodeName: "node0",
+						Status:   pod_status.Running,
+					},
+					"pending_shared_dra_job": {
+						NodeName: "node0",
+						Status:   pod_status.Running,
+					},
+				},
+				Mocks: &test_utils.TestMock{
+					CacheRequirements: &test_utils.CacheMocking{
+						NumberOfCacheBinds: 1,
+					},
+				},
+			},
+			RoundsUntilMatch:   1,
+			RoundsAfterMatch:   1,
+			SchedulingDuration: time.Millisecond,
+		},
+	})
+}

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)
+	}
+}