Skip to content

Commit 95e496b

Browse files
KaiPilotBotRaya Solanoenoodle
authored
fix(scheduler): count shared DRA GPU device once per node (#2024)
Signed-off-by: Raya Solano <raya@mbinf.de> Signed-off-by: Erez Freiberger <enoodle@gmail.com> Signed-off-by: Raya Solano <raya.solano@mbinf.de> Co-authored-by: Raya Solano <raya@mbinf.de> Co-authored-by: Erez Freiberger <enoodle@gmail.com>
1 parent 3767c11 commit 95e496b

6 files changed

Lines changed: 638 additions & 0 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Fixed
2+
body: Count a GPU shared by multiple pods through one DRA ResourceClaim once per node, preventing negative idle GPUs.
3+
time: 2026-07-17T23:34:28.999387137Z
4+
custom:
5+
Author: TensorRaya
6+
Issue: "1930"

‎pkg/scheduler/actions/integration_tests/allocate/allocate_shared_dra_test.go‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
// Copyright 2026 NVIDIA CORPORATION
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package allocate
5+
6+
import (
7+
"testing"
8+
"time"
9+
10+
resourceapi "k8s.io/api/resource/v1"
11+
12+
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
13+
featuregates "github.com/kai-scheduler/KAI-scheduler/pkg/common/feature_gates"
14+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/integration_tests/integration_tests_utils"
15+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
16+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
17+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
18+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/dra_fake"
19+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
20+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
21+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/tasks_fake"
22+
)
23+
24+
func TestSharedDRADeviceDoesNotBlockCPUOnlyPod(t *testing.T) {
25+
featuregates.SetDynamicResourcesEnabledForTest(true)
26+
t.Cleanup(func() {
27+
featuregates.SetDynamicResourcesEnabledForTest(false)
28+
})
29+
30+
integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
31+
{
32+
Name: "shared DRA device does not block CPU-only pod",
33+
TestTopologyBasic: test_utils.TestTopologyBasic{
34+
Name: "shared DRA device does not block CPU-only pod",
35+
Jobs: []*jobs_fake.TestJobBasic{
36+
{
37+
Name: "shared_dra_job0",
38+
Namespace: "test",
39+
Priority: constants.PriorityTrainNumber,
40+
QueueName: "queue1",
41+
Tasks: []*tasks_fake.TestTaskBasic{
42+
{
43+
NodeName: "node0",
44+
State: pod_status.Running,
45+
ResourceClaimNames: []string{"shared-claim"},
46+
},
47+
},
48+
},
49+
{
50+
Name: "shared_dra_job1",
51+
Namespace: "test",
52+
Priority: constants.PriorityTrainNumber,
53+
QueueName: "queue1",
54+
Tasks: []*tasks_fake.TestTaskBasic{
55+
{
56+
NodeName: "node0",
57+
State: pod_status.Running,
58+
ResourceClaimNames: []string{"shared-claim"},
59+
},
60+
},
61+
},
62+
{
63+
Name: "cpu_only_job",
64+
Namespace: "test",
65+
Priority: constants.PriorityTrainNumber,
66+
QueueName: "queue1",
67+
Tasks: []*tasks_fake.TestTaskBasic{
68+
{
69+
State: pod_status.Pending,
70+
NodeAffinityNames: []string{"node0"},
71+
},
72+
},
73+
},
74+
},
75+
TestDRAObjects: dra_fake.TestDRAObjects{
76+
DeviceClasses: []string{"nvidia.com/gpu"},
77+
ResourceSlices: []*dra_fake.TestResourceSlice{
78+
{
79+
Name: "node0-gpu",
80+
DeviceClassName: "nvidia.com/gpu",
81+
NodeName: "node0",
82+
Count: 1,
83+
},
84+
},
85+
ResourceClaims: []*dra_fake.TestResourceClaim{
86+
{
87+
Name: "shared-claim",
88+
Namespace: "test",
89+
DeviceClassName: "nvidia.com/gpu",
90+
Count: 1,
91+
Labels: map[string]string{
92+
commonconstants.DefaultQueueLabel: "queue1",
93+
},
94+
ClaimStatus: &resourceapi.ResourceClaimStatus{
95+
Allocation: &resourceapi.AllocationResult{
96+
Devices: resourceapi.DeviceAllocationResult{
97+
Results: []resourceapi.DeviceRequestAllocationResult{
98+
{
99+
Request: "request",
100+
Driver: "nvidia.com/gpu",
101+
Pool: "node0",
102+
Device: "0",
103+
},
104+
},
105+
},
106+
},
107+
ReservedFor: []resourceapi.ResourceClaimConsumerReference{
108+
{Resource: "pods", Name: "shared_dra_job0-0", UID: "shared_dra_job0-0"},
109+
{Resource: "pods", Name: "shared_dra_job1-0", UID: "shared_dra_job1-0"},
110+
},
111+
},
112+
},
113+
},
114+
},
115+
Nodes: map[string]nodes_fake.TestNodeBasic{
116+
"node0": {},
117+
},
118+
Queues: []test_utils.TestQueueBasic{
119+
{
120+
Name: "queue1",
121+
DeservedGPUs: 1,
122+
},
123+
},
124+
JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
125+
"shared_dra_job0": {
126+
NodeName: "node0",
127+
Status: pod_status.Running,
128+
},
129+
"shared_dra_job1": {
130+
NodeName: "node0",
131+
Status: pod_status.Running,
132+
},
133+
"cpu_only_job": {
134+
NodeName: "node0",
135+
Status: pod_status.Running,
136+
},
137+
},
138+
Mocks: &test_utils.TestMock{
139+
CacheRequirements: &test_utils.CacheMocking{
140+
NumberOfCacheBinds: 1,
141+
},
142+
},
143+
},
144+
RoundsUntilMatch: 1,
145+
RoundsAfterMatch: 1,
146+
SchedulingDuration: time.Millisecond,
147+
},
148+
})
149+
}
150+
151+
func TestPendingPodCanUseSharedDRADevice(t *testing.T) {
152+
featuregates.SetDynamicResourcesEnabledForTest(true)
153+
t.Cleanup(func() {
154+
featuregates.SetDynamicResourcesEnabledForTest(false)
155+
})
156+
157+
integration_tests_utils.RunTests(t, []integration_tests_utils.TestTopologyMetadata{
158+
{
159+
Name: "pending pod can use shared DRA device",
160+
TestTopologyBasic: test_utils.TestTopologyBasic{
161+
Name: "pending pod can use shared DRA device",
162+
Jobs: []*jobs_fake.TestJobBasic{
163+
{
164+
Name: "running_shared_dra_job",
165+
Namespace: "test",
166+
Priority: constants.PriorityTrainNumber,
167+
QueueName: "queue1",
168+
Tasks: []*tasks_fake.TestTaskBasic{
169+
{
170+
NodeName: "node0",
171+
State: pod_status.Running,
172+
ResourceClaimNames: []string{"shared-claim"},
173+
},
174+
},
175+
},
176+
{
177+
Name: "pending_shared_dra_job",
178+
Namespace: "test",
179+
Priority: constants.PriorityTrainNumber,
180+
QueueName: "queue1",
181+
Tasks: []*tasks_fake.TestTaskBasic{
182+
{
183+
State: pod_status.Pending,
184+
NodeAffinityNames: []string{"node0"},
185+
ResourceClaimNames: []string{"shared-claim"},
186+
},
187+
},
188+
},
189+
},
190+
TestDRAObjects: dra_fake.TestDRAObjects{
191+
DeviceClasses: []string{"nvidia.com/gpu"},
192+
ResourceSlices: []*dra_fake.TestResourceSlice{
193+
{
194+
Name: "node0-gpu",
195+
DeviceClassName: "nvidia.com/gpu",
196+
NodeName: "node0",
197+
Count: 1,
198+
},
199+
},
200+
ResourceClaims: []*dra_fake.TestResourceClaim{
201+
{
202+
Name: "shared-claim",
203+
Namespace: "test",
204+
DeviceClassName: "nvidia.com/gpu",
205+
Count: 1,
206+
Labels: map[string]string{
207+
commonconstants.DefaultQueueLabel: "queue1",
208+
},
209+
ClaimStatus: &resourceapi.ResourceClaimStatus{
210+
Allocation: &resourceapi.AllocationResult{
211+
Devices: resourceapi.DeviceAllocationResult{
212+
Results: []resourceapi.DeviceRequestAllocationResult{
213+
{
214+
Request: "request",
215+
Driver: "nvidia.com/gpu",
216+
Pool: "node0",
217+
Device: "0",
218+
},
219+
},
220+
},
221+
},
222+
ReservedFor: []resourceapi.ResourceClaimConsumerReference{
223+
{Resource: "pods", Name: "running_shared_dra_job-0", UID: "running_shared_dra_job-0"},
224+
{Resource: "pods", Name: "pending_shared_dra_job-0", UID: "pending_shared_dra_job-0"},
225+
},
226+
},
227+
},
228+
},
229+
},
230+
Nodes: map[string]nodes_fake.TestNodeBasic{
231+
"node0": {},
232+
},
233+
Queues: []test_utils.TestQueueBasic{
234+
{
235+
Name: "queue1",
236+
DeservedGPUs: 1,
237+
},
238+
},
239+
JobExpectedResults: map[string]test_utils.TestExpectedResultBasic{
240+
"running_shared_dra_job": {
241+
NodeName: "node0",
242+
Status: pod_status.Running,
243+
},
244+
"pending_shared_dra_job": {
245+
NodeName: "node0",
246+
Status: pod_status.Running,
247+
},
248+
},
249+
Mocks: &test_utils.TestMock{
250+
CacheRequirements: &test_utils.CacheMocking{
251+
NumberOfCacheBinds: 1,
252+
},
253+
},
254+
},
255+
RoundsUntilMatch: 1,
256+
RoundsAfterMatch: 1,
257+
SchedulingDuration: time.Millisecond,
258+
},
259+
})
260+
}

‎pkg/scheduler/api/node_info/dra_shared_device_info.go‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
// Copyright 2025 NVIDIA CORPORATION
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package node_info
5+
6+
import (
7+
resourceapi "k8s.io/api/resource/v1"
8+
9+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
10+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
11+
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
12+
)
13+
14+
// draDeviceKey uniquely identifies a physical DRA device on the node.
15+
func draDeviceKey(result resourceapi.DeviceRequestAllocationResult) string {
16+
return result.Driver + "/" + result.Pool + "/" + result.Device
17+
}
18+
19+
// allocatedGPUDeviceKeys returns the keys of all GPU devices allocated to the
20+
// task via DRA ResourceClaims. Non-GPU devices are ignored: they are not part
21+
// of the GPU accounting that this dedup protects.
22+
func (ni *NodeInfo) allocatedGPUDeviceKeys(task *pod_info.PodInfo) []string {
23+
var keys []string
24+
for _, claimAllocation := range task.ResourceClaimInfo {
25+
if claimAllocation == nil || claimAllocation.Allocation == nil {
26+
continue
27+
}
28+
for _, result := range claimAllocation.Allocation.Devices.Results {
29+
if !resources.IsGPUDeviceClass(result.Driver) {
30+
continue
31+
}
32+
keys = append(keys, draDeviceKey(result))
33+
}
34+
}
35+
return keys
36+
}
37+
38+
// sharedDRAGpuDiscount returns the number of GPU devices the task requests via
39+
// DRA claims that are already counted on this node for other pods. A task that
40+
// shares an allocated device with a running pod does not need additional GPU
41+
// capacity for that device.
42+
func (ni *NodeInfo) sharedDRAGpuDiscount(task *pod_info.PodInfo) float64 {
43+
discount := 0.0
44+
for _, key := range ni.allocatedGPUDeviceKeys(task) {
45+
if ni.DRASharedDeviceRefCount[key] > 0 {
46+
discount++
47+
}
48+
}
49+
return discount
50+
}
51+
52+
// dedupSharedDRAGpus removes from resourcesToTrack the GPU count that would
53+
// double-count physical DRA devices already referenced by other pods on the
54+
// node. It also updates the node's per-device reference count. It must be
55+
// called once per addTaskResources, before the vector is added to UsedVector.
56+
func (ni *NodeInfo) dedupSharedDRAGpus(task *pod_info.PodInfo, resourcesToTrack resource_info.ResourceVector) {
57+
current := resourcesToTrack.Get(resource_info.GPUIndex)
58+
if current <= 0 {
59+
// The task contributes no GPUs to the used vector (e.g. a resource
60+
// reservation task whose GPU index was zeroed). Tracking its devices
61+
// would both risk a negative deduction below and mask the reference
62+
// count of the real consuming pods, so leave the accounting untouched.
63+
return
64+
}
65+
66+
alreadyCounted := 0.0
67+
for _, key := range ni.allocatedGPUDeviceKeys(task) {
68+
if ni.DRASharedDeviceRefCount[key] > 0 {
69+
// Another pod on this node already contributed this physical
70+
// device to the used vector: do not count it again.
71+
alreadyCounted++
72+
}
73+
ni.DRASharedDeviceRefCount[key]++
74+
}
75+
76+
if alreadyCounted > current {
77+
// Never deduct more than the task's own GPU contribution.
78+
alreadyCounted = current
79+
}
80+
if alreadyCounted > 0 {
81+
resourcesToTrack.Set(resource_info.GPUIndex, current-alreadyCounted)
82+
}
83+
}
84+
85+
// releaseSharedDRAGpus is the inverse of dedupSharedDRAGpus: it decrements the
86+
// per-device reference count and adds back the GPU count for devices that
87+
// remain referenced by other pods (and were therefore never subtracted on this
88+
// task's removal path). It must be called once per removeTaskResources.
89+
func (ni *NodeInfo) releaseSharedDRAGpus(task *pod_info.PodInfo, resourcesToTrack resource_info.ResourceVector) {
90+
current := resourcesToTrack.Get(resource_info.GPUIndex)
91+
if current <= 0 {
92+
// Mirror of dedupSharedDRAGpus: a task that contributed no GPUs never
93+
// incremented the reference count, so it must not decrement it here.
94+
return
95+
}
96+
97+
stillShared := 0.0
98+
for _, key := range ni.allocatedGPUDeviceKeys(task) {
99+
if ni.DRASharedDeviceRefCount[key] > 1 {
100+
// The device stays referenced by another pod after this removal:
101+
// it must remain in the used vector, so this task's removal must
102+
// not subtract it.
103+
stillShared++
104+
}
105+
if ni.DRASharedDeviceRefCount[key] > 0 {
106+
ni.DRASharedDeviceRefCount[key]--
107+
}
108+
if ni.DRASharedDeviceRefCount[key] == 0 {
109+
delete(ni.DRASharedDeviceRefCount, key)
110+
}
111+
}
112+
113+
if stillShared > current {
114+
// Never add back more than the task's own GPU contribution.
115+
stillShared = current
116+
}
117+
if stillShared > 0 {
118+
resourcesToTrack.Set(resource_info.GPUIndex, current-stillShared)
119+
}
120+
}

0 commit comments

Comments
 (0)