From ab12ab5317f5775f9832be79dff59e2221328763 Mon Sep 17 00:00:00 2001 From: Shivkumar Ople Date: Tue, 3 Nov 2026 19:35:58 +0530 Subject: [PATCH] Respect driver upgrade policy during orphan pod migration Signed-off-by: Shivkumar Ople --- controllers/nodelabeling_controller.go | 33 +++- controllers/nodelabeling_controller_test.go | 19 ++ controllers/nodelabeling_migration_test.go | 186 ++++++++++++++++++++ 3 files changed, 231 insertions(+), 7 deletions(-) create mode 100644 controllers/nodelabeling_migration_test.go diff --git a/controllers/nodelabeling_controller.go b/controllers/nodelabeling_controller.go index 54a581bda..651bfe27a 100644 --- a/controllers/nodelabeling_controller.go +++ b/controllers/nodelabeling_controller.go @@ -28,6 +28,7 @@ import ( resourcev1 "k8s.io/api/resource/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/tools/events" "k8s.io/client-go/util/workqueue" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" @@ -58,6 +59,7 @@ type NodeLabelingReconciler struct { Scheme *runtime.Scheme Namespace string Log logr.Logger + recorder events.EventRecorder } // nodeLabelingController holds per-reconcile state so that helper methods don't need to @@ -70,6 +72,7 @@ type nodeLabelingController struct { clusterPolicy *gpuv1.ClusterPolicy gpuCluster *nvidiav1alpha1.GPUCluster logger logr.Logger + recorder events.EventRecorder // draPluginRemovalDeferred records that gpu.deploy.dra-driver removal was skipped on // at least one node because pods holding gpu.nvidia.com claims are still present; the @@ -213,6 +216,7 @@ func (r *NodeLabelingReconciler) Reconcile(ctx context.Context, req ctrl.Request clusterPolicy: clusterPolicy, gpuCluster: gpuCluster, logger: logger, + recorder: r.recorder, } gpuLabelUpdateResult, err := nlc.labelGPUNodes(ctx) @@ -607,7 +611,19 @@ func (nlc *nodeLabelingController) labelNodesWithOrphanedDriverPods(ctx context. nlc.logger.Error(err, "failed to get node for orphaned driver pod", "pod", pod.Name, "node", pod.Spec.NodeName) continue } - if !nodeOwnedByNVIDIADriver(node, nvidiaDrivers.Items) { + owner := nodeNVIDIADriverOwner(node, nvidiaDrivers.Items) + if owner == nil { + continue + } + + // Migration must use the same per-driver policy as the upgrade controller, + // which removes upgrade-state labels when automatic upgrades are disabled. + if !owner.Spec.GetUpgradePolicyWithDefaults().AutoUpgrade { + nlc.recorder.Eventf(owner, &pod, corev1.EventTypeWarning, "ManualDriverPodReplacementRequired", "SkipMigration", + "Orphaned driver pod %s/%s on node %s requires manual replacement because NVIDIADriver %s has autoUpgrade disabled", + pod.Namespace, pod.Name, node.Name, owner.Name) + nlc.logger.Info("Orphaned driver pod requires manual replacement because automatic upgrades are disabled", + "nvidiaDriver", owner.Name, "node", node.Name, "pod", pod.Name) continue } @@ -632,20 +648,21 @@ func (nlc *nodeLabelingController) labelNodesWithOrphanedDriverPods(ctx context. return nil } -// nodeOwnedByNVIDIADriver returns true when the node has an owner label matching a live NVIDIADriver. -func nodeOwnedByNVIDIADriver(node *corev1.Node, nvidiaDrivers []nvidiav1alpha1.NVIDIADriver) bool { +// nodeNVIDIADriverOwner returns the live NVIDIADriver selected by the node's owner label. +func nodeNVIDIADriverOwner(node *corev1.Node, nvidiaDrivers []nvidiav1alpha1.NVIDIADriver) *nvidiav1alpha1.NVIDIADriver { if node.Labels == nil || node.Labels[consts.NVIDIADriverOwnerLabel] == "" { - return false + return nil } - for _, nvidiaDriver := range nvidiaDrivers { + for i := range nvidiaDrivers { + nvidiaDriver := &nvidiaDrivers[i] if nvidiaDriver.HasDeletionTimestamp() { continue } if node.Labels[consts.NVIDIADriverOwnerLabel] == nvidiaDriver.Name { - return true + return nvidiaDriver } } - return false + return nil } // isDriverUpgradeRequestAllowed returns true when migration can request a driver upgrade @@ -677,6 +694,8 @@ func singletonCRPredicate[T interface { // SetupWithManager registers the NodeLabelingReconciler with the controller-runtime manager. func (r *NodeLabelingReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error { + r.recorder = mgr.GetEventRecorder("nvidia-gpu-operator") + mapToSingleton := func(_ context.Context, _ client.Object) []reconcile.Request { return []reconcile.Request{{NamespacedName: types.NamespacedName{Name: nodeLabelingControllerSingletonName}}} } diff --git a/controllers/nodelabeling_controller_test.go b/controllers/nodelabeling_controller_test.go index 852058fd5..65c490e6e 100644 --- a/controllers/nodelabeling_controller_test.go +++ b/controllers/nodelabeling_controller_test.go @@ -31,6 +31,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/tools/events" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/client/fake" "sigs.k8s.io/controller-runtime/pkg/client/interceptor" @@ -1051,6 +1052,7 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { const namespace = "test-ns" const driverName = "gpu-driver" + upgrade.SetDriverName("gpu") upgradeStateLabel := upgrade.GetUpgradeStateLabelKey() // liveDriver returns a NVIDIADriver with no deletion timestamp. @@ -1170,6 +1172,18 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, expectedUpgradeState: map[string]string{"node-1": ""}, }, + { + name: "orphaned pod on node with auto-upgrade disabled → not labeled", + nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{{ + ObjectMeta: metav1.ObjectMeta{Name: driverName}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + UpgradePolicy: &nvidiav1alpha1.DriverUpgradePolicySpec{AutoUpgrade: false}, + }, + }}, + nodes: []*corev1.Node{ownedNode("node-1", "")}, + pods: []*corev1.Pod{orphanedPod("pod-1", "node-1")}, + expectedUpgradeState: map[string]string{"node-1": ""}, + }, { name: "NVIDIADriver with deletion timestamp → node treated as unowned, not labeled", nvidiaDrivers: []*nvidiav1alpha1.NVIDIADriver{{ @@ -1207,6 +1221,7 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { client: fakeClient, namespace: namespace, logger: logr.Discard(), + recorder: events.NewFakeRecorder(10), } require.NoError(t, nlc.labelNodesWithOrphanedDriverPods(context.Background())) @@ -1214,6 +1229,10 @@ func TestLabelNodesWithOrphanedDriverPods(t *testing.T) { for nodeName, expectedUpgradeState := range tc.expectedUpgradeState { node := &corev1.Node{} require.NoError(t, fakeClient.Get(context.Background(), types.NamespacedName{Name: nodeName}, node)) + if expectedUpgradeState == "" { + assert.NotContains(t, node.Labels, upgradeStateLabel, "node %q upgrade state label", nodeName) + continue + } assert.Equal(t, expectedUpgradeState, node.Labels[upgradeStateLabel], "node %q upgrade state", nodeName) } }) diff --git a/controllers/nodelabeling_migration_test.go b/controllers/nodelabeling_migration_test.go new file mode 100644 index 000000000..2616147fa --- /dev/null +++ b/controllers/nodelabeling_migration_test.go @@ -0,0 +1,186 @@ +/** +# 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 controllers + +import ( + "context" + "testing" + + "github.com/NVIDIA/k8s-operator-libs/pkg/upgrade" + "github.com/go-logr/logr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" + "k8s.io/client-go/tools/events" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + gpuv1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1" + nvidiav1alpha1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1alpha1" + "github.com/NVIDIA/gpu-operator/internal/consts" +) + +func TestNodeLabelingMigrationRespectsPerDriverPolicies(t *testing.T) { + for _, stack := range []string{"ClusterPolicy", "GPUCluster"} { + t.Run(stack, func(t *testing.T) { + ctx := context.Background() + disabled := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "disabled-driver"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + Default: true, + UpgradePolicy: &nvidiav1alpha1.DriverUpgradePolicySpec{AutoUpgrade: false}, + }, + } + enabled := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "enabled-driver"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + NodeSelector: map[string]string{"driver-pool": "enabled"}, + UpgradePolicy: &nvidiav1alpha1.DriverUpgradePolicySpec{AutoUpgrade: true}, + }, + } + defaulted := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "defaulted-driver"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + NodeSelector: map[string]string{"driver-pool": "defaulted"}, + }, + } + objects := []client.Object{disabled, enabled, defaulted} + if stack == "ClusterPolicy" { + objects = append(objects, &gpuv1.ClusterPolicy{ + ObjectMeta: metav1.ObjectMeta{Name: "cluster-policy"}, + Spec: gpuv1.ClusterPolicySpec{Driver: gpuv1.DriverSpec{ + UseNvidiaDriverCRD: new(true), + }}, + }) + } else { + objects = append(objects, &nvidiav1alpha1.GPUCluster{ + ObjectMeta: metav1.ObjectMeta{Name: "gpu-cluster"}, + }) + } + for _, pool := range []string{"disabled", "enabled", "defaulted"} { + node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{ + Name: pool + "-node", + Labels: map[string]string{ + gpuPCILabelKey: "true", commonGPULabelKey: commonGPULabelValue, + "driver-pool": pool, + }, + }} + if pool == "disabled" { + node.Annotations = map[string]string{driverAutoUpgradeAnnotationKey: "true"} + } + objects = append(objects, node, orphanedDriverPod(pool+"-pod", testOperatorNamespace, node.Name)) + } + upgrader, _ := newTestUpgradeReconciler(t, objects...) + recorder := events.NewFakeRecorder(10) + labeler := &NodeLabelingReconciler{ + Client: upgrader.Client, Namespace: testOperatorNamespace, recorder: recorder, + } + stateKey := upgrade.GetUpgradeStateLabelKey() + + for cycle := range 3 { + _, err := labeler.Reconcile(ctx, reconcile.Request{}) + require.NoError(t, err) + for _, pool := range []string{"disabled", "enabled", "defaulted"} { + node := &corev1.Node{} + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: pool + "-node"}, node)) + assert.Equal(t, pool+"-driver", node.Labels[consts.NVIDIADriverOwnerLabel]) + if pool == "disabled" { + assert.NotContains(t, node.Labels, stateKey, "cycle %d", cycle) + assert.NotContains(t, node.Annotations, driverAutoUpgradeAnnotationKey) + } else { + assert.Equal(t, upgrade.UpgradeStateUpgradeRequired, node.Labels[stateKey]) + assert.Equal(t, "true", node.Annotations[driverAutoUpgradeAnnotationKey]) + } + } + require.Len(t, recorder.Events, 1) + event := <-recorder.Events + assert.Contains(t, event, "Warning ManualDriverPodReplacementRequired") + assert.Contains(t, event, "disabled-driver") + assert.Contains(t, event, "disabled-node") + assert.Contains(t, event, "disabled-pod") + + _, err = upgrader.Reconcile(ctx, upgradeSingletonRequest()) + require.NoError(t, err) + node := &corev1.Node{} + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: "disabled-node"}, node)) + assert.NotContains(t, node.Labels, stateKey) + } + + // Enabling automatic upgrades must allow migration to resume without changing ownership. + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: disabled.Name}, disabled)) + disabled.Spec.UpgradePolicy.AutoUpgrade = true + require.NoError(t, upgrader.Update(ctx, disabled)) + _, err := labeler.Reconcile(ctx, reconcile.Request{}) + require.NoError(t, err) + node := &corev1.Node{} + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: "disabled-node"}, node)) + assert.Equal(t, upgrade.UpgradeStateUpgradeRequired, node.Labels[stateKey]) + assert.Equal(t, "true", node.Annotations[driverAutoUpgradeAnnotationKey]) + assert.Empty(t, recorder.Events) + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{ + Name: "disabled-pod", Namespace: testOperatorNamespace, + }, &corev1.Pod{})) + }) + } +} + +func TestOrphanedDriverPodDisabledPolicyPreservesExistingState(t *testing.T) { + upgrade.SetDriverName("gpu") + stateKey := upgrade.GetUpgradeStateLabelKey() + for _, state := range []string{ + upgrade.UpgradeStateUnknown, upgrade.UpgradeStateDone, upgrade.UpgradeStateUpgradeRequired, + upgrade.UpgradeStatePodRestartRequired, upgrade.UpgradeStateFailed, + } { + t.Run("state="+state, func(t *testing.T) { + ctx := context.Background() + driver := &nvidiav1alpha1.NVIDIADriver{ + ObjectMeta: metav1.ObjectMeta{Name: "disabled-driver"}, + Spec: nvidiav1alpha1.NVIDIADriverSpec{ + UpgradePolicy: &nvidiav1alpha1.DriverUpgradePolicySpec{AutoUpgrade: false}, + }, + } + node := nodeWithUpgradeState("disabled-node", driver.Name) + if state == upgrade.UpgradeStateUnknown { + delete(node.Labels, stateKey) + } else { + node.Labels[stateKey] = state + } + pod := orphanedDriverPod("legacy-driver", testOperatorNamespace, node.Name) + upgrader, manager := newTestUpgradeReconciler(t, driver, node, pod) + recorder := events.NewFakeRecorder(10) + labeler := &nodeLabelingController{ + client: upgrader.Client, namespace: testOperatorNamespace, logger: logr.Discard(), recorder: recorder, + } + require.NoError(t, labeler.labelNodesWithOrphanedDriverPods(ctx)) + updated := &corev1.Node{} + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: node.Name}, updated)) + assert.Equal(t, node.Labels, updated.Labels) + require.Len(t, recorder.Events, 1) + assert.Contains(t, <-recorder.Events, "ManualDriverPodReplacementRequired") + + // The upgrade reconciler remains responsible for clearing stale upgrade state. + _, err := upgrader.Reconcile(ctx, upgradeSingletonRequest()) + require.NoError(t, err) + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: node.Name}, updated)) + assert.NotContains(t, updated.Labels, stateKey) + assert.Zero(t, manager.applyCalls) + require.NoError(t, upgrader.Get(ctx, types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}, &corev1.Pod{})) + }) + } +}