Skip to content

fix(util): make same-pod node lock acquisition idempotent #2256

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
shivv23 wants to merge 1 commit into from
Closed

fix(util): make same-pod node lock acquisition idempotent #2256

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions pkg/scheduler/scheduler_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2091,6 +2091,74 @@ func (m *bindLockMockDevice) ReleaseNodeLock(_ *corev1.Node, _ *corev1.Pod) erro
return nil
}

// sharedLockMockDevice mirrors backends (e.g. nvidia, ascend, hygon, metax) that
// all acquire the same node-lock annotation key "hami.io/mutex.lock" through the
// shared pkg/util/nodelock helper.
type sharedLockMockDevice struct {
registerMockDevice
vendor string
}

func (m *sharedLockMockDevice) CommonWord() string { return m.vendor }
func (m *sharedLockMockDevice) LockNode(n *corev1.Node, p *corev1.Pod) error {
return nodelockutil.LockNode(n.Name, nodelockutil.NodeLockKey, p)
}
func (m *sharedLockMockDevice) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error {
return nodelockutil.ReleaseNodeLock(n.Name, nodelockutil.NodeLockKey, p, false)
}

// Test_Bind_MultiDeviceBackendsSharingNodeLock reproduces #2243: a pod requesting
// two device types whose backends both lock the same node annotation must not
// contend with itself during Bind.
func Test_Bind_MultiDeviceBackendsSharingNodeLock(t *testing.T) {
pod := &corev1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "pod-multi-device", Namespace: "default", UID: types.UID("uid-multi-device"),
},
}
mockA := &sharedLockMockDevice{vendor: "shared-lock-a"}
mockB := &sharedLockMockDevice{vendor: "shared-lock-b"}

oldRetry := config.NodeLockRetryTimeout
config.NodeLockRetryTimeout = 500 * time.Millisecond
oldDevicesMap := device.DevicesMap
device.DevicesMap = map[string]device.Devices{mockA.vendor: mockA, mockB.vendor: mockB}
t.Cleanup(func() {
config.NodeLockRetryTimeout = oldRetry
device.DevicesMap = oldDevicesMap
})

s := NewScheduler()
t.Cleanup(func() { close(s.stopCh) })
scheme := runtime.NewScheme()
_ = corev1.AddToScheme(scheme)
s.eventRecorder = record.NewBroadcaster().NewRecorder(scheme, corev1.EventSource{})

node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node1"}}
fakeClient := fake.NewSimpleClientset(pod, node)
s.kubeClient = fakeClient
client.KubeClient = fakeClient

@coderabbitai coderabbitai Bot Aug 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.

Suggested change
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.


// lockAllDevices must succeed when both backends lock the same node for the
// same pod, and the resulting annotation must reference the locking pod.
if err := s.lockAllDevices(node, pod); err != nil {
t.Fatalf("lockAllDevices failed for same-pod multi-backend lock: %v", err)
}
nodeAfter, err := fakeClient.CoreV1().Nodes().Get(context.TODO(), node.Name, metav1.GetOptions{})
require.NoError(t, err)
lockValue, ok := nodeAfter.Annotations[nodelockutil.NodeLockKey]
require.True(t, ok, "node lock annotation must be set")
if !strings.HasSuffix(lockValue, nodelockutil.NodeLockSep+nodelockutil.GeneratePodNamespaceName(pod, nodelockutil.NodeLockSep)) {
t.Fatalf("node lock %q does not reference the locking pod", lockValue)
}

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")
Comment on lines +2155 to +2159

@coderabbitai coderabbitai Bot Aug 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.go

Repository: 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.go

Repository: 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.

}

var errContention = fmt.Errorf("contended: %w", nodelockutil.ErrNodeLockContention)

func setupBindLockRetryTest(t *testing.T, retryTimeout time.Duration, pod *corev1.Pod, mock *bindLockMockDevice) (*Scheduler, extenderv1.ExtenderBindingArgs, func()) {
Expand Down
8 changes: 8 additions & 0 deletions pkg/util/nodelock/nodelock.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,14 @@ func LockNode(nodeName string, lockname string, pods *corev1.Pod) error {
return err
}

// 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
}
Comment on lines +232 to +238

@coderabbitai coderabbitai Bot Aug 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 parse pods.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.


var skipOwnerCheck = false
if time.Since(lockTime) > NodeLockTimeout {
klog.InfoS("Node lock expired", "node", nodeName, "lockTime", lockTime, "timeout", NodeLockTimeout)
Expand Down
31 changes: 30 additions & 1 deletion pkg/util/nodelock/nodelock_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,35 @@ func Test_LockNode(t *testing.T) {
args: args{
nodeName: func() string {
name := "worker-1"
client.KubeClient.CoreV1().Pods("hami-ns").Create(context.TODO(), &corev1.Pod{
ObjectMeta: metav1.ObjectMeta{Name: "hami", Namespace: "hami-ns"},
}, metav1.CreateOptions{})
client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
ObjectMeta: metav1.ObjectMeta{
Name: name,
Annotations: map[string]string{
NodeLockKey: GenerateNodeLockKeyByPod(&corev1.Pod{
ObjectMeta: metav1.ObjectMeta{Name: "hami", Namespace: "hami-ns"},
}),
},
},
}, metav1.CreateOptions{})
return name
},
pods: &corev1.Pod{
ObjectMeta: metav1.ObjectMeta{
Name: "other",
Namespace: "other-ns",
},
},
},
wantErr: true,
},
{
name: "node lock is idempotent for the same pod",
args: args{
nodeName: func() string {
name := "worker-1b"
client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
ObjectMeta: metav1.ObjectMeta{
Name: name,
Expand All @@ -82,7 +111,7 @@ func Test_LockNode(t *testing.T) {
},
},
},
wantErr: true,
wantErr: false,
},
{
name: "node lock is invalid",
Expand Down
Loading