kubernetes/pkg/scheduler/framework/plugins/dynamicresources/extendeddynamicresources_test.go
Alay Patel f8ccc4c4d7 dra scheduler plugin: refactor extendeddynamicresources.go for readibility
Signed-off-by: Alay Patel <alayp@nvidia.com>
2025-11-06 15:49:33 -05:00

565 lines
19 KiB
Go

/*
Copyright 2025 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package dynamicresources
import (
"sort"
"testing"
v1 "k8s.io/api/core/v1"
resourceapi "k8s.io/api/resource/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
fwk "k8s.io/kube-scheduler/framework"
st "k8s.io/kubernetes/pkg/scheduler/testing"
"k8s.io/kubernetes/test/utils/ktesting"
)
func Test_createRequestsAndMappings_requests(t *testing.T) {
pod1 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
v1.ResourceName(extendedResourceName + "1"): "2",
}).
Obj()
pod2 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "1"): "2",
}).
Obj()
podInit := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
podInit2 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
podInit3 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
SidecarReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
res := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
}
res2 := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
v1.ResourceName(extendedResourceName + "1"): 2,
}
resInit := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
v1.ResourceName(extendedResourceName + "init"): 2,
}
devMap := map[v1.ResourceName]*resourceapi.DeviceClass{
v1.ResourceName(extendedResourceName): {
ObjectMeta: metav1.ObjectMeta{
Name: "class",
},
},
}
devMap2 := map[v1.ResourceName]*resourceapi.DeviceClass{
v1.ResourceName(extendedResourceName): {
ObjectMeta: metav1.ObjectMeta{
Name: "class",
},
},
v1.ResourceName(extendedResourceName + "1"): {
ObjectMeta: metav1.ObjectMeta{
Name: "class1",
},
},
}
devMapInit := map[v1.ResourceName]*resourceapi.DeviceClass{
v1.ResourceName(extendedResourceName): {
ObjectMeta: metav1.ObjectMeta{
Name: "class",
},
},
v1.ResourceName(extendedResourceName + "init"): {
ObjectMeta: metav1.ObjectMeta{
Name: "classInit",
},
},
}
devReq := resourceapi.DeviceRequest{
Name: "container-0-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 1,
},
}
devReq2 := resourceapi.DeviceRequest{
Name: "container-0-request-1",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class1",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 2,
},
}
devReq3 := resourceapi.DeviceRequest{
Name: "container-1-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class1",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 2,
},
}
devReqInit := resourceapi.DeviceRequest{
Name: "container-1-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 1,
},
}
devReqSidecar := resourceapi.DeviceRequest{
Name: "container-1-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "classInit",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 1,
},
}
devReq2Init := resourceapi.DeviceRequest{
Name: "container-1-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "classInit",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 2,
},
}
devReq6Init := resourceapi.DeviceRequest{
Name: "container-0-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "classInit",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 2,
},
}
devReq3Init := resourceapi.DeviceRequest{
Name: "container-2-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 1,
},
}
devReq4Init := resourceapi.DeviceRequest{
Name: "container-3-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "class",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 1,
},
}
devReq5Init := resourceapi.DeviceRequest{
Name: "container-2-request-0",
Exactly: &resourceapi.ExactDeviceRequest{
DeviceClassName: "classInit",
AllocationMode: resourceapi.DeviceAllocationModeExactCount,
Count: 2,
},
}
testcases := map[string]struct {
pod *v1.Pod
extendedResources map[v1.ResourceName]int64
cache fwk.DeviceClassResolver
wantDeviceRequests []resourceapi.DeviceRequest
}{
"nil": {
pod: pod1,
wantDeviceRequests: nil,
},
"one resource match": {
pod: pod1,
extendedResources: res,
cache: &mockDeviceClassResolver{mapping: devMap},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq},
},
"one resource match, one resource not match": {
pod: pod1,
extendedResources: res2,
cache: &mockDeviceClassResolver{mapping: devMap},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq},
},
"two resources match": {
pod: pod1,
extendedResources: res2,
cache: &mockDeviceClassResolver{mapping: devMap2},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq, devReq2},
},
"two containers match": {
pod: pod2,
extendedResources: res2,
cache: &mockDeviceClassResolver{mapping: devMap2},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq, devReq3},
},
"one init container, one regular container": {
pod: podInit,
extendedResources: resInit,
cache: &mockDeviceClassResolver{mapping: devMapInit},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq6Init, devReqInit},
},
"two init containers, one regular container": {
pod: podInit2,
extendedResources: resInit,
cache: &mockDeviceClassResolver{mapping: devMapInit},
wantDeviceRequests: []resourceapi.DeviceRequest{devReq2Init, devReq3Init},
},
"three init containers, one sidecar, one regular container": {
pod: podInit3,
extendedResources: resInit,
cache: &mockDeviceClassResolver{mapping: devMapInit},
wantDeviceRequests: []resourceapi.DeviceRequest{devReqSidecar, devReq5Init, devReq4Init},
},
}
for name, tc := range testcases {
t.Run(name, func(t *testing.T) {
logger, _ := ktesting.NewTestContext(t)
gotDeviceRequests, _ := createRequestsAndMappings(tc.pod, tc.extendedResources, logger, tc.cache)
if len(tc.wantDeviceRequests) != len(gotDeviceRequests) {
t.Fatalf("different length, want %#v, len=%v, got %#v, len=%v", tc.wantDeviceRequests, len(tc.wantDeviceRequests), gotDeviceRequests, len(gotDeviceRequests))
}
// gotDeviceRequests should already be sorted by createRequestsAndMappings
for i, r := range tc.wantDeviceRequests {
if r.Name != gotDeviceRequests[i].Name {
t.Errorf("different name, want %#v, got %#v", r, gotDeviceRequests[i])
}
if r.Exactly.DeviceClassName != gotDeviceRequests[i].Exactly.DeviceClassName {
t.Errorf("different deviceClassName, want %#v, got %#v", r, gotDeviceRequests[i])
}
if r.Exactly.AllocationMode != gotDeviceRequests[i].Exactly.AllocationMode {
t.Errorf("different allocationMode, want %#v, got %#v", r, gotDeviceRequests[i])
}
if r.Exactly.Count != gotDeviceRequests[i].Exactly.Count {
t.Errorf("different count, want %#v, got %#v", r.Exactly.Count, gotDeviceRequests[i].Exactly.Count)
}
}
})
}
}
func Test_createRequestsAndMappings_mappings(t *testing.T) {
pod1 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
v1.ResourceName(extendedResourceName + "1"): "2",
}).
Obj()
pod1InitImplicit := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
v1.ResourceName(resourceapi.ResourceDeviceClassPrefix + "classInit"): "2",
}).
Obj()
pod2 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "1"): "2",
}).
Obj()
podInit := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
podInit2 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
podInitImplicit := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(resourceapi.ResourceDeviceClassPrefix + "classInit"): "2",
}).
Obj()
podInit3 := st.MakePod().Name(podName).Namespace(namespace).
UID(podUID).
Res(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
SidecarReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "1",
}).
InitReq(map[v1.ResourceName]string{
v1.ResourceName(extendedResourceName + "init"): "2",
}).
Obj()
res := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
v1.ResourceName(extendedResourceName + "1"): 2,
}
resInit := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
v1.ResourceName(extendedResourceName + "init"): 2,
}
resInitImplicit := map[v1.ResourceName]int64{
v1.ResourceName(extendedResourceName): 1,
v1.ResourceName(extendedResourceName + "init"): 2,
v1.ResourceName(resourceapi.ResourceDeviceClassPrefix + "classInit"): 2,
}
devMap := map[v1.ResourceName]*resourceapi.DeviceClass{
v1.ResourceName(extendedResourceName): {
ObjectMeta: metav1.ObjectMeta{
Name: "class",
},
},
v1.ResourceName(extendedResourceName + "1"): {
ObjectMeta: metav1.ObjectMeta{
Name: "class1",
},
},
}
devMapInit := map[v1.ResourceName]*resourceapi.DeviceClass{
v1.ResourceName(extendedResourceName): {
ObjectMeta: metav1.ObjectMeta{
Name: "class",
},
},
v1.ResourceName(extendedResourceName + "init"): {
ObjectMeta: metav1.ObjectMeta{
Name: "classInit",
},
},
v1.ResourceName(resourceapi.ResourceDeviceClassPrefix + "classInit"): {
ObjectMeta: metav1.ObjectMeta{
Name: "classInit",
},
},
}
cer := v1.ContainerExtendedResourceRequest{
ContainerName: "con0",
ResourceName: extendedResourceName,
RequestName: "container-0-request-0",
}
cer1 := v1.ContainerExtendedResourceRequest{
ContainerName: "con0",
ResourceName: extendedResourceName + "1",
RequestName: "container-0-request-1",
}
cer2 := v1.ContainerExtendedResourceRequest{
ContainerName: "con1",
ResourceName: extendedResourceName + "1",
RequestName: "container-1-request-0",
}
cer3 := v1.ContainerExtendedResourceRequest{
ContainerName: "con0",
ResourceName: extendedResourceName,
RequestName: "container-1-request-0",
}
cer4 := v1.ContainerExtendedResourceRequest{
ContainerName: "con0",
ResourceName: extendedResourceName,
RequestName: "container-2-request-0",
}
cer5 := v1.ContainerExtendedResourceRequest{
ContainerName: "con0",
ResourceName: extendedResourceName,
RequestName: "container-3-request-0",
}
cerInit := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: extendedResourceName + "init",
RequestName: "container-1-request-0",
}
cerInit0 := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: extendedResourceName + "init",
RequestName: "container-0-request-0",
}
cerInit1 := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: extendedResourceName + "init",
RequestName: "container-1-request-0",
}
cerInit2 := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con1",
ResourceName: extendedResourceName + "init",
RequestName: "container-1-request-0",
}
cerInit3 := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con2",
ResourceName: extendedResourceName + "init",
RequestName: "container-2-request-0",
}
cerSidecar := v1.ContainerExtendedResourceRequest{
ContainerName: "sidecar-con1",
ResourceName: extendedResourceName + "init",
RequestName: "container-1-request-0",
}
cerInitImplicit := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: extendedResourceName + "init",
RequestName: "container-0-request-0",
}
cerInit4Implicit := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: extendedResourceName + "init",
RequestName: "container-0-request-1",
}
cerInit2Implicit := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con1",
ResourceName: resourceapi.ResourceDeviceClassPrefix + "classInit",
RequestName: "container-1-request-0",
}
cerInit3Implicit := v1.ContainerExtendedResourceRequest{
ContainerName: "init-con0",
ResourceName: resourceapi.ResourceDeviceClassPrefix + "classInit",
RequestName: "container-0-request-0",
}
testcases := map[string]struct {
pod *v1.Pod
extnededResources map[v1.ResourceName]int64
deviceClassMapping fwk.DeviceClassResolver
wantReqMappings []v1.ContainerExtendedResourceRequest
}{
"one container, two requests": {
pod: pod1,
extnededResources: res,
deviceClassMapping: &mockDeviceClassResolver{devMap},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cer, cer1},
},
"one container, one explicit and one implicit request": {
pod: pod1InitImplicit,
extnededResources: resInitImplicit,
deviceClassMapping: &mockDeviceClassResolver{devMapInit},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cerInit3Implicit, cerInit4Implicit},
},
"two containers, two requests": {
pod: pod2,
extnededResources: res,
deviceClassMapping: &mockDeviceClassResolver{devMap},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cer, cer2},
},
"one init container, one regular container, one request": {
pod: podInit,
extnededResources: resInit,
deviceClassMapping: &mockDeviceClassResolver{devMapInit},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cerInit0, cer3},
},
"three containers (two are init container), two requests": {
pod: podInit2,
extnededResources: resInit,
deviceClassMapping: &mockDeviceClassResolver{devMapInit},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cerInit, cerInit2, cer4},
},
"three containers (two are init container), both explicit and implicit resources": {
pod: podInitImplicit,
extnededResources: resInitImplicit,
deviceClassMapping: &mockDeviceClassResolver{devMapInit},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cerInitImplicit, cerInit2Implicit, cer4},
},
"four containers (two are init container, one sidecar), three requests": {
pod: podInit3,
extnededResources: resInit,
deviceClassMapping: &mockDeviceClassResolver{devMapInit},
wantReqMappings: []v1.ContainerExtendedResourceRequest{cerInit1, cerSidecar, cerInit3, cer5},
},
}
for name, tc := range testcases {
t.Run(name, func(t *testing.T) {
logger, _ := ktesting.NewTestContext(t)
_, gotReqMappings := createRequestsAndMappings(tc.pod, tc.extnededResources, logger, tc.deviceClassMapping)
if len(tc.wantReqMappings) != len(gotReqMappings) {
t.Fatalf("different length, want %#v, got %#v", tc.wantReqMappings, gotReqMappings)
}
sort.Slice(gotReqMappings, func(i, j int) bool {
if gotReqMappings[i].RequestName < gotReqMappings[j].RequestName {
return true
}
if gotReqMappings[i].RequestName > gotReqMappings[j].RequestName {
return false
}
return gotReqMappings[i].ContainerName < gotReqMappings[j].ContainerName
})
for i, r := range tc.wantReqMappings {
if r.RequestName != gotReqMappings[i].RequestName {
t.Errorf("different request name, want %#v, got %#v", r, gotReqMappings[i])
}
if r.ContainerName != gotReqMappings[i].ContainerName {
t.Errorf("different container name, want %#v, got %#v", r, gotReqMappings[i])
}
if r.ResourceName != gotReqMappings[i].ResourceName {
t.Errorf("different resource name, want %#v, got %#v", r, gotReqMappings[i])
}
}
})
}
}