Add optional restart-only predicate to inplace upgrade flow

rajathagasthya · rajathagasthya · commit 80df34917c4f · 2026-06-16T15:01:19.000-05:00
Consumers occasionally need a DaemonSet's pods rolled to a new revision
without the disruptive pod-eviction and drain sequence -- for example
when a pod-template change does not affect the managed software (a
label-only change) yet still changes the controller revision hash, so
the node is otherwise driven through the full upgrade flow.

Add an optional RestartOnlyPredicate hook, registered via
WithRestartOnlyPredicate. In the inplace ProcessUpgradeRequiredNodes,
when the predicate reports that a full upgrade is not needed, cordon the
node and route it straight to pod-restart-required, skipping
wait-for-jobs, pod-deletion, and drain. The pod is still restarted so
the DaemonSet converges to the new revision. Cordoning keeps the node
unschedulable if the restart fails, matching the full flow; the
uncordon-required state uncordons it on success.

The change is additive and backward compatible: a nil predicate (the
default, and every existing consumer) preserves current behavior, and
podInSyncWithDS is unchanged. If the predicate returns an error or the
cordon fails, the node is kept in upgrade-required and retried on a
later reconcile, with a Warning event recorded -- an upgrade is never
started on an unknown answer. The predicate is not consulted for
orphaned pods, nodes with an explicit upgrade-requested annotation, or
nodes waiting for safe driver load -- the safe-load handshake relies on
the full flow to evict workloads before the driver load is unblocked at
pod-restart-required. The maxParallelUpgrades throttle applies to both
paths. The upgrade-requested state is captured before the annotation is
cleared, because the node provider re-fetches and overwrites the
in-memory node object.

Signed-off-by: Rajath Agasthya &lt;ragasthya@nvidia.com&gt;
diff --git a/pkg/upgrade/common_manager.go b/pkg/upgrade/common_manager.go
@@ -79,6 +79,17 @@ func NewClusterUpgradeState() ClusterUpgradeState {
 	return ClusterUpgradeState{NodeStates: make(map[string][]*NodeUpgradeState)}
 }
 
+// RestartOnlyPredicate is used for a node whose driver pod is out-of-sync with
+// its DaemonSet. Returning true means the difference does not affect the installed
+// driver, so the node is cordoned and the driver pod restarted in place, skipping
+// pod-deletion (workload eviction) and drain; the consumer guarantees the running
+// driver does not need to change across the restart. Returning false (the default
+// when unset) routes the node through the full upgrade flow. Returning an error
+// keeps the node in upgrade-required to be retried on a later reconcile. It is
+// never called for orphaned pods, upgrade-requested nodes, or nodes waiting for
+// safe driver load.
+type RestartOnlyPredicate func(ctx context.Context, pod *corev1.Pod, ds *appsv1.DaemonSet) (bool, error)
+
 // CommonUpgradeManagerImpl is an implementation of the CommonUpgradeStateManager interface.
 // It facilitates common logic implementation for both upgrade modes: in-place and requestor (e.g. maintenance OP).
 type CommonUpgradeManagerImpl struct {
@@ -97,6 +108,8 @@ type CommonUpgradeManagerImpl struct {
 	// optional states
 	podDeletionStateEnabled bool
 	validationStateEnabled  bool
+	// optional: when set, route immaterial pod-template changes to a restart-only path
+	restartOnlyPredicate RestartOnlyPredicate
 }
 
 // NewCommonUpgradeStateManager creates a new instance of CommonUpgradeManagerImpl
diff --git a/pkg/upgrade/upgrade_inplace.go b/pkg/upgrade/upgrade_inplace.go
@@ -19,6 +19,7 @@ package upgrade
 import (
 	"context"
 
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 
 	"github.com/NVIDIA/k8s-operator-libs/api/upgrade/v1alpha1"
@@ -69,7 +70,8 @@ func (m *InplaceNodeStateManagerImpl) ProcessUpgradeRequiredNodes(
 		"maximum nodes that can be unavailable", maxUnavailable)
 
 	for _, nodeState := range currentClusterState.NodeStates[UpgradeStateUpgradeRequired] {
-		if m.IsUpgradeRequested(nodeState.Node) {
+		upgradeRequested := m.IsUpgradeRequested(nodeState.Node)
+		if upgradeRequested {
 			// Make sure to remove the upgrade-requested annotation
 			err := m.NodeUpgradeStateProvider.ChangeNodeUpgradeAnnotation(ctx, nodeState.Node,
 				GetUpgradeRequestedAnnotationKey(), "null")
@@ -96,14 +98,62 @@ func (m *InplaceNodeStateManagerImpl) ProcessUpgradeRequiredNodes(
 			}
 		}
 
-		err := m.NodeUpgradeStateProvider.ChangeNodeUpgradeState(ctx, nodeState.Node, UpgradeStateCordonRequired)
+		// Default to the full upgrade flow. A registered restart-only predicate can route an
+		// out-of-sync node to an in-place driver pod restart; it is not used for orphaned pods,
+		// nodes that explicitly requested an upgrade, or nodes waiting for safe driver load.
+		targetState := UpgradeStateCordonRequired
+		if m.restartOnlyPredicate != nil && !upgradeRequested && !nodeState.IsOrphanedPod() {
+			// A node waiting for safe driver load must take the full flow, which evicts workloads
+			// before the load is unblocked at pod-restart-required. Restart-only would skip that eviction.
+			waitingForSafeLoad, serr := m.SafeDriverLoadManager.IsWaitingForSafeDriverLoad(ctx, nodeState.Node)
+			if serr != nil {
+				m.Log.V(consts.LogLevelError).Error(serr,
+					"failed to check safe driver load status; node kept in upgrade-required for retry",
+					"node", nodeState.Node.Name)
+				logEventf(m.EventRecorder, nodeState.Node, corev1.EventTypeWarning, GetEventReason(),
+					"Failed to check safe driver load status, will retry: %v", serr)
+				continue
+			}
+			if !waitingForSafeLoad {
+				restartOnly, perr := m.restartOnlyPredicate(ctx, nodeState.DriverPod, nodeState.DriverDaemonSet)
+				if perr != nil {
+					// On error, keep the node in upgrade-required and retry on the next
+					// reconcile instead of starting a full upgrade.
+					m.Log.V(consts.LogLevelError).Error(perr,
+						"restart-only predicate failed; node kept in upgrade-required for retry",
+						"node", nodeState.Node.Name)
+					logEventf(m.EventRecorder, nodeState.Node, corev1.EventTypeWarning, GetEventReason(),
+						"Failed to evaluate restart-only predicate, will retry: %v", perr)
+					continue
+				}
+				if restartOnly {
+					// Cordon the node so it stays unschedulable if the pod restart fails, as in
+					// the full upgrade flow; the uncordon-required state uncordons it on success.
+					if cerr := m.CordonManager.Cordon(ctx, nodeState.Node); cerr != nil {
+						m.Log.V(consts.LogLevelError).Error(cerr,
+							"failed to cordon node; node kept in upgrade-required for retry",
+							"node", nodeState.Node.Name)
+						logEventf(m.EventRecorder, nodeState.Node, corev1.EventTypeWarning, GetEventReason(),
+							"Failed to cordon node for restart-only upgrade, will retry: %v", cerr)
+						continue
+					}
+					targetState = UpgradeStatePodRestartRequired
+					m.Log.V(consts.LogLevelInfo).Info(
+						"Restart-only change detected; cordoning node and restarting driver pod "+
+							"in place, skipping pod-deletion and drain",
+						"node", nodeState.Node.Name)
+				}
+			}
+		}
+
+		err := m.NodeUpgradeStateProvider.ChangeNodeUpgradeState(ctx, nodeState.Node, targetState)
 		if err == nil {
 			upgradesAvailable--
-			m.Log.V(consts.LogLevelInfo).Info("Node waiting for cordon",
-				"node", nodeState.Node.Name)
+			m.Log.V(consts.LogLevelInfo).Info("Node moving to next upgrade state",
+				"node", nodeState.Node.Name, "state", targetState)
 		} else {
 			m.Log.V(consts.LogLevelError).Error(
-				err, "Failed to change node upgrade state", "state", UpgradeStateCordonRequired)
+				err, "Failed to change node upgrade state", "state", targetState)
 			return err
 		}
 	}
diff --git a/pkg/upgrade/upgrade_state.go b/pkg/upgrade/upgrade_state.go
@@ -40,6 +40,9 @@ type ClusterUpgradeStateManager interface {
 	// WithValidationEnabled provides an option to enable the optional 'validation' state
 	// and pass a podSelector to specify which pods are performing the validation
 	WithValidationEnabled(podSelector string) ClusterUpgradeStateManager
+	// WithRestartOnlyPredicate registers an optional predicate (see RestartOnlyPredicate);
+	// a nil predicate, the default, keeps the full upgrade flow for every out-of-sync node.
+	WithRestartOnlyPredicate(predicate RestartOnlyPredicate) ClusterUpgradeStateManager
 	// BuildState builds a point-in-time snapshot of the driver upgrade state in the cluster.
 	BuildState(ctx context.Context, namespace string,
 		driverLabels map[string]string) (*ClusterUpgradeState, error)
@@ -349,6 +352,14 @@ func (m *ClusterUpgradeStateManagerImpl) WithValidationEnabled(podSelector strin
 	return m
 }
 
+// WithRestartOnlyPredicate registers an optional restart-only predicate; a nil predicate
+// preserves the default full upgrade flow for every out-of-sync node.
+func (m *ClusterUpgradeStateManagerImpl) WithRestartOnlyPredicate(
+	predicate RestartOnlyPredicate) ClusterUpgradeStateManager {
+	m.restartOnlyPredicate = predicate
+	return m
+}
+
 // buildNodeUpgradeState creates a mapping between a node,
 // the driver POD running on them and the daemon set, controlling this pod
 func (m *ClusterUpgradeStateManagerImpl) buildNodeUpgradeState(
diff --git a/pkg/upgrade/upgrade_state_test.go b/pkg/upgrade/upgrade_state_test.go
@@ -37,8 +37,8 @@ import (
 	maintenancev1alpha1 "github.com/Mellanox/maintenance-operator/api/v1alpha1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 
-	v1alpha1 "github.com/NVIDIA/k8s-operator-libs/api/upgrade/v1alpha1"
-	upgrade "github.com/NVIDIA/k8s-operator-libs/pkg/upgrade"
+	"github.com/NVIDIA/k8s-operator-libs/api/upgrade/v1alpha1"
+	"github.com/NVIDIA/k8s-operator-libs/pkg/upgrade"
 	"github.com/NVIDIA/k8s-operator-libs/pkg/upgrade/mocks"
 )
 
@@ -347,6 +347,179 @@ var _ = Describe("UpgradeStateManager tests", func() {
 			Expect(stateCount[upgrade.UpgradeStateUpgradeRequired]).To(Equal(2))
 			Expect(stateCount[upgrade.UpgradeStateCordonRequired]).To(Equal(maxParallelUpgrades))
 		})
+		When("a RestartOnlyPredicate is registered", func() {
+			var (
+				daemonSet   *appsv1.DaemonSet
+				outdatedPod *corev1.Pod
+			)
+			BeforeEach(func() {
+				daemonSet = &appsv1.DaemonSet{ObjectMeta: v1.ObjectMeta{}}
+				// pod revision hash differs from the mocked DS hash ("test-hash-12345") -> out of sync
+				outdatedPod = &corev1.Pod{ObjectMeta: v1.ObjectMeta{
+					Labels: map[string]string{upgrade.PodControllerRevisionHashLabelKey: "test-hash-outdated"}}}
+			})
+
+			It("cordons the node and routes it to PodRestartRequired when the predicate returns true", func() {
+				var gotPod *corev1.Pod
+				var gotDS *appsv1.DaemonSet
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, pod *corev1.Pod, ds *appsv1.DaemonSet) (bool, error) {
+						gotPod, gotDS = pod, ds
+						return true, nil
+					})
+				localCordonManager := mocks.CordonManager{}
+				localCordonManager.On("Cordon", mock.Anything, mock.Anything).Return(nil)
+				stateManager.CordonManager = &localCordonManager
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStatePodRestartRequired))
+				localCordonManager.AssertCalled(GinkgoT(), "Cordon", mock.Anything, node)
+				// predicate received this node's driver pod and owning DaemonSet
+				Expect(gotPod).To(Equal(outdatedPod))
+				Expect(gotDS).To(Equal(daemonSet))
+			})
+
+			It("leaves the node in UpgradeRequired for retry when the cordon fails", func() {
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						return true, nil
+					})
+				localCordonManager := mocks.CordonManager{}
+				localCordonManager.On("Cordon", mock.Anything, mock.Anything).Return(fmt.Errorf("cordon failure"))
+				stateManager.CordonManager = &localCordonManager
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateUpgradeRequired))
+			})
+
+			It("routes the node to CordonRequired when the predicate returns false", func() {
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						return false, nil
+					})
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateCordonRequired))
+			})
+
+			It("leaves the node in UpgradeRequired for retry when the predicate returns an error", func() {
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						return true, fmt.Errorf("predicate failure")
+					})
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateUpgradeRequired))
+			})
+
+			It("does not invoke the predicate for an orphaned pod", func() {
+				called := false
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						called = true
+						return true, nil
+					})
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: &corev1.Pod{}, DriverDaemonSet: nil}, // orphaned
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(called).To(BeFalse())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateCordonRequired))
+			})
+
+			It("does not invoke the predicate for an upgrade-requested node", func() {
+				called := false
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						called = true
+						return true, nil
+					})
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				node.Annotations[upgrade.GetUpgradeRequestedAnnotationKey()] = "true"
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(called).To(BeFalse())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateCordonRequired))
+			})
+
+			It("does not invoke the predicate for a node waiting for safe driver load", func() {
+				called := false
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						called = true
+						return true, nil
+					})
+				node := nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired)
+				node.Annotations[upgrade.GetUpgradeDriverWaitForSafeLoadAnnotationKey()] = "true"
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = []*upgrade.NodeUpgradeState{
+					{Node: node, DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState,
+					&v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true})).To(Succeed())
+				Expect(called).To(BeFalse())
+				Expect(getNodeUpgradeState(node)).To(Equal(upgrade.UpgradeStateCordonRequired))
+			})
+
+			It("honors the maxParallelUpgrades throttle for restart-only nodes", func() {
+				const maxParallelUpgrades = 1
+				stateManagerInterface.WithRestartOnlyPredicate(
+					func(_ context.Context, _ *corev1.Pod, _ *appsv1.DaemonSet) (bool, error) {
+						return true, nil
+					})
+				nodeStates := []*upgrade.NodeUpgradeState{
+					{Node: nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired),
+						DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+					{Node: nodeWithUpgradeState(upgrade.UpgradeStateUpgradeRequired),
+						DriverPod: outdatedPod, DriverDaemonSet: daemonSet},
+				}
+				clusterState := upgrade.NewClusterUpgradeState()
+				clusterState.NodeStates[upgrade.UpgradeStateUpgradeRequired] = nodeStates
+				policy := &v1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true, MaxParallelUpgrades: maxParallelUpgrades}
+
+				Expect(stateManagerInterface.ApplyState(testCtx, &clusterState, policy)).To(Succeed())
+				stateCount := make(map[string]int)
+				for i := range nodeStates {
+					stateCount[getNodeUpgradeState(nodeStates[i].Node)]++
+				}
+				Expect(stateCount[upgrade.UpgradeStatePodRestartRequired]).To(Equal(maxParallelUpgrades))
+				Expect(stateCount[upgrade.UpgradeStateUpgradeRequired]).To(Equal(1))
+			})
+		})
 		It("UpgradeStateManager should start additional upgrades if maxParallelUpgrades limit is not reached", func() {
 			const maxParallelUpgrades = 4
 
