diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index 0e785c363b..d3baa4cdd9 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -122,6 +122,23 @@ spec: - name: ENABLE_TOPOLOGY_SCORE value: "true" {{- end }} + {{- with .Values.devicePlugin.numaRefit }} + {{- if .enabled }} + - name: HAMI_SCHEDULER_ENDPOINT + value: {{ .schedulerEndpoint | default (printf "https://%s.%s.svc:%v" (include "hami-vgpu.scheduler" $) (include "hami-vgpu.namespace" $) ($.Values.scheduler.service.httpPort | default 443)) | quote }} + {{- if .tlsInsecure }} + - name: HAMI_SCHEDULER_TLS_INSECURE + value: "true" + {{- end }} + {{- if .caFile }} + - name: HAMI_SCHEDULER_CA_FILE + value: {{ .caFile | quote }} + {{- else if .caSecret }} + - name: HAMI_SCHEDULER_CA_FILE + value: "/etc/hami/numa-refit-ca/ca.crt" + {{- end }} + {{- end }} + {{- end }} {{- with .Values.devicePlugin.extraEnvs }} {{- . | toYaml | nindent 12 }} {{- end }} @@ -147,6 +164,11 @@ spec: subPath: device-config.yaml - name: cdi-root mountPath: /var/run/cdi + {{- if and .Values.devicePlugin.numaRefit .Values.devicePlugin.numaRefit.enabled .Values.devicePlugin.numaRefit.caSecret }} + - name: numa-refit-ca + mountPath: /etc/hami/numa-refit-ca + readOnly: true + {{- end }} {{- if typeIs "string" .Values.devicePlugin.nvidiaDriverRoot }} # We always mount the driver root at /driver-root in the container. # This is required for CDI detection to work correctly. @@ -205,6 +227,11 @@ spec: - name: hosttmp mountPath: /tmp volumes: + {{- if and .Values.devicePlugin.numaRefit .Values.devicePlugin.numaRefit.enabled .Values.devicePlugin.numaRefit.caSecret }} + - name: numa-refit-ca + secret: + secretName: {{ .Values.devicePlugin.numaRefit.caSecret }} + {{- end }} - name: ctrs hostPath: path: {{ .Values.devicePlugin.monitor.ctrPath }} diff --git a/charts/hami/values.yaml b/charts/hami/values.yaml index 14ad58f60c..9d031ee346 100644 --- a/charts/hami/values.yaml +++ b/charts/hami/values.yaml @@ -351,6 +351,23 @@ devicePlugin: # TopologyManager can align CPU and GPU NUMA nodes. Opt-in because it changes # admission behavior when topologyManagerPolicy is single-numa-node. enableNumaTopology: false + # NUMA alignment refit (#2080): when enabled, the device plugin asks the + # scheduler to re-run its fit over the NUMA-restricted device set kubelet + # allows, keeping scheduler accounting authoritative. Requires + # enableNumaTopology plus the per-node enablegetpreferredallocation node + # configuration; pods opt in with the hami.io/numa-alignment annotation + # (best-effort or strict). Disabled by default: mismatches are only logged. + numaRefit: + enabled: false + # Scheduler base URL override. Defaults to the in-cluster scheduler + # service, https://-scheduler..svc:. + schedulerEndpoint: "" + # The scheduler serves the admission webhook's self-signed certificate, + # so verification is skipped by default with the same posture as the + # extender configmap (tlsConfig.insecure). Set false to verify, and use + # caFile to point at a CA bundle mounted into the device plugin. + tlsInsecure: true + caFile: "" # Pre-configured device memory in MB for GPUs that don't support memory query (e.g., unified memory architecture GPUs like NVIDIA GB10/DGX Spark). # Set to 0 to use auto-detection (default). For unified memory GPUs, set to the total GPU memory (e.g., 131072 for 128GB). # Can be overridden per-node via nodeConfiguration.config. diff --git a/cmd/scheduler/main.go b/cmd/scheduler/main.go index f7c68d1585..1cc94944de 100644 --- a/cmd/scheduler/main.go +++ b/cmd/scheduler/main.go @@ -145,6 +145,7 @@ func start() error { router := httprouter.New() router.POST("/filter", routes.PredicateRoute(sher)) router.POST("/bind", routes.Bind(sher)) + router.POST("/refit", routes.NumaRefit(sher)) router.POST("/webhook", routes.WebHookRoute()) router.GET("/healthz", routes.HealthzRoute()) router.GET("/readyz", routes.ReadyzRoute(sher)) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go new file mode 100644 index 0000000000..e1b63a00ff --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go @@ -0,0 +1,221 @@ +/* +Copyright 2026 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. +*/ + +package plugin + +import ( + "bytes" + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "strconv" + "strings" + "time" + + corev1 "k8s.io/api/core/v1" + "k8s.io/klog/v2" + kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" + "github.com/Project-HAMi/HAMi/pkg/util" +) + +const ( + // SchedulerEndpointEnvName holds the HAMi scheduler base URL used for + // the NUMA refit, for example https://hami-scheduler.kube-system.svc:443. + // Empty disables the refit: mismatches are then only logged, exactly as + // before the refit existed. + SchedulerEndpointEnvName = "HAMI_SCHEDULER_ENDPOINT" + // SchedulerCAFileEnvName optionally points at a PEM bundle used to + // verify the scheduler endpoint's TLS certificate. + SchedulerCAFileEnvName = "HAMI_SCHEDULER_CA_FILE" + // SchedulerTLSInsecureEnvName set to true skips TLS verification of the + // scheduler endpoint. The scheduler serves the admission webhook's + // self-signed certificate, so the chart enables this by default with the + // same posture as the extender configmap (tlsConfig.insecure: true). + SchedulerTLSInsecureEnvName = "HAMI_SCHEDULER_TLS_INSECURE" + + numaRefitPath = "/refit" + + // numaRefitTimeout bounds one refit round trip. Kubelet applies no + // deadline of its own to GetPreferredAllocation and admits pods on a + // single serialized loop, so this client timeout is the node's only + // protection against a slow or unreachable scheduler. + numaRefitTimeout = 2 * time.Second +) + +// numaRefitTLSConfig verifies the scheduler certificate by default, against +// SchedulerCAFileEnvName when provided; SchedulerTLSInsecureEnvName is an +// explicit operator opt-out for the self-signed webhook certificate. +func numaRefitTLSConfig() *tls.Config { + config := &tls.Config{MinVersion: tls.VersionTLS12} + if caFile := os.Getenv(SchedulerCAFileEnvName); caFile != "" { + pem, err := os.ReadFile(caFile) + if err != nil { + klog.ErrorS(err, "cannot read scheduler CA bundle", "path", caFile) + } else if pool := x509.NewCertPool(); pool.AppendCertsFromPEM(pem) { + config.RootCAs = pool + } else { + klog.ErrorS(nil, "scheduler CA bundle contains no usable certificates", "path", caFile) + } + } + if insecure, err := strconv.ParseBool(os.Getenv(SchedulerTLSInsecureEnvName)); err == nil { + config.InsecureSkipVerify = insecure + } + return config +} + +// numaRefitHTTPClient reaches the scheduler service. +var numaRefitHTTPClient = &http.Client{ + Timeout: numaRefitTimeout, + Transport: &http.Transport{ + TLSClientConfig: numaRefitTLSConfig(), + }, +} + +// tryNumaRefit asks the scheduler to move this container's pending +// allocation onto kubelet's allowed device set. It returns the preferred +// replica IDs on success. A nil slice with a nil error means the refit did +// not apply (disabled, pod not opted in, or best-effort failure); a non-nil +// error means strict mode failed and the allocation must fail. +func (plugin *NvidiaDevicePlugin) tryNumaRefit(ctx context.Context, pod *corev1.Pod, containerIndex int, req *kubeletdevicepluginv1beta1.ContainerPreferredAllocationRequest, cause error) ([]string, error) { + if pod == nil || plugin.operatingMode == nvidia.MigMode || !errors.Is(cause, errAnnotatedDeviceUnavailable) { + return nil, nil + } + mode, parseErr := util.GetNumaAlignmentModeByPod(pod) + if parseErr != nil || mode == util.NumaAlignmentNone { + return nil, nil + } + if os.Getenv(SchedulerEndpointEnvName) == "" { + return nil, nil + } + + // When kubelet pins replicas via MustIncludeDeviceIDs, only their + // physical devices can satisfy the allocation, so restrict the refit to + // them; otherwise any available physical device is eligible. + allowedUUIDs := allowedPhysicalDeviceIDs(req.AvailableDeviceIDs) + if len(req.MustIncludeDeviceIDs) > 0 { + allowedUUIDs = allowedPhysicalDeviceIDs(req.MustIncludeDeviceIDs) + } + newDevices, err := plugin.requestNumaRefit(ctx, pod, containerIndex, allowedUUIDs) + if err == nil { + replicas, selectErr := plugin.selectPreferredDeviceIDsFromAnnotatedDevices(req.AvailableDeviceIDs, req.MustIncludeDeviceIDs, newDevices, int(req.AllocationSize)) + if selectErr == nil { + klog.InfoS("NUMA refit succeeded", "pod", klog.KObj(pod), "container", containerIndex, "devices", replicas) + return replicas, nil + } + // The scheduler has already committed the move at this point. + // Falling back to kubelet's own selection would leave runtime and + // accounting divergent, so fail the allocation in both modes. + return nil, fmt.Errorf("numa refit committed but kubelet cannot honor the selection: %w", selectErr) + } + + if mode == util.NumaAlignmentStrict { + return nil, fmt.Errorf("numa-alignment strict: %w", err) + } + klog.InfoS("NUMA refit failed; best-effort keeps kubelet's own selection", "pod", klog.KObj(pod), "container", containerIndex, "err", err) + return nil, nil +} + +// requestNumaRefit performs one refit round trip against the scheduler. +func (plugin *NvidiaDevicePlugin) requestNumaRefit(ctx context.Context, pod *corev1.Pod, containerIndex int, allowedUUIDs []string) (device.ContainerDevices, error) { + payload, err := json.Marshal(device.NumaRefitRequest{ + PodUID: string(pod.UID), + PodNamespace: pod.Namespace, + PodName: pod.Name, + NodeName: os.Getenv(util.NodeNameEnvName), + ContainerIndex: containerIndex, + ContainerName: podContainerNameAt(pod, containerIndex), + DeviceType: nvidia.NvidiaGPUDevice, + AllowedDeviceUUIDs: allowedUUIDs, + }) + if err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, numaRefitTimeout) + defer cancel() + url := strings.TrimSuffix(os.Getenv(SchedulerEndpointEnvName), "/") + numaRefitPath + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + + httpResp, err := numaRefitHTTPClient.Do(httpReq) + if err != nil { + return nil, err + } + defer httpResp.Body.Close() + if httpResp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("scheduler refit returned status %d", httpResp.StatusCode) + } + + var response device.NumaRefitResponse + if err := json.NewDecoder(io.LimitReader(httpResp.Body, 1<<20)).Decode(&response); err != nil { + return nil, err + } + if !response.Succeeded { + return nil, fmt.Errorf("scheduler refused refit: %s", response.FailureReason) + } + devices, err := device.DecodeContainerDevices(response.ContainerDevices) + if err != nil { + return nil, fmt.Errorf("cannot decode refit devices: %w", err) + } + if len(devices) == 0 { + return nil, errors.New("scheduler refit returned no devices") + } + return devices, nil +} + +// podContainerNameAt returns the pod's container name at the PodDevices +// position, counting init containers first, for the scheduler's cross-check. +func podContainerNameAt(pod *corev1.Pod, index int) string { + if index < 0 { + return "" + } + if index < len(pod.Spec.InitContainers) { + return pod.Spec.InitContainers[index].Name + } + index -= len(pod.Spec.InitContainers) + if index >= 0 && index < len(pod.Spec.Containers) { + return pod.Spec.Containers[index].Name + } + return "" +} + +// allowedPhysicalDeviceIDs maps kubelet's replica IDs to their unique +// physical device UUIDs, preserving first-seen order. +func allowedPhysicalDeviceIDs(available []string) []string { + seen := make(map[string]struct{}, len(available)) + physical := make([]string, 0, len(available)) + for _, id := range available { + p := physicalDeviceID(id) + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + physical = append(physical, p) + } + return physical +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go new file mode 100644 index 0000000000..69dfb126c7 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go @@ -0,0 +1,242 @@ +/* +Copyright 2026 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. +*/ + +package plugin + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" + "github.com/Project-HAMi/HAMi/pkg/util" +) + +func numaRefitTestPod(mode string) *corev1.Pod { + annotations := map[string]string{ + "hami.io/vgpu-devices-to-allocate": device.EncodePodSingleDevice(device.PodSingleDevice{ + device.ContainerDevices{{UUID: numaTestGPUA, Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}}, + }), + } + if mode != "" { + annotations[util.NumaAlignmentAnnotationKey] = mode + } + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: "refit-pod-uid", Name: "numa-pod", Namespace: "default", + Annotations: annotations, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}, + } +} + +// numaRefitTestServer serves /refit with the given response and captures the +// last request. +func numaRefitTestServer(t *testing.T, response device.NumaRefitResponse) (*httptest.Server, *device.NumaRefitRequest) { + t.Helper() + var lastRequest device.NumaRefitRequest + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodPost, r.Method) + require.Equal(t, numaRefitPath, r.URL.Path) + require.NoError(t, json.NewDecoder(r.Body).Decode(&lastRequest)) + w.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(w).Encode(response)) + })) + t.Cleanup(server.Close) + return server, &lastRequest +} + +func TestRequestNumaRefitRoundTrip(t *testing.T) { + refitted := device.ContainerDevices{{UUID: numaTestGPUB, Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}} + server, lastRequest := numaRefitTestServer(t, device.NumaRefitResponse{ + Succeeded: true, + ContainerDevices: device.EncodeContainerDevices(refitted), + }) + t.Setenv(SchedulerEndpointEnvName, server.URL) + t.Setenv(util.NodeNameEnvName, "node-a") + + plugin := &NvidiaDevicePlugin{} + devices, err := plugin.requestNumaRefit(context.Background(), numaRefitTestPod("best-effort"), 0, []string{numaTestGPUB}) + + require.NoError(t, err) + require.Len(t, devices, 1) + require.Equal(t, numaTestGPUB, devices[0].UUID) + require.Equal(t, int32(20000), devices[0].Usedmem) + + require.Equal(t, "refit-pod-uid", lastRequest.PodUID) + require.Equal(t, "default", lastRequest.PodNamespace) + require.Equal(t, "numa-pod", lastRequest.PodName) + require.Equal(t, "node-a", lastRequest.NodeName) + require.Equal(t, 0, lastRequest.ContainerIndex) + require.Equal(t, "main", lastRequest.ContainerName) + require.Equal(t, nvidia.NvidiaGPUDevice, lastRequest.DeviceType) + require.Equal(t, []string{numaTestGPUB}, lastRequest.AllowedDeviceUUIDs) +} + +func TestRequestNumaRefitRefused(t *testing.T) { + server, _ := numaRefitTestServer(t, device.NumaRefitResponse{Succeeded: false, FailureReason: "no allowed device fits"}) + t.Setenv(SchedulerEndpointEnvName, server.URL) + + plugin := &NvidiaDevicePlugin{} + _, err := plugin.requestNumaRefit(context.Background(), numaRefitTestPod("strict"), 0, []string{numaTestGPUB}) + + require.Error(t, err) + require.Contains(t, err.Error(), "no allowed device fits") +} + +func TestRequestNumaRefitBadStatus(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(server.Close) + t.Setenv(SchedulerEndpointEnvName, server.URL) + + plugin := &NvidiaDevicePlugin{} + _, err := plugin.requestNumaRefit(context.Background(), numaRefitTestPod("strict"), 0, []string{numaTestGPUB}) + + require.Error(t, err) + require.Contains(t, err.Error(), "status 500") +} + +func TestAllowedPhysicalDeviceIDs(t *testing.T) { + got := allowedPhysicalDeviceIDs([]string{ + numaTestGPUB + "-0", numaTestGPUB + "-1", numaTestGPUA + "-0", + }) + require.Equal(t, []string{numaTestGPUB, numaTestGPUA}, got) +} + +// End-to-end through GetPreferredAllocation: best-effort pod, mismatch, and a +// scheduler that refits onto GPU-B — kubelet receives GPU-B replicas. +func TestGetPreferredAllocationRefitBestEffort(t *testing.T) { + refitted := device.ContainerDevices{{UUID: numaTestGPUB, Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}} + server, lastRequest := numaRefitTestServer(t, device.NumaRefitResponse{ + Succeeded: true, + ContainerDevices: device.EncodeContainerDevices(refitted), + }) + t.Setenv(SchedulerEndpointEnvName, server.URL) + t.Setenv(util.NodeNameEnvName, "node-a") + setupInRequestDevices(t) + + pod := numaRefitTestPod("best-effort") + mockAllocateGlobals(t, pod) + + plugin := &NvidiaDevicePlugin{} + response, err := plugin.GetPreferredAllocation(context.Background(), &kubeletdevicepluginv1beta1.PreferredAllocationRequest{ + ContainerRequests: []*kubeletdevicepluginv1beta1.ContainerPreferredAllocationRequest{{ + AvailableDeviceIDs: []string{numaTestGPUB + "-0", numaTestGPUB + "-1"}, + AllocationSize: 1, + }}, + }) + + require.NoError(t, err) + require.Len(t, response.ContainerResponses, 1) + require.Equal(t, []string{numaTestGPUB + "-0"}, response.ContainerResponses[0].DeviceIDs) + require.Equal(t, []string{numaTestGPUB}, lastRequest.AllowedDeviceUUIDs) +} + +// Strict pod with a scheduler that refuses: the allocation must fail. +func TestGetPreferredAllocationRefitStrictFailure(t *testing.T) { + server, _ := numaRefitTestServer(t, device.NumaRefitResponse{Succeeded: false, FailureReason: "no allowed device fits"}) + t.Setenv(SchedulerEndpointEnvName, server.URL) + t.Setenv(util.NodeNameEnvName, "node-a") + setupInRequestDevices(t) + + pod := numaRefitTestPod("strict") + mockAllocateGlobals(t, pod) + + plugin := &NvidiaDevicePlugin{} + _, err := plugin.GetPreferredAllocation(context.Background(), &kubeletdevicepluginv1beta1.PreferredAllocationRequest{ + ContainerRequests: []*kubeletdevicepluginv1beta1.ContainerPreferredAllocationRequest{{ + AvailableDeviceIDs: []string{numaTestGPUB + "-0"}, + AllocationSize: 1, + }}, + }) + + require.Error(t, err) + require.Contains(t, err.Error(), "numa-alignment strict") +} + +// Without the endpoint configured even strict pods keep detection-only +// behavior: empty response, no error. +func TestGetPreferredAllocationRefitDisabled(t *testing.T) { + t.Setenv(SchedulerEndpointEnvName, "") + t.Setenv(util.NodeNameEnvName, "node-a") + setupInRequestDevices(t) + + pod := numaRefitTestPod("strict") + mockAllocateGlobals(t, pod) + + plugin := &NvidiaDevicePlugin{} + response, err := plugin.GetPreferredAllocation(context.Background(), &kubeletdevicepluginv1beta1.PreferredAllocationRequest{ + ContainerRequests: []*kubeletdevicepluginv1beta1.ContainerPreferredAllocationRequest{{ + AvailableDeviceIDs: []string{numaTestGPUB + "-0"}, + AllocationSize: 1, + }}, + }) + + require.NoError(t, err) + require.Len(t, response.ContainerResponses, 0) +} + +func TestNumaRefitTLSConfigVerifiesByDefault(t *testing.T) { + t.Setenv(SchedulerTLSInsecureEnvName, "") + t.Setenv(SchedulerCAFileEnvName, "") + config := numaRefitTLSConfig() + require.False(t, config.InsecureSkipVerify) + require.Nil(t, config.RootCAs) +} + +func TestNumaRefitTLSConfigInsecureOptOut(t *testing.T) { + t.Setenv(SchedulerTLSInsecureEnvName, "true") + config := numaRefitTLSConfig() + require.True(t, config.InsecureSkipVerify) +} + +// A committed refit whose devices kubelet cannot honor must fail the +// allocation even in best-effort mode: falling back would leave scheduler +// accounting and runtime divergent. +func TestGetPreferredAllocationRefitCommittedButUnmappable(t *testing.T) { + refitted := device.ContainerDevices{{UUID: numaTestGPUA, Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}} + server, _ := numaRefitTestServer(t, device.NumaRefitResponse{ + Succeeded: true, + ContainerDevices: device.EncodeContainerDevices(refitted), + }) + t.Setenv(SchedulerEndpointEnvName, server.URL) + t.Setenv(util.NodeNameEnvName, "node-a") + setupInRequestDevices(t) + + pod := numaRefitTestPod("best-effort") + mockAllocateGlobals(t, pod) + + plugin := &NvidiaDevicePlugin{} + _, err := plugin.GetPreferredAllocation(context.Background(), &kubeletdevicepluginv1beta1.PreferredAllocationRequest{ + ContainerRequests: []*kubeletdevicepluginv1beta1.ContainerPreferredAllocationRequest{{ + AvailableDeviceIDs: []string{numaTestGPUB + "-0"}, + AllocationSize: 1, + }}, + }) + + require.Error(t, err) + require.Contains(t, err.Error(), "committed") +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index e0480a6584..bc23113f44 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -647,10 +647,14 @@ func (plugin *NvidiaDevicePlugin) GetPreferredAllocation(ctx context.Context, r // Filter out empty annotations to match kubelet's ContainerRequests order. // Kubelet only sends requests for containers that need GPUs, but annotations // include all containers (init + regular), some of which may be empty. + // annotationIndices keeps each entry's original PodDevices container + // position, which the NUMA refit protocol requires. var nonEmptyAnnotations []device.ContainerDevices - for _, ann := range annotatedRequests { + var annotationIndices []int + for i, ann := range annotatedRequests { if len(ann) > 0 { nonEmptyAnnotations = append(nonEmptyAnnotations, ann) + annotationIndices = append(annotationIndices, i) } } @@ -664,7 +668,25 @@ func (plugin *NvidiaDevicePlugin) GetPreferredAllocation(ctx context.Context, r }) } else { klog.Warningf("err: %v", err) - plugin.reportAnnotatedDeviceMismatch(pendingPod, idx, err) + refitDevices, refitErr := plugin.tryNumaRefit(ctx, pendingPod, annotationIndices[idx], req, err) + switch { + case refitErr != nil: + // Strict mode, or a committed refit kubelet cannot + // honor: fail the pod's admission rather than running + // misaligned or diverging from accounting. Mark the + // bind failed so the node lock is released instead of + // blocking the node until the lock timeout. + if nodename != "" && pendingPod != nil { + PodAllocationFailed(nodename, pendingPod, NodeLockNvidia) + } + return nil, refitErr + case len(refitDevices) > 0: + response.ContainerResponses = append(response.ContainerResponses, &kubeletdevicepluginv1beta1.ContainerPreferredAllocationResponse{ + DeviceIDs: refitDevices, + }) + default: + plugin.reportAnnotatedDeviceMismatch(pendingPod, idx, err) + } } } } @@ -695,8 +717,12 @@ func (plugin *NvidiaDevicePlugin) reportAnnotatedDeviceMismatch(pod *corev1.Pod, return } - klog.InfoS("scheduler-annotated GPU has no replica in kubelet's available devices; NUMA refit is not implemented yet, so kubelet will select a device on its own", - "err", err, "pod", klog.KObj(pod), "containerRequest", containerRequest, "numaAlignment", string(mode)) + const message = "scheduler-annotated GPU has no replica in kubelet's available devices and the NUMA refit did not handle it, so kubelet will select a device on its own" + if mode == util.NumaAlignmentStrict { + klog.ErrorS(err, message, "pod", klog.KObj(pod), "containerRequest", containerRequest, "numaAlignment", string(mode)) + return + } + klog.InfoS(message, "err", err, "pod", klog.KObj(pod), "containerRequest", containerRequest, "numaAlignment", string(mode)) } func (plugin *NvidiaDevicePlugin) selectPreferredDeviceIDsFromAnnotatedDevices(available, required []string, desired device.ContainerDevices, allocationSize int) ([]string, error) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_numa_alignment_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_numa_alignment_test.go index f138fd00d3..df0562afb8 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_numa_alignment_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_numa_alignment_test.go @@ -87,7 +87,7 @@ func TestGetPreferredAllocationMismatchResponseUnchanged(t *testing.T) { availableIDs: []string{numaTestGPUB + "-0", numaTestGPUB + "-1"}, }, { - name: "strict is not accepted yet", + name: "strict without the refit endpoint", annotations: map[string]string{util.NumaAlignmentAnnotationKey: "strict"}, availableIDs: []string{numaTestGPUB + "-0", numaTestGPUB + "-1"}, }, @@ -216,12 +216,13 @@ func TestReportAnnotatedDeviceMismatch(t *testing.T) { wantFields: []string{`numaAlignment="best-effort"`, "default/numa-pod", "containerRequest=0"}, }, { - name: "strict warns as invalid until the refit lands", + name: "strict mismatch is reported at error severity", plugin: &NvidiaDevicePlugin{}, pod: podWithMode("strict"), err: mismatch, - wantLogged: "ignoring invalid numa-alignment annotation", - wantPrefix: "W", + wantLogged: mismatchMessage, + wantPrefix: "E", + wantFields: []string{`numaAlignment="strict"`, "default/numa-pod"}, }, } diff --git a/pkg/device/pods.go b/pkg/device/pods.go index 0fb40e242a..1a953707d6 100644 --- a/pkg/device/pods.go +++ b/pkg/device/pods.go @@ -166,6 +166,23 @@ func (m *PodManager) TakeAndDeletePod(pod *corev1.Pod) (*PodInfo, bool) { return pi, ok } +// ReplacePodDevices swaps the tracked devices for pod and returns the +// previous set. Unlike UpdatePodDevice it leaves +// InitContainerResourceReleased untouched, so a device swap (for example the +// NUMA refit from #2080) does not disable the init-container usage shrink. +func (m *PodManager) ReplacePodDevices(pod *corev1.Pod, newDevices PodDevices) (oldDevices PodDevices, ok bool) { + m.mutex.Lock() + defer m.mutex.Unlock() + + pi, exists := m.pods[pod.UID] + if !exists { + return nil, false + } + oldDevices = pi.Devices + pi.Devices = newDevices + return oldDevices, true +} + func (m *PodManager) UpdatePodDevice(pod *corev1.Pod, newDevices PodDevices) (oldDevices PodDevices, ok bool) { m.mutex.Lock() defer m.mutex.Unlock() diff --git a/pkg/device/pods_replace_test.go b/pkg/device/pods_replace_test.go new file mode 100644 index 0000000000..802075956d --- /dev/null +++ b/pkg/device/pods_replace_test.go @@ -0,0 +1,47 @@ +/* +Copyright 2026 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. +*/ + +package device + +import ( + "testing" + + "gotest.tools/v3/assert" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +func TestReplacePodDevicesPreservesInitFlag(t *testing.T) { + m := NewPodManager() + pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Name: "pod", Namespace: "ns"}} + m.AddPod(pod, "node-1", PodDevices{"NVIDIA": {{{UUID: "GPU-a", Usedmem: 100}}}}) + + old, ok := m.ReplacePodDevices(pod, PodDevices{"NVIDIA": {{{UUID: "GPU-b", Usedmem: 100}}}}) + assert.Equal(t, ok, true) + assert.Equal(t, old["NVIDIA"][0][0].UUID, "GPU-a") + + pi, found := m.GetPod(pod) + assert.Equal(t, found, true) + assert.Equal(t, pi.Devices["NVIDIA"][0][0].UUID, "GPU-b") + // Unlike UpdatePodDevice, the init-container shrink must stay armed. + assert.Equal(t, pi.InitContainerResourceReleased, false) +} + +func TestReplacePodDevicesUnknownPod(t *testing.T) { + m := NewPodManager() + _, ok := m.ReplacePodDevices(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "missing"}}, nil) + assert.Equal(t, ok, false) +} diff --git a/pkg/scheduler/event.go b/pkg/scheduler/event.go index 21b4075310..601a84d89c 100644 --- a/pkg/scheduler/event.go +++ b/pkg/scheduler/event.go @@ -40,6 +40,11 @@ const ( EventReasonBindingFailed = "BindingFailed" // EventReasonBindingSucceed indicates that binding succeed. EventReasonBindingSucceed = "BindingSucceed" + + // EventReasonNumaRefitFailed indicates that a NUMA refit failed. + EventReasonNumaRefitFailed = "NumaRefitFailed" + // EventReasonNumaRefitSucceed indicates that a NUMA refit succeeded. + EventReasonNumaRefitSucceed = "NumaRefitSucceed" ) func (s *Scheduler) addAllEventHandlers() { @@ -65,6 +70,19 @@ func (s *Scheduler) recordScheduleBindingResultEvent(pod *corev1.Pod, eventReaso } } +// recordNumaRefitResultEvent emits the outcome of a NUMA refit on the pod. +func (s *Scheduler) recordNumaRefitResultEvent(pod *corev1.Pod, successMsg string, refitErr error) { + // eventRecorder maybe nil + if pod == nil || s.eventRecorder == nil { + return + } + if refitErr == nil { + s.eventRecorder.Event(pod, corev1.EventTypeNormal, EventReasonNumaRefitSucceed, successMsg) + } else { + s.eventRecorder.Event(pod, corev1.EventTypeWarning, EventReasonNumaRefitFailed, refitErr.Error()) + } +} + func (s *Scheduler) recordScheduleFilterResultEvent(pod *corev1.Pod, eventReason string, successMsg string, schedulerErr error) { // eventRecorder maybe nil if pod == nil || s.eventRecorder == nil { diff --git a/pkg/scheduler/numa_refit_handler.go b/pkg/scheduler/numa_refit_handler.go new file mode 100644 index 0000000000..3643018490 --- /dev/null +++ b/pkg/scheduler/numa_refit_handler.go @@ -0,0 +1,337 @@ +/* +Copyright 2026 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. +*/ + +package scheduler + +import ( + "context" + "encoding/json" + "fmt" + "maps" + "strings" + + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + k8stypes "k8s.io/apimachinery/pkg/types" + "k8s.io/klog/v2" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" + "github.com/Project-HAMi/HAMi/pkg/util" + "github.com/Project-HAMi/HAMi/pkg/util/client" +) + +// patchPodAnnotations patches only the given annotations, with the pod's +// resourceVersion as a precondition so a stale refit cannot overwrite a +// newer update to the same annotations (for example Allocate consuming a +// to-allocate entry). Conflicts fail the refit; there is no retry with +// cached values. Also a test seam. +var patchPodAnnotations = func(pod *corev1.Pod, annotations map[string]string) error { + metadata := map[string]any{"annotations": annotations} + if pod.ResourceVersion != "" { + metadata["resourceVersion"] = pod.ResourceVersion + } + payload, err := json.Marshal(map[string]any{"metadata": metadata}) + if err != nil { + return err + } + _, err = client.GetClient().CoreV1().Pods(pod.Namespace). + Patch(context.Background(), pod.Name, k8stypes.MergePatchType, payload, metav1.PatchOptions{}) + return err +} + +// maxAllowedDeviceUUIDs bounds the allowed set a refit request may carry. +const maxAllowedDeviceUUIDs = 512 + +// RefitNumaAllocation moves one container's device reservation onto a device +// from the caller-supplied allowed set, re-running the pod's normal +// policy-chain fit restricted to that set. The device plugin calls it (via +// the /refit route) when kubelet's Topology Manager restricted an allocation +// to replicas of devices the scheduler did not annotate; see issue #2080. +// +// The scheduler stays authoritative: the refit runs the same fit and +// capacity checks as scheduling, patches hami.io/vgpu-devices-to-allocate +// and hami.io/vgpu-devices-allocated together in one merge patch, and only +// then moves the in-memory reservation. Failures are reported in-band and +// leave both annotations and accounting untouched. +func (s *Scheduler) RefitNumaAllocation(req device.NumaRefitRequest) device.NumaRefitResponse { + if req.PodUID == "" || req.PodNamespace == "" || req.PodName == "" || req.NodeName == "" { + return numaRefitFailure(nil, "incomplete refit request: pod UID, namespace, name, and node are required") + } + if len(req.AllowedDeviceUUIDs) == 0 { + return numaRefitFailure(nil, "refit request carries an empty allowed device set") + } + if len(req.AllowedDeviceUUIDs) > maxAllowedDeviceUUIDs { + return numaRefitFailure(nil, "refit request carries %d allowed devices, limit is %d", len(req.AllowedDeviceUUIDs), maxAllowedDeviceUUIDs) + } + if _, ok := device.GetDevices()[req.DeviceType]; !ok { + return numaRefitFailure(nil, "unknown device type %q", req.DeviceType) + } + // HAMi's type matching is substring based, so a refit for one type could + // restrict sibling types whose names contain it. Only the NVIDIA plugin + // sends refits today; widening this needs exact type identity first. + if req.DeviceType != nvidia.NvidiaGPUDevice { + return numaRefitFailure(nil, "device type %q is not supported by the NUMA refit yet", req.DeviceType) + } + + s.allocLock.Lock() + defer s.allocLock.Unlock() + + key := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: k8stypes.UID(req.PodUID)}} + pi, ok := s.podManager.GetPod(key) + if !ok { + return numaRefitFailure(nil, "pod %s is not tracked by the scheduler", req.PodUID) + } + pod := pi.Pod + if pod.Namespace != req.PodNamespace || pod.Name != req.PodName { + return numaRefitFailure(nil, "pod %s does not match %s/%s", req.PodUID, req.PodNamespace, req.PodName) + } + if pi.NodeID != req.NodeName { + return s.numaRefitFailureEvent(pod, "pod %s/%s is tracked on node %s, not %s", pod.Namespace, pod.Name, pi.NodeID, req.NodeName) + } + + name, ok := containerNameAt(pod, req.ContainerIndex) + if !ok { + return s.numaRefitFailureEvent(pod, "container index %d is outside the pod's containers", req.ContainerIndex) + } + if req.ContainerName != "" && name != req.ContainerName { + return s.numaRefitFailureEvent(pod, "container index %d is %q, not %q", req.ContainerIndex, name, req.ContainerName) + } + + // Per-container state comes from the annotations: podManager stores + // init-collapsed aggregates whose entries do not map to container + // positions. The refit only applies while this container's allocation is + // still pending: a blank to-allocate entry means Allocate consumed it. + toAllocate, err := device.DecodePodDevices(device.InRequestDevices, pod.Annotations) + if err != nil { + return s.numaRefitFailureEvent(pod, "cannot decode pending allocation annotation: %v", err) + } + allocated, err := device.DecodePodDevices(device.SupportDevices, pod.Annotations) + if err != nil { + return s.numaRefitFailureEvent(pod, "cannot decode allocated annotation: %v", err) + } + pendingSingle := toAllocate[req.DeviceType] + allocatedSingle := allocated[req.DeviceType] + if req.ContainerIndex < 0 || req.ContainerIndex >= len(pendingSingle) || len(pendingSingle[req.ContainerIndex]) == 0 { + return s.numaRefitFailureEvent(pod, "container %d has no pending %s allocation to refit", req.ContainerIndex, req.DeviceType) + } + if req.ContainerIndex >= len(allocatedSingle) || len(allocatedSingle[req.ContainerIndex]) == 0 { + return s.numaRefitFailureEvent(pod, "container %d has no recorded %s allocation", req.ContainerIndex, req.DeviceType) + } + current := allocatedSingle[req.ContainerIndex] + + allowed := make(map[string]struct{}, len(req.AllowedDeviceUUIDs)) + for _, id := range req.AllowedDeviceUUIDs { + allowed[id] = struct{}{} + } + alreadyAllowed := true + for _, d := range current { + if _, ok := allowed[d.UUID]; !ok { + alreadyAllowed = false + break + } + } + if alreadyAllowed { + return device.NumaRefitResponse{Succeeded: true, ContainerDevices: device.EncodeContainerDevices(current)} + } + // A fit request carries one memory/core amount for all requested + // devices, so a reservation with differing per-device amounts (possible + // with percentage requests on mixed GPUs) cannot be re-fit faithfully: + // replaying current[0] would rewrite the other devices' accounting. + for _, d := range current[1:] { + if d.Usedmem != current[0].Usedmem || d.Usedcores != current[0].Usedcores { + return s.numaRefitFailureEvent(pod, "container %d reserves differing amounts per device; the refit does not support heterogeneous reservations", req.ContainerIndex) + } + } + + nodeUsageMap, _, failedNodes, err := s.getNodesUsage(&[]string{req.NodeName}, pod) + if err != nil { + return s.numaRefitFailureEvent(pod, "cannot compute node usage: %v", err) + } + nodeUsage, ok := (*nodeUsageMap)[req.NodeName] + if !ok { + return s.numaRefitFailureEvent(pod, "node %s unavailable: %s", req.NodeName, failedNodes[req.NodeName]) + } + for _, deviceList := range nodeUsage.Devices.DeviceLists { + if _, ok := allowed[deviceList.Device.ID]; ok && deviceList.Device.Mode == nvidia.MigMode { + return s.numaRefitFailureEvent(pod, "allowed device %s is in MIG mode; the NUMA refit does not support MIG", deviceList.Device.ID) + } + } + // The snapshot includes this container's own reservation; release it so + // capacity checks do not double count the pod against itself. + releaseContainerUsage(nodeUsage, current) + + weights, err := util.GetDeviceScoringWeightsByPod(pod) + if err != nil { + return s.numaRefitFailureEvent(pod, "invalid device scoring weights: %v", err) + } + + // Seed the fit with the pod's other container allocations so exclusivity + // and custom filter rules see them, and release the pod's quota usage so + // the namespace quota check does not count the pod against itself. + devinput := device.PodDevices{} + for deviceType, single := range allocated { + for containerIndex, containerDevices := range single { + if deviceType == req.DeviceType && containerIndex == req.ContainerIndex { + continue + } + if len(containerDevices) == 0 { + continue + } + devinput[deviceType] = append(devinput[deviceType], containerDevices) + } + } + seeded := len(devinput[req.DeviceType]) + s.quotaManager.RmUsage(pod, pi.Devices) + failWithQuotaRestore := func(format string, args ...any) device.NumaRefitResponse { + s.quotaManager.AddUsage(pod, pi.Devices) + return s.numaRefitFailureEvent(pod, format, args...) + } + + // Preserve the reservation's accounted amounts; only the device moves. + requests := device.ContainerDeviceRequests{req.DeviceType: { + Nums: int32(len(current)), + Type: req.DeviceType, + Memreq: current[0].Usedmem, + Coresreq: current[0].Usedcores, + }} + fit, reason := fitInRestrictedDevices(nodeUsage, requests, req.DeviceType, req.AllowedDeviceUUIDs, pod, nodeUsage.NodeInfo, &devinput, weights) + if !fit { + return failWithQuotaRestore("no allowed device fits: %s", reason) + } + selected := devinput[req.DeviceType] + if len(selected) != seeded+1 || len(selected[seeded]) != len(current) { + return failWithQuotaRestore("restricted fit selected %d container sets, want %d with %d devices", len(selected), seeded+1, len(current)) + } + newDevices := selected[seeded] + + // Patch both annotations by replacing only this container's entry inside + // the current raw values: entries Allocate already consumed stay blank, + // the separator layout survives byte for byte, and a scheduler restart + // rebuilds accounting onto the refitted device. + pendingValue, err := replaceContainerDeviceEntry(pod.Annotations[device.InRequestDevices[req.DeviceType]], req.ContainerIndex, newDevices) + if err != nil { + return failWithQuotaRestore("cannot rewrite pending allocation annotation: %v", err) + } + allocatedValue, err := replaceContainerDeviceEntry(pod.Annotations[device.SupportDevices[req.DeviceType]], req.ContainerIndex, newDevices) + if err != nil { + return failWithQuotaRestore("cannot rewrite allocated annotation: %v", err) + } + annotations := map[string]string{ + device.InRequestDevices[req.DeviceType]: pendingValue, + device.SupportDevices[req.DeviceType]: allocatedValue, + } + if err := patchPodAnnotations(pod, annotations); err != nil { + return failWithQuotaRestore("cannot patch pod annotations: %v", err) + } + + // Rebuild the in-memory reservation from the patched annotations exactly + // like the informer's add path does, so cached accounting and the + // durable record cannot drift apart. + patchedAnnotations := make(map[string]string, len(pod.Annotations)+len(annotations)) + maps.Copy(patchedAnnotations, pod.Annotations) + maps.Copy(patchedAnnotations, annotations) + if rawDevices, decodeErr := device.DecodePodDevices(device.SupportDevices, patchedAnnotations); decodeErr == nil { + effective := device.CollapseInitContainerUsage(pod, rawDevices) + if _, ok := s.podManager.ReplacePodDevices(key, effective); ok { + s.quotaManager.AddUsage(pod, effective) + } else { + // The pod left the cache between lookup and update; the informer + // rebuilds accounting from the patched annotations on re-add. + klog.InfoS("pod left the scheduler cache during NUMA refit; annotations remain authoritative", "pod", klog.KObj(pod)) + } + } else { + s.quotaManager.AddUsage(pod, pi.Devices) + klog.ErrorS(decodeErr, "cannot rebuild accounting from patched annotations; keeping previous usage", "pod", klog.KObj(pod)) + } + + message := fmt.Sprintf("NUMA refit moved container %d from %s to %s", req.ContainerIndex, containerDeviceIDs(current), containerDeviceIDs(newDevices)) + klog.InfoS(message, "pod", klog.KObj(pod), "node", req.NodeName) + s.recordNumaRefitResultEvent(pod, message, nil) + return device.NumaRefitResponse{Succeeded: true, ContainerDevices: device.EncodeContainerDevices(newDevices)} +} + +// replaceContainerDeviceEntry swaps one container's entry inside an encoded +// pod-device annotation value, preserving every other byte, so blanked +// entries and the separator layout survive the round trip unchanged. +func replaceContainerDeviceEntry(annotation string, index int, devices device.ContainerDevices) (string, error) { + parts := strings.Split(annotation, device.OnePodMultiContainerSplitSymbol) + if index < 0 || index >= len(parts) || parts[index] == "" { + return "", fmt.Errorf("annotation has no container entry at index %d", index) + } + parts[index] = device.EncodeContainerDevices(devices) + return strings.Join(parts, device.OnePodMultiContainerSplitSymbol), nil +} + +// containerNameAt returns the pod's container name at the PodDevices +// position, counting init containers first. +func containerNameAt(pod *corev1.Pod, index int) (string, bool) { + if index < 0 { + return "", false + } + if index < len(pod.Spec.InitContainers) { + return pod.Spec.InitContainers[index].Name, true + } + index -= len(pod.Spec.InitContainers) + if index < len(pod.Spec.Containers) { + return pod.Spec.Containers[index].Name, true + } + return "", false +} + +// releaseContainerUsage subtracts one container's reserved usage from the +// node usage snapshot in place. +func releaseContainerUsage(node *NodeUsage, reserved device.ContainerDevices) { + for _, r := range reserved { + for _, deviceList := range node.Devices.DeviceLists { + if deviceList.Device.ID != r.UUID { + continue + } + if deviceList.Device.Used > 0 { + deviceList.Device.Used-- + } + deviceList.Device.Usedmem = max(deviceList.Device.Usedmem-r.Usedmem, 0) + deviceList.Device.Usedcores = max(deviceList.Device.Usedcores-r.Usedcores, 0) + break + } + } +} + +func containerDeviceIDs(devices device.ContainerDevices) []string { + ids := make([]string, 0, len(devices)) + for _, d := range devices { + ids = append(ids, d.UUID) + } + return ids +} + +// numaRefitFailure logs and wraps a refit refusal for the wire. +func numaRefitFailure(pod *corev1.Pod, format string, args ...any) device.NumaRefitResponse { + reason := fmt.Sprintf(format, args...) + if pod != nil { + klog.InfoS("NUMA refit refused", "pod", klog.KObj(pod), "reason", reason) + } else { + klog.InfoS("NUMA refit refused", "reason", reason) + } + return device.NumaRefitResponse{Succeeded: false, FailureReason: reason} +} + +// numaRefitFailureEvent additionally records a warning event on the pod. +func (s *Scheduler) numaRefitFailureEvent(pod *corev1.Pod, format string, args ...any) device.NumaRefitResponse { + response := numaRefitFailure(pod, format, args...) + s.recordNumaRefitResultEvent(pod, "", fmt.Errorf("%s", response.FailureReason)) + return response +} diff --git a/pkg/scheduler/numa_refit_handler_test.go b/pkg/scheduler/numa_refit_handler_test.go new file mode 100644 index 0000000000..84ff8f552b --- /dev/null +++ b/pkg/scheduler/numa_refit_handler_test.go @@ -0,0 +1,468 @@ +/* +Copyright 2026 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. +*/ + +package scheduler + +import ( + "errors" + "maps" + "strings" + "testing" + + "gotest.tools/v3/assert" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + k8stypes "k8s.io/apimachinery/pkg/types" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/device/nvidia" +) + +const ( + refitPodUID = "refit-pod-uid" + refitPodName = "refit-pod" + refitNode = "node-1" +) + +// refitFixture builds a cache-only scheduler tracking one pod that reserves +// GPU-a, on a node carrying GPU-a (NUMA 1) and GPU-b (NUMA 0). +func refitFixture(t *testing.T, gpuBDevmem int32) (*Scheduler, *corev1.Pod) { + t.Helper() + nodes := newNodeManager() + nodes.addNode(refitNode, &device.NodeInfo{ + ID: refitNode, Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: refitNode}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: { + {ID: "GPU-a", Count: 10, Devmem: 40000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-b", Count: 10, Devmem: gpuBDevmem, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Health: true}, + }}, + }) + + reserved := device.PodSingleDevice{{{UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}}} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: refitPodUID, Name: refitPodName, Namespace: "default", + Annotations: map[string]string{ + device.InRequestDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + device.SupportDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + }, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}, + } + pods := device.NewPodManager() + pods.AddPod(pod, refitNode, device.PodDevices{nvidia.NvidiaGPUDevice: reserved}) + + s := &Scheduler{nodeManager: nodes, podManager: pods, quotaManager: device.NewQuotaManager()} + s.quotaManager.Quotas = map[string]*device.DeviceQuota{} + return s, pod +} + +// stubRefitPatch replaces the annotation patch with a capture; returns the +// captured map and a call counter. +func stubRefitPatch(t *testing.T, fail error) (map[string]string, *int) { + t.Helper() + captured := map[string]string{} + calls := 0 + previous := patchPodAnnotations + patchPodAnnotations = func(_ *corev1.Pod, annotations map[string]string) error { + calls++ + if fail != nil { + return fail + } + maps.Copy(captured, annotations) + return nil + } + t.Cleanup(func() { patchPodAnnotations = previous }) + return captured, &calls +} + +func refitTestRequestFor(allowed ...string) device.NumaRefitRequest { + return device.NumaRefitRequest{ + PodUID: refitPodUID, + PodNamespace: "default", + PodName: refitPodName, + NodeName: refitNode, + ContainerIndex: 0, + DeviceType: nvidia.NvidiaGPUDevice, + AllowedDeviceUUIDs: allowed, + } +} + +func trackedUUID(t *testing.T, s *Scheduler) string { + t.Helper() + pi, ok := s.podManager.GetPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: refitPodUID}}) + assert.Equal(t, ok, true) + return pi.Devices[nvidia.NvidiaGPUDevice][0][0].UUID +} + +func TestRefitNumaAllocationMovesReservation(t *testing.T) { + s, _ := refitFixture(t, 40000) + captured, calls := stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-b")) + + assert.Equal(t, response.Succeeded, true, "refit failed: %s", response.FailureReason) + devices, err := device.DecodeContainerDevices(response.ContainerDevices) + assert.NilError(t, err) + assert.Equal(t, len(devices), 1) + assert.Equal(t, devices[0].UUID, "GPU-b") + assert.Equal(t, devices[0].Usedmem, int32(20000)) + assert.Equal(t, devices[0].Usedcores, int32(30)) + + // Both annotations were patched together onto the new device, with the + // value shape preserved byte for byte (no phantom container entries). + assert.Equal(t, *calls, 1) + expected := device.EncodePodSingleDevice(device.PodSingleDevice{{{UUID: "GPU-b", Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}}}) + for _, key := range []string{device.InRequestDevices[nvidia.NvidiaGPUDevice], device.SupportDevices[nvidia.NvidiaGPUDevice]} { + value, ok := captured[key] + assert.Equal(t, ok, true, "annotation %s not patched", key) + assert.Equal(t, value, expected, "annotation %s", key) + } + + // In-memory accounting moved with it, without arming the init shrink. + assert.Equal(t, trackedUUID(t, s), "GPU-b") + pi, _ := s.podManager.GetPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: refitPodUID}}) + assert.Equal(t, pi.InitContainerResourceReleased, false) +} + +func TestRefitNumaAllocationNoOpWhenAlreadyAllowed(t *testing.T) { + s, _ := refitFixture(t, 40000) + _, calls := stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-a", "GPU-b")) + + assert.Equal(t, response.Succeeded, true) + devices, err := device.DecodeContainerDevices(response.ContainerDevices) + assert.NilError(t, err) + assert.Equal(t, devices[0].UUID, "GPU-a") + assert.Equal(t, *calls, 0) + assert.Equal(t, trackedUUID(t, s), "GPU-a") +} + +func TestRefitNumaAllocationInsufficientCapacity(t *testing.T) { + s, _ := refitFixture(t, 16000) + _, calls := stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-b")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, "no allowed device fits"), "reason: %s", response.FailureReason) + assert.Equal(t, *calls, 0) + assert.Equal(t, trackedUUID(t, s), "GPU-a") +} + +func TestRefitNumaAllocationUnmatchedAllowedSet(t *testing.T) { + s, _ := refitFixture(t, 40000) + stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-a-0", "GPU-a-1")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, AllowedSetUnmatched), "reason: %s", response.FailureReason) +} + +func TestRefitNumaAllocationPatchFailureLeavesStateUntouched(t *testing.T) { + s, _ := refitFixture(t, 40000) + _, calls := stubRefitPatch(t, errors.New("apiserver unavailable")) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-b")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, "apiserver unavailable")) + assert.Equal(t, *calls, 1) + assert.Equal(t, trackedUUID(t, s), "GPU-a") +} + +func TestRefitNumaAllocationConsumedAllocation(t *testing.T) { + s, pod := refitFixture(t, 40000) + stubRefitPatch(t, nil) + // Simulate Allocate having consumed the container's pending entry. + pod.Annotations[device.InRequestDevices[nvidia.NvidiaGPUDevice]] = device.EncodePodSingleDevice(device.PodSingleDevice{{}}) + s.podManager.UpdatePod(pod) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-b")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, "no pending"), "reason: %s", response.FailureReason) + assert.Equal(t, trackedUUID(t, s), "GPU-a") +} + +func TestRefitNumaAllocationValidation(t *testing.T) { + tests := []struct { + name string + mutate func(*device.NumaRefitRequest) + wantReason string + }{ + { + name: "unknown pod", + mutate: func(r *device.NumaRefitRequest) { r.PodUID = "other-uid" }, + wantReason: "not tracked", + }, + { + name: "pod identity mismatch", + mutate: func(r *device.NumaRefitRequest) { r.PodName = "other-name" }, + wantReason: "does not match", + }, + { + name: "wrong node", + mutate: func(r *device.NumaRefitRequest) { r.NodeName = "node-2" }, + wantReason: "tracked on node", + }, + { + name: "unknown device type", + mutate: func(r *device.NumaRefitRequest) { r.DeviceType = "NoSuchVendor" }, + wantReason: "unknown device type", + }, + { + name: "empty allowed set", + mutate: func(r *device.NumaRefitRequest) { r.AllowedDeviceUUIDs = nil }, + wantReason: "empty allowed device set", + }, + { + name: "container index out of range", + mutate: func(r *device.NumaRefitRequest) { r.ContainerIndex = 3 }, + wantReason: "outside the pod's containers", + }, + { + name: "negative container index", + mutate: func(r *device.NumaRefitRequest) { r.ContainerIndex = -1 }, + wantReason: "outside the pod's containers", + }, + { + name: "incomplete request", + mutate: func(r *device.NumaRefitRequest) { r.PodNamespace = "" }, + wantReason: "incomplete refit request", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + s, _ := refitFixture(t, 40000) + _, calls := stubRefitPatch(t, nil) + + request := refitTestRequestFor("GPU-b") + test.mutate(&request) + response := s.RefitNumaAllocation(request) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, test.wantReason), + "reason %q does not contain %q", response.FailureReason, test.wantReason) + assert.Equal(t, *calls, 0) + assert.Equal(t, trackedUUID(t, s), "GPU-a") + }) + } +} + +func TestRefitNumaAllocationSeedsExistingAllocations(t *testing.T) { + // Two GPU containers: container 0 already consumed its to-allocate + // entry, container 1 is pending. The refit of container 1 must keep + // container 0's entries intact, including the blank to-allocate slot. + nodes := newNodeManager() + nodes.addNode(refitNode, &device.NodeInfo{ + ID: refitNode, Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: refitNode}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: { + {ID: "GPU-a", Count: 10, Devmem: 40000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-b", Count: 10, Devmem: 40000, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Health: true}, + }}, + }) + + first := device.ContainerDevices{{UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 5000, Usedcores: 10}} + second := device.ContainerDevices{{UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: refitPodUID, Name: refitPodName, Namespace: "default", + Annotations: map[string]string{ + device.InRequestDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(device.PodSingleDevice{{}, second}), + device.SupportDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(device.PodSingleDevice{first, second}), + }, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "first"}, {Name: "second"}}}, + } + pods := device.NewPodManager() + pods.AddPod(pod, refitNode, device.PodDevices{nvidia.NvidiaGPUDevice: {first, second}}) + s := &Scheduler{nodeManager: nodes, podManager: pods, quotaManager: device.NewQuotaManager()} + s.quotaManager.Quotas = map[string]*device.DeviceQuota{} + captured, _ := stubRefitPatch(t, nil) + + request := refitTestRequestFor("GPU-b") + request.ContainerIndex = 1 + response := s.RefitNumaAllocation(request) + + assert.Equal(t, response.Succeeded, true, "refit failed: %s", response.FailureReason) + toAllocate := captured[device.InRequestDevices[nvidia.NvidiaGPUDevice]] + allocatedAnno := captured[device.SupportDevices[nvidia.NvidiaGPUDevice]] + // Container 0's consumed to-allocate slot stays blank; its allocated + // record stays on GPU-a; only container 1 moves to GPU-b. + assert.Assert(t, strings.HasPrefix(toAllocate, ";"), "to-allocate: %q", toAllocate) + assert.Assert(t, strings.Contains(toAllocate, "GPU-b"), "to-allocate: %q", toAllocate) + assert.Assert(t, strings.Contains(allocatedAnno, "GPU-a,NVIDIA,5000,10"), "allocated: %q", allocatedAnno) + assert.Assert(t, strings.Contains(allocatedAnno, "GPU-b,NVIDIA,20000,30"), "allocated: %q", allocatedAnno) + + // In-memory accounting is rebuilt collapsed (one aggregate entry), the + // same shape Filter and the informer store: GPU-a keeps container 0's + // usage, GPU-b carries the refitted container 1. + pi, _ := s.podManager.GetPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: refitPodUID}}) + collapsed := pi.Devices[nvidia.NvidiaGPUDevice] + assert.Equal(t, len(collapsed), 1) + byUUID := map[string]device.ContainerDevice{} + for _, d := range collapsed[0] { + byUUID[d.UUID] = d + } + assert.Equal(t, byUUID["GPU-a"].Usedmem, int32(5000)) + assert.Equal(t, byUUID["GPU-b"].Usedmem, int32(20000)) +} + +func TestRefitNumaAllocationCompetingRefits(t *testing.T) { + // GPU-b only has room for one of the two reservations: after the first + // refit moves onto it, the second must be refused on capacity. + nodes := newNodeManager() + nodes.addNode(refitNode, &device.NodeInfo{ + ID: refitNode, Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: refitNode}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: { + {ID: "GPU-a", Count: 10, Devmem: 40000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-b", Count: 10, Devmem: 20000, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Health: true}, + }}, + }) + pods := device.NewPodManager() + makePod := func(uid, name string) *corev1.Pod { + reserved := device.PodSingleDevice{{{UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 15000, Usedcores: 10}}} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: k8stypes.UID(uid), Name: name, Namespace: "default", + Annotations: map[string]string{ + device.InRequestDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + device.SupportDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + }, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}, + } + pods.AddPod(pod, refitNode, device.PodDevices{nvidia.NvidiaGPUDevice: reserved}) + return pod + } + makePod("pod-1-uid", "pod-1") + makePod("pod-2-uid", "pod-2") + s := &Scheduler{nodeManager: nodes, podManager: pods, quotaManager: device.NewQuotaManager()} + s.quotaManager.Quotas = map[string]*device.DeviceQuota{} + stubRefitPatch(t, nil) + + requestFor := func(uid, name string) device.NumaRefitRequest { + return device.NumaRefitRequest{ + PodUID: uid, PodNamespace: "default", PodName: name, NodeName: refitNode, + DeviceType: nvidia.NvidiaGPUDevice, AllowedDeviceUUIDs: []string{"GPU-b"}, + } + } + + firstResponse := s.RefitNumaAllocation(requestFor("pod-1-uid", "pod-1")) + assert.Equal(t, firstResponse.Succeeded, true, "first refit failed: %s", firstResponse.FailureReason) + + secondResponse := s.RefitNumaAllocation(requestFor("pod-2-uid", "pod-2")) + assert.Equal(t, secondResponse.Succeeded, false) + assert.Assert(t, strings.Contains(secondResponse.FailureReason, "no allowed device fits"), + "reason: %s", secondResponse.FailureReason) + + // The loser keeps its original reservation. + pi, _ := s.podManager.GetPod(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-2-uid"}}) + assert.Equal(t, pi.Devices[nvidia.NvidiaGPUDevice][0][0].UUID, "GPU-a") +} + +func TestRefitNumaAllocationContainerNameMismatch(t *testing.T) { + s, _ := refitFixture(t, 40000) + _, calls := stubRefitPatch(t, nil) + + request := refitTestRequestFor("GPU-b") + request.ContainerName = "sidecar" + response := s.RefitNumaAllocation(request) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, `is "main", not "sidecar"`), "reason: %s", response.FailureReason) + assert.Equal(t, *calls, 0) + + // The matching name is accepted. + request.ContainerName = "main" + response = s.RefitNumaAllocation(request) + assert.Equal(t, response.Succeeded, true, "refit failed: %s", response.FailureReason) +} + +func TestRefitNumaAllocationRejectsMigDevice(t *testing.T) { + nodes := newNodeManager() + nodes.addNode(refitNode, &device.NodeInfo{ + ID: refitNode, Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: refitNode}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: { + {ID: "GPU-a", Count: 10, Devmem: 40000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-b", Count: 7, Devmem: 40000, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Mode: nvidia.MigMode, Health: true}, + }}, + }) + reserved := device.PodSingleDevice{{{UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}}} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: refitPodUID, Name: refitPodName, Namespace: "default", + Annotations: map[string]string{ + device.InRequestDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + device.SupportDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + }, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}, + } + pods := device.NewPodManager() + pods.AddPod(pod, refitNode, device.PodDevices{nvidia.NvidiaGPUDevice: reserved}) + s := &Scheduler{nodeManager: nodes, podManager: pods, quotaManager: device.NewQuotaManager()} + s.quotaManager.Quotas = map[string]*device.DeviceQuota{} + stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-b")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, "MIG"), "reason: %s", response.FailureReason) +} + +func TestRefitNumaAllocationHeterogeneousReservation(t *testing.T) { + // Two devices in one container with differing reserved amounts (possible + // with percentage requests on mixed GPUs) cannot be re-fit faithfully. + nodes := newNodeManager() + nodes.addNode(refitNode, &device.NodeInfo{ + ID: refitNode, Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: refitNode}}, + Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: { + {ID: "GPU-a", Count: 10, Devmem: 40000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-b", Count: 10, Devmem: 20000, Devcore: 100, Numa: 1, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-c", Count: 10, Devmem: 40000, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Health: true}, + {ID: "GPU-d", Count: 10, Devmem: 40000, Devcore: 100, Numa: 0, Type: nvidia.NvidiaGPUDevice, Health: true}, + }}, + }) + reserved := device.PodSingleDevice{{ + {UUID: "GPU-a", Type: nvidia.NvidiaGPUDevice, Usedmem: 40000, Usedcores: 30}, + {UUID: "GPU-b", Type: nvidia.NvidiaGPUDevice, Usedmem: 20000, Usedcores: 30}, + }} + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + UID: refitPodUID, Name: refitPodName, Namespace: "default", + Annotations: map[string]string{ + device.InRequestDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + device.SupportDevices[nvidia.NvidiaGPUDevice]: device.EncodePodSingleDevice(reserved), + }, + }, + Spec: corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}, + } + pods := device.NewPodManager() + pods.AddPod(pod, refitNode, device.PodDevices{nvidia.NvidiaGPUDevice: reserved}) + s := &Scheduler{nodeManager: nodes, podManager: pods, quotaManager: device.NewQuotaManager()} + s.quotaManager.Quotas = map[string]*device.DeviceQuota{} + _, calls := stubRefitPatch(t, nil) + + response := s.RefitNumaAllocation(refitTestRequestFor("GPU-c", "GPU-d")) + + assert.Equal(t, response.Succeeded, false) + assert.Assert(t, strings.Contains(response.FailureReason, "heterogeneous"), "reason: %s", response.FailureReason) + assert.Equal(t, *calls, 0) +} diff --git a/pkg/scheduler/routes/numa_refit_route_test.go b/pkg/scheduler/routes/numa_refit_route_test.go new file mode 100644 index 0000000000..97fd6bdd22 --- /dev/null +++ b/pkg/scheduler/routes/numa_refit_route_test.go @@ -0,0 +1,87 @@ +/* +Copyright 2026 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. +*/ + +package routes + +import ( + "context" + "encoding/json" + "net/http/httptest" + "strings" + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device" + "github.com/Project-HAMi/HAMi/pkg/scheduler" +) + +func TestNumaRefitRoute_NilBody(t *testing.T) { + req := httptest.NewRequest("POST", "/refit", nil) + req.Body = nil + w := httptest.NewRecorder() + + handler := NumaRefit(&scheduler.Scheduler{}) + handler(w, req, nil) + + if w.Code != 400 { + t.Errorf("expected 400 for nil body, got %d", w.Code) + } +} + +func TestNumaRefitRoute_DecodeError(t *testing.T) { + req := httptest.NewRequest("POST", "/refit", strings.NewReader("{not-json")) + w := httptest.NewRecorder() + + handler := NumaRefit(&scheduler.Scheduler{}) + handler(w, req, nil) + + if w.Code != 200 { + t.Errorf("expected 200, got %d", w.Code) + } + var response device.NumaRefitResponse + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("failed to unmarshal response: %v", err) + } + if response.Succeeded || response.FailureReason == "" { + t.Errorf("expected in-band decode failure, got %+v", response) + } +} + +func TestNumaRefitRoute_CacheNotSynced(t *testing.T) { + body, err := json.Marshal(device.NumaRefitRequest{PodUID: "uid"}) + if err != nil { + t.Fatalf("failed to marshal request: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // cancel immediately so WaitForCacheSync fails fast instead of polling forever + + req := httptest.NewRequest("POST", "/refit", strings.NewReader(string(body))).WithContext(ctx) + w := httptest.NewRecorder() + + handler := NumaRefit(&scheduler.Scheduler{}) // zero value: synced defaults to false + handler(w, req, nil) + + if w.Code != 200 { + t.Errorf("expected 200, got %d", w.Code) + } + var response device.NumaRefitResponse + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("failed to unmarshal response: %v", err) + } + if response.Succeeded || !strings.Contains(response.FailureReason, "context cancelled") { + t.Errorf("expected cache-not-synced failure, got %+v", response) + } +} diff --git a/pkg/scheduler/routes/route.go b/pkg/scheduler/routes/route.go index c409f53a8d..743b1d5630 100644 --- a/pkg/scheduler/routes/route.go +++ b/pkg/scheduler/routes/route.go @@ -27,6 +27,7 @@ import ( "k8s.io/klog/v2" extenderv1 "k8s.io/kube-scheduler/extender/v1" + "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/scheduler" ) @@ -184,3 +185,40 @@ func ReadyzRoute(s *scheduler.Scheduler) httprouter.Handle { w.WriteHeader(http.StatusOK) } } + +// NumaRefit handles device-plugin requests to move a pending allocation onto +// kubelet's NUMA-restricted device set. See issue #2080. +func NumaRefit(s *scheduler.Scheduler) httprouter.Handle { + klog.Infoln("Initializing NumaRefit Route") + return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) { + if !checkBody(w, r) { + return + } + + // Limit the body size to prevent deep nesting/resource exhaustion attacks + body := io.LimitReader(r.Body, maxRequestSize) + + var response device.NumaRefitResponse + var request device.NumaRefitRequest + if err := json.NewDecoder(body).Decode(&request); err != nil { + klog.ErrorS(err, "Failed to decode NUMA refit request") + response = device.NumaRefitResponse{FailureReason: err.Error()} + } else if !s.WaitForCacheSync(r.Context()) { + // Poll may return false when context is cancelled + err := fmt.Errorf("context cancelled") + klog.ErrorS(err, "Cache not synced, cannot refit") + response = device.NumaRefitResponse{FailureReason: err.Error()} + } else { + response = s.RefitNumaAllocation(request) + } + + resultBody, err := json.Marshal(response) + if err != nil { + klog.ErrorS(err, "Failed to marshal NUMA refit response", "response", response) + resultBody, _ = json.Marshal(device.NumaRefitResponse{FailureReason: err.Error()}) + writeResponse(w, http.StatusInternalServerError, resultBody) + return + } + writeResponse(w, http.StatusOK, resultBody) + } +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 6fcf559275..12cba75558 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -79,6 +79,13 @@ type Scheduler struct { lock sync.RWMutex synced bool + + // allocLock serializes reservation mutations between Filter and the + // NUMA refit (RefitNumaAllocation). kube-scheduler already serializes + // Filter calls per scheduling cycle, so in the common path this adds no + // contention; it exists so a refit cannot interleave with Filter's + // take-fit-readd span and observe or produce half-applied accounting. + allocLock sync.Mutex } func NewScheduler() *Scheduler { @@ -1089,6 +1096,9 @@ func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFi return s.filterSimulation(args, resourceReqs) } + s.allocLock.Lock() + defer s.allocLock.Unlock() + if pi, ok := s.podManager.TakeAndDeletePod(args.Pod); ok { s.quotaManager.RmUsage(args.Pod, pi.Devices) } diff --git a/pkg/util/numa_alignment.go b/pkg/util/numa_alignment.go index 48c7440a9d..9793d66706 100644 --- a/pkg/util/numa_alignment.go +++ b/pkg/util/numa_alignment.go @@ -28,20 +28,21 @@ import ( // Manager restricts an allocation to. It is distinct from the // nvidia.com/numa-bind annotation, which requests GPU-to-GPU co-location on // one NUMA node at scheduling time. -// -// Only best-effort exists today. A strict mode that fails the allocation on -// an unreconcilable mismatch is introduced together with the NUMA refit that -// can actually enforce it (issue #2080), so the annotation never promises -// semantics that are not implemented yet. type NumaAlignmentMode string const ( // NumaAlignmentNone means the Pod does not opt into NUMA alignment // handling and mismatches are treated exactly as before. NumaAlignmentNone NumaAlignmentMode = "" - // NumaAlignmentBestEffort surfaces a mismatch but never fails the - // allocation because of it. + // NumaAlignmentBestEffort asks for the NUMA refit when it is available + // but never fails the allocation: on a refit failure or with the refit + // disabled, the mismatch is only surfaced and kubelet's own selection + // stands. NumaAlignmentBestEffort NumaAlignmentMode = "best-effort" + // NumaAlignmentStrict fails the allocation when the NUMA refit is + // enabled on the cluster and cannot reconcile the mismatch. With the + // refit disabled, strict only logs the mismatch at error severity. + NumaAlignmentStrict NumaAlignmentMode = "strict" ) // GetNumaAlignmentModeByPod returns the Pod's NUMA alignment mode, or @@ -63,10 +64,10 @@ func GetNumaAlignmentModeByPod(pod *corev1.Pod) (NumaAlignmentMode, error) { // value. Values are case-insensitive and surrounding whitespace is ignored. func ParseNumaAlignmentMode(value string) (NumaAlignmentMode, error) { switch mode := NumaAlignmentMode(strings.ToLower(strings.TrimSpace(value))); mode { - case NumaAlignmentBestEffort: + case NumaAlignmentBestEffort, NumaAlignmentStrict: return mode, nil default: - return NumaAlignmentNone, fmt.Errorf("invalid %s annotation %q: expected %q", - NumaAlignmentAnnotationKey, value, NumaAlignmentBestEffort) + return NumaAlignmentNone, fmt.Errorf("invalid %s annotation %q: expected %q or %q", + NumaAlignmentAnnotationKey, value, NumaAlignmentBestEffort, NumaAlignmentStrict) } } diff --git a/pkg/util/numa_alignment_test.go b/pkg/util/numa_alignment_test.go index 88798fd7d4..7138bb6f36 100644 --- a/pkg/util/numa_alignment_test.go +++ b/pkg/util/numa_alignment_test.go @@ -70,12 +70,18 @@ func TestGetNumaAlignmentModeByPod(t *testing.T) { want: NumaAlignmentBestEffort, }, { - name: "strict is rejected until the refit can enforce it", + name: "strict", pod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ Annotations: map[string]string{NumaAlignmentAnnotationKey: "strict"}, }}, - want: NumaAlignmentNone, - wantErr: true, + want: NumaAlignmentStrict, + }, + { + name: "mixed case strict is accepted", + pod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{ + Annotations: map[string]string{NumaAlignmentAnnotationKey: "Strict"}, + }}, + want: NumaAlignmentStrict, }, { name: "explicit empty value is rejected",