Fix node drain deadlock when device plugin holds kernel module

When draining a node running a KMM kernel module with an associated device
plugin, a race condition between the Module reconciler and NMC reconciler
can cause a deadlock:

1. Module reconciler removes the module from NMC spec (node unschedulable)
2. NMC reconciler creates an unloader pod (orphan status path)
3. The kmm-ready label stays on the node (removeOrphanedLabels requires
   both spec AND status to be absent, but status persists until unload
   succeeds)
4. Device plugin DaemonSet pod stays running (nodeSelector still matches
   and DaemonSet pods tolerate the unschedulable taint)
5. Unloader's modprobe -r fails indefinitely (device file held open)

Fix: after processing orphan statuses, check if the node is unschedulable
for each orphan module. If so, add the kmm-ready labels to the removal set.
This ensures the device plugin pod is evicted, releasing the device file so
the unloader can succeed.

This fix was written by bug buddy - ai workflow

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: TomerNewman

internal/controllers/nmc_reconciler.go

@@ -141,6 +141,19 @@ func (r *NMCReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 		}
 	}
 
+	// Fix drain deadlock: when the Module reconciler removes the NMC spec before the NMC
+	// reconciler processes the node-taint change, orphan statuses get unloader pods but the
+	// kmm-ready label stays (removeOrphanedLabels requires status to be absent). This keeps
+	// the device plugin DaemonSet pod running and holding the device file, deadlocking the
+	// unloader. Remove kmm-ready labels for orphan statuses on unschedulable nodes so the
+	// device plugin terminates and the unloader can eventually succeed.
+	for _, status := range statusMap {
+		if !r.nodeAPI.IsNodeSchedulable(&node, status.Tolerations) {
+			readyLabelsToRemove[utils.GetKernelModuleReadyNodeLabel(status.Namespace, status.Name)] = ""
+			readyLabelsToRemove[utils.GetKernelModuleVersionReadyNodeLabel(status.Namespace, status.Name)] = ""
+		}
+	}
+
 	// removing label of loaded kmods
 	if len(readyLabelsToRemove) != 0 {
 		if err := r.nodeAPI.UpdateLabels(ctx, &node, nil, readyLabelsToRemove); err != nil {

internal/controllers/nmc_reconciler_test.go

@@ -236,6 +236,75 @@ var _ = Describe("NodeModulesConfigReconciler_Reconcile", func() {
 		Expect(node).ToNot(Equal(expectedNode))
 	})
 
+	It("should remove kmod labels for orphan statuses on unschedulable node to prevent drain deadlock", func() {
+		const (
+			modNamespace = "test-ns"
+			modName      = "test-module"
+		)
+
+		orphanStatus := kmmv1beta1.NodeModuleStatus{
+			ModuleItem: kmmv1beta1.ModuleItem{
+				Namespace: modNamespace,
+				Name:      modName,
+			},
+			Config: kmmv1beta1.ModuleConfig{
+				KernelVersion: "5.14.0",
+			},
+		}
+
+		nmc := &kmmv1beta1.NodeModulesConfig{
+			ObjectMeta: metav1.ObjectMeta{Name: nmcName},
+			Spec:       kmmv1beta1.NodeModulesConfigSpec{},
+			Status: kmmv1beta1.NodeModulesConfigStatus{
+				Modules: []kmmv1beta1.NodeModuleStatus{orphanStatus},
+			},
+		}
+
+		node := v1.Node{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: map[string]string{
+					kmodReadyLabel:   "",
+					labelNotToRemove: "",
+				},
+			},
+			Spec: v1.NodeSpec{
+				Taints: []v1.Taint{
+					{
+						Key:    "node.kubernetes.io/unschedulable",
+						Effect: v1.TaintEffectNoSchedule,
+					},
+				},
+			},
+		}
+
+		contextWithValueMatch := gomock.AssignableToTypeOf(
+			reflect.TypeOf((*context.Context)(nil)).Elem(),
+		)
+
+		gomock.InOrder(
+			kubeClient.
+				EXPECT().
+				Get(ctx, nmcNsn, &kmmv1beta1.NodeModulesConfig{}).
+				Do(func(_ context.Context, _ types.NamespacedName, kubeNmc ctrlclient.Object, _ ...ctrlclient.Options) {
+					*kubeNmc.(*kmmv1beta1.NodeModulesConfig) = *nmc
+				}),
+			kubeClient.EXPECT().Get(ctx, types.NamespacedName{Name: nmc.Name}, &v1.Node{}).DoAndReturn(
+				func(_ context.Context, _ types.NamespacedName, fetchedNode *v1.Node, _ ...ctrlclient.Options) error {
+					*fetchedNode = node
+					return nil
+				},
+			),
+			wh.EXPECT().SyncStatus(ctx, nmc, &node),
+			wh.EXPECT().ProcessUnconfiguredModuleStatus(contextWithValueMatch, nmc, &orphanStatus, &node),
+			nm.EXPECT().IsNodeSchedulable(&node, nil).Return(false),
+			nm.EXPECT().UpdateLabels(ctx, &node, nil, map[string]string{kmodReadyLabel: "", kmodVersionReadyLabel: ""}).Return(nil),
+		)
+
+		res, err := r.Reconcile(ctx, req)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(res).To(BeZero())
+	})
+
 	It("should process spec entries and orphan statuses", func() {
 		const (
 			mod0Name = "mod0"
@@ -304,6 +373,7 @@ var _ = Describe("NodeModulesConfigReconciler_Reconcile", func() {
 			nm.EXPECT().IsNodeSchedulable(&node, nil).Return(true),
 			wh.EXPECT().ProcessModuleSpec(contextWithValueMatch, nmc, &spec1, nil, &node),
 			wh.EXPECT().ProcessUnconfiguredModuleStatus(contextWithValueMatch, nmc, &status2, &node),
+			nm.EXPECT().IsNodeSchedulable(&node, nil).Return(true),
 			wh.EXPECT().GarbageCollectInUseLabels(ctx, nmc),
 			wh.EXPECT().GarbageCollectWorkerPods(ctx, nmc),
 			wh.EXPECT().UpdateNodeLabels(ctx, nmc, &node).Return(loaded, unloaded, err),
@@ -384,6 +454,7 @@ var _ = Describe("NodeModulesConfigReconciler_Reconcile", func() {
 			nm.EXPECT().IsNodeSchedulable(&node, nil).Return(true),
 			wh.EXPECT().ProcessModuleSpec(contextWithValueMatch, nmc, &spec0, &status0, &node).Return(errors.New(errorMeassge)),
 			wh.EXPECT().ProcessUnconfiguredModuleStatus(contextWithValueMatch, nmc, &status2, &node).Return(errors.New(errorMeassge)),
+			nm.EXPECT().IsNodeSchedulable(&node, nil).Return(true),
 			wh.EXPECT().GarbageCollectInUseLabels(ctx, nmc).Return(errors.New(errorMeassge)),
 			wh.EXPECT().GarbageCollectWorkerPods(ctx, nmc).Return(errors.New(errorMeassge)),
 			wh.EXPECT().UpdateNodeLabels(ctx, nmc, &node).Return(nil, nil, errors.New(errorMeassge)),