Skip to content

Commit

Permalink
poc with some tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@kannon92
kannon92 committed Sep 20, 2024
1 parent ca7265f commit de0e60d
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 46 deletions.
16 changes: 16 additions & 0 deletions pkg/util/testing/wrappers.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,22 @@ func (p *PodSetWrapper) Limit(r corev1.ResourceName, q string) *PodSetWrapper {
return p
}

func (p *PodSetWrapper) Claim(claim corev1.ResourceClaim) *PodSetWrapper {
if p.Template.Spec.Containers[0].Resources.Claims == nil {
p.Template.Spec.Containers[0].Resources.Claims = []corev1.ResourceClaim{}
}
p.Template.Spec.Containers[0].Resources.Claims = append(p.Template.Spec.Containers[0].Resources.Claims, claim)
return p
}

func (p *PodSetWrapper) ResourceClaim(resourceClaim corev1.PodResourceClaim) *PodSetWrapper {
if p.Template.Spec.ResourceClaims == nil {
p.Template.Spec.ResourceClaims = []corev1.PodResourceClaim{}
}
p.Template.Spec.ResourceClaims = append(p.Template.Spec.ResourceClaims, resourceClaim)
return p
}

func (p *PodSetWrapper) Image(image string) *PodSetWrapper {
p.Template.Spec.Containers[0].Image = image
return p
Expand Down
136 changes: 90 additions & 46 deletions pkg/workload/resources_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,12 @@ import (
nodev1 "k8s.io/api/node/v1"
dra "k8s.io/api/resource/v1alpha3"
"k8s.io/apimachinery/pkg/api/resource"
v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/utils/ptr"

kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
"sigs.k8s.io/kueue/pkg/controller/core/indexer"
"sigs.k8s.io/kueue/pkg/features"
"sigs.k8s.io/kueue/pkg/resources"
utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
)
Expand Down Expand Up @@ -486,67 +489,108 @@ func TestAdjustResources(t *testing.T) {
func TestAddDeviceClassesToContainerRequests(t *testing.T) {
cases := map[string]struct {
wl *kueue.Workload
resourceClaimTemplate dra.ResourceClaimTemplate
enableDRAGate bool
resourceClaimTemplate []dra.ResourceClaimTemplate
wantWl *kueue.Workload
}{
"single device class request": {
"dra feature gate off; ignore devices": {
enableDRAGate: false,
resourceClaimTemplate: []dra.ResourceClaimTemplate{
{
ObjectMeta: v1.ObjectMeta{
Name: "single-gpu",
Namespace: "",
},
Spec: dra.ResourceClaimTemplateSpec{
Spec: dra.ResourceClaimSpec{
Devices: dra.DeviceClaim{
Requests: []dra.DeviceRequest{{
Name: "single-gpu",
DeviceClassName: "gpu.example.com",
}},
},
},
},
},
},
wl: utiltesting.MakeWorkload("foo", "").
PodSets(
*utiltesting.MakePodSet("a", 1).
InitContainers(corev1.Container{}).
Obj(),
*utiltesting.MakePodSet("b", 1).
InitContainers(corev1.Container{}).
Limit(corev1.ResourceCPU, "6").
Obj(),
*utiltesting.MakePodSet("c", 1).
InitContainers(corev1.Container{}).
Limit(corev1.ResourceCPU, "2").
Request(corev1.ResourceCPU, "1").
Claim(corev1.ResourceClaim{
Name: "gpu",
}).
ResourceClaim(corev1.PodResourceClaim{
Name: "gpu",
ResourceClaimTemplateName: ptr.To("single-gpu"),
}).
Obj(),
).
Obj(),
wantWl: utiltesting.MakeWorkload("foo", "").
PodSets(
*utiltesting.MakePodSet("a", 1).
Limit(corev1.ResourceCPU, "4").
Request(corev1.ResourceCPU, "3").
InitContainers(corev1.Container{
Resources: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(4, resource.DecimalSI),
},
Requests: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(3, resource.DecimalSI),
},
},
Limit(corev1.ResourceCPU, "2").
Request(corev1.ResourceCPU, "1").
Claim(corev1.ResourceClaim{
Name: "gpu",
}).
ResourceClaim(corev1.PodResourceClaim{
Name: "gpu",
ResourceClaimTemplateName: ptr.To("single-gpu"),
}).
Obj(),
*utiltesting.MakePodSet("b", 1).
Limit(corev1.ResourceCPU, "6").
Request(corev1.ResourceCPU, "3").
InitContainers(corev1.Container{
Resources: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(4, resource.DecimalSI),
},
Requests: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(3, resource.DecimalSI),
},
).
Obj(),
},
"single device class request in a container": {
enableDRAGate: true,
resourceClaimTemplate: []dra.ResourceClaimTemplate{
{
ObjectMeta: v1.ObjectMeta{
Name: "single-gpu",
Namespace: "",
},
Spec: dra.ResourceClaimTemplateSpec{
Spec: dra.ResourceClaimSpec{
Devices: dra.DeviceClaim{
Requests: []dra.DeviceRequest{{
Name: "single-gpu",
DeviceClassName: "gpu.example.com",
}},
},
},
},
},
},
wl: utiltesting.MakeWorkload("foo", "").
PodSets(
*utiltesting.MakePodSet("a", 1).
Limit(corev1.ResourceCPU, "2").
Request(corev1.ResourceCPU, "1").
Claim(corev1.ResourceClaim{
Name: "gpu",
}).
ResourceClaim(corev1.PodResourceClaim{
Name: "gpu",
ResourceClaimTemplateName: ptr.To("single-gpu"),
}).
Obj(),
*utiltesting.MakePodSet("c", 1).
Limit(corev1.ResourceCPU, "4").
).
Obj(),
wantWl: utiltesting.MakeWorkload("foo", "").
PodSets(
*utiltesting.MakePodSet("a", 1).
Limit(corev1.ResourceCPU, "2").
Request(corev1.ResourceCPU, "1").
InitContainers(corev1.Container{
Resources: corev1.ResourceRequirements{
Limits: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(4, resource.DecimalSI),
},
Requests: corev1.ResourceList{
corev1.ResourceCPU: *resource.NewQuantity(3, resource.DecimalSI),
},
},
Request("gpu.example.com", "1").
Claim(corev1.ResourceClaim{
Name: "gpu",
}).
ResourceClaim(corev1.PodResourceClaim{
Name: "gpu",
ResourceClaimTemplateName: ptr.To("single-gpu"),
}).
Obj(),
).
Expand All @@ -556,11 +600,11 @@ func TestAddDeviceClassesToContainerRequests(t *testing.T) {
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
cl := utiltesting.NewClientBuilder().WithLists(
&nodev1.RuntimeClassList{Items: tc.runtimeClasses},
&corev1.LimitRangeList{Items: tc.limitranges},
).WithIndex(&corev1.LimitRange{}, indexer.LimitRangeHasContainerType, indexer.IndexLimitRangeHasContainerType).
&dra.ResourceClaimTemplateList{Items: tc.resourceClaimTemplate},
).
Build()
ctx, _ := utiltesting.ContextWithLog(t)
features.SetFeatureGateDuringTest(t, features.DynamicResourceStructuredParameters, tc.enableDRAGate)
AddDeviceClassesToContainerRequests(ctx, cl, tc.wl)
if diff := cmp.Diff(tc.wl, tc.wantWl); diff != "" {
t.Errorf("Unexpected resources after adjusting (-want,+got): %s", diff)
Expand Down

0 comments on commit de0e60d

Please sign in to comment.