feat: implement stale lock detection and enhance node lock management

haitwang-cloud · haitwang-cloud · commit d50cb7002090 · 2025-03-25T16:14:29.000+08:00
diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go
@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -189,10 +190,39 @@ func (dev *NvidiaGPUDevices) LockNode(n *corev1.Node, p *corev1.Pod) error {
 	if !found {
 		return nil
 	}
+
+	if stale, err := dev.checkStaleLock(n); err == nil && stale {
+		klog.Warningf("Detected stale lock on node %s, forcing release", n.Name)
+		_ = nodelock.ReleaseNodeLock(n.Name, NodeLockNvidia, p, true)
+	}
+
 	return nodelock.LockNode(n.Name, NodeLockNvidia, p)
 }
 
+func (dev *NvidiaGPUDevices) checkStaleLock(n *corev1.Node) (bool, error) {
+	lockStr := n.Annotations[nodelock.NodeLockKey]
+	if lockStr == "" {
+		return false, nil
+	}
+
+	parts := strings.Split(lockStr, "|")
+	if len(parts) < 4 {
+		return true, nil // stale lock
+	}
+
+	createTime, err := time.Parse(time.RFC3339Nano, parts[0])
+	if err != nil {
+		return true, nil
+	}
+
+	return time.Since(createTime) > nodelock.LockTTL, nil
+}
+
 func (dev *NvidiaGPUDevices) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error {
+	defer func() {
+		klog.Errorf("Final fallback release for node %s", n.Name)
+		_ = nodelock.ReleaseNodeLock(n.Name, NodeLockNvidia, p, true)
+	}()
 	found := false
 	for _, val := range p.Spec.Containers {
 		if (dev.GenerateResourceRequests(&val).Nums) > 0 {
diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go
@@ -19,13 +19,15 @@ package nvidia
 import (
 	"errors"
 	"testing"
+	"time"
 
 	"gotest.tools/v3/assert"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/Project-HAMi/HAMi/pkg/util"
+	"github.com/Project-HAMi/HAMi/pkg/util/nodelock"
 )
 
 func Test_DefaultResourceNum(t *testing.T) {
@@ -614,3 +616,84 @@ func Test_GetNodeDevices(t *testing.T) {
 		})
 	}
 }
+func Test_checkStaleLock(t *testing.T) {
+	tests := []struct {
+		name      string
+		node      corev1.Node
+		wantStale bool
+		wantErr   bool
+	}{
+		{
+			name: "no lock annotation",
+			node: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{},
+				},
+			},
+			wantStale: false,
+			wantErr:   false,
+		},
+		{
+			name: "stale lock with insufficient parts",
+			node: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						nodelock.NodeLockKey: "2023-10-10T10:10:10Z|part2|part3",
+					},
+				},
+			},
+			wantStale: true,
+			wantErr:   false,
+		},
+		{
+			name: "stale lock with invalid time format",
+			node: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						nodelock.NodeLockKey: "invalid-time|part2|part3|part4",
+					},
+				},
+			},
+			wantStale: true,
+			wantErr:   false,
+		},
+		{
+			name: "stale lock with expired time",
+			node: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						nodelock.NodeLockKey: time.Now().Add(-2*nodelock.LockTTL).Format(time.RFC3339Nano) + "|part2|part3|part4",
+					},
+				},
+			},
+			wantStale: true,
+			wantErr:   false,
+		},
+		{
+			name: "valid lock",
+			node: corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{
+						nodelock.NodeLockKey: time.Now().Format(time.RFC3339Nano) + "|part2|part3|part4",
+					},
+				},
+			},
+			wantStale: false,
+			wantErr:   false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			gpuDevices := &NvidiaGPUDevices{}
+			gotStale, err := gpuDevices.checkStaleLock(&test.node)
+			if (err != nil) != test.wantErr {
+				t.Errorf("checkStaleLock() error = %v, wantErr %v", err, test.wantErr)
+				return
+			}
+			if gotStale != test.wantStale {
+				t.Errorf("checkStaleLock() gotStale = %v, want %v", gotStale, test.wantStale)
+			}
+		})
+	}
+}
diff --git a/pkg/util/nodelock/nodelock.go b/pkg/util/nodelock/nodelock.go
@@ -19,21 +19,25 @@ package nodelock
 import (
 	"context"
 	"fmt"
+	"math"
 	"strings"
 	"sync"
 	"time"
 
-	"github.com/Project-HAMi/HAMi/pkg/util/client"
-
 	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/klog/v2"
+
+	"github.com/Project-HAMi/HAMi/pkg/util/client"
 )
 
 const (
-	NodeLockKey  = "hami.io/mutex.lock"
-	MaxLockRetry = 5
-	NodeLockSep  = ","
+	NodeLockKey       = "hami.io/mutex.lock"
+	MaxLockRetry      = 8
+	NodeLockSep       = ","
+	LockTTL           = 1 * time.Minute
+	RetryBaseInterval = 50 * time.Millisecond
 )
 
 var lock sync.Mutex
@@ -53,9 +57,14 @@ func SetNodeLock(nodeName string, lockname string, pods *corev1.Pod) error {
 	newNode.ObjectMeta.Annotations[NodeLockKey] = GenerateNodeLockKeyByPod(pods)
 	_, err = client.GetClient().CoreV1().Nodes().Update(ctx, newNode, metav1.UpdateOptions{})
 	for i := 0; i < MaxLockRetry && err != nil; i++ {
-		klog.ErrorS(err, "Failed to update node", "node", nodeName, "retry", i)
-		time.Sleep(100 * time.Millisecond)
+		sleepDuration := time.Duration(math.Pow(2, float64(i))) * RetryBaseInterval
+		time.Sleep(sleepDuration)
+
 		node, err = client.GetClient().CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+		if apierrors.IsConflict(err) {
+			klog.Warningf("Version conflict on node %s, retrying...", nodeName)
+			continue
+		}
 		if err != nil {
 			klog.ErrorS(err, "Failed to get node when retry to update", "node", nodeName)
 			continue
@@ -152,6 +161,9 @@ func GenerateNodeLockKeyByPod(pods *corev1.Pod) string {
 	if pods == nil {
 		return time.Now().Format(time.RFC3339)
 	}
-	ns, name := pods.Namespace, pods.Name
-	return fmt.Sprintf("%s%s%s%s%s", time.Now().Format(time.RFC3339), NodeLockSep, ns, NodeLockSep, name)
+	return fmt.Sprintf("%s|%s|%s|%s",
+		time.Now().Format(time.RFC3339Nano), // use nano second as lock time
+		string(pods.UID),                    // use pod UID as lock namespace
+		pods.Namespace,
+		pods.Name)
 }
diff --git a/pkg/util/nodelock/nodelock_test.go b/pkg/util/nodelock/nodelock_test.go
@@ -18,7 +18,9 @@ package nodelock
 
 import (
 	"context"
+	"strings"
 	"testing"
+	"time"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -235,3 +237,51 @@ func TestReleaseNodeLock(t *testing.T) {
 		})
 	}
 }
+func TestGenerateNodeLockKeyByPod(t *testing.T) {
+	tests := []struct {
+		name string
+		pods *corev1.Pod
+	}{
+		{
+			name: "nil pod",
+			pods: nil,
+		},
+		{
+			name: "valid pod",
+			pods: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					UID:       "12345",
+					Name:      "hami",
+					Namespace: "hami-ns",
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GenerateNodeLockKeyByPod(tt.pods)
+			if tt.pods == nil {
+				if _, err := time.Parse(time.RFC3339, got); err != nil {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want valid RFC3339 time", got)
+				}
+			} else {
+				parts := strings.Split(got, "|")
+				if len(parts) != 4 {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want 4 parts", got)
+				}
+				if _, err := time.Parse(time.RFC3339Nano, parts[0]); err != nil {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want valid RFC3339Nano time", got)
+				}
+				if parts[1] != string(tt.pods.UID) {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want UID %v", got, tt.pods.UID)
+				}
+				if parts[2] != tt.pods.Namespace {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want Namespace %v", got, tt.pods.Namespace)
+				}
+				if parts[3] != tt.pods.Name {
+					t.Errorf("GenerateNodeLockKeyByPod() = %v, want Name %v", got, tt.pods.Name)
+				}
+			}
+		})
+	}
+}
