Fix slinky-drainer to wait for fully drained state before deleting (#919)

Signed-off-by: Fagani Hajizada <fagani.hajizada@gmail.com>

Author: faganihajizada

plugins/mock-slurm-operator/pkg/controller/node_controller.go

@@ -32,6 +32,7 @@ import (
 const (
 	annotationKey           = "nodeset.slinky.slurm.net/node-cordon-reason"
 	slurmDrainConditionType = "SlurmNodeStateDrain"
+	slurmIdleConditionType  = "SlurmNodeStateIdle"
 )
 
 type NodeReconciler struct {
@@ -127,13 +128,22 @@ func (r *NodeReconciler) listSlinkyPodsOnNode(ctx context.Context, nodeName stri
 }
 
 func (r *NodeReconciler) updatePodDrainCondition(ctx context.Context, pod *corev1.Pod, reason string) error {
-	pod.Status.Conditions = append(pod.Status.Conditions, corev1.PodCondition{
-		Type:               slurmDrainConditionType,
-		Status:             corev1.ConditionTrue,
-		LastTransitionTime: metav1.Now(),
-		Reason:             "NodeCordonedBySlinkyDrainer",
-		Message:            reason,
-	})
+	pod.Status.Conditions = append(pod.Status.Conditions,
+		corev1.PodCondition{
+			Type:               slurmDrainConditionType,
+			Status:             corev1.ConditionTrue,
+			LastTransitionTime: metav1.Now(),
+			Reason:             "NodeCordonedBySlinkyDrainer",
+			Message:            reason,
+		},
+		corev1.PodCondition{
+			Type:               slurmIdleConditionType,
+			Status:             corev1.ConditionTrue,
+			LastTransitionTime: metav1.Now(),
+			Reason:             "NodeDrained",
+			Message:            "Node is idle after drain",
+		},
+	)
 
 	return r.Status().Update(ctx, pod)
 }

plugins/slinky-drainer/pkg/controller/drainrequest_controller.go

@@ -45,6 +45,15 @@ const (
 	annotationPrefix                 = "[J] [NVSentinel]"
 	nvsentinelStateLabelKey          = "dgxc.nvidia.com/nvsentinel-state"
 	drainRequestFinalizer            = "nvsentinel.nvidia.com/slinky-drainer"
+
+	// Slurm base-state conditions that indicate the node still has running work.
+	// A pod is only considered fully drained when SlurmNodeStateDrain is True
+	// and none of these busy-state conditions are True.
+	// Ref: https://github.com/SlinkyProject/slurm-operator/blob/main/pkg/conditions/conditions.go
+	slurmNodeStateAllocatedConditionType  = "SlurmNodeStateAllocated"
+	slurmNodeStateMixedConditionType      = "SlurmNodeStateMixed"
+	slurmNodeStateCompletingConditionType = "SlurmNodeStateCompleting"
+	slurmNodeStateUndrainConditionType    = "SlurmNodeStateUndrain"
 )
 
 type DrainRequestReconciler struct {
@@ -100,9 +109,9 @@ func (r *DrainRequestReconciler) Reconcile(ctx context.Context, req ctrl.Request
 		return r.markDrainComplete(ctx, drainReq, "NoPods", "No Slinky pods found on node")
 	}
 
-	allReady, notReadyPods := r.checkPodsReadyForDrain(pods)
-	if !allReady {
-		slog.Info("Waiting for pods to be ready for drain",
+	allDrained, notReadyPods := r.checkPodsFullyDrained(pods)
+	if !allDrained {
+		slog.Info("Waiting for pods to be fully drained",
 			"drainrequest", req.NamespacedName,
 			"total", len(pods),
 			"notReady", len(notReadyPods),
@@ -289,15 +298,15 @@ func (r *DrainRequestReconciler) getSlinkyPods(ctx context.Context, nodeName str
 	return podList.Items, nil
 }
 
-func (r *DrainRequestReconciler) checkPodsReadyForDrain(pods []corev1.Pod) (bool, []string) {
+func (r *DrainRequestReconciler) checkPodsFullyDrained(pods []corev1.Pod) (bool, []string) {
 	var notReady []string
 
 	for _, pod := range pods {
 		if !isPodReady(&pod) {
 			continue
 		}
 
-		if !hasSlurmDrainCondition(&pod) {
+		if !isPodDrained(&pod) {
 			notReady = append(notReady, pod.Name)
 		}
 	}
@@ -315,9 +324,28 @@ func isPodReady(pod *corev1.Pod) bool {
 	return false
 }
 
-func hasSlurmDrainCondition(pod *corev1.Pod) bool {
+// isPodDrained returns true when the Slurm node behind this pod has the DRAIN
+// flag set, the UNDRAIN flag is NOT set, and no work is running. This mirrors
+// the Slinky operator's IsNodeDrained logic so we only delete pods after all
+// Slurm jobs have finished.
+func isPodDrained(pod *corev1.Pod) bool {
+	return isNodeDrain(pod) && !isPodBusy(pod)
+}
+
+func isNodeDrain(pod *corev1.Pod) bool {
+	return hasPodCondition(pod, slurmNodeStateDrainConditionType) &&
+		!hasPodCondition(pod, slurmNodeStateUndrainConditionType)
+}
+
+func isPodBusy(pod *corev1.Pod) bool {
+	return hasPodCondition(pod, slurmNodeStateAllocatedConditionType) ||
+		hasPodCondition(pod, slurmNodeStateMixedConditionType) ||
+		hasPodCondition(pod, slurmNodeStateCompletingConditionType)
+}
+
+func hasPodCondition(pod *corev1.Pod, condType corev1.PodConditionType) bool {
 	for _, cond := range pod.Status.Conditions {
-		if cond.Type == slurmNodeStateDrainConditionType && cond.Status == corev1.ConditionTrue {
+		if cond.Type == condType && cond.Status == corev1.ConditionTrue {
 			return true
 		}
 	}

plugins/slinky-drainer/pkg/controller/drainrequest_controller_test.go

@@ -67,7 +67,14 @@ func TestReconcile_FullDrainCycle(t *testing.T) {
 
 	assertNodeAnnotation(t, tc, node.Name, "[J] [NVSentinel] 79 GPU:0 - GPU has fallen off the bus")
 
-	markPodDrainReady(t, tc, pod.Name, pod.Namespace)
+	// Simulate DRAINING state: Slurm accepted drain but jobs still running.
+	markPodDraining(t, tc, pod.Name, pod.Namespace)
+
+	// Pod must NOT be deleted while in DRAINING state (busy).
+	assertPodNotDeleted(t, tc, pod.Name, pod.Namespace, 3*time.Second)
+
+	// Simulate transition to DRAINED state: jobs finished, node is idle.
+	markPodDrained(t, tc, pod.Name, pod.Namespace)
 
 	waitForDrainComplete(t, tc, "drain-full-cycle", "default")
 	waitForPodDeletion(t, tc, pod.Name, pod.Namespace)
@@ -97,6 +104,30 @@ func TestReconcile_PreExistingAnnotationPreserved(t *testing.T) {
 	assertNodeAnnotation(t, tc, node.Name, "Manual drain by operator")
 }
 
+func TestReconcile_DrainingPodNotDeleted(t *testing.T) {
+	tc := setupTestEnv(t, "drain-still-draining")
+
+	node := createNode(t, tc, "test-node-draining", nil, map[string]string{
+		nvsentinelStateLabelKey: "draining",
+	})
+	pod := createSlinkyPod(t, tc, node.Name)
+	markPodReady(t, tc, pod.Name, pod.Namespace)
+	createDrainRequest(t, tc, "drain-still-draining", drainv1alpha1.DrainRequestSpec{
+		NodeName:  node.Name,
+		ErrorCode: []string{"79"},
+		Reason:    "GPU has fallen off the bus",
+	})
+
+	assertNodeAnnotation(t, tc, node.Name, "[J] [NVSentinel] 79 - GPU has fallen off the bus")
+
+	// Set pod to DRAINING: Drain flag set but node is still busy (Allocated).
+	markPodDraining(t, tc, pod.Name, pod.Namespace)
+
+	// Verify pod is NOT deleted and DrainRequest is NOT completed while draining.
+	assertPodNotDeleted(t, tc, pod.Name, pod.Namespace, 5*time.Second)
+	assertDrainNotComplete(t, tc, "drain-still-draining", "default")
+}
+
 // ---------------------------------------------------------------------------
 // Test setup
 // ---------------------------------------------------------------------------
@@ -235,24 +266,43 @@ func createFailedPod(t *testing.T, tc *testEnvContext, nodeName string) {
 func markPodReady(t *testing.T, tc *testEnvContext, podName, podNamespace string) {
 	t.Helper()
 
+	var pod corev1.Pod
+
+	require.Eventually(t, func() bool {
+		return tc.client.Get(tc.ctx, types.NamespacedName{Name: podName, Namespace: podNamespace}, &pod) == nil
+	}, testTimeout, testPollInterval, "Pod %s/%s should exist", podNamespace, podName)
+
+	pod.Status.Phase = corev1.PodRunning
+	pod.Status.Conditions = []corev1.PodCondition{
+		{Type: corev1.PodReady, Status: corev1.ConditionTrue},
+	}
+	require.NoError(t, tc.client.Status().Update(tc.ctx, &pod))
+}
+
+func markPodDraining(t *testing.T, tc *testEnvContext, podName, podNamespace string) {
+	t.Helper()
+
 	pod := &corev1.Pod{}
 	require.NoError(t, tc.client.Get(tc.ctx, types.NamespacedName{Name: podName, Namespace: podNamespace}, pod))
 
-	pod.Status.Phase = corev1.PodRunning
 	pod.Status.Conditions = []corev1.PodCondition{
 		{Type: corev1.PodReady, Status: corev1.ConditionTrue},
+		{Type: slurmNodeStateDrainConditionType, Status: corev1.ConditionTrue},
+		{Type: slurmNodeStateAllocatedConditionType, Status: corev1.ConditionTrue},
 	}
 	require.NoError(t, tc.client.Status().Update(tc.ctx, pod))
 }
 
-func markPodDrainReady(t *testing.T, tc *testEnvContext, podName, podNamespace string) {
+func markPodDrained(t *testing.T, tc *testEnvContext, podName, podNamespace string) {
 	t.Helper()
 
 	pod := &corev1.Pod{}
 	require.NoError(t, tc.client.Get(tc.ctx, types.NamespacedName{Name: podName, Namespace: podNamespace}, pod))
 
 	pod.Status.Conditions = []corev1.PodCondition{
+		{Type: corev1.PodReady, Status: corev1.ConditionTrue},
 		{Type: slurmNodeStateDrainConditionType, Status: corev1.ConditionTrue},
+		{Type: "SlurmNodeStateIdle", Status: corev1.ConditionTrue},
 	}
 	require.NoError(t, tc.client.Status().Update(tc.ctx, pod))
 }
@@ -328,6 +378,32 @@ func waitForAnnotationRemoved(t *testing.T, tc *testEnvContext, nodeName string)
 	}, testTimeout, testPollInterval, "Annotation on node %s should be removed", nodeName)
 }
 
+func assertPodNotDeleted(t *testing.T, tc *testEnvContext, podName, podNamespace string, waitDuration time.Duration) {
+	t.Helper()
+
+	assert.Never(t, func() bool {
+		p := &corev1.Pod{}
+		if err := tc.client.Get(tc.ctx, types.NamespacedName{Name: podName, Namespace: podNamespace}, p); err != nil {
+			return apierrors.IsNotFound(err)
+		}
+
+		return p.DeletionTimestamp != nil
+	}, waitDuration, testPollInterval, "Pod %s/%s should NOT be deleted while draining", podNamespace, podName)
+}
+
+func assertDrainNotComplete(t *testing.T, tc *testEnvContext, drName, drNamespace string) {
+	t.Helper()
+
+	dr := &drainv1alpha1.DrainRequest{}
+	require.NoError(t, tc.client.Get(tc.ctx, types.NamespacedName{Name: drName, Namespace: drNamespace}, dr))
+
+	for _, c := range dr.Status.Conditions {
+		if c.Type == drainCompleteConditionType && c.Status == metav1.ConditionTrue {
+			t.Fatalf("DrainRequest %s/%s should NOT have DrainComplete=True while pods are still draining", drNamespace, drName)
+		}
+	}
+}
+
 func assertNodeAnnotation(t *testing.T, tc *testEnvContext, nodeName, expectedValue string) {
 	t.Helper()