diff --git a/CHANGELOG.md b/CHANGELOG.md index 42ae5a6824..ff4a6f3b82 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -408,5 +408,7 @@ Add "NVIDIA_VISIBLE_DEVICES=none" to none-gpu tasks - Fix initialization error when using tensor parallelism on vLLM above 0.18 - Fix multiple device typos - Add unit test coverage for node discovery handshake parsing with malformed or empty annotations +- Fix webhook mutating requests for pods with initContainers and multiple application containers + diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/cdi/api_mock.go b/pkg/device-plugin/nvidiadevice/nvinternal/cdi/api_mock.go index fd1d2d6fb0..b1d28e7745 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/cdi/api_mock.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/cdi/api_mock.go @@ -1,3 +1,19 @@ +/* +Copyright 2024 The HAMi 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. +*/ + // Code generated by moq; DO NOT EDIT. // github.com/matryer/moq diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go index 77465e444e..6bdda71e6a 100644 --- a/pkg/device/nvidia/device_test.go +++ b/pkg/device/nvidia/device_test.go @@ -1467,6 +1467,9 @@ func TestFitQuota(t *testing.T) { device.DevicesMap[NvidiaGPUDevice] = dev qm := device.NewQuotaManager() + t.Cleanup(func() { + delete(qm.Quotas, "default") + }) qm.AddQuota(&corev1.ResourceQuota{ TypeMeta: metav1.TypeMeta{ APIVersion: "v1", diff --git a/pkg/device/quota.go b/pkg/device/quota.go index 9b12f5f3f4..355110d556 100644 --- a/pkg/device/quota.go +++ b/pkg/device/quota.go @@ -45,7 +45,7 @@ type QuotaManager struct { var localCache QuotaManager func GetLocalCache() *QuotaManager { - return &localCache + return NewQuotaManager() } var once sync.Once @@ -123,6 +123,9 @@ func (q *QuotaManager) AddUsage(pod *corev1.Pod, podDev PodDevices) { } q.mutex.Lock() defer q.mutex.Unlock() + if q.Quotas == nil { + q.Quotas = make(map[string]*DeviceQuota) + } if q.Quotas[pod.Namespace] == nil { q.Quotas[pod.Namespace] = &DeviceQuota{} } @@ -232,6 +235,9 @@ func (q *QuotaManager) addQuotaLocked(quota *corev1.ResourceQuota) { if !ok { continue } + if q.Quotas == nil { + q.Quotas = make(map[string]*DeviceQuota) + } if q.Quotas[quota.Namespace] == nil { q.Quotas[quota.Namespace] = &DeviceQuota{} } diff --git a/pkg/scheduler/webhook.go b/pkg/scheduler/webhook.go index 6cc1163303..47a59e6884 100644 --- a/pkg/scheduler/webhook.go +++ b/pkg/scheduler/webhook.go @@ -70,6 +70,17 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res klog.V(5).Infof(template, pod.Namespace, pod.Name, pod.UID) privilegedName, hasPrivileged := privilegedContainerName(pod) hasResource := false + for idx := range pod.Spec.InitContainers { + c := &pod.Spec.InitContainers[idx] + for _, val := range device.GetDevices() { + found, err := val.MutateAdmission(c, pod) + if err != nil { + klog.Errorf("validating pod failed:%s", err.Error()) + return admission.Errored(http.StatusInternalServerError, err) + } + hasResource = hasResource || found + } + } for idx := range pod.Spec.Containers { c := &pod.Spec.Containers[idx] for _, val := range device.GetDevices() { @@ -139,6 +150,21 @@ func fitResourceQuota(pod *corev1.Pod) bool { // container spec here. It applies its own memory factor, defaults and // template rounding, which is what the scheduler later records as used, // so this keeps admission and the scheduler on the same numbers. + var initMemoryReqMax, initCoresReqMax int64 + for i := range pod.Spec.InitContainers { + req := dev.GenerateResourceRequests(&pod.Spec.InitContainers[i]) + if req.Nums == 0 { + continue + } + mem := int64(req.Memreq) * int64(req.Nums) + cores := int64(req.Coresreq) * int64(req.Nums) + if mem > initMemoryReqMax { + initMemoryReqMax = mem + } + if cores > initCoresReqMax { + initCoresReqMax = cores + } + } var memoryReq, coresReq int64 for i := range pod.Spec.Containers { req := dev.GenerateResourceRequests(&pod.Spec.Containers[i]) @@ -148,6 +174,12 @@ func fitResourceQuota(pod *corev1.Pod) bool { memoryReq += int64(req.Memreq) * int64(req.Nums) coresReq += int64(req.Coresreq) * int64(req.Nums) } + if initMemoryReqMax > memoryReq { + memoryReq = initMemoryReqMax + } + if initCoresReqMax > coresReq { + coresReq = initCoresReqMax + } if memoryReq == 0 && coresReq == 0 { continue } diff --git a/pkg/scheduler/webhook_test.go b/pkg/scheduler/webhook_test.go index 10762221b4..5607032f74 100644 --- a/pkg/scheduler/webhook_test.go +++ b/pkg/scheduler/webhook_test.go @@ -276,6 +276,9 @@ func TestFitResourceQuota(t *testing.T) { memName: &device.Quota{Used: 1000, Limit: 2000, LimitSet: true}, coreName: &device.Quota{Used: 200, Limit: 400, LimitSet: true}, } + t.Cleanup(func() { + delete(qm.Quotas, ns) + }) testCases := []struct { name string @@ -1070,3 +1073,371 @@ func TestPrivilegedContainerDenied(t *testing.T) { }) } } + +func TestMutateAdmissionMultiContainerAndInitContainers(t *testing.T) { + config.SchedulerName = "hami-scheduler" + sConfig := &config.Config{ + NvidiaConfig: nvidia.NvidiaConfig{ + ResourceCountName: "nvidia.com/gpu", + ResourceMemoryName: "nvidia.com/gpumem", + ResourceMemoryPercentageName: "nvidia.com/gpumem-percentage", + ResourceCoreName: "nvidia.com/gpucores", + DefaultMemory: 0, + DefaultCores: 0, + DefaultGPUNum: 1, + }, + } + + if err := config.InitDevicesWithConfig(sConfig); err != nil { + t.Fatalf("Failed to initialize devices with config: %v", err) + } + + wh, err := NewWebHook() + if err != nil { + t.Fatalf("Error creating WebHook: %v", err) + } + + scheme := runtime.NewScheme() + corev1.AddToScheme(scheme) + codec := serializer.NewCodecFactory(scheme).LegacyCodec(corev1.SchemeGroupVersion) + + t.Run("initContainer and multi-container pod patch validation", func(t *testing.T) { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "multi-container-init-pod", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-downloader", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-main", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + }, + }, + }, + { + Name: "sidecar-logging", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{}, + }, + }, + }, + }, + } + + podBytes, err := runtime.Encode(codec, pod) + if err != nil { + t.Fatalf("Error encoding pod: %v", err) + } + + req := admission.Request{ + AdmissionRequest: admissionv1.AdmissionRequest{ + UID: "multi-container-uid", + Namespace: pod.Namespace, + Name: pod.Name, + Object: runtime.RawExtension{ + Raw: podBytes, + }, + }, + } + + resp := wh.Handle(context.Background(), req) + if !resp.Allowed { + t.Fatalf("Expected allowed response for multi-container pod with initContainers, got denied: %+v", resp.Result) + } + + if len(resp.Patches) == 0 { + t.Fatalf("Expected JSON patches to be generated for multi-container pod, but got 0 patches") + } + + // Verify JSON patches target both initContainers and containers appropriately + initContainerPatched := false + mainContainerPatched := false + for _, patch := range resp.Patches { + if strings.HasPrefix(patch.Path, "/spec/initContainers/0") { + initContainerPatched = true + } + if strings.HasPrefix(patch.Path, "/spec/containers/0") { + mainContainerPatched = true + } + } + + if !initContainerPatched { + t.Errorf("Expected JSON patches for initContainers[0], but none found in patches: %+v", resp.Patches) + } + if !mainContainerPatched { + t.Errorf("Expected JSON patches for containers[0], but none found in patches: %+v", resp.Patches) + } + }) +} + +func TestFitResourceQuotaInitContainers(t *testing.T) { + sConfig := &config.Config{ + NvidiaConfig: nvidia.NvidiaConfig{ + ResourceCountName: "nvidia.com/gpu", + ResourceMemoryName: "nvidia.com/gpumem", + ResourceMemoryPercentageName: "nvidia.com/gpumem-percentage", + ResourceCoreName: "nvidia.com/gpucores", + DefaultMemory: 0, + DefaultCores: 0, + DefaultGPUNum: 1, + }, + } + + if err := config.InitDevicesWithConfig(sConfig); err != nil { + t.Fatalf("Failed to initialize devices with config: %v", err) + } + + t.Run("init container requests larger than container sum", func(t *testing.T) { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "init-larger-pod", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-large", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("2000"), + "nvidia.com/gpucores": resource.MustParse("50"), + }, + }, + }, + { + Name: "init-small", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("1000"), + "nvidia.com/gpucores": resource.MustParse("20"), + }, + }, + }, + { + Name: "init-no-gpu", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{}, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-main", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("500"), + "nvidia.com/gpucores": resource.MustParse("10"), + }, + }, + }, + { + Name: "app-no-gpu", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{}, + }, + }, + }, + }, + } + + if !fitResourceQuota(pod) { + t.Errorf("Expected fitResourceQuota to return true for pod under no quota limit") + } + }) + + t.Run("container sum larger than init container max", func(t *testing.T) { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "container-larger-pod", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-small", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("500"), + "nvidia.com/gpucores": resource.MustParse("10"), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-1", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("1000"), + "nvidia.com/gpucores": resource.MustParse("30"), + }, + }, + }, + { + Name: "app-2", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("1000"), + "nvidia.com/gpucores": resource.MustParse("30"), + }, + }, + }, + }, + }, + } + + if !fitResourceQuota(pod) { + t.Errorf("Expected fitResourceQuota to return true for container-larger pod under no quota limit") + } + }) + + t.Run("quota denial via init container request", func(t *testing.T) { + ns := "quota-test-init-deny" + cache := device.GetLocalCache() + cache.Quotas[ns] = &device.DeviceQuota{ + "nvidia.com/gpumem": &device.Quota{ + Limit: 1500, + LimitSet: true, + Used: 0, + }, + } + defer delete(cache.Quotas, ns) + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "init-deny-pod", + Namespace: ns, + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-heavy", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("2000"), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-light", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpumem": resource.MustParse("500"), + }, + }, + }, + }, + }, + } + + if fitResourceQuota(pod) { + t.Errorf("Expected fitResourceQuota to return false when init container exceeds memory quota") + } + }) + + t.Run("quota denial via core limit in init container", func(t *testing.T) { + ns := "quota-test-core-deny" + cache := device.GetLocalCache() + cache.Quotas[ns] = &device.DeviceQuota{ + "nvidia.com/gpucores": &device.Quota{ + Limit: 30, + LimitSet: true, + Used: 0, + }, + } + defer delete(cache.Quotas, ns) + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "core-deny-pod", + Namespace: ns, + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-heavy-cores", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpucores": resource.MustParse("50"), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-light-cores", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + "nvidia.com/gpucores": resource.MustParse("10"), + }, + }, + }, + }, + }, + } + + if fitResourceQuota(pod) { + t.Errorf("Expected fitResourceQuota to return false when init container exceeds core quota") + } + }) + + t.Run("pod with zero GPU requests", func(t *testing.T) { + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "no-gpu-pod", + Namespace: "default", + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{ + { + Name: "init-cpu", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "cpu": resource.MustParse("1"), + }, + }, + }, + }, + Containers: []corev1.Container{ + { + Name: "app-cpu", + Resources: corev1.ResourceRequirements{ + Limits: corev1.ResourceList{ + "cpu": resource.MustParse("1"), + }, + }, + }, + }, + }, + } + + if !fitResourceQuota(pod) { + t.Errorf("Expected fitResourceQuota to return true for pod requesting no GPU resources") + } + }) +} diff --git a/pkg/util/client/client.go b/pkg/util/client/client.go index 759e50d32a..219dacd09c 100644 --- a/pkg/util/client/client.go +++ b/pkg/util/client/client.go @@ -35,8 +35,10 @@ type Client struct { } var ( - KubeClient kubernetes.Interface - once sync.Once + KubeClient kubernetes.Interface + once sync.Once + buildConfigFromFlags = clientcmd.BuildConfigFromFlags + inClusterConfig = rest.InClusterConfig ) func init() { @@ -92,10 +94,10 @@ func loadKubeConfig() (*rest.Config, error) { kubeConfigPath = filepath.Join(os.Getenv("HOME"), ".kube", "config") } - config, err := clientcmd.BuildConfigFromFlags("", kubeConfigPath) + config, err := buildConfigFromFlags("", kubeConfigPath) if err != nil { klog.Infof("BuildConfigFromFlags failed for file %s: %v. Using in-cluster config.", kubeConfigPath, err) - return rest.InClusterConfig() + return inClusterConfig() } return config, nil } diff --git a/pkg/util/client/client_test.go b/pkg/util/client/client_test.go index 8995fc27c3..d8125253f9 100644 --- a/pkg/util/client/client_test.go +++ b/pkg/util/client/client_test.go @@ -29,18 +29,10 @@ import ( "gotest.tools/v3/assert" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/rest" - "k8s.io/client-go/tools/clientcmd" -) - -// Mock functions for testing. -var ( - buildConfigFromFlags = clientcmd.BuildConfigFromFlags - inClusterConfig = rest.InClusterConfig ) // TestGetClient tests the GetClient function. func TestGetClient(t *testing.T) { - InitGlobalClient() tests := []struct { name string kubeConfig string @@ -91,6 +83,10 @@ func TestGetClient(t *testing.T) { os.Setenv("KUBECONFIG", tt.kubeConfig) defer os.Setenv("KUBECONFIG", oldKubeConfig) + once = sync.Once{} + KubeClient = nil + InitGlobalClient() + // Call GetClient and check the result. client := GetClient() if tt.expectError { @@ -108,6 +104,12 @@ func TestGetClient(t *testing.T) { // TestClientWithOptions tests client initialization with options. func TestClientWithOptions(t *testing.T) { + oldBuildConfigFromFlags := buildConfigFromFlags + buildConfigFromFlags = func(masterUrl, kubeconfigPath string) (*rest.Config, error) { + return &rest.Config{Host: "https://example.com"}, nil + } + defer func() { buildConfigFromFlags = oldBuildConfigFromFlags }() + KubeClient = nil once = sync.Once{}