fix(recovery): use pod names for node error recovery

pkg/recovery/node_recovery_test.go

@@ -0,0 +1,155 @@
+package recovery
+
+import (
+	"context"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/baizeai/kcover/pkg/constants"
+	"github.com/baizeai/kcover/pkg/events"
+
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/fields"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/fake"
+	ktesting "k8s.io/client-go/testing"
+)
+
+func TestNodeErrorRestartsJobsOnFaultyNode(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		pods []*corev1.Pod
+		want map[string]int
+	}{
+		{
+			name: "pod names differ from job names",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train", "job-a-worker-0", "job-a", "node-a", true, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "job-b-worker-0", "job-b", "node-b", true, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "ordinary-pod", "", "node-a", true, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{"train/job-a": 1},
+		},
+		{
+			name: "multiple pods restart the same job once",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train", "job-a-worker-0", "job-a", "node-a", true, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "job-a-worker-1", "job-a", "node-a", true, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{"train/job-a": 1},
+		},
+		{
+			name: "same job name in different namespaces",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train-a", "job-worker-0", "job", "node-a", true, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train-b", "job-worker-0", "job", "node-a", true, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{"train-a/job": 1, "train-b/job": 1},
+		},
+		{
+			name: "eligible pod follows disabled and never pods",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train", "job-disabled", "job", "node-a", false, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "job-never", "job", "node-a", true, corev1.RestartPolicyNever),
+				nodeRecoveryPod("train", "job-worker", "job", "node-a", true, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{"train/job": 1},
+		},
+		{
+			name: "eligible pod precedes never and disabled pods",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train", "job-worker", "job", "node-a", true, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "job-never", "job", "node-a", true, corev1.RestartPolicyNever),
+				nodeRecoveryPod("train", "job-disabled", "job", "node-a", false, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{"train/job": 1},
+		},
+		{
+			name: "only ineligible pods",
+			pods: []*corev1.Pod{
+				nodeRecoveryPod("train", "job-disabled", "job", "node-a", false, corev1.RestartPolicyAlways),
+				nodeRecoveryPod("train", "job-never", "job", "node-a", true, corev1.RestartPolicyNever),
+				nodeRecoveryPod("train", "other-worker", "other", "node-b", true, corev1.RestartPolicyAlways),
+			},
+			want: map[string]int{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			objects := []runtime.Object{
+				&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node-a"}},
+				&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node-b"}},
+			}
+			for _, pod := range tt.pods {
+				objects = append(objects, pod)
+			}
+			client := fake.NewSimpleClientset(objects...)
+			// The fake client does not apply field selectors itself.
+			client.PrependReactor("list", "pods", func(action ktesting.Action) (bool, runtime.Object, error) {
+				list := action.(ktesting.ListAction)
+				restrictions := list.GetListRestrictions()
+				if restrictions.Fields.String() != "spec.nodeName=node-a" {
+					t.Fatalf("pod field selector = %q, want spec.nodeName=node-a", restrictions.Fields)
+				}
+				pods := &corev1.PodList{}
+				for _, pod := range tt.pods {
+					if list.GetNamespace() != "" && pod.Namespace != list.GetNamespace() {
+						continue
+					}
+					if restrictions.Fields.Matches(fields.Set{"spec.nodeName": pod.Spec.NodeName}) &&
+						restrictions.Labels.Matches(labels.Set(pod.Labels)) {
+						pods.Items = append(pods.Items, *pod.DeepCopy())
+					}
+				}
+				return true, pods, nil
+			})
+			restarts := map[string]int{}
+			client.PrependReactor("delete-collection", "pods", func(action ktesting.Action) (bool, runtime.Object, error) {
+				selector := action.(ktesting.DeleteCollectionAction).GetListRestrictions().Labels
+				job, ok := selector.RequiresExactMatch(constants.KubeflowJobLabel)
+				if !ok {
+					t.Fatalf("pod deletion selector %q does not select a job", selector)
+				}
+				restarts[action.GetNamespace()+"/"+job]++
+				// Keep pods available so the restart ledger must deduplicate subsequent pods.
+				return true, nil, nil
+			})
+
+			controller := NewController(client, nil, nil, 0, 0)
+			controller.jobRestartLedger.now = func() time.Time { return time.Unix(100, 0) }
+			controller.onEvent(context.Background(), events.Event{
+				ResourceType: events.Node,
+				Name:         "node-a",
+				EventType:    events.Error,
+			})
+
+			if !reflect.DeepEqual(restarts, tt.want) {
+				t.Errorf("job restarts = %v, want %v", restarts, tt.want)
+			}
+			assertNodeUnschedulable(t, client, "node-a", true)
+			assertNodeUnschedulable(t, client, "node-b", false)
+		})
+	}
+}
+
+func nodeRecoveryPod(namespace, name, job, node string, enabled bool, policy corev1.RestartPolicy) *corev1.Pod {
+	labels := map[string]string{}
+	if job != "" {
+		labels[constants.KubeflowJobLabel] = job
+	}
+	if enabled {
+		labels[constants.EnabledRecoveryLabel] = constants.True
+	}
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: name, Labels: labels},
+		Spec:       corev1.PodSpec{NodeName: node, RestartPolicy: policy},
+	}
+}

pkg/recovery/recovery.go

@@ -180,20 +180,21 @@ func (r *RecoveryController) onNodeError(ctx context.Context, e events.Event) {
 		return
 	}
 
-	jobs, err := r.listJobsOnNode(ctx, e.Name)
+	pods, err := r.listJobPodsOnNode(ctx, e.Name)
 	if err != nil {
-		klog.ErrorS(err, "failed to list jobs on node", "node", e.Name)
+		klog.ErrorS(err, "failed to list job pods on node", "node", e.Name)
 		return
 	}
 
-	for _, job := range jobs {
-		r.onPodError(ctx, job.ns, job.name)
+	// Check each pod's recovery policy; the restart ledger deduplicates jobs.
+	for _, pod := range pods {
+		r.onPodError(ctx, pod.ns, pod.name)
 	}
 
 	r.ensureNodeUnschedulable(ctx, e.Name)
 }
 
-func (r *RecoveryController) listJobsOnNode(ctx context.Context, nodeName string) ([]nsName, error) {
+func (r *RecoveryController) listJobPodsOnNode(ctx context.Context, nodeName string) ([]nsName, error) {
 	// TODO: traner v2 & lws?
 	ctx, cancel := kube.WithRequestTimeout(ctx)
 	defer cancel()
@@ -206,19 +207,12 @@ func (r *RecoveryController) listJobsOnNode(ctx context.Context, nodeName string
 		return nil, err
 	}
 
-	jobs := make(map[nsName]struct{})
+	items := make([]nsName, 0, len(pods.Items))
 	for _, pod := range pods.Items {
-		jobLabel, ok := pod.Labels[constants.KubeflowJobLabel]
-		if !ok {
+		if _, ok := pod.Labels[constants.KubeflowJobLabel]; !ok {
 			continue
 		}
-
-		jobs[nsName{ns: pod.Namespace, name: jobLabel}] = struct{}{}
-	}
-
-	items := make([]nsName, 0, len(jobs))
-	for job := range jobs {
-		items = append(items, job)
+		items = append(items, nsName{ns: pod.Namespace, name: pod.Name})
 	}
 
 	return items, nil