fix(cambricon): use shared nodelock package for node-level locking

pkg/device/cambricon/device.go

@@ -17,24 +17,18 @@ limitations under the License.
 package cambricon
 
 import (
-	"context"
-	"encoding/json"
 	"flag"
 	"fmt"
-	"math/rand"
 	"slices"
 	"strings"
-	"time"
 
 	"github.com/Project-HAMi/HAMi/pkg/device"
 	"github.com/Project-HAMi/HAMi/pkg/device/common"
 	"github.com/Project-HAMi/HAMi/pkg/util"
-	"github.com/Project-HAMi/HAMi/pkg/util/client"
+	"github.com/Project-HAMi/HAMi/pkg/util/nodelock"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
 )
 
@@ -50,10 +44,12 @@ const (
 	MLUUseUUID = "cambricon.com/use-gpuuuid"
 	// MLUNoUseUUID annotation specifies a comma-separated list of MLU UUIDs to exclude.
 	MLUNoUseUUID          = "cambricon.com/nouse-gpuuuid"
-	DsmluLockTime         = "cambricon.com/dsmlu.lock"
 	DsmluProfile          = "CAMBRICON_DSMLU_PROFILE"
 	DsmluResourceAssigned = "CAMBRICON_DSMLU_ASSIGNED"
-	retry                 = 5
+	// NodeLockCambricon should be the same as the node lock name used by the
+	// device plugin; nodelock hard-codes the annotation key regardless of the
+	// name passed in, so this only needs to be distinct for readability.
+	NodeLockCambricon = "hami.io/mutex.lock"
 	// MemoryFactor converts the vmemory unit used in the pod spec into the MiB
 	// HAMi accounts internally. One cambricon.com/mlu.smlu.vmemory is 256 MiB.
 	MemoryFactor = 256
@@ -96,31 +92,6 @@ func (dev *CambriconDevices) CommonWord() string {
 	return CambriconMLUCommonWord
 }
 
-func (dev *CambriconDevices) setNodeLock(node *corev1.Node) error {
-	ctx := context.Background()
-	patchedAnnotation, err := json.Marshal(
-		map[string]any{
-			"metadata": map[string]map[string]string{"annotations": {
-				DsmluLockTime: time.Now().Format(time.RFC3339),
-			}}})
-	if err != nil {
-		klog.ErrorS(err, "Failed to marshal node annotation", "node", node.Name)
-		return fmt.Errorf("marshal node annotation: %w", err)
-	}
-
-	_, err = client.GetClient().CoreV1().Nodes().Patch(ctx, node.Name, types.MergePatchType, patchedAnnotation, metav1.PatchOptions{})
-	for i := 0; i < retry && err != nil; i++ {
-		klog.ErrorS(err, "Failed to patch node annotation", "node", node.Name, "retry", i)
-		time.Sleep(time.Duration(rand.Intn(i+1)) * 10 * time.Millisecond)
-		_, err = client.GetClient().CoreV1().Nodes().Patch(ctx, node.Name, types.MergePatchType, patchedAnnotation, metav1.PatchOptions{})
-	}
-	if err != nil {
-		return fmt.Errorf("setNodeLock exceeds retry count %d", retry)
-	}
-	klog.InfoS("Node lock set", "node", node.Name)
-	return nil
-}
-
 func (dev *CambriconDevices) LockNode(n *corev1.Node, p *corev1.Pod) error {
 	found := false
 	for _, val := range p.Spec.Containers {
@@ -132,69 +103,21 @@ func (dev *CambriconDevices) LockNode(n *corev1.Node, p *corev1.Pod) error {
 	if !found {
 		return nil
 	}
-	if _, ok := n.Annotations[DsmluLockTime]; !ok {
-		return dev.setNodeLock(n)
-	}
-	lockTime, err := time.Parse(time.RFC3339, n.Annotations[DsmluLockTime])
-	if err != nil {
-		return err
-	}
-	if time.Since(lockTime) > time.Minute*2 {
-		klog.InfoS("Node lock expired", "node", n.Name, "lockTime", lockTime)
-		err = dev.ReleaseNodeLock(n, p)
-		if err != nil {
-			klog.ErrorS(err, "Failed to release node lock", "node", n.Name)
-			return err
-		}
-		return dev.setNodeLock(n)
-	}
-	return fmt.Errorf("node %s has been locked within 2 minutes", n.Name)
+	return nodelock.LockNode(n.Name, NodeLockCambricon, p)
 }
 
 func (dev *CambriconDevices) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error {
-	ctx := context.Background()
-	nodeName := n.Name
-
-	current, err := client.GetClient().CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
-	if err != nil {
-		return fmt.Errorf("failed to get node for lock release: %w", err)
-	}
-	if current.Annotations == nil {
-		return nil
-	}
-	if _, ok := current.Annotations[DsmluLockTime]; !ok {
-		klog.InfoS("Node lock not set", "node", nodeName)
-		return nil
-	}
-
-	patchData, err := json.Marshal(
-		map[string]any{
-			"metadata": map[string]map[string]any{"annotations": {
-				DsmluLockTime: nil,
-			}}})
-	if err != nil {
-		return fmt.Errorf("marshal patch for lock release: %w", err)
-	}
-
-	_, err = client.GetClient().CoreV1().Nodes().Patch(ctx, nodeName, types.MergePatchType, patchData, metav1.PatchOptions{})
-	for i := 0; i < retry && err != nil; i++ {
-		klog.ErrorS(err, "Failed to release node lock", "node", nodeName, "retry", i)
-		time.Sleep(time.Duration(rand.Intn(i+1)) * 10 * time.Millisecond)
-
-		current, err = client.GetClient().CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
-		if err != nil {
-			return fmt.Errorf("failed to re-get node during release retry: %w", err)
-		}
-		if current.Annotations == nil || current.Annotations[DsmluLockTime] == "" {
-			return nil
+	found := false
+	for _, val := range p.Spec.Containers {
+		if (dev.GenerateResourceRequests(&val).Nums) > 0 {
+			found = true
+			break
 		}
-		_, err = client.GetClient().CoreV1().Nodes().Patch(ctx, nodeName, types.MergePatchType, patchData, metav1.PatchOptions{})
 	}
-	if err != nil {
-		return fmt.Errorf("releaseNodeLock exceeds retry count %d", retry)
+	if !found {
+		return nil
 	}
-	klog.InfoS("Node lock released", "node", nodeName)
-	return nil
+	return nodelock.ReleaseNodeLock(n.Name, NodeLockCambricon, p, false)
 }
 
 func (dev *CambriconDevices) NodeCleanUp(nn string) error {