Preserve pre-existing node cordons across driver upgrades

cmd/driver-manager/main.go

@@ -390,6 +390,14 @@ func (dm *DriverManager) uninstallDriver() error {
 		// Remove stale PID file from previous container
 		dm.removePIDFile()
 
+		// Release a cordon left behind by an interrupted previous cycle. A no-op when
+		// the previous cycle never cordoned.
+		if dm.isGPUPodEvictionEnabled() || dm.isAutoDrainEnabled() {
+			if err := dm.kubeClient.UncordonNode(dm.config.nodeName); err != nil {
+				dm.log.Warnf("Failed to uncordon node: %v", err)
+			}
+		}
+
 		if err := dm.rescheduleGPUOperatorComponents(); err != nil {
 			return fmt.Errorf("failed to reschedule GPU operator components: %w", err)
 		}

internal/kubernetes/client.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -46,6 +47,8 @@ const (
 	nvidiaMigResourcePrefix  = nvidiaDomainPrefix + "/" + "mig-"
 	nvidiaDRADriverName      = "gpu." + nvidiaDomainPrefix
 
+	nodeInitialStateAnnotation = nvidiaDomainPrefix + "/" + "driver-manager.node-initial-state.unschedulable"
+
 	kubeClientPollInterval = 5 * time.Second
 )
 
@@ -54,7 +57,7 @@ type Client struct {
 	ctx context.Context
 	log *logrus.Logger
 
-	clientset *kubernetes.Clientset
+	clientset kubernetes.Interface
 }
 
 // DrainOptions represents the option parameters that can passed to the drain.Helper struct
@@ -103,33 +106,11 @@ func (c *Client) GetNodeLabelValue(nodeName, label string) (string, error) {
 // UpdateNodeLabels updates the labels on a Node given a Node name and a string map of label key-value pairs
 // This method uses a strategic merge patch to avoid conflicts with concurrent updates
 func (c *Client) UpdateNodeLabels(nodeName string, nodeLabels map[string]string) error {
-	patch := map[string]interface{}{
-		"metadata": map[string]interface{}{
-			"labels": nodeLabels,
-		},
-	}
-
-	patchBytes, err := json.Marshal(patch)
-	if err != nil {
-		return fmt.Errorf("failed to marshal patch: %w", err)
+	labels := make(map[string]interface{}, len(nodeLabels))
+	for key, value := range nodeLabels {
+		labels[key] = value
 	}
-
-	backoff := wait.Backoff{
-		Duration: time.Second,
-		Factor:   2.0,
-		Jitter:   0.2,
-		Steps:    7,
-	}
-
-	return retry.OnError(backoff, func(err error) bool {
-		return true
-	}, func() error {
-		_, err := c.clientset.CoreV1().Nodes().Patch(c.ctx, nodeName, types.StrategicMergePatchType, patchBytes, metav1.PatchOptions{})
-		if err != nil {
-			c.log.Warnf("Failed to update labels on node %s, retrying: %v", nodeName, err)
-		}
-		return err
-	})
+	return c.patchNodeMetadata(nodeName, "labels", labels)
 }
 
 // GetNodeAnnotationValue returns the annotation value given a node name and annotation key
@@ -144,30 +125,102 @@ func (c *Client) GetNodeAnnotationValue(nodeName, annotation string) (string, er
 	return node.Annotations[annotation], nil
 }
 
-// CordonNode cordons a Node given a Node name marking it as Unschedulable
+// CordonNode cordons a Node given a Node name marking it as Unschedulable. The node's
+// current schedulable state is first recorded in the nodeInitialStateAnnotation so that
+// UncordonNode can restore it. The annotation is written before the cordon, and a
+// recording is kept only while the node is Unschedulable: a restart after the cordon
+// must not overwrite it, whereas on a schedulable node any recording is stale, since a
+// cycle interrupted after the cordon would have left the node Unschedulable.
 func (c *Client) CordonNode(nodeName string) error {
-	c.log.Infof("Cordoning node %s", nodeName)
-
 	node, err := c.clientset.CoreV1().Nodes().Get(c.ctx, nodeName, metav1.GetOptions{})
 	if err != nil {
 		return fmt.Errorf("failed to get node %s: %w", nodeName, err)
 	}
 
+	value, ok := node.Annotations[nodeInitialStateAnnotation]
+	initialState := strconv.FormatBool(node.Spec.Unschedulable)
+	if !ok || (!node.Spec.Unschedulable && value != initialState) {
+		c.log.Infof("Recording initial state of node %s in annotation %s=%s", nodeName, nodeInitialStateAnnotation, initialState)
+		if err := c.setNodeAnnotation(nodeName, nodeInitialStateAnnotation, initialState); err != nil {
+			return fmt.Errorf("failed to record initial state of node %s: %w", nodeName, err)
+		}
+	}
+
+	c.log.Infof("Cordoning node %s", nodeName)
 	drainHelper := &drain.Helper{Ctx: c.ctx, Client: c.clientset}
 	return drain.RunCordonOrUncordon(drainHelper, node, true)
 }
 
-// UncordonNode uncordons a Node given a Node name marking it as Schedulable
+// UncordonNode restores the schedulable state recorded by CordonNode and removes the
+// nodeInitialStateAnnotation. If the annotation records that the node was already
+// Unschedulable, or if the annotation is absent (CordonNode never acted on the node),
+// the cordon is left in place.
 func (c *Client) UncordonNode(nodeName string) error {
-	c.log.Infof("Uncordoning node %s", nodeName)
-
 	node, err := c.clientset.CoreV1().Nodes().Get(c.ctx, nodeName, metav1.GetOptions{})
 	if err != nil {
 		return fmt.Errorf("failed to get node %s: %w", nodeName, err)
 	}
 
-	drainHelper := &drain.Helper{Ctx: c.ctx, Client: c.clientset}
-	return drain.RunCordonOrUncordon(drainHelper, node, false)
+	value, ok := node.Annotations[nodeInitialStateAnnotation]
+	if !ok {
+		c.log.Infof("Annotation %s not present on node %s, node was not cordoned by driver-manager, skipping uncordon", nodeInitialStateAnnotation, nodeName)
+		return nil
+	}
+
+	if value == "true" {
+		c.log.Infof("Node %s was already cordoned before the driver upgrade, skipping uncordon", nodeName)
+	} else {
+		c.log.Infof("Uncordoning node %s", nodeName)
+		drainHelper := &drain.Helper{Ctx: c.ctx, Client: c.clientset}
+		if err := drain.RunCordonOrUncordon(drainHelper, node, false); err != nil {
+			return err
+		}
+	}
+
+	return c.removeNodeAnnotation(nodeName, nodeInitialStateAnnotation)
+}
+
+// setNodeAnnotation sets an annotation on a Node given a Node name and an annotation key-value pair
+func (c *Client) setNodeAnnotation(nodeName, annotation, value string) error {
+	return c.patchNodeMetadata(nodeName, "annotations", map[string]interface{}{annotation: value})
+}
+
+// removeNodeAnnotation removes an annotation from a Node given a Node name and an annotation key
+func (c *Client) removeNodeAnnotation(nodeName, annotation string) error {
+	return c.patchNodeMetadata(nodeName, "annotations", map[string]interface{}{annotation: nil})
+}
+
+// patchNodeMetadata updates a metadata field (labels or annotations) on a Node, where a
+// nil value removes the key. This method uses a strategic merge patch to avoid conflicts
+// with concurrent updates
+func (c *Client) patchNodeMetadata(nodeName, field string, values map[string]interface{}) error {
+	patch := map[string]interface{}{
+		"metadata": map[string]interface{}{
+			field: values,
+		},
+	}
+
+	patchBytes, err := json.Marshal(patch)
+	if err != nil {
+		return fmt.Errorf("failed to marshal patch: %w", err)
+	}
+
+	backoff := wait.Backoff{
+		Duration: time.Second,
+		Factor:   2.0,
+		Jitter:   0.2,
+		Steps:    7,
+	}
+
+	return retry.OnError(backoff, func(err error) bool {
+		return true
+	}, func() error {
+		_, err := c.clientset.CoreV1().Nodes().Patch(c.ctx, nodeName, types.StrategicMergePatchType, patchBytes, metav1.PatchOptions{})
+		if err != nil {
+			c.log.Warnf("Failed to update %s on node %s, retrying: %v", field, nodeName, err)
+		}
+		return err
+	})
 }
 
 // WaitForPodTermination will wait for the termination of pods matching labels from the selectorMap on the node with the specified namespace.

internal/kubernetes/client_test.go

@@ -0,0 +1,167 @@
+/*
+ * Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kubernetes
+
+import (
+	"context"
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+func newTestClient(node *corev1.Node) *Client {
+	return &Client{
+		ctx:       context.Background(),
+		log:       logrus.New(),
+		clientset: fake.NewSimpleClientset(node),
+	}
+}
+
+func getTestNode(t *testing.T, c *Client, nodeName string) *corev1.Node {
+	node, err := c.clientset.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
+	require.NoError(t, err)
+	return node
+}
+
+func TestCordonUncordonNode(t *testing.T) {
+	testCases := []struct {
+		description           string
+		unschedulable         bool
+		annotations           map[string]string
+		cordon                bool
+		expectedUnschedulable bool
+	}{
+		{
+			description:           "schedulable node is cordoned and uncordoned",
+			unschedulable:         false,
+			cordon:                true,
+			expectedUnschedulable: false,
+		},
+		{
+			description:           "already cordoned node stays cordoned after uncordon",
+			unschedulable:         true,
+			cordon:                true,
+			expectedUnschedulable: true,
+		},
+		{
+			description:           "restart after cordon keeps the recorded initial state and uncordons",
+			unschedulable:         true,
+			annotations:           map[string]string{nodeInitialStateAnnotation: "false"},
+			cordon:                true,
+			expectedUnschedulable: false,
+		},
+		{
+			description:           "stale recording on a schedulable node is reconciled before cordon",
+			unschedulable:         false,
+			annotations:           map[string]string{nodeInitialStateAnnotation: "true"},
+			cordon:                true,
+			expectedUnschedulable: false,
+		},
+		{
+			description:           "uncordon without prior cordon leaves an external cordon in place",
+			unschedulable:         true,
+			cordon:                false,
+			expectedUnschedulable: true,
+		},
+		{
+			description:           "uncordon without prior cordon leaves a schedulable node schedulable",
+			unschedulable:         false,
+			cordon:                false,
+			expectedUnschedulable: false,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			node := &corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:        "test-node",
+					Annotations: tc.annotations,
+				},
+				Spec: corev1.NodeSpec{
+					Unschedulable: tc.unschedulable,
+				},
+			}
+			c := newTestClient(node)
+
+			if tc.cordon {
+				require.NoError(t, c.CordonNode("test-node"))
+				cordoned := getTestNode(t, c, "test-node")
+				require.True(t, cordoned.Spec.Unschedulable)
+			}
+
+			require.NoError(t, c.UncordonNode("test-node"))
+			uncordoned := getTestNode(t, c, "test-node")
+			require.Equal(t, tc.expectedUnschedulable, uncordoned.Spec.Unschedulable)
+			require.NotContains(t, uncordoned.Annotations, nodeInitialStateAnnotation)
+		})
+	}
+}
+
+func TestCordonNodeRecordsInitialState(t *testing.T) {
+	testCases := []struct {
+		description        string
+		unschedulable      bool
+		annotations        map[string]string
+		expectedAnnotation string
+	}{
+		{
+			description:        "schedulable node is recorded as schedulable",
+			unschedulable:      false,
+			expectedAnnotation: "false",
+		},
+		{
+			description:        "already cordoned node is recorded as unschedulable",
+			unschedulable:      true,
+			expectedAnnotation: "true",
+		},
+		{
+			description:        "an existing recording is not overwritten while the node is unschedulable",
+			unschedulable:      true,
+			annotations:        map[string]string{nodeInitialStateAnnotation: "false"},
+			expectedAnnotation: "false",
+		},
+		{
+			description:        "a stale recording on a schedulable node is overwritten",
+			unschedulable:      false,
+			annotations:        map[string]string{nodeInitialStateAnnotation: "true"},
+			expectedAnnotation: "false",
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			node := &corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:        "test-node",
+					Annotations: tc.annotations,
+				},
+				Spec: corev1.NodeSpec{
+					Unschedulable: tc.unschedulable,
+				},
+			}
+			c := newTestClient(node)
+
+			require.NoError(t, c.CordonNode("test-node"))
+			cordoned := getTestNode(t, c, "test-node")
+			require.True(t, cordoned.Spec.Unschedulable)
+			require.Equal(t, tc.expectedAnnotation, cordoned.Annotations[nodeInitialStateAnnotation])
+		})
+	}
+}