fix(scheduler): add prioritize extender and defer allocation to bind - #2273
fix(scheduler): add prioritize extender and defer allocation to bind#2273blackdragoon26 wants to merge 8 commits into
Conversation
Make Filter side-effect free, introduce Prioritize for policy-based scoring, and defer concrete device allocation to the Bind phase for safe revalidation. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
hami-robot
Bot
added
kind/bug
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: blackdragoon26 The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe scheduler now exposes extender prioritization, retains all feasible nodes during filtering, normalizes node scores, and defers device allocation until bind-time validation. Tests cover endpoint validation, scoring, allocation, rollback, retries, and stale pod protection. ChangesScheduler prioritization and binding
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant KubeScheduler
participant SchedulerHTTP
participant Scheduler
participant PodManager
KubeScheduler->>SchedulerHTTP: POST /prioritize
SchedulerHTTP->>Scheduler: Prioritize(extenderArgs)
Scheduler-->>SchedulerHTTP: normalized host priorities
SchedulerHTTP-->>KubeScheduler: JSON priorities
KubeScheduler->>Scheduler: Bind(selected node)
Scheduler->>PodManager: reserve devices
Scheduler->>Scheduler: patch allocation annotations
Scheduler->>Scheduler: bind pod
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
pkg/scheduler/routes/route_test.go (1)
79-90: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the payload prove truncation at
maxRequestSize.The body contains only spaces. The decoder returns
io.EOFfor an empty document, so the 400 status does not show thatio.LimitReadertruncated the request. Use an oversized but otherwise valid JSON document. Truncation then produces "unexpected EOF", which is caused by the limit.♻️ Proposed change to the payload
- hugePayload := strings.Repeat(" ", maxRequestSize+100) + // Valid JSON larger than maxRequestSize; truncation must surface as a decode error. + hugePayload := `{"Pod":{"metadata":{"name":"` + strings.Repeat("a", maxRequestSize+100) + `"}}}` req := httptest.NewRequest("POST", "/prioritize", strings.NewReader(hugePayload))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/scheduler/routes/route_test.go` around lines 79 - 90, The TestMaxRequestSizePrioritize test currently uses whitespace, which only verifies empty-document rejection rather than request truncation. Replace hugePayload with an oversized, otherwise valid JSON document so truncation at maxRequestSize causes an unexpected EOF while preserving the expected 400 response.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/scheduler/routes/route.go`:
- Around line 101-117: Validate that decoded extenderArgs.Pod is non-nil
immediately after decoding and before WaitForCacheSync or s.Prioritize,
returning an HTTP 400 error for missing Pod requests. Ensure the validation also
prevents the Prioritize error path from dereferencing a nil Pod.
In `@pkg/scheduler/scheduler_test.go`:
- Around line 955-962: Update setupBindAllocationTest around the
client.KubeClient assignment to save the existing package-level client and
register cleanup that restores it after the test. Keep fake.NewSimpleClientset
unchanged, and ensure PatchPodAnnotations still uses the test fake client during
the test.
In `@pkg/scheduler/scheduler.go`:
- Line 1020: Update the failedNodes declaration in the surrounding scheduler
flow to declare the variable without initializing it, preserving the existing
assignments in both branches and eliminating the unused-value lint error.
- Around line 906-918: Update the patch-failure rollback in the allocation flow
around podManager.AddPod and util.PatchPodAnnotations to remove the added guard
from quotaManager.RmUsage, ensuring quota usage is rolled back whenever
allocation is non-nil. Keep podManager.DelPod guarded by allocation as currently
implemented.
---
Nitpick comments:
In `@pkg/scheduler/routes/route_test.go`:
- Around line 79-90: The TestMaxRequestSizePrioritize test currently uses
whitespace, which only verifies empty-document rejection rather than request
truncation. Replace hugePayload with an oversized, otherwise valid JSON document
so truncation at maxRequestSize causes an unexpected EOF while preserving the
expected 400 response.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: af43d6f5-cf10-42fb-bed5-b89670a9e3f6
📒 Files selected for processing (6)
charts/hami/templates/scheduler/configmap.yamlcmd/scheduler/main.gopkg/scheduler/routes/route.gopkg/scheduler/routes/route_test.gopkg/scheduler/scheduler.gopkg/scheduler/scheduler_test.go
Return 400 Bad Request when extender args lack a Pod. Also fix the failedNodes lint error, strengthen request-size coverage, and restore global client state after tests. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
Codecov Report❌ Patch coverage is
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Refactor prioritizeRoute to accept dependencies for better testability. Add unit tests covering cache sync, success, scheduler errors, and zero device requests. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
Reject stale Pod UID binding requests and avoid replacing existing allocation records. Roll back only allocations created by the current attempt while preserving reservations across Bind API failures for scheduler retries. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/device/pods.go`:
- Around line 82-103: Protect reservations created by AddPodIfAbsent from later
AddPod informer updates by adding an ownership or generation marker to PodInfo
and preserving bind-owned entries during informer reconciliation. Update Bind
rollback to remove the entry only when its marker still matches the reservation
it created, preventing deletion of a replacement reservation. Add a regression
test covering bind insertion, a subsequent AddPod update, and rollback.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e7265ec6-5840-4b4e-98bf-46c0e3015e75
📒 Files selected for processing (6)
pkg/device/pod_test.gopkg/device/pods.gopkg/scheduler/routes/route.gopkg/scheduler/routes/route_test.gopkg/scheduler/scheduler.gopkg/scheduler/scheduler_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/scheduler/scheduler.go
Track bind-owned allocations with reservation tokens so informer reconciliation and stale rollbacks cannot remove persisted or replacement allocations. Compare decoded informer state with the canonical Bind allocation and cover ambiguous patch failures. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
Also, one major design question before I move this out of draft: with reservation shifted to Bind, PodGroup members can pass Filter/Permit before HAMi records usage. Bind revalidates under the node lock, but a permitted pod may need a scheduler retry. Is that retry behaviour acceptable for the HAMi?? |
Bind had the same retry gap, so I have fixed that here and left Filter recovery to #2044 if it is to be merged first. |
Keep Bind reservations owned until informer confirmation. Revalidate retries without dropping the previous allocation, and restore it atomically when the pod annotation patch fails. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
|
@CodeRabbit review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
pkg/device/pods.go (1)
146-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the device comparison with
AddPod, and log the rejection.
AddPodcompares informer devices withreservationMatchesInformerDevices, which normalizes trailing empty container entries.ReplacePodReservationcompares with rawreflect.DeepEqual. If an informer update rewritescurrent.Devicesinto the decoded representation between theGetPodcall inBindand this call, the comparison fails andBindreports "pod allocation already exists for UID". The outcome is a conservative bind failure, not data corruption, but the message does not describe the real cause.Reuse the normalized comparison, and add a log line on the rejection path so the compare-and-swap miss is diagnosable.
♻️ Proposed change
current, exists := m.pods[pod.UID] - if !exists || expected == nil || current.NodeID != expected.NodeID || !reflect.DeepEqual(current.Devices, expected.Devices) { + if !exists || expected == nil || current.NodeID != expected.NodeID || + !reservationMatchesInformerDevices(current.Devices, expected.Devices) { + klog.V(5).InfoS("Pod reservation replacement rejected", + "pod", klog.KRef(pod.Namespace, pod.Name), + "exists", exists, + ) return 0, false }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/device/pods.go` around lines 146 - 158, Update ReplacePodReservation to use reservationMatchesInformerDevices for the device comparison, matching AddPod’s normalization of trailing empty container entries instead of reflect.DeepEqual. On the compare-and-swap rejection path, add a diagnostic log identifying the pod reservation mismatch before returning 0, false.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pkg/device/pods.go`:
- Around line 146-158: Update ReplacePodReservation to use
reservationMatchesInformerDevices for the device comparison, matching AddPod’s
normalization of trailing empty container entries instead of reflect.DeepEqual.
On the compare-and-swap rejection path, add a diagnostic log identifying the pod
reservation mismatch before returning 0, false.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e48644b5-a030-47a3-9d05-2dc662643487
📒 Files selected for processing (4)
pkg/device/pod_test.gopkg/device/pods.gopkg/scheduler/scheduler.gopkg/scheduler/scheduler_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/scheduler/scheduler.go
Return allocation snapshots under the PodManager lock and retain bind-owned state when Pod annotation patch results are ambiguous. Reconcile delayed informer updates and clean unconfirmed reservations on Pod deletion. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
|
@CodeRabbit review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (2)
pkg/device/pods.go (1)
151-155: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the compared node identifiers to the rejection log.
ReplacePodReservationreturnsfalsefor four distinct causes: missing entry, nilexpected, node mismatch, and device mismatch. The log records onlyexists. Bind converts the rejection intopod allocation already exists for UID ..., which does not identify the cause. Include the compared node identifiers so operators can separate a node mismatch from a device mismatch.♻️ Proposed logging improvement
current, exists := m.pods[pod.UID] if !exists || expected == nil || current.NodeID != expected.NodeID || !reservationMatchesInformerDevices(expected.Devices, current.Devices) { klog.V(5).InfoS("Pod reservation replacement rejected", "pod", klog.KRef(pod.Namespace, pod.Name), "exists", exists, + "currentNodeID", currentNodeID, + "expectedNodeID", expectedNodeID, ) return false }Compute the two identifiers before the check so that a missing entry or a nil
expectedstays safe:currentNodeID := "" if exists { currentNodeID = current.NodeID } expectedNodeID := "" if expected != nil { expectedNodeID = expected.NodeID }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/device/pods.go` around lines 151 - 155, Update ReplacePodReservation’s rejection log to include the compared current and expected node identifiers. Compute them safely before the rejection check, using empty values when the reservation is missing or expected is nil, then add both identifiers alongside the existing exists field in the klog.V(5) message.pkg/device/pod_test.go (1)
438-479: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd coverage for the
ReplacePodReservationrejection path.Both tests cover only the success path.
ReplacePodReservationreturnsfalsefor a node mismatch, a device mismatch, a nilexpected, and a missing entry. Bind converts thatfalseinto a binding failure, so the rejection branch changes user-visible behavior. Add a case that passes a staleexpectedwith a different node identifier and asserts that the stored allocation is unchanged.♻️ Proposed additional test
func TestReplacePodReservationRejectsStaleExpectation(t *testing.T) { manager := NewPodManager() pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{UID: "pod-uid", Namespace: "default", Name: "pod"}} currentDevices := PodDevices{"device": {{{UUID: "device-current"}}}} manager.AddPod(pod, "node-current", currentDevices) stale := &PodInfo{Pod: pod, NodeID: "node-stale", Devices: PodDevices{"device": {{{UUID: "device-stale"}}}}} assert.Equal(t, false, manager.ReplacePodReservation(pod, stale, "node-new", PodDevices{"device": {{{UUID: "device-new"}}}})) assert.Equal(t, false, manager.ReplacePodReservation(pod, nil, "node-new", PodDevices{"device": {{{UUID: "device-new"}}}})) allocation, ok := manager.GetPod(pod) assert.Equal(t, true, ok) assert.Equal(t, "node-current", allocation.NodeID) _, owned := manager.reservations[pod.UID] assert.Equal(t, false, owned) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/device/pod_test.go` around lines 438 - 479, Add rejection-path coverage alongside TestReplacementReservationRetainsOwnershipUntilObserved: create a current reservation, call ReplacePodReservation with a stale PodInfo whose NodeID differs, and assert it returns false while the stored allocation remains unchanged. Also cover a nil expected value if supported by the existing method contract, verifying the reservation is still present with its original ownership state.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pkg/device/pod_test.go`:
- Around line 438-479: Add rejection-path coverage alongside
TestReplacementReservationRetainsOwnershipUntilObserved: create a current
reservation, call ReplacePodReservation with a stale PodInfo whose NodeID
differs, and assert it returns false while the stored allocation remains
unchanged. Also cover a nil expected value if supported by the existing method
contract, verifying the reservation is still present with its original ownership
state.
In `@pkg/device/pods.go`:
- Around line 151-155: Update ReplacePodReservation’s rejection log to include
the compared current and expected node identifiers. Compute them safely before
the rejection check, using empty values when the reservation is missing or
expected is nil, then add both identifiers alongside the existing exists field
in the klog.V(5) message.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7cd3eaea-be8c-417c-ab30-e91282819c24
📒 Files selected for processing (4)
pkg/device/pod_test.gopkg/device/pods.gopkg/scheduler/scheduler.gopkg/scheduler/scheduler_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/scheduler/scheduler.go
Verify stale and missing expected allocations cannot replace or disown the active reservation. Include current and expected node IDs in rejection diagnostics. Signed-off-by: blackdragoon26 <sankalp.jha9643@gmail.com>
|
@CodeRabbit review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
pkg/scheduler/scheduler_test.go (1)
1086-1089: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract a shared quota-assertion helper.
The block that fetches the namespace quota and asserts
mock-memory/mock-coresUsedvalues repeats verbatim six times: Lines 1086-1089, 1127-1130, 1151-1154, 1182-1185, 1223-1226, and 1244-1247. Only the two expected values change between occurrences.Extract a helper, for example
assertQuotaUsage(t, s, namespace, wantMem, wantCores int64), and call it from each test. This reduces duplication and keeps future bind-retry test additions consistent.♻️ Proposed helper extraction
func assertQuotaUsage(t *testing.T, s *Scheduler, namespace string, wantMem, wantCores int64) { t.Helper() quota := s.quotaManager.GetResourceQuota()[namespace] assert.Assert(t, quota != nil) assert.Equal(t, wantMem, (*quota)["example.com/mock-memory"].Used) assert.Equal(t, wantCores, (*quota)["example.com/mock-cores"].Used) }- quota := s.quotaManager.GetResourceQuota()[pod.Namespace] - assert.Assert(t, quota != nil) - assert.Equal(t, int64(1), (*quota)["example.com/mock-memory"].Used) - assert.Equal(t, int64(1), (*quota)["example.com/mock-cores"].Used) + assertQuotaUsage(t, s, pod.Namespace, 1, 1)Also applies to: 1127-1130, 1137-1158, 1151-1154, 1160-1186, 1182-1185, 1188-1227, 1223-1226, 1229-1248, 1244-1247
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/scheduler/scheduler_test.go` around lines 1086 - 1089, Extract the repeated quota lookup and mock resource usage assertions into an assertQuotaUsage helper near the scheduler tests, accepting the test, scheduler, namespace, and expected memory/core values and marking itself as a helper. Replace all six duplicated assertion blocks with calls to this helper, preserving each test’s existing expected values.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pkg/scheduler/scheduler_test.go`:
- Around line 1086-1089: Extract the repeated quota lookup and mock resource
usage assertions into an assertQuotaUsage helper near the scheduler tests,
accepting the test, scheduler, namespace, and expected memory/core values and
marking itself as a helper. Replace all six duplicated assertion blocks with
calls to this helper, preserving each test’s existing expected values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 09b8434b-c985-4077-88b6-b41c4ddd0438
📒 Files selected for processing (4)
pkg/device/pod_test.gopkg/device/pods.gopkg/scheduler/scheduler.gopkg/scheduler/scheduler_test.go
🚧 Files skipped from review as they are similar to previous changes (3)
- pkg/scheduler/scheduler.go
- pkg/device/pods.go
- pkg/device/pod_test.go
|
closing for now, an architecture change like moving reservation from filter to bind needs design agreement in #2264 first, not a 1100 line pr. pls keep the discussion there w/ a short design note (bind time contention, old configmaps w/o prioritizeVerb, perf on large clusters) and once maintainers agree on the direction a fresh pr is welcome. |
|
Yeah, fair point. I honestly started this thinking it was a small extender fix, but once reservation moved from Filter to Bind, the retry and state ownership paths kept expanding, and the PR became much bigger and bigggerr than it should have before design agreement. Well, I learned a lot while working through it though.
|
2 participants
What type of PR is this?
/kind bug
What this PR does / why we need it:
HAMi's scheduler extender currently selects a single node during
Filter. This prevents kube-scheduler scoring plugins, including preferred node affinity, topology spread, and other soft placement policies, from comparing all HAMi-feasible nodes.This change separates the extender phases:
Filterreturns every node on which the requested devices can fit, without patching the Pod or reserving devices.Prioritizeexposes HAMi's binpack or spread preference as normalized Kubernetes extender scores in the0..10range.Bindlocks and revalidates the node selected by kube-scheduler, then commits the concrete device allocation and Pod annotations immediately before binding.The Helm scheduler configuration now enables the
prioritizeverb for both current and legacy kube-scheduler configuration formats.Which issue(s) this PR fixes:
Fixes #2264
Special notes for your reviewer:
This is opened as a draft because the change moves device reservation from Filter to Bind and I would like maintainer feedback on that phase boundary and retry semantics.
Verification performed:
go test ./pkg/scheduler ./pkg/scheduler/routes -count=1go test -race ./pkg/scheduler ./pkg/scheduler/routes -count=1go vet ./pkg/scheduler/... ./cmd/schedulergofmt -don every modified Go filehelm template hami-2264 ./charts/hami --kube-version 1.36.1The kind reproduction used two GPU-capable workers. HAMi's binpack calculation preferred
hami-2264-repro-worker2(2.25) while preferred node affinity selectedhami-2264-repro-worker. Filter returned both workers, kube-scheduler selected the affinity-preferred worker, and Bind revalidated and allocatedGPU-MOCK-WORKER-0on that selected worker. kube-scheduler reportedfeasibleNodes=2and completed the binding successfully.The full local Go package run passed for ordinary packages. Cluster-dependent E2E packages could not connect to the local kind API from the restricted test process; the scheduler behavior was instead verified directly against that cluster as described above.
Does this PR introduce a user-facing change?:
Yes. HAMi-managed Pods can now participate in kube-scheduler scoring across all HAMi-feasible nodes. Placement can therefore differ when Kubernetes soft placement preferences conflict with HAMi's binpack or spread preference.
AI assistance disclosure: I used Codex for codebase analysis, implementation support, test generation, reproduction design, and drafting this PR description. I manually executed and reviewed the unit, race, static, Helm, image-build, and kind-cluster verification described above, inspected the resulting diff, and authored the signed commit under my own identity.
Summary by CodeRabbit
New Features
/prioritizeendpoint.Bug Fixes
Tests