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