diff --git a/pkg/device/cambricon/device.go b/pkg/device/cambricon/device.go index 222b4cf3e0..b5ef776ae7 100644 --- a/pkg/device/cambricon/device.go +++ b/pkg/device/cambricon/device.go @@ -98,25 +98,21 @@ func (dev *CambriconDevices) CommonWord() string { func (dev *CambriconDevices) setNodeLock(node *corev1.Node) error { ctx := context.Background() - if _, ok := node.Annotations[DsmluLockTime]; ok { - return fmt.Errorf("node %s is locked", node.Name) - } - 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 patch node annotation", "node", node.Name) - return fmt.Errorf("patch node annotation %v", err) + 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.StrategicMergePatchType, patchedAnnotation, metav1.PatchOptions{}) + _, 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.StrategicMergePatchType, patchedAnnotation, metav1.PatchOptions{}) + _, 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) @@ -156,27 +152,48 @@ func (dev *CambriconDevices) LockNode(n *corev1.Node, p *corev1.Pod) error { } func (dev *CambriconDevices) ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error { - if n.Annotations == nil { + 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 := n.Annotations[DsmluLockTime]; !ok { - klog.InfoS("Node lock not set", "node", n.Name) + if _, ok := current.Annotations[DsmluLockTime]; !ok { + klog.InfoS("Node lock not set", "node", nodeName) return nil } - newNode := n.DeepCopy() - delete(newNode.Annotations, DsmluLockTime) - _, err := client.GetClient().CoreV1().Nodes().Update(context.Background(), newNode, metav1.UpdateOptions{}) + 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 patch node annotation", "node", n.Name, "retry", i) + klog.ErrorS(err, "Failed to release node lock", "node", nodeName, "retry", i) time.Sleep(time.Duration(rand.Intn(i+1)) * 10 * time.Millisecond) - _, err = client.GetClient().CoreV1().Nodes().Update(context.Background(), newNode, metav1.UpdateOptions{}) + + 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 + } + _, 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) } - delete(n.Annotations, DsmluLockTime) - klog.InfoS("Node lock released", "node", n.Name) + klog.InfoS("Node lock released", "node", nodeName) return nil } diff --git a/pkg/device/cambricon/device_test.go b/pkg/device/cambricon/device_test.go index eb3ffa007e..69fcda91c0 100644 --- a/pkg/device/cambricon/device_test.go +++ b/pkg/device/cambricon/device_test.go @@ -19,15 +19,19 @@ package cambricon import ( "context" "flag" - "strings" + "fmt" "testing" "time" "github.com/stretchr/testify/assert" corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/util/client" @@ -395,76 +399,39 @@ func Test_PatchAnnotations(t *testing.T) { } func Test_setNodeLock(t *testing.T) { - tests := []struct { - name string - node corev1.Node - expectErr bool - expectMsg string - }{ - { - name: "node is locked", - node: corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node-01", - Annotations: map[string]string{ - "cambricon.com/dsmlu.lock": "test123", - }, - }, - }, - expectErr: true, - expectMsg: "node node-01 is locked", - }, - { - name: "set node lock", - node: corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node-02", - Annotations: map[string]string{}, - }, - }, - expectErr: false, - }, - } - client.KubeClient = fake.NewClientset() k8sClient := client.GetClient() - if k8sClient != nil { - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() + if k8sClient == nil { + t.Skip("no k8s client available") + } - defer func() { - if tt.node.Name != "" { - // Delete the node to clean up - err := k8sClient.CoreV1().Nodes().Delete(ctx, tt.node.Name, metav1.DeleteOptions{}) - if err != nil { - t.Errorf("failed to delete node %s: %v", tt.node.Name, err) - } - } - }() + node := corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-set-lock", + Annotations: map[string]string{}, + }, + } - _, err := k8sClient.CoreV1().Nodes().Create(ctx, &tt.node, metav1.CreateOptions{}) - if err != nil { - t.Fatalf("failed to create node %s: %v", tt.node.Name, err) - } + ctx := context.Background() + _, err := k8sClient.CoreV1().Nodes().Create(ctx, &node, metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create node: %v", err) + } + defer k8sClient.CoreV1().Nodes().Delete(ctx, node.Name, metav1.DeleteOptions{}) - dev := CambriconDevices{} - err = dev.setNodeLock(&tt.node) + dev := CambriconDevices{} + err = dev.setNodeLock(&node) + if err != nil { + t.Errorf("did not expect error but got %v", err) + } - if tt.expectErr { - if err == nil { - t.Errorf("expected error but got none") - } else if !strings.Contains(err.Error(), tt.expectMsg) { - t.Errorf("expected error to contain '%s' but got '%s'", tt.expectMsg, err.Error()) - } - } else { - if err != nil { - t.Errorf("did not expect error but got %v", err) - } - } - }) - } + // Verify the annotation was set on the apiserver. + fetchedNode, err := k8sClient.CoreV1().Nodes().Get(ctx, node.Name, metav1.GetOptions{}) + if err != nil { + t.Fatalf("failed to fetch node: %v", err) + } + if _, ok := fetchedNode.Annotations[DsmluLockTime]; !ok { + t.Error("Expected node to be locked but it wasn't") } } @@ -589,55 +556,226 @@ func Test_LockNode(t *testing.T) { } func Test_ReleaseNodeLock(t *testing.T) { + clientset := fake.NewClientset() + client.KubeClient = clientset + ctx := context.Background() + tests := []struct { - name string - args struct { - node corev1.Node - pod corev1.Pod - } - err error + name string + nodeName string + annotations map[string]string + callerAnnotations map[string]string + skipCreate bool + wantErr bool }{ { - name: "no annation", - args: struct { - node corev1.Node - pod corev1.Pod - }{ - node: corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node-01", - }, - }, - pod: corev1.Pod{}, + name: "no annotation — lock not present", + nodeName: "node-rel-01", + annotations: nil, + wantErr: false, + }, + { + name: "annotation without lock key", + nodeName: "node-rel-02", + annotations: map[string]string{ + "test": "test123", }, - err: nil, + wantErr: false, }, { - name: "annation no lock value", - args: struct { - node corev1.Node - pod corev1.Pod - }{ - node: corev1.Node{ - ObjectMeta: metav1.ObjectMeta{ - Name: "node-02", - Annotations: map[string]string{ - "test": "test123", - }, - }, - }, - pod: corev1.Pod{}, + name: "lock present — successfully released", + nodeName: "node-rel-03", + annotations: map[string]string{ + DsmluLockTime: time.Now().Format(time.RFC3339), + }, + wantErr: false, + }, + { + name: "stale caller — lock on server but not on caller node", + nodeName: "node-rel-04", + annotations: map[string]string{ + DsmluLockTime: time.Now().Format(time.RFC3339), }, - err: nil, + callerAnnotations: map[string]string{}, + wantErr: false, + }, + { + name: "node not found", + nodeName: "node-rel-nonexistent", + skipCreate: true, + wantErr: true, }, } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if !tt.skipCreate { + node := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: tt.nodeName, + Annotations: tt.annotations, + }, + } + _, err := clientset.CoreV1().Nodes().Create(ctx, node, metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create node %s: %v", tt.nodeName, err) + } + defer clientset.CoreV1().Nodes().Delete(ctx, tt.nodeName, metav1.DeleteOptions{}) + } + + callerNode := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: tt.nodeName, + }, + } + if tt.callerAnnotations != nil { + callerNode.Annotations = tt.callerAnnotations + } else if tt.annotations != nil { + callerNode.Annotations = tt.annotations + } + dev := CambriconDevices{} - result := dev.ReleaseNodeLock(&test.args.node, &test.args.pod) - assert.Equal(t, test.err, result) + err := dev.ReleaseNodeLock(callerNode, &corev1.Pod{}) + if (err != nil) != tt.wantErr { + t.Errorf("ReleaseNodeLock() error = %v, wantErr %v", err, tt.wantErr) + } + + if !tt.wantErr && !tt.skipCreate && tt.annotations != nil { + if _, ok := tt.annotations[DsmluLockTime]; ok { + fetched, ferr := clientset.CoreV1().Nodes().Get(ctx, tt.nodeName, metav1.GetOptions{}) + if ferr != nil { + t.Fatalf("failed to fetch node: %v", ferr) + } + if _, ok := fetched.Annotations[DsmluLockTime]; ok { + t.Error("Expected lock annotation to be removed but it's still present") + } + } + } }) } + + // Retry/conflict tests using PrependReactor. + t.Run("patch error — retry succeeds", func(t *testing.T) { + cs := fake.NewClientset() + client.KubeClient = cs + + node := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-retry-01", + Annotations: map[string]string{ + DsmluLockTime: time.Now().Format(time.RFC3339), + }, + }, + } + _, err := cs.CoreV1().Nodes().Create(ctx, node, metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create node: %v", err) + } + defer cs.CoreV1().Nodes().Delete(ctx, "node-retry-01", metav1.DeleteOptions{}) + + patchFailCount := 0 + cs.PrependReactor("patch", "nodes", func(action k8stesting.Action) (bool, runtime.Object, error) { + patchFailCount++ + if patchFailCount <= 2 { + return true, nil, apierrors.NewInternalError(fmt.Errorf("patch error")) + } + return false, nil, nil + }) + + dev := CambriconDevices{} + err = dev.ReleaseNodeLock(node, &corev1.Pod{}) + if err != nil { + t.Errorf("ReleaseNodeLock() unexpected error: %v", err) + } + if patchFailCount < 2 { + t.Errorf("expected at least 2 patch errors before success, got %d", patchFailCount) + } + + fetched, _ := cs.CoreV1().Nodes().Get(ctx, "node-retry-01", metav1.GetOptions{}) + if _, ok := fetched.Annotations[DsmluLockTime]; ok { + t.Error("Expected lock annotation to be removed after retries") + } + }) + + t.Run("lock already removed during retry", func(t *testing.T) { + cs := fake.NewClientset() + client.KubeClient = cs + + node := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-retry-02", + Annotations: map[string]string{ + DsmluLockTime: time.Now().Format(time.RFC3339), + }, + }, + } + _, err := cs.CoreV1().Nodes().Create(ctx, node, metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create node: %v", err) + } + defer cs.CoreV1().Nodes().Delete(ctx, "node-retry-02", metav1.DeleteOptions{}) + + cs.PrependReactor("patch", "nodes", func(action k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, apierrors.NewInternalError(fmt.Errorf("patch error")) + }) + + getCount := 0 + cs.PrependReactor("get", "nodes", func(action k8stesting.Action) (bool, runtime.Object, error) { + getCount++ + if getCount == 2 { + return true, &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-retry-02", + Annotations: map[string]string{}, + }, + }, nil + } + return false, nil, nil + }) + + dev := CambriconDevices{} + err = dev.ReleaseNodeLock(node, &corev1.Pod{}) + if err != nil { + t.Errorf("ReleaseNodeLock() unexpected error: %v", err) + } + }) + + t.Run("re-get fails during retry", func(t *testing.T) { + cs := fake.NewClientset() + client.KubeClient = cs + + node := &corev1.Node{ + ObjectMeta: metav1.ObjectMeta{ + Name: "node-retry-03", + Annotations: map[string]string{ + DsmluLockTime: time.Now().Format(time.RFC3339), + }, + }, + } + _, err := cs.CoreV1().Nodes().Create(ctx, node, metav1.CreateOptions{}) + if err != nil { + t.Fatalf("failed to create node: %v", err) + } + defer cs.CoreV1().Nodes().Delete(ctx, "node-retry-03", metav1.DeleteOptions{}) + + cs.PrependReactor("patch", "nodes", func(action k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, apierrors.NewInternalError(fmt.Errorf("patch error")) + }) + getCount := 0 + cs.PrependReactor("get", "nodes", func(action k8stesting.Action) (bool, runtime.Object, error) { + getCount++ + if getCount == 2 { + return true, nil, apierrors.NewNotFound(schema.GroupResource{Resource: "nodes"}, "node-retry-03") + } + return false, nil, nil + }) + + dev := CambriconDevices{} + err = dev.ReleaseNodeLock(node, &corev1.Pod{}) + if err == nil { + t.Error("ReleaseNodeLock() expected error but got nil") + } + }) } func TestDevices_Fit(t *testing.T) {