116 lines
3.8 KiB
Go
116 lines
3.8 KiB
Go
/*
|
|
Copyright 2018 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 tpu
|
|
|
|
import (
|
|
"reflect"
|
|
"testing"
|
|
|
|
"k8s.io/apimachinery/pkg/api/resource"
|
|
. "k8s.io/autoscaler/cluster-autoscaler/utils/test"
|
|
|
|
apiv1 "k8s.io/api/core/v1"
|
|
)
|
|
|
|
var (
|
|
ResourceTPUV2 = ResourceTPUPrefix + "v2"
|
|
ResourceTPUPreemptibleV2 = ResourceTPUPrefix + "preemptible-v2"
|
|
)
|
|
|
|
type requests map[apiv1.ResourceName]int64
|
|
type containerSpecs []requests
|
|
|
|
func testContainer(requests requests) apiv1.Container {
|
|
container := apiv1.Container{
|
|
Resources: apiv1.ResourceRequirements{
|
|
Requests: apiv1.ResourceList{},
|
|
},
|
|
}
|
|
for res, request := range requests {
|
|
container.Resources.Requests[res] = *resource.NewQuantity(request, resource.DecimalSI)
|
|
}
|
|
return container
|
|
}
|
|
|
|
func testPod(name string, containers ...requests) *apiv1.Pod {
|
|
pod := BuildTestPod(name, 0, 0)
|
|
for _, requests := range containers[1:] {
|
|
pod.Spec.Containers = append(pod.Spec.Containers, testContainer(requests))
|
|
}
|
|
return pod
|
|
}
|
|
|
|
func TestClearTPURequests(t *testing.T) {
|
|
cpuPod := testPod("cpuPod", requests{apiv1.ResourceCPU: 10})
|
|
memoryPod := testPod("memoryPod", requests{apiv1.ResourceMemory: 100})
|
|
cpuMemoryPod := testPod("cpuMemoryPod", requests{apiv1.ResourceCPU: 10, apiv1.ResourceMemory: 30}, requests{apiv1.ResourceMemory: 20})
|
|
tpuPod := testPod("tpuPod", requests{apiv1.ResourceName(ResourceTPUV2): 1})
|
|
sanitizedTPUPod := testPod("tpuPod", requests{})
|
|
preemptibleTPUPod := testPod("preemptibleTPUPod", requests{apiv1.ResourceName(ResourceTPUPreemptibleV2): 1})
|
|
sanitizedPreemptibleTPUPod := testPod("preemptibleTPUPod", requests{})
|
|
tpuMemoryPod := testPod("tpuMemoryPod", requests{apiv1.ResourceName(ResourceTPUV2): 1, apiv1.ResourceMemory: 30}, requests{apiv1.ResourceName(ResourceTPUV2): 2, apiv1.ResourceMemory: 13})
|
|
sanitizedTPUMemoryPod := testPod("tpuMemoryPod", requests{apiv1.ResourceMemory: 30}, requests{apiv1.ResourceMemory: 13})
|
|
|
|
podsWithoutTPUs := []*apiv1.Pod{cpuPod, memoryPod, cpuMemoryPod}
|
|
mixedPods := []*apiv1.Pod{cpuPod, tpuPod, preemptibleTPUPod, memoryPod}
|
|
sanitizedMixedPods := []*apiv1.Pod{cpuPod, sanitizedTPUPod, sanitizedPreemptibleTPUPod, memoryPod}
|
|
podsWithTPUs := []*apiv1.Pod{tpuPod, preemptibleTPUPod, tpuMemoryPod}
|
|
sanitizedPodsWithTPUs := []*apiv1.Pod{sanitizedTPUPod, sanitizedPreemptibleTPUPod, sanitizedTPUMemoryPod}
|
|
|
|
testCases := []struct {
|
|
desc string
|
|
pods []*apiv1.Pod
|
|
expected []*apiv1.Pod
|
|
}{
|
|
{"Empty pod list", []*apiv1.Pod{}, []*apiv1.Pod{}},
|
|
{
|
|
desc: "Pods without TPU request",
|
|
pods: podsWithoutTPUs,
|
|
expected: podsWithoutTPUs,
|
|
},
|
|
{
|
|
desc: "Pods with & without TPU reqest",
|
|
pods: mixedPods,
|
|
expected: sanitizedMixedPods,
|
|
},
|
|
{
|
|
desc: "Just TPU pods",
|
|
pods: podsWithTPUs,
|
|
expected: sanitizedPodsWithTPUs,
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.desc, func(t *testing.T) {
|
|
copied := make([]*apiv1.Pod, len(tc.pods))
|
|
for i, pod := range tc.pods {
|
|
copied[i] = pod.DeepCopy()
|
|
}
|
|
actual := ClearTPURequests(tc.pods)
|
|
// Verify that output is as expected.
|
|
if !reflect.DeepEqual(actual, tc.expected) {
|
|
t.Errorf("Pod list should be as expected, got: %v, want: %v", actual, tc.expected)
|
|
}
|
|
// Verify that original pods weren't modified.
|
|
if !reflect.DeepEqual(tc.pods, copied) {
|
|
t.Errorf("Original pod list shouldn't be modified, got: %v, want: %v", tc.pods, copied)
|
|
}
|
|
})
|
|
}
|
|
|
|
}
|