From 8dc2871286de9c1ba4b0511a6b018a594ded9895 Mon Sep 17 00:00:00 2001 From: Raya Solano Date: Thu, 16 Jul 2026 07:19:09 +0000 Subject: [PATCH 1/2] 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 --- .../unreleased/Fixed-20260717-233428.yaml | 6 + .../api/node_info/dra_shared_device_info.go | 120 ++++++++++ .../node_info/dra_shared_device_info_test.go | 226 ++++++++++++++++++ .../api/node_info/gpu_sharing_node_info.go | 11 + pkg/scheduler/api/node_info/node_info.go | 15 ++ 5 files changed, 378 insertions(+) create mode 100644 .changes/unreleased/Fixed-20260717-233428.yaml create mode 100644 pkg/scheduler/api/node_info/dra_shared_device_info.go create mode 100644 pkg/scheduler/api/node_info/dra_shared_device_info_test.go diff --git a/.changes/unreleased/Fixed-20260717-233428.yaml b/.changes/unreleased/Fixed-20260717-233428.yaml new file mode 100644 index 000000000..a5f9b10a4 --- /dev/null +++ b/.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" diff --git a/pkg/scheduler/api/node_info/dra_shared_device_info.go b/pkg/scheduler/api/node_info/dra_shared_device_info.go new file mode 100644 index 000000000..3876eb086 --- /dev/null +++ b/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) + } +} diff --git a/pkg/scheduler/api/node_info/dra_shared_device_info_test.go b/pkg/scheduler/api/node_info/dra_shared_device_info_test.go new file mode 100644 index 000000000..1b418e68d --- /dev/null +++ b/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") +} diff --git a/pkg/scheduler/api/node_info/gpu_sharing_node_info.go b/pkg/scheduler/api/node_info/gpu_sharing_node_info.go index e0eff252c..4bf166b76 100644 --- a/pkg/scheduler/api/node_info/gpu_sharing_node_info.go +++ b/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 } diff --git a/pkg/scheduler/api/node_info/node_info.go b/pkg/scheduler/api/node_info/node_info.go index 447616a0a..582e6414f 100644 --- a/pkg/scheduler/api/node_info/node_info.go +++ b/pkg/scheduler/api/node_info/node_info.go @@ -472,6 +472,10 @@ func (ni *NodeInfo) addTaskResources(task *pod_info.PodInfo) { resourcesToTrackVector.Set(resource_info.GPUIndex, 0) } + // A physical DRA device shared by several pods (one ResourceClaim with + // multiple reservedFor entries) must be counted once, not once per pod. + ni.dedupSharedDRAGpus(task, resourcesToTrackVector) + // DRA-backed extended resources are absent from node.Status.Allocatable and must // not be charged against the node's vector — the DRA allocator handles them. for i := range len(resourcesToTrackVector) { @@ -530,6 +534,10 @@ func (ni *NodeInfo) removeTaskResources(task *pod_info.PodInfo) { resourcesToTrackVector.Set(resource_info.GPUIndex, 0) } + // Mirror of dedupSharedDRAGpus: keep a shared physical DRA device in the + // used vector as long as another pod on the node still references it. + ni.releaseSharedDRAGpus(task, resourcesToTrackVector) + // Mirror the zeroing done in addTaskResources so vectors stay consistent. for i := range len(resourcesToTrackVector) { if ni.AllocatableVector.Get(i) == 0 { @@ -758,6 +766,13 @@ func (ni *NodeInfo) lessEqualTaskToNodeResources( if !ni.isValidGpuPortion(&task.GpuRequirement) { return false } + // A task sharing an already-counted DRA device does not need additional + // GPU capacity for that device. + if discount := ni.sharedDRAGpuDiscount(task); discount > 0 { + adjusted := nodeResourcesVector.Clone() + adjusted.Set(resource_info.GPUIndex, adjusted.Get(resource_info.GPUIndex)+discount) + return task.ResReqVector.LessEqual(adjusted) + } return task.ResReqVector.LessEqual(nodeResourcesVector) } From 4e9cf1e3cbe1916955f50857bfffc18c55e53fed Mon Sep 17 00:00:00 2001 From: Erez Freiberger Date: Sat, 18 Jul 2026 00:57:53 +0200 Subject: [PATCH 2/2] add test cases Signed-off-by: Erez Freiberger Signed-off-by: Raya Solano --- .../allocate/allocate_shared_dra_test.go | 260 ++++++++++++++++++ 1 file changed, 260 insertions(+) create mode 100644 pkg/scheduler/actions/integration_tests/allocate/allocate_shared_dra_test.go diff --git a/pkg/scheduler/actions/integration_tests/allocate/allocate_shared_dra_test.go b/pkg/scheduler/actions/integration_tests/allocate/allocate_shared_dra_test.go new file mode 100644 index 000000000..8a53b5a4b --- /dev/null +++ b/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, + }, + }) +}