From 1c81382c6c432f4c99bc4100899a86ed9bfe3ef2 Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 11:26:36 +0530 Subject: [PATCH 01/11] feat(util): support strict NUMA alignment Adds strict as a valid hami.io/numa-alignment mode. When NUMA refit is enabled, strict mode fails the allocation if the mismatch cannot be fixed. Without refit enabled, it only logs the mismatch as an error. Signed-off-by: Saiyam Pathak --- pkg/util/numa_alignment.go | 21 +++++++++++---------- pkg/util/numa_alignment_test.go | 12 +++++++++--- 2 files changed, 20 insertions(+), 13 deletions(-) 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", From fa84afe3593961c2b3de25167a0b28d7af708e9d Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 11:26:36 +0530 Subject: [PATCH 02/11] feat(device): add ReplacePodDevices Adds ReplacePodDevices for replacing a pod's tracked devices without triggering init-container resource release. This is used by NUMA refit when moving an existing reservation to different devices. Signed-off-by: Saiyam Pathak --- pkg/device/pods.go | 17 ++++++++++++ pkg/device/pods_replace_test.go | 47 +++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+) create mode 100644 pkg/device/pods_replace_test.go 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) +} From 7f8b1602fc3b9e48a0d847bc28c65f84ce71bf22 Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 11:26:36 +0530 Subject: [PATCH 03/11] feat(scheduler): add NUMA refit endpoint Adds the scheduler-side NUMA refit handler and /refit route. The handler re-runs device fitting using the devices allowed by kubelet, updates the affected container's allocation annotations, and rebuilds the pod's reservation from the updated annotations. It also serializes refit with normal scheduling and rejects unsupported or invalid requests without changing the existing allocation. Signed-off-by: Saiyam Pathak --- cmd/scheduler/main.go | 1 + pkg/scheduler/event.go | 18 + pkg/scheduler/numa_refit_handler.go | 308 +++++++++++++ pkg/scheduler/numa_refit_handler_test.go | 421 ++++++++++++++++++ pkg/scheduler/routes/numa_refit_route_test.go | 87 ++++ pkg/scheduler/routes/route.go | 40 ++ pkg/scheduler/scheduler.go | 10 + 7 files changed, 885 insertions(+) create mode 100644 pkg/scheduler/numa_refit_handler.go create mode 100644 pkg/scheduler/numa_refit_handler_test.go create mode 100644 pkg/scheduler/routes/numa_refit_route_test.go 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/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..9f407427a7 --- /dev/null +++ b/pkg/scheduler/numa_refit_handler.go @@ -0,0 +1,308 @@ +/* +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 ( + "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" +) + +// patchPodAnnotations is a test seam around util.PatchPodAnnotations. +var patchPodAnnotations = util.PatchPodAnnotations + +// 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) + } + + if req.ContainerIndex < 0 { + return s.numaRefitFailureEvent(pod, "container index %d is negative", req.ContainerIndex) + } + if name, ok := containerNameAt(pod, req.ContainerIndex); req.ContainerName != "" && ok && 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)} + } + + 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..fc045146c2 --- /dev/null +++ b/pkg/scheduler/numa_refit_handler_test.go @@ -0,0 +1,421 @@ +/* +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), + }, + }} + 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: "no pending", + }, + { + name: "negative container index", + mutate: func(r *device.NumaRefitRequest) { r.ContainerIndex = -1 }, + wantReason: "negative", + }, + { + 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), + }, + }} + 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, pod := refitFixture(t, 40000) + pod.Spec.Containers = []corev1.Container{{Name: "main"}} + s.podManager.UpdatePod(pod) + _, 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), + }, + }} + 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) +} 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..cbfdbcac51 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,42 @@ 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 + } + + var buf bytes.Buffer + // Limit the body size to prevent deep nesting/resource exhaustion attacks + limitedReader := io.LimitReader(r.Body, maxRequestSize) + body := io.TeeReader(limitedReader, &buf) + + 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) } From dfc6ceafb612729f4ac692a5ccbe6f25967a199b Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 11:26:36 +0530 Subject: [PATCH 04/11] feat(device-plugin): refit on NUMA allocation mismatch When the GPU selected by the scheduler is not available to kubelet, ask the scheduler to refit the allocation onto one of kubelet's allowed devices. Best-effort mode keeps the existing fallback behavior if refit fails, while strict mode fails the allocation. The refit client is only enabled when HAMI_SCHEDULER_ENDPOINT is configured. Signed-off-by: Saiyam Pathak --- .../nvinternal/plugin/numa_refit_client.go | 215 ++++++++++++++++++ .../plugin/numa_refit_client_test.go | 214 +++++++++++++++++ .../nvidiadevice/nvinternal/plugin/server.go | 33 ++- .../plugin/server_numa_alignment_test.go | 9 +- 4 files changed, 463 insertions(+), 8 deletions(-) create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go create mode 100644 pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go 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..75f8ebd915 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go @@ -0,0 +1,215 @@ +/* +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 + } + err = fmt.Errorf("refit-selected devices are not allocatable: %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 < 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..7cb14c741d --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go @@ -0,0 +1,214 @@ +/* +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) +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index e0480a6584..4d0f6f7967 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,24 @@ 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: failing the call terminally fails the + // pod's admission rather than running misaligned. 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 +716,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"}, }, } From ae759f19ea0d76d447c8c1e150a6a80f6ade384f Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 11:26:36 +0530 Subject: [PATCH 05/11] feat(helm): add NUMA refit configuration Adds devicePlugin.numaRefit Helm configuration. When enabled, the device plugin is given the scheduler endpoint and TLS settings needed for NUMA refit. The feature remains disabled by default. Signed-off-by: Saiyam Pathak --- .../device-plugin/daemonsetnvidia.yaml | 14 ++++++++++++++ charts/hami/values.yaml | 17 +++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index 0e785c363b..679924ca51 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -122,6 +122,20 @@ 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 or (not (hasKey . "tlsInsecure")) .tlsInsecure }} + - name: HAMI_SCHEDULER_TLS_INSECURE + value: "true" + {{- end }} + {{- if .caFile }} + - name: HAMI_SCHEDULER_CA_FILE + value: {{ .caFile | quote }} + {{- end }} + {{- end }} + {{- end }} {{- with .Values.devicePlugin.extraEnvs }} {{- . | toYaml | nindent 12 }} {{- end }} 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. From e323b5cd1c29fc34624859b0d04ab14328d6032b Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 20:18:20 +0530 Subject: [PATCH 06/11] fix(scheduler): refuse heterogeneous reservations in refit A fit request carries one memory/core amount for all devices, so a reservation with differing per-device amounts (possible with percentage requests on mixed GPUs) cannot be re-fit faithfully. Refuse it instead of rewriting the other devices' accounting. Signed-off-by: Saiyam Pathak --- pkg/scheduler/numa_refit_handler.go | 9 ++++++ pkg/scheduler/numa_refit_handler_test.go | 37 ++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/pkg/scheduler/numa_refit_handler.go b/pkg/scheduler/numa_refit_handler.go index 9f407427a7..74e4f4357b 100644 --- a/pkg/scheduler/numa_refit_handler.go +++ b/pkg/scheduler/numa_refit_handler.go @@ -127,6 +127,15 @@ func (s *Scheduler) RefitNumaAllocation(req device.NumaRefitRequest) device.Numa 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 { diff --git a/pkg/scheduler/numa_refit_handler_test.go b/pkg/scheduler/numa_refit_handler_test.go index fc045146c2..53d8ba26ba 100644 --- a/pkg/scheduler/numa_refit_handler_test.go +++ b/pkg/scheduler/numa_refit_handler_test.go @@ -419,3 +419,40 @@ func TestRefitNumaAllocationRejectsMigDevice(t *testing.T) { 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), + }, + }} + 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) +} From b3662c7c033493b3daa46fcba5f831eb315abde5 Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 20:18:20 +0530 Subject: [PATCH 07/11] fix(helm): require explicit tlsInsecure for the refit Only set HAMI_SCHEDULER_TLS_INSECURE when tlsInsecure is explicitly true. A values file that omits the key now gets certificate verification instead of silently skipping it. Signed-off-by: Saiyam Pathak --- charts/hami/templates/device-plugin/daemonsetnvidia.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index 679924ca51..9fc201430f 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -126,7 +126,7 @@ spec: {{- 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 or (not (hasKey . "tlsInsecure")) .tlsInsecure }} + {{- if .tlsInsecure }} - name: HAMI_SCHEDULER_TLS_INSECURE value: "true" {{- end }} From f4494ae028064fa7e2e755be2133dff95674be2b Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Thu, 20 Aug 2026 20:18:20 +0530 Subject: [PATCH 08/11] chore: address review nitpicks in the refit path Drops the unused request body copy in the refit route and guards podContainerNameAt against a negative index, matching its scheduler-side twin. Signed-off-by: Saiyam Pathak --- .../nvidiadevice/nvinternal/plugin/numa_refit_client.go | 3 +++ pkg/scheduler/routes/route.go | 4 +--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go index 75f8ebd915..3269a8ceaa 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go @@ -188,6 +188,9 @@ func (plugin *NvidiaDevicePlugin) requestNumaRefit(ctx context.Context, pod *cor // 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 } diff --git a/pkg/scheduler/routes/route.go b/pkg/scheduler/routes/route.go index cbfdbcac51..743b1d5630 100644 --- a/pkg/scheduler/routes/route.go +++ b/pkg/scheduler/routes/route.go @@ -195,10 +195,8 @@ func NumaRefit(s *scheduler.Scheduler) httprouter.Handle { return } - var buf bytes.Buffer // Limit the body size to prevent deep nesting/resource exhaustion attacks - limitedReader := io.LimitReader(r.Body, maxRequestSize) - body := io.TeeReader(limitedReader, &buf) + body := io.LimitReader(r.Body, maxRequestSize) var response device.NumaRefitResponse var request device.NumaRefitRequest From 5447a3a4dbf63c96057ad3155fb344fcb26f831d Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Fri, 21 Aug 2026 13:37:43 +0530 Subject: [PATCH 09/11] fix(scheduler): validate refit container index against the pod spec Rejects indexes with no matching init or regular container, and adds a resource-version precondition to the annotation patch so a stale refit cannot overwrite a newer update; conflicts fail instead of retrying. Signed-off-by: Saiyam Pathak --- pkg/scheduler/numa_refit_handler.go | 30 +++++++++-- pkg/scheduler/numa_refit_handler_test.go | 68 ++++++++++++++---------- 2 files changed, 64 insertions(+), 34 deletions(-) diff --git a/pkg/scheduler/numa_refit_handler.go b/pkg/scheduler/numa_refit_handler.go index 74e4f4357b..3643018490 100644 --- a/pkg/scheduler/numa_refit_handler.go +++ b/pkg/scheduler/numa_refit_handler.go @@ -17,6 +17,8 @@ limitations under the License. package scheduler import ( + "context" + "encoding/json" "fmt" "maps" "strings" @@ -29,10 +31,27 @@ import ( "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 is a test seam around util.PatchPodAnnotations. -var patchPodAnnotations = util.PatchPodAnnotations +// 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 @@ -84,10 +103,11 @@ func (s *Scheduler) RefitNumaAllocation(req device.NumaRefitRequest) device.Numa return s.numaRefitFailureEvent(pod, "pod %s/%s is tracked on node %s, not %s", pod.Namespace, pod.Name, pi.NodeID, req.NodeName) } - if req.ContainerIndex < 0 { - return s.numaRefitFailureEvent(pod, "container index %d is negative", req.ContainerIndex) + name, ok := containerNameAt(pod, req.ContainerIndex) + if !ok { + return s.numaRefitFailureEvent(pod, "container index %d is outside the pod's containers", req.ContainerIndex) } - if name, ok := containerNameAt(pod, req.ContainerIndex); req.ContainerName != "" && ok && name != req.ContainerName { + if req.ContainerName != "" && name != req.ContainerName { return s.numaRefitFailureEvent(pod, "container index %d is %q, not %q", req.ContainerIndex, name, req.ContainerName) } diff --git a/pkg/scheduler/numa_refit_handler_test.go b/pkg/scheduler/numa_refit_handler_test.go index 53d8ba26ba..84ff8f552b 100644 --- a/pkg/scheduler/numa_refit_handler_test.go +++ b/pkg/scheduler/numa_refit_handler_test.go @@ -51,13 +51,16 @@ func refitFixture(t *testing.T, gpuBDevmem int32) (*Scheduler, *corev1.Pod) { }) 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), + 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}) @@ -230,12 +233,12 @@ func TestRefitNumaAllocationValidation(t *testing.T) { { name: "container index out of range", mutate: func(r *device.NumaRefitRequest) { r.ContainerIndex = 3 }, - wantReason: "no pending", + wantReason: "outside the pod's containers", }, { name: "negative container index", mutate: func(r *device.NumaRefitRequest) { r.ContainerIndex = -1 }, - wantReason: "negative", + wantReason: "outside the pod's containers", }, { name: "incomplete request", @@ -335,13 +338,16 @@ func TestRefitNumaAllocationCompetingRefits(t *testing.T) { 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), + 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 } @@ -372,9 +378,7 @@ func TestRefitNumaAllocationCompetingRefits(t *testing.T) { } func TestRefitNumaAllocationContainerNameMismatch(t *testing.T) { - s, pod := refitFixture(t, 40000) - pod.Spec.Containers = []corev1.Container{{Name: "main"}} - s.podManager.UpdatePod(pod) + s, _ := refitFixture(t, 40000) _, calls := stubRefitPatch(t, nil) request := refitTestRequestFor("GPU-b") @@ -401,13 +405,16 @@ func TestRefitNumaAllocationRejectsMigDevice(t *testing.T) { }}, }) 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), + 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()} @@ -437,13 +444,16 @@ func TestRefitNumaAllocationHeterogeneousReservation(t *testing.T) { {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), + 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()} From 515e2c3f3c7623bd66f11e8c0c4a8bc0c4124568 Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Fri, 21 Aug 2026 13:37:43 +0530 Subject: [PATCH 10/11] fix(device-plugin): fail allocation when a committed refit cannot be honored Once the scheduler has moved the reservation, falling back to kubelet's own selection would leave runtime and accounting divergent, so an unmappable refit response now fails the allocation in both modes. Signed-off-by: Saiyam Pathak --- .../nvinternal/plugin/numa_refit_client.go | 5 +++- .../plugin/numa_refit_client_test.go | 28 +++++++++++++++++++ .../nvidiadevice/nvinternal/plugin/server.go | 9 +++--- 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go index 3269a8ceaa..e1b63a00ff 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client.go @@ -124,7 +124,10 @@ func (plugin *NvidiaDevicePlugin) tryNumaRefit(ctx context.Context, pod *corev1. klog.InfoS("NUMA refit succeeded", "pod", klog.KObj(pod), "container", containerIndex, "devices", replicas) return replicas, nil } - err = fmt.Errorf("refit-selected devices are not allocatable: %w", selectErr) + // 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 { 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 index 7cb14c741d..69dfb126c7 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/numa_refit_client_test.go @@ -212,3 +212,31 @@ func TestNumaRefitTLSConfigInsecureOptOut(t *testing.T) { 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 4d0f6f7967..bc23113f44 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -671,10 +671,11 @@ func (plugin *NvidiaDevicePlugin) GetPreferredAllocation(ctx context.Context, r refitDevices, refitErr := plugin.tryNumaRefit(ctx, pendingPod, annotationIndices[idx], req, err) switch { case refitErr != nil: - // Strict mode: failing the call terminally fails the - // pod's admission rather than running misaligned. Mark - // the bind failed so the node lock is released instead - // of blocking the node until the lock timeout. + // 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) } From 7e54be343a1771fd03754b976837e498fd2f8cd2 Mon Sep 17 00:00:00 2001 From: Saiyam Pathak Date: Fri, 21 Aug 2026 13:37:43 +0530 Subject: [PATCH 11/11] feat(helm): mount an optional CA secret for the refit devicePlugin.numaRefit.caSecret mounts a Secret (key ca.crt) read-only and points HAMI_SCHEDULER_CA_FILE at it, so verified TLS needs no manual file placement. caFile keeps working for pre-provisioned paths. Signed-off-by: Saiyam Pathak --- .../templates/device-plugin/daemonsetnvidia.yaml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index 9fc201430f..d3baa4cdd9 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -133,6 +133,9 @@ spec: {{- 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 }} @@ -161,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. @@ -219,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 }}