fix(util): make same-pod node lock acquisition idempotent - #2256
Conversation
Signed-off-by: shivv23 <shivamkumar0423@gmail.com>
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: shivv23 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 |
📝 WalkthroughWalkthrough
ChangesShared node-lock binding
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/scheduler_test.go`:
- Line 2140: Update the test cleanup around the assignment to client.KubeClient
to save the original global client before replacing it, then restore that value
during cleanup alongside config.NodeLockRetryTimeout and device.DevicesMap.
Ensure later tests cannot retain the fake client or its mutated objects.
- Around line 2155-2159: Update the release flow around releaseAllDevices and
ReleaseNodeLock so releasing the first backend does not clear the shared
node-lock annotation while the pod still has other backends. Add an intermediate
assertion in the scheduler test after the first backend release confirming the
lock remains, then retain the existing lock-release behavior only after all
backends finish.
In `@pkg/util/nodelock/nodelock.go`:
- Around line 232-238: Use Pod UID to distinguish ownership in the idempotent
acquisition branch around the node-lock ownership parsing and comparison in
pkg/util/nodelock/nodelock.go:232-238; persist and parse pods.UID in
GenerateNodeLockKeyByPod, require a non-empty matching UID alongside namespace
and name, and treat legacy annotations without a UID as non-idempotent so
stale-lock handling runs. Add coverage in
pkg/util/nodelock/nodelock_test.go:91-114 for identical namespace/name with a
different UID, verifying it does not succeed as a same-pod re-acquisition.
🪄 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: 14942620-f253-4c93-9c1a-0f7c79a1d97e
📒 Files selected for processing (3)
pkg/scheduler/scheduler_test.gopkg/util/nodelock/nodelock.gopkg/util/nodelock/nodelock_test.go
| node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}} | ||
| fakeClient := fake.NewSimpleClientset(pod, node) | ||
| s.kubeClient = fakeClient | ||
| client.KubeClient = fakeClient |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Restore client.KubeClient during cleanup.
Line 2140 overwrites the package-global Kubernetes client. The cleanup restores config.NodeLockRetryTimeout and device.DevicesMap, but not this client. Later tests can use this fake client and its mutated objects.
Proposed fix
oldDevicesMap := device.DevicesMap
+oldKubeClient := client.KubeClient
device.DevicesMap = map[string]device.Devices{mockA.vendor: mockA, mockB.vendor: mockB}
t.Cleanup(func() {
config.NodeLockRetryTimeout = oldRetry
device.DevicesMap = oldDevicesMap
+ client.KubeClient = oldKubeClient
})📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| client.KubeClient = fakeClient | |
| oldDevicesMap := device.DevicesMap | |
| oldKubeClient := client.KubeClient | |
| device.DevicesMap = map[string]device.Devices{mockA.vendor: mockA, mockB.vendor: mockB} | |
| t.Cleanup(func() { | |
| config.NodeLockRetryTimeout = oldRetry | |
| device.DevicesMap = oldDevicesMap | |
| client.KubeClient = oldKubeClient | |
| }) |
🤖 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` at line 2140, Update the test cleanup around
the assignment to client.KubeClient to save the original global client before
replacing it, then restore that value during cleanup alongside
config.NodeLockRetryTimeout and device.DevicesMap. Ensure later tests cannot
retain the fake client or its mutated objects.
| s.releaseAllDevices(node, pod) | ||
| nodeAfter, err = fakeClient.CoreV1().Nodes().Get(context.TODO(), node.Name, metav1.GetOptions{}) | ||
| require.NoError(t, err) | ||
| _, ok = nodeAfter.Annotations[nodelockutil.NodeLockKey] | ||
| require.False(t, ok, "node lock annotation must be released after all backends finish") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline pkg/scheduler/scheduler.go --match releaseAllDevices --view expanded
rg -n -A60 -B5 '\bfunc \(.*\) releaseAllDevices\b' pkg/scheduler/scheduler.go
rg -n -A90 -B5 '\bfunc ReleaseNodeLock\b' pkg/util/nodelock/nodelock.goRepository: Project-HAMi/HAMi
Length of output: 6813
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' pkg/util/nodelock/nodelock.go
printf '\n--- lock/release node usages ---\n'
rg -n -A20 -B10 '\bReleaseNodeLock\(|\bSetNodeLock\(|NodeLockKey|NodeLockSep|GeneratePodNamespaceName|ParseNodeLock' pkg -g '*.go'
printf '\n--- relevant tests ---\n'
rg -n -A35 -B15 'releaseAllDevices|Node lock annotation must be released|Node lock released|Multiple.*node lock|Multi-device|node lock' pkg/scheduler pkg/util/nodelock -g '*_test.go'Repository: Project-HAMi/HAMi
Length of output: 50374
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- nodelock helpers ---'
rg -n -A40 -B5 'func GenerateNodeLockKeyByPod|func GeneratePodNamespaceName|func ParseNodeLock|func ' pkg/util/nodelock/nodelock.go
printf '%s\n' '--- locked release behavior ---'
python3 - <<'PY'
from pathlib import Path
text = Path('pkg/util/nodelock/nodelock.go').read_text()
idx = text.find('func ReleaseNodeLock')
print(text[idx:idx+1800])
PY
printf '%s\n' '--- release lifecycle comment context ---'
sed -n '2130,2160p' pkg/scheduler/scheduler_test.goRepository: Project-HAMi/HAMi
Length of output: 13008
Guard shared node-lock release until the pod has no remaining backends.
releaseAllDevices calls ReleaseNodeLock for every device, and ReleaseNodeLock clears hami.io/mutex.lock for any matching pod annotation. Add an intermediate assertion after the first backend release, then fix the release path to keep the lock until all backends for that pod have finished.
🤖 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 2155 - 2159, Update the release
flow around releaseAllDevices and ReleaseNodeLock so releasing the first backend
does not clear the shared node-lock annotation while the pod still has other
backends. Add an intermediate assertion in the scheduler test after the first
backend release confirming the lock remains, then retain the existing
lock-release behavior only after all backends finish.
| // A lock already held by this same pod (e.g., a multi-device pod whose | ||
| // earlier backend locked the node through the shared annotation key) is an | ||
| // idempotent re-acquisition, not contention. | ||
| if ns != "" && previousPodName != "" && ns == pods.Namespace && previousPodName == pods.Name { | ||
| klog.InfoS("Node lock already held by this pod", "node", nodeName, "podName", pods.Name, "podNamespace", pods.Namespace) | ||
| return nil | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Use Pod UID for same-pod ownership.
A namespace and name match can identify a replacement Pod after the original Pod is deleted. GenerateNodeLockKeyByPod stores only those fields, so this branch can return success for a different Pod and bypass stale-lock handling.
pkg/util/nodelock/nodelock.go#L232-L238: persist and parsepods.UID, then require UID equality for idempotent acquisition. Treat legacy annotations without a UID as non-idempotent.pkg/util/nodelock/nodelock_test.go#L91-L114: add a case with the same namespace and name but a different UID. It must not succeed as a same-pod re-acquisition.
📍 Affects 2 files
pkg/util/nodelock/nodelock.go#L232-L238(this comment)pkg/util/nodelock/nodelock_test.go#L91-L114
🤖 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/util/nodelock/nodelock.go` around lines 232 - 238, Use Pod UID to
distinguish ownership in the idempotent acquisition branch around the node-lock
ownership parsing and comparison in pkg/util/nodelock/nodelock.go:232-238;
persist and parse pods.UID in GenerateNodeLockKeyByPod, require a non-empty
matching UID alongside namespace and name, and treat legacy annotations without
a UID as non-idempotent so stale-lock handling runs. Add coverage in
pkg/util/nodelock/nodelock_test.go:91-114 for identical namespace/name with a
different UID, verifying it does not succeed as a same-pod re-acquisition.
|
closing, #2255 was opened earlier and handles the expired own lock case correctly. pls check existing prs before opening a new one. |
|
@mesutoezdil Sure sure. I won't repeat that. Thanks! |
2 participants
/kind bug
What this PR does / why we need it:
Makes node lock acquisition idempotent for the same pod. lockAllDevices acquires a node lock for every registered device backend, and several backends (nvidia, ascend, hygon, metax, kunlun, etc.) share the same hami.io/mutex.lock annotation through pkg/util/nodelock. When a pod requests two of those device types, the second backend sees the lock created by the first backend, which is held by the same pod, but LockNode treats it as contention and fails the Bind.
Which issue(s) this PR fixes:
Fixes #2243
Special notes for your reviewer:
LockNode previously only released and re-acquired the lock when it was expired or dangling (owner pod no longer exists). A lock owned by the same pod fell through to the contention error, so a single Bind that locked the node through multiple backends would fail with
ode ... has been locked within 5m0s: node lock contention.
The fix treats a lock already held by the same pod as an idempotent re-acquisition and returns success. Release behavior is unchanged: ReleaseNodeLock already skips locks not owned by the calling pod, so the first backend to finish does not remove the shared lock.
Validation:
Does this PR introduce a user-facing change?:
No.
AI assistance disclosure:
This PR description was generated with the assistance of an AI coding tool. The author reviewed and verified all changes.
Summary by CodeRabbit
Bug Fixes
Tests