diff --git a/executor/pkg/plugin/k8s/child_pod_gpu_fault_test.go b/executor/pkg/plugin/k8s/child_pod_gpu_fault_test.go new file mode 100644 index 00000000000..26f83f7c310 --- /dev/null +++ b/executor/pkg/plugin/k8s/child_pod_gpu_fault_test.go @@ -0,0 +1,524 @@ +package k8s + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + v1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + k8stypes "k8s.io/apimachinery/pkg/types" + k8sscheme "k8s.io/client-go/kubernetes/scheme" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + pluginsCoreMock "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/gpufault" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/k8s" + k8sMocks "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/k8s/mocks" + "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core" +) + +// childPodPlugin is a plugin that tracks a CRD and can name the pods an operator expanded +// from it. Only ChildPods is ever called on it, so the rest of k8s.Plugin is left nil. +type childPodPlugin struct { + k8s.Plugin + selector labels.Selector + err error +} + +func (p childPodPlugin) ChildPods(_ context.Context, taskCtx pluginsCore.TaskExecutionMetadata, _ client.Object) (labels.Selector, error) { + if p.err != nil { + return nil, p.err + } + if p.selector != nil { + return p.selector, nil + } + return flytek8s.AttemptPodSelector(taskCtx), nil +} + +var _ k8s.ChildPodDiscovery = childPodPlugin{} + +const childPodNamespace = "ns" + +// attemptLabels are the labels the framework stamps on the attempt and every plugin merges +// into the pod templates it builds, so the operator's pods carry them. +func attemptLabels() map[string]string { + return map[string]string{ + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", + } +} + +func childPodTaskContext(t *testing.T, podLabels map[string]string) pluginsCore.TaskExecutionContext { + t.Helper() + meta := &pluginsCoreMock.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(podLabels) + tCtx := &pluginsCoreMock.TaskExecutionContext{} + tCtx.EXPECT().TaskExecutionMetadata().Return(meta) + return tCtx +} + +// trackedJobSet stands in for any CRD a plugin tracks instead of a Pod. +func trackedJobSet() *jobsetv1alpha2.JobSet { + return &jobsetv1alpha2.JobSet{ + TypeMeta: metav1.TypeMeta{Kind: "JobSet", APIVersion: jobsetv1alpha2.SchemeGroupVersion.String()}, + ObjectMeta: metav1.ObjectMeta{Namespace: childPodNamespace, Name: "job", UID: "jobset-uid"}, + } +} + +type workerPodOpts struct { + name string + uid k8stypes.UID + podLabels map[string]string + finishedAt time.Time + exitCode int32 + phase v1.PodPhase +} + +func workerPod(opts workerPodOpts) *v1.Pod { + podLabels := opts.podLabels + if podLabels == nil { + podLabels = attemptLabels() + } + pod := &v1.Pod{ + TypeMeta: metav1.TypeMeta{Kind: "Pod", APIVersion: "v1"}, + ObjectMeta: metav1.ObjectMeta{ + Namespace: childPodNamespace, + Name: opts.name, + UID: opts.uid, + Labels: podLabels, + }, + Status: v1.PodStatus{Phase: opts.phase}, + } + if !opts.finishedAt.IsZero() { + exitCode := opts.exitCode + if exitCode == 0 && opts.phase != v1.PodSucceeded { + exitCode = 1 + } + pod.Status.ContainerStatuses = []v1.ContainerStatus{{ + Name: "primary", + State: v1.ContainerState{Terminated: &v1.ContainerStateTerminated{ + ExitCode: exitCode, + FinishedAt: metav1.NewTime(opts.finishedAt), + }}, + }} + } + return pod +} + +func childPodManager(t *testing.T, plugin k8s.Plugin, events map[watchedObjectKey][]*eventInfo, pods ...*v1.Pod) *PluginManager { + t.Helper() + builder := fake.NewClientBuilder().WithScheme(k8sscheme.Scheme) + for _, pod := range pods { + builder = builder.WithObjects(pod) + } + kubeClient := &pluginsCoreMock.KubeClient{} + kubeClient.EXPECT().GetClient().Return(builder.Build()).Maybe() + + pm := NewPluginManager("test-plugin", plugin, kubeClient) + pm.eventWatcher = &fakeEventWatcher{events: events} + return pm +} + +func podEventKey(name string) watchedObjectKey { + return watchedObjectKey{Namespace: childPodNamespace, Name: name, Kind: "Pod"} +} + +func crdFailure() pluginsCore.PhaseInfo { + // Every CRD plugin stamps the failure with the time of the reconcile that noticed it, + // which is why a child pod's own termination is the better anchor. + now := time.Now() + return pluginsCore.PhaseInfoRetryableFailure("UnknownError", "JobSet failed", &pluginsCore.TaskInfo{OccurredAt: &now}) +} + +func TestClassifyGpuFailureOnChildPods(t *testing.T) { + t.Run("the earliest critical fault across workers explains the failure", func(t *testing.T) { + // Two workers fault. The later one is the one whose exit the operator reported, + // but the fault that explains the job is the first one the hardware produced. + early := time.Now().Add(-3 * time.Minute) + late := time.Now().Add(-1 * time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-1"): {gpuFaultEventFor(48, gpufault.SeverityCritical, late, late, "worker-1-uid")}, + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, early, early, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: early}), + workerPod(workerPodOpts{name: "job-worker-1", uid: "worker-1-uid", finishedAt: late}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + require.NotNil(t, got.Err()) + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + assert.Equal(t, core.ExecutionError_SYSTEM, got.Err().GetKind()) + assert.Equal(t, pluginsCore.PhaseRetryableFailure, got.Phase()) + require.NotNil(t, got.Err().GetGpuFault()) + assert.EqualValues(t, 79, got.Err().GetGpuFault().GetCode()) + }) + + t.Run("the pods a fault was seen on are named for the user", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + workerPod(workerPodOpts{name: "job-worker-1", uid: "worker-1-uid", finishedAt: observed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + require.NotNil(t, got.Info()) + reasons := make([]string, 0, len(got.Info().AdditionalReasons)) + for _, reason := range got.Info().AdditionalReasons { + reasons = append(reasons, reason.Reason) + } + // Only the worker that actually faulted is named, so the user is not sent to + // read the logs of every pod in the job. + assert.Equal(t, []string{"GPU fault recorded on pod job-worker-0"}, reasons) + }) + + t.Run("a plugin that cannot name its child pods classifies nothing", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + // The run label did not survive sanitization, so the attempt cannot be identified. + podLabels := attemptLabels() + delete(podLabels, flytek8s.RunLabel) + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + + phase := crdFailure() + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, podLabels), trackedJobSet(), phase) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + // A plugin that tracks a CRD and does not implement ChildPodDiscovery has to behave + // exactly as it did before the interface existed. Everything here is real: a plugin + // that genuinely lacks the method, a task context that is present, and fault events + // well inside the relevance window keyed to a pod that genuinely exists and matches + // the attempt, so the only reason nothing is classified is the missing interface. + t.Run("a plugin that does not implement the interface is left alone", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + plugin := k8sMocks.NewPlugin(t) + require.NotImplements(t, (*k8s.ChildPodDiscovery)(nil), plugin) + + pm := childPodManager(t, plugin, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + + // The same fixture reclassifies when the plugin can name its pods, so the + // assertions below are about the interface and nothing else. + withDiscovery := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + reclassified := withDiscovery.classifyGpuFailure( + context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + require.Equal(t, gpufault.CodeGpuFallenOffBus, reclassified.Err().GetCode()) + + phase := crdFailure() + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), phase) + + assert.Equal(t, phase.Phase(), got.Phase()) + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + assert.Empty(t, got.Info().AdditionalReasons) + }) + + t.Run("a plugin that fails to name its child pods classifies nothing", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{err: assert.AnError}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("a pod from another attempt of the same action is not consulted", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + // The pod left behind by the previous attempt still carries the run and action, + // but not this attempt. + previousAttempt := attemptLabels() + previousAttempt[flytek8s.AttemptLabel] = "1" + current := attemptLabels() + current[flytek8s.AttemptLabel] = "2" + + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", podLabels: previousAttempt, finishedAt: observed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, current), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("a pod that is already gone yields nothing rather than a wrong answer", func(t *testing.T) { + // The operator tore the pods down with the job. The fault is still cached under + // the pod's name, but nothing can name that pod any more, and an operator's child + // pod name is not derivable because it carries a random suffix. + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("a fault recorded against an earlier incarnation of the pod is rejected", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "an-older-pod-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("relevance is anchored on the pod that terminated, not on the reconcile", func(t *testing.T) { + // The worker died and its fault was recorded two hours ago; the operator only + // reported the job as failed on this reconcile. Anchoring on the reconcile would + // age the fault out, anchoring on the pod's own termination keeps it. + diedAt := time.Now().Add(-2 * time.Hour) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, diedAt, diedAt, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: diedAt}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + }) + + t.Run("a fault long before the pod terminated does not explain the failure", func(t *testing.T) { + diedAt := time.Now().Add(-time.Minute) + longBefore := diedAt.Add(-2 * gpuFaultRelevanceWindow) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, longBefore, longBefore, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: diedAt}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("a critical fault on one worker outranks a user fault on another", func(t *testing.T) { + // The user fault came first in time, but severity decides which fault names the + // failure, and a critical fault is not the workload's doing. + userAt := time.Now().Add(-3 * time.Minute) + criticalAt := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(31, gpufault.SeverityUser, userAt, userAt, "worker-0-uid")}, + podEventKey("job-worker-1"): {gpuFaultEventFor(79, gpufault.SeverityCritical, criticalAt, criticalAt, "worker-1-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: userAt}), + workerPod(workerPodOpts{name: "job-worker-1", uid: "worker-1-uid", finishedAt: criticalAt}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + assert.Equal(t, core.ExecutionError_SYSTEM, got.Err().GetKind()) + }) +} + +func TestClassifyGpuFailureOnChildPodsSelection(t *testing.T) { + t.Run("a worker that finished its work does not contribute its faults", func(t *testing.T) { + // The classic shape of this: an MPI worker exits 0 after logging a critical Xid + // mid-run, while the launcher fails on the user's own code. Crediting the worker's + // fault would turn a user error into a system one and stop charging the retry. + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{ + name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed, + exitCode: 0, phase: v1.PodSucceeded, + }), + workerPod(workerPodOpts{name: "job-launcher", uid: "launcher-uid", finishedAt: observed, phase: v1.PodFailed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, "UnknownError", got.Err().GetCode()) + assert.Nil(t, got.Err().GetGpuFault()) + }) + + t.Run("a worker still running does contribute its faults", func(t *testing.T) { + // A worker wedged on a GPU that fell off the bus never terminates, and it is + // exactly the case this path exists for. + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", phase: v1.PodRunning}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + }) + + t.Run("the fault that happened first wins, not the one last seen", func(t *testing.T) { + // The root cause is still repeating, so its last observation is newer than that of + // the one-shot fault that followed it. Ordering on the last observation would let + // the downstream symptom name the failure. + rootCauseAt := time.Now().Add(-10 * time.Minute) + stillRepeatingAt := time.Now().Add(-time.Minute) + downstreamAt := time.Now().Add(-5 * time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): { + gpuFaultEventFor(79, gpufault.SeverityCritical, rootCauseAt, stillRepeatingAt, "worker-0-uid"), + }, + podEventKey("job-worker-1"): { + gpuFaultEventFor(48, gpufault.SeverityCritical, downstreamAt, downstreamAt, "worker-1-uid"), + }, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", phase: v1.PodRunning}), + workerPod(workerPodOpts{name: "job-worker-1", uid: "worker-1-uid", finishedAt: downstreamAt}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + require.Len(t, got.Info().AdditionalReasons, 1) + assert.Equal(t, "GPU fault recorded on pod job-worker-0", got.Info().AdditionalReasons[0].Reason) + }) +} + +func TestAttachFaultingPod(t *testing.T) { + t.Run("names only the pod whose fault settled the verdict", func(t *testing.T) { + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(92, gpufault.SeverityWarn, observed, observed, "worker-0-uid")}, + podEventKey("job-worker-1"): {gpuFaultEventFor(79, gpufault.SeverityCritical, observed, observed, "worker-1-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + workerPod(workerPodOpts{name: "job-worker-1", uid: "worker-1-uid", finishedAt: observed}), + ) + + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), crdFailure()) + + // The warning on worker-0 is in window and rode along, but the critical on + // worker-1 is what named the failure, so worker-0 must not be pointed at. + require.Len(t, got.Info().AdditionalReasons, 1) + assert.Equal(t, "GPU fault recorded on pod job-worker-1", got.Info().AdditionalReasons[0].Reason) + }) + + t.Run("names nothing when the faults did not explain the failure", func(t *testing.T) { + // A task that ran out of memory keeps that verdict, and a warning that happened to + // coincide rides along as data only. Pointing at a pod here would send the user + // looking for hardware trouble that did not cause anything. + observed := time.Now().Add(-time.Minute) + events := map[watchedObjectKey][]*eventInfo{ + podEventKey("job-worker-0"): {gpuFaultEventFor(92, gpufault.SeverityWarn, observed, observed, "worker-0-uid")}, + } + pm := childPodManager(t, childPodPlugin{}, events, + workerPod(workerPodOpts{name: "job-worker-0", uid: "worker-0-uid", finishedAt: observed}), + ) + + now := time.Now() + phase := pluginsCore.PhaseInfoFailure("OOMKilled", "out of memory", &pluginsCore.TaskInfo{OccurredAt: &now}) + got := pm.classifyGpuFailure(context.Background(), childPodTaskContext(t, attemptLabels()), trackedJobSet(), phase) + + assert.Equal(t, "OOMKilled", got.Err().GetCode()) + assert.NotNil(t, got.Err().GetGpuFault()) + assert.Empty(t, got.Info().AdditionalReasons) + }) + + t.Run("a pod that names itself is not named again", func(t *testing.T) { + // On the single pod path the failure is already reported against that pod, so a + // reason pointing at it would be self-referential noise. + base := time.Now().Add(-time.Minute) + watcher := &fakeEventWatcher{events: map[watchedObjectKey][]*eventInfo{ + {Namespace: "ns", Name: "pod", Kind: "Pod"}: {gpuFaultEvent(79, gpufault.SeverityCritical, base)}, + }} + pm := NewPluginManager("test-plugin", nil, nil) + pm.eventWatcher = watcher + + phase := pluginsCore.PhaseInfoRetryableFailure("UnknownError", "Pod failed", &pluginsCore.TaskInfo{}) + got := pm.classifyGpuFailure(context.Background(), nil, failedPod(), phase) + + assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) + assert.Empty(t, got.Info().AdditionalReasons) + }) +} + +// The anchor itself is podFailureTime, covered in plugin_manager_test.go. What is specific +// to a child pod is the bound against the attempt's own failure. +func TestChildPodFailureTime(t *testing.T) { + failureAt := time.Date(2026, 8, 25, 12, 0, 0, 0, time.UTC) + + t.Run("anchors the pod on itself when it died before the job failed", func(t *testing.T) { + diedAt := failureAt.Add(-40 * time.Minute) + pod := &v1.Pod{Status: v1.PodStatus{ContainerStatuses: []v1.ContainerStatus{ + {State: v1.ContainerState{Terminated: &v1.ContainerStateTerminated{FinishedAt: metav1.NewTime(diedAt)}}}, + }}} + assert.Equal(t, diedAt, childPodFailureTime(pod, failureAt)) + }) + + t.Run("never anchors later than the attempt's own failure", func(t *testing.T) { + // A pod torn down twenty minutes after the job failed must not let faults recorded + // in the meantime explain it; the slack past the failure is deliberately small. + deletion := metav1.NewTime(failureAt.Add(20 * time.Minute)) + pod := &v1.Pod{ObjectMeta: metav1.ObjectMeta{DeletionTimestamp: &deletion}} + assert.Equal(t, failureAt, childPodFailureTime(pod, failureAt)) + + late := metav1.NewTime(failureAt.Add(20 * time.Minute)) + pod = &v1.Pod{Status: v1.PodStatus{ContainerStatuses: []v1.ContainerStatus{ + {State: v1.ContainerState{Terminated: &v1.ContainerStateTerminated{FinishedAt: late}}}, + }}} + assert.Equal(t, failureAt, childPodFailureTime(pod, failureAt)) + }) + + t.Run("does not bound against a failure time the plugin never reported", func(t *testing.T) { + // With no reported time podFailureTime falls back to now, and bounding that + // against the zero time would anchor every pod in 1970. + assert.WithinDuration(t, time.Now(), childPodFailureTime(&v1.Pod{}, time.Time{}), time.Minute) + }) +} diff --git a/executor/pkg/plugin/k8s/plugin_manager.go b/executor/pkg/plugin/k8s/plugin_manager.go index aff845d9bb9..cc478e99672 100644 --- a/executor/pkg/plugin/k8s/plugin_manager.go +++ b/executor/pkg/plugin/k8s/plugin_manager.go @@ -3,6 +3,7 @@ package k8s import ( "context" "fmt" + "sort" "sync" "time" @@ -13,6 +14,7 @@ import ( k8serrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" k8stypes "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/sets" "k8s.io/apimachinery/pkg/util/validation" "sigs.k8s.io/controller-runtime/pkg/client" @@ -272,7 +274,7 @@ func (pm *PluginManager) Handle(ctx context.Context, tCtx pluginsCore.TaskExecut lastEventUpdate, lastEventRecordedAt, ) - phaseInfo = pm.classifyGpuFailure(resource, phaseInfo) + phaseInfo = pm.classifyGpuFailure(ctx, tCtx, resource, phaseInfo) transition.SetInfo(phaseInfo) } @@ -389,32 +391,248 @@ const gpuFaultRelevanceWindow = 30 * time.Minute // says nothing about whether it caused the failure. See faultOverlapsFailure. const gpuFaultAfterFailureSlack = 2 * time.Minute -// classifyGpuFailure folds the GPU faults recorded against a failed attempt's pod into -// the failure the plugin reported, so that a fault the node saw becomes the code and -// the message the user reads. Anything that is not a failed pod is left alone. +// observedFault is a fault recorded against one pod, kept with the times that order it and +// the pod it was observed on. The pod name is what tells the user which worker of a +// distributed job the fault happened on. +type observedFault struct { + fault *core.GpuFault + // createdAt is when the fault was first recorded, which is when it happened. It is + // what orders faults gathered from several pods into the single sequence + // ClassifyFailure expects, so that the first fault of a severity is the earliest. + createdAt time.Time + podName string +} + +// classifyGpuFailure folds the GPU faults recorded against a failed attempt's pods into +// the failure the plugin reported, so that a fault the node saw becomes the code and the +// message the user reads. Anything that is not a failure is left alone. +// +// The pods it looks at are the tracked resource itself when the plugin tracks a Pod, and +// otherwise the child pods the plugin names through k8s.ChildPodDiscovery. A plugin that +// tracks a CRD and does not implement that interface contributes no pods and so no faults, +// which is exactly what it did before the interface existed. func (pm *PluginManager) classifyGpuFailure( + ctx context.Context, + tCtx pluginsCore.TaskExecutionContext, resource client.Object, phaseInfo pluginsCore.PhaseInfo, ) pluginsCore.PhaseInfo { if pm.eventWatcher == nil || resource == nil || !phaseInfo.Phase().IsFailure() { return phaseInfo } - if _, isPod := resource.(*v1.Pod); !isPod { + + // Every event cached for a pod, not only the ones since the last watermark: the Xid + // that killed the task is usually recorded rounds before the pod's status catches up + // with it, and by then the watermark has moved past it. What bounds the search is the + // identity and the recency of each event. + // + // The attempt's own failure time is what a child pod's anchor is bounded by, and what + // stands in for a pod that offers nothing better. Every CRD plugin stamps it with the + // time of the reconcile that noticed the failure rather than with anything the kubelet + // recorded, which is why each pod is anchored on itself below. + attemptFailedAt := phaseInfoOccurredAt(phaseInfo) + + // A plugin that tracks the Pod is looking at the pod the failure is already reported + // against, so naming it again would tell the user nothing. Only a CRD's child pods + // need naming, which is what namesTheFaultingPod tracks. + var observed []observedFault + namesTheFaultingPod := false + if pod, isPod := resource.(*v1.Pod); isPod { + observed = pm.faultsOnPod(ctx, objectKeyFor(resource), resource.GetUID(), podFailureTime(pod, attemptFailedAt)) + } else { + observed = pm.faultsOnChildPods(ctx, tCtx, resource, attemptFailedAt) + namesTheFaultingPod = true + } + if len(observed) == 0 { return phaseInfo } - // Every event cached for the pod, not only the ones since the last watermark: the - // Xid that killed the task is usually recorded rounds before the pod's status - // catches up with it, and by then the watermark has moved past it. What bounds the - // search is the identity and the recency of each event, checked below. - failureAt := podFailureTime(resource.(*v1.Pod), phaseInfoOccurredAt(phaseInfo)) + // ClassifyFailure reads the list as a sequence in time, taking the first fault of the + // severity it settles on. Faults gathered from several pods arrive grouped by pod, so + // they have to be put back in order for first to mean earliest. The order is on when + // each fault was first recorded, not on when it was last seen: a fault that is still + // repeating has a later last observation than a one-shot fault that followed it, and + // ordering on that would let a downstream symptom outrank the root cause. + sort.SliceStable(observed, func(i, j int) bool { + return observed[i].createdAt.Before(observed[j].createdAt) + }) - events := pm.eventWatcher.List(objectKeyFor(resource), time.Time{}, time.Time{}) - if len(events) == 0 { + faults := make([]*core.GpuFault, 0, len(observed)) + for _, o := range observed { + faults = append(faults, o.fault) + } + + classified := gpufault.ClassifyFailure(phaseInfo, faults) + if namesTheFaultingPod { + classified = attachFaultingPod(classified, observed) + } + return classified +} + +// gpuFaultCodes are the codes ClassifyFailure puts on a failure it settled with a fault. +// Their presence is how the caller tells a failure the fault explained from one the fault +// only rode along with. +var gpuFaultCodes = sets.NewString( + gpufault.CodeGpuXidError, + gpufault.CodeGpuFallenOffBus, + gpufault.CodeGpuEccUncorrectable, + gpufault.CodeGpuRowRemapPending, + gpufault.CodeGpuNvlinkError, + gpufault.CodeGpuGspError, +) + +// attachFaultingPod names the one pod whose fault the classification settled on. The fault +// sentence carries the Xid, the GPU and the node but not the pod, which on a job with many +// workers is not enough to act on, and the pod is not part of the fault contract itself +// because a fault reported as an event on a pod already says which pod it is about to +// anyone reading the event. +// +// Only the fault that decided the verdict is named. Naming every pod that saw any fault +// would put a misleading reason on a failure the faults did not explain, for example a +// plain OOMKilled that a warning happened to coincide with, and would emit one cluster +// event per pod when a whole node's worth of GPUs faults at once. +func attachFaultingPod(phaseInfo pluginsCore.PhaseInfo, observed []observedFault) pluginsCore.PhaseInfo { + info := phaseInfo.Info() + if info == nil || !gpuFaultCodes.Has(phaseInfo.Err().GetCode()) { return phaseInfo } - faults := make([]*core.GpuFault, 0, len(events)) + // The fault the classification kept is carried on the error, so matching it back by + // identity finds the pod it came from without repeating the precedence rules. + settled := phaseInfo.Err().GetGpuFault() + for _, o := range observed { + if o.fault != settled || o.podName == "" { + continue + } + occurredAt := o.createdAt + info.AdditionalReasons = append(info.AdditionalReasons, pluginsCore.ReasonInfo{ + Reason: fmt.Sprintf("GPU fault recorded on pod %s", o.podName), + OccurredAt: &occurredAt, + }) + return phaseInfo + } + + return phaseInfo +} + +// faultsOnChildPods gathers the faults recorded against the pods an operator expanded from +// the resource this plugin tracks. A plugin that cannot name its child pods, either because +// it does not implement the interface or because it declined for this resource, contributes +// nothing: there is no safe wider search, since an operator's child pod names carry a random +// suffix and the events are cached under those names. +func (pm *PluginManager) faultsOnChildPods( + ctx context.Context, + tCtx pluginsCore.TaskExecutionContext, + resource client.Object, + failureAt time.Time, +) []observedFault { + discovery, ok := pm.plugin.(k8s.ChildPodDiscovery) + if !ok || tCtx == nil { + return nil + } + + tracked := fmt.Sprintf("%s/%s", resource.GetNamespace(), resource.GetName()) + + selector, err := discovery.ChildPods(ctx, tCtx.TaskExecutionMetadata(), resource) + if err != nil { + logger.Warnf(ctx, "plugin [%s] failed to name the child pods of %s: %v", pm.GetID(), tracked, err) + return nil + } + if selector == nil { + logger.Debugf(ctx, "plugin [%s] did not name the child pods of %s, so its GPU faults are not classified", + pm.GetID(), tracked) + return nil + } + + podList := &v1.PodList{} + listOptions := []client.ListOption{ + client.InNamespace(resource.GetNamespace()), + client.MatchingLabelsSelector{Selector: selector}, + } + if err := pm.kubeClient.GetClient().List(ctx, podList, listOptions...); err != nil { + logger.Warnf(ctx, "failed to list the child pods of %s for GPU fault classification: %v", tracked, err) + return nil + } + if len(podList.Items) == 0 { + // The pods are gone by the time the failure is classified, most often because the + // operator tore them down with the job. Nothing can be recovered: the events are + // cached under pod names that are not derivable without the pods. + logger.Debugf(ctx, "no child pods left for %s matching %s, so any GPU fault on them is not classified", + tracked, selector.String()) + return nil + } + + var observed []observedFault + for i := range podList.Items { + pod := &podList.Items[i] + // A worker that finished its work cannot be the reason the job failed. It may + // still have logged a fault mid-run, and crediting that to another pod's failure + // would turn a plain user error into a hardware one. A pod that is still running + // is kept: a worker wedged on a GPU that fell off the bus is exactly the case + // this path exists for. + if podSucceeded(pod) { + continue + } + podKey := watchedObjectKey{Namespace: pod.Namespace, Name: pod.Name, Kind: "Pod"} + observed = append(observed, pm.faultsOnPod(ctx, podKey, pod.UID, childPodFailureTime(pod, failureAt))...) + } + return observed +} + +// podSucceeded reports whether this pod finished its work. The phase is the kubelet's own +// verdict; the container check catches the pod whose containers have all exited cleanly +// but whose phase has not caught up yet. +func podSucceeded(pod *v1.Pod) bool { + if pod.Status.Phase == v1.PodSucceeded { + return true + } + if pod.Status.Phase != v1.PodRunning || len(pod.Status.ContainerStatuses) == 0 { + return false + } + for _, status := range pod.Status.ContainerStatuses { + terminated := status.State.Terminated + if terminated == nil || terminated.ExitCode != 0 { + return false + } + } + return true +} + +// childPodFailureTime is podFailureTime bounded by the attempt's own failure. +// +// A child pod is anchored on itself, so a worker that died an hour before the operator +// admitted the job had failed is judged against its own death, which is the whole point of +// anchoring per pod. What it must not do is anchor later than the failure: an operator +// tearing its pods down long after the job failed would otherwise let faults recorded in +// the meantime explain it, and the slack after the failure is deliberately small. The risk +// the earlier direction leaves, a fault from a pod's earlier life explaining an unrelated +// failure, is what excluding succeeded pods above bounds. +func childPodFailureTime(pod *v1.Pod, failureAt time.Time) time.Time { + anchor := podFailureTime(pod, failureAt) + if !failureAt.IsZero() && anchor.After(failureAt) { + return failureAt + } + return anchor +} + +// faultsOnPod reads the faults recorded against one pod. +// +// It takes every event cached for the pod, not only the ones since the last watermark: the +// Xid that killed the task is usually recorded rounds before the pod's status catches up +// with it, and by then the watermark has moved past it. What bounds the search is the +// identity and the recency of each event, checked below. +func (pm *PluginManager) faultsOnPod( + ctx context.Context, + podKey watchedObjectKey, + podUID k8stypes.UID, + failureAt time.Time, +) []observedFault { + events := pm.eventWatcher.List(podKey, time.Time{}, time.Time{}) + if len(events) == 0 { + return nil + } + + observed := make([]observedFault, 0, len(events)) for _, event := range events { // Only events the GPU fault emitter wrote, recognized by their reason, are // parsed; the message prefix alone is free text anyone can put in an event. @@ -429,25 +647,29 @@ func (pm *PluginManager) classifyGpuFailure( // when the pod was deleted before this round reached it; the name match the // cache is keyed on is then all there is, and it is used knowingly: a same-name // replacement pod's faults could be credited here, a deliberate trade against - // losing every fault on the path where the hardware most clearly failed. + // losing every fault on the path where the hardware most clearly failed. A child + // pod always arrives with its UID, because it was read out of the cache as an + // object rather than named from the task, so it never takes that trade. if event.RegardingUID == "" { continue } - if resource.GetUID() != "" && event.RegardingUID != resource.GetUID() { + if podUID != "" && event.RegardingUID != podUID { continue } if !faultOverlapsFailure(event, failureAt) { - logger.Debugf(context.TODO(), + logger.Debugf(ctx, "ignoring GPU fault event %q on %s: active %s to %s, which does not reach the failure at %s", - event.Reason, objectKeyFor(resource).Name, event.CreatedAt, event.LastObservedAt, failureAt) + event.Reason, podKey.Name, event.CreatedAt, event.LastObservedAt, failureAt) continue } if fault := gpufault.FromEventMessage(event.Message); fault != nil { - faults = append(faults, fault) + // The last observation decided relevance above; what orders the fault against + // the others is when it was first recorded, which is when it happened. + observed = append(observed, observedFault{fault: fault, createdAt: event.CreatedAt, podName: podKey.Name}) } } - return gpufault.ClassifyFailure(phaseInfo, faults) + return observed } // faultOverlapsFailure reports whether a fault event was active close enough to the diff --git a/executor/pkg/plugin/k8s/plugin_manager_test.go b/executor/pkg/plugin/k8s/plugin_manager_test.go index 54feaa17ea2..4e7f6567e55 100644 --- a/executor/pkg/plugin/k8s/plugin_manager_test.go +++ b/executor/pkg/plugin/k8s/plugin_manager_test.go @@ -466,7 +466,7 @@ func TestClassifyGpuFailure(t *testing.T) { pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - got := pm.classifyGpuFailure(failedPod(), tt.phaseInfo) + got := pm.classifyGpuFailure(context.Background(), nil, failedPod(), tt.phaseInfo) assert.Equal(t, tt.wantPhase, got.Phase()) require.NotNil(t, got.Err()) @@ -497,7 +497,7 @@ func TestClassifyGpuFailureRelevanceIsAnchoredOnTheFailure(t *testing.T) { }} pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - got := pm.classifyGpuFailure(failedPod(), phase) + got := pm.classifyGpuFailure(context.Background(), nil, failedPod(), phase) assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) }) @@ -508,7 +508,7 @@ func TestClassifyGpuFailureRelevanceIsAnchoredOnTheFailure(t *testing.T) { }} pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - got := pm.classifyGpuFailure(failedPod(), phase) + got := pm.classifyGpuFailure(context.Background(), nil, failedPod(), phase) assert.Equal(t, "UnknownError", got.Err().GetCode()) assert.Nil(t, got.Err().GetGpuFault()) }) @@ -531,7 +531,7 @@ func TestClassifyGpuFailureRelevanceIsAnInterval(t *testing.T) { }} pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - return pm.classifyGpuFailure(failedPod(), phase) + return pm.classifyGpuFailure(context.Background(), nil, failedPod(), phase) } t.Run("a fault that started before the failure and kept firing after it explains it", func(t *testing.T) { @@ -639,7 +639,7 @@ func TestClassifyGpuFailureAnchorsOnThePodNotItsStartTime(t *testing.T) { pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - got := pm.classifyGpuFailure(pod, phase) + got := pm.classifyGpuFailure(context.Background(), nil, pod, phase) assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) assert.Equal(t, core.ExecutionError_SYSTEM, got.Err().GetKind()) @@ -657,7 +657,7 @@ func TestClassifyGpuFailureIdentity(t *testing.T) { }} pm := NewPluginManager("test-plugin", nil, nil) pm.eventWatcher = watcher - got := pm.classifyGpuFailure(failedPod(), phase) + got := pm.classifyGpuFailure(context.Background(), nil, failedPod(), phase) assert.Nil(t, got.Err().GetGpuFault()) assert.Equal(t, "UnknownError", got.Err().GetCode()) }) @@ -673,7 +673,7 @@ func TestClassifyGpuFailureIdentity(t *testing.T) { pm.eventWatcher = watcher pod := failedPod() pod.UID = "" - got := pm.classifyGpuFailure(pod, phase) + got := pm.classifyGpuFailure(context.Background(), nil, pod, phase) assert.Equal(t, gpufault.CodeGpuFallenOffBus, got.Err().GetCode()) assert.NotNil(t, got.Err().GetGpuFault()) }) @@ -700,14 +700,6 @@ func TestClassifyGpuFailureSkips(t *testing.T) { resource: failedPod(), phaseInfo: pluginsCore.PhaseInfoSuccess(nil), }, - { - name: "the resource is not a pod", - resource: &v1.Service{ - TypeMeta: metav1.TypeMeta{Kind: "Pod", APIVersion: "v1"}, - ObjectMeta: metav1.ObjectMeta{Namespace: "ns", Name: "pod"}, - }, - phaseInfo: pluginsCore.PhaseInfoRetryableFailure("UnknownError", "Pod failed", nil), - }, { name: "there is no event watcher", resource: failedPod(), @@ -728,7 +720,7 @@ func TestClassifyGpuFailureSkips(t *testing.T) { pm.eventWatcher = &fakeEventWatcher{events: events} } - got := pm.classifyGpuFailure(tt.resource, tt.phaseInfo) + got := pm.classifyGpuFailure(context.Background(), nil, tt.resource, tt.phaseInfo) assert.Equal(t, tt.phaseInfo.Phase(), got.Phase()) if tt.phaseInfo.Err() != nil { diff --git a/executor/pkg/plugin/task_exec_metadata.go b/executor/pkg/plugin/task_exec_metadata.go index ec22f21a825..c8397d7fdaa 100644 --- a/executor/pkg/plugin/task_exec_metadata.go +++ b/executor/pkg/plugin/task_exec_metadata.go @@ -24,22 +24,16 @@ import ( var _ pluginsCore.TaskExecutionMetadata = &taskExecutionMetadata{} var _ pluginsCore.TaskExecutionID = &taskExecutionID{} +// Labels stamped on task pods so external per-pod metrics (DCGM, cAdvisor, +// kube-state-metrics) can be joined back to Flyte semantics, and so the framework can +// find the pods belonging to one attempt of one action. They are defined in flytek8s +// alongside the managed label, because the plugins that build pod templates have to keep +// them intact for the pods an operator derives from those templates to carry them. const ( - // Labels stamped on task pods so external per-pod metrics (DCGM, cAdvisor, - // kube-state-metrics) can be joined back to Flyte semantics. They are bare - // names, matching the project/domain/organization labels already injected, - // which kube-state-metrics surfaces as `label_run`, `label_action`, - // `label_attempt` and `label_task_name`. - - // RunLabel carries the name of the run that owns the action. - RunLabel = "run" - // ActionLabel carries the name of the action the pod is executing. - ActionLabel = "action" - // AttemptLabel carries the 1-based attempt number, so the pod of a retried - // action can be told apart from the pods of its earlier attempts. - AttemptLabel = "attempt" - // TaskNameLabel carries the registered task name from the task template. - TaskNameLabel = "task-name" + RunLabel = flytek8s.RunLabel + ActionLabel = flytek8s.ActionLabel + AttemptLabel = flytek8s.AttemptLabel + TaskNameLabel = flytek8s.TaskNameLabel ) type taskExecutionID struct { diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/attempt_pod_selector_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/attempt_pod_selector_test.go new file mode 100644 index 00000000000..743f264f7b3 --- /dev/null +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/attempt_pod_selector_test.go @@ -0,0 +1,169 @@ +package flytek8s + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "k8s.io/apimachinery/pkg/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + pluginsCoreMock "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" +) + +// completeAttemptLabels mirrors what NewTaskExecutionMetadata injects into every task's +// labels and what every plugin then merges into the Pod templates it builds. +func completeAttemptLabels() map[string]string { + return map[string]string{ + ManagedLabelKey: ManagedLabelValue, + RunLabel: "run-abc", + ActionLabel: "a0", + AttemptLabel: "1", + TaskNameLabel: "train", + } +} + +func metadataWithLabels(podLabels map[string]string) pluginsCore.TaskExecutionMetadata { + meta := &pluginsCoreMock.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(podLabels) + return meta +} + +func TestAttemptPodSelector(t *testing.T) { + t.Run("matches a pod carrying the attempt's labels", func(t *testing.T) { + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + assert.True(t, selector.Matches(labels.Set(completeAttemptLabels()))) + }) + + t.Run("ignores labels it does not select on", func(t *testing.T) { + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + // A worker pod carries the operator's own labels alongside the attempt's; they + // must not stop it matching. + podLabels := completeAttemptLabels() + podLabels["ray.io/cluster"] = "cluster-abcde" + podLabels["ray.io/node-type"] = "worker" + assert.True(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("does not match another attempt of the same action", func(t *testing.T) { + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + podLabels := completeAttemptLabels() + podLabels[AttemptLabel] = "2" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("does not match another action of the same run", func(t *testing.T) { + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + podLabels := completeAttemptLabels() + podLabels[ActionLabel] = "a1" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("does not match an unmanaged pod", func(t *testing.T) { + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + podLabels := completeAttemptLabels() + delete(podLabels, ManagedLabelKey) + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + // A run or action name that sanitizes to nothing is left off the pod rather than + // stamped blank, so the attempt cannot be identified and no selector may be built. + // Selecting on what is left would reach another action's pods. + for _, missing := range []string{ManagedLabelKey, RunLabel, ActionLabel, AttemptLabel} { + t.Run("refuses to select without "+missing, func(t *testing.T) { + podLabels := completeAttemptLabels() + delete(podLabels, missing) + assert.Nil(t, AttemptPodSelector(metadataWithLabels(podLabels))) + }) + + t.Run("refuses to select on a blank "+missing, func(t *testing.T) { + podLabels := completeAttemptLabels() + podLabels[missing] = "" + assert.Nil(t, AttemptPodSelector(metadataWithLabels(podLabels))) + }) + } + + t.Run("refuses to select without metadata", func(t *testing.T) { + assert.Nil(t, AttemptPodSelector(nil)) + assert.Nil(t, AttemptPodSelector(metadataWithLabels(nil))) + }) +} + +func TestPreservedPodLabels(t *testing.T) { + t.Run("carries the identity the selector looks for", func(t *testing.T) { + preserved := PreservedPodLabels(metadataWithLabels(completeAttemptLabels())) + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + + // The point of the helper: what a plugin re-applies last is exactly what the + // selector goes looking for, so the two can never drift apart. + assert.True(t, selector.Matches(labels.Set(preserved))) + }) + + t.Run("drops the labels a user supplied over the identity", func(t *testing.T) { + preserved := PreservedPodLabels(metadataWithLabels(completeAttemptLabels())) + + // This is what a plugin's union does: the user's labels first, the preserved set + // applied over them last. + podLabels := map[string]string{ + RunLabel: "a-run-the-user-made-up", + ActionLabel: "not-this-action", + AttemptLabel: "99", + ManagedLabelKey: "false", + "their-label": "kept", + } + for key, value := range preserved { + podLabels[key] = value + } + + selector := AttemptPodSelector(metadataWithLabels(completeAttemptLabels())) + require.NotNil(t, selector) + assert.True(t, selector.Matches(labels.Set(podLabels))) + assert.Equal(t, "kept", podLabels["their-label"]) + }) + + t.Run("always carries the managed label", func(t *testing.T) { + // Even when the attempt cannot be identified, the label the Pod cache selects on + // has to survive, or the executor cannot see its own pod at all. + incomplete := completeAttemptLabels() + delete(incomplete, RunLabel) + + preserved := PreservedPodLabels(metadataWithLabels(incomplete)) + assert.Equal(t, ManagedLabelValue, preserved[ManagedLabelKey]) + assert.NotContains(t, preserved, ActionLabel) + + assert.Equal(t, map[string]string{ManagedLabelKey: ManagedLabelValue}, PreservedPodLabels(nil)) + }) +} + +func TestAttemptIdentityLabels(t *testing.T) { + t.Run("returns only the identity, not every execution label", func(t *testing.T) { + podLabels := completeAttemptLabels() + podLabels["unrelated"] = "value" + + identity := AttemptIdentityLabels(metadataWithLabels(podLabels)) + + assert.Equal(t, map[string]string{ + ManagedLabelKey: ManagedLabelValue, + RunLabel: "run-abc", + ActionLabel: "a0", + AttemptLabel: "1", + }, identity) + }) + + t.Run("refuses an incomplete identity", func(t *testing.T) { + podLabels := completeAttemptLabels() + delete(podLabels, AttemptLabel) + assert.Nil(t, AttemptIdentityLabels(metadataWithLabels(podLabels))) + }) +} diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go index 627f45a9d2a..1fda638f6a4 100644 --- a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go +++ b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go @@ -16,6 +16,7 @@ import ( v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/util/sets" "sigs.k8s.io/controller-runtime/pkg/client" @@ -72,6 +73,90 @@ const ( ManagedLabelValue = "true" ) +// Labels identifying which attempt of which action a Pod belongs to. They are stamped on +// every task's execution metadata and every plugin merges that metadata into the Pod +// templates it builds, so an operator that expands one of those templates carries them +// onto the Pods it creates. They live here for the same reason ManagedLabelKey does: the +// plugins that build Pod templates have to keep them intact, and the framework has to be +// able to select on them afterwards. +// +// They are bare names, matching the project/domain/organization labels already injected, +// which kube-state-metrics surfaces as label_run, label_action and label_attempt. +const ( + // RunLabel carries the name of the run that owns the action. + RunLabel = "run" + // ActionLabel carries the name of the action the pod is executing. + ActionLabel = "action" + // AttemptLabel carries the 1-based attempt number, so the pods of a retried action + // can be told apart from the pods of its earlier attempts. + AttemptLabel = "attempt" + // TaskNameLabel carries the registered task name from the task template. + TaskNameLabel = "task-name" +) + +// attemptPodLabelKeys are the labels that together identify one attempt of one action. +// All of them have to be present for a selector to be exact. +var attemptPodLabelKeys = []string{ManagedLabelKey, RunLabel, ActionLabel, AttemptLabel} + +// AttemptIdentityLabels returns the subset of the task's execution labels that identifies +// this attempt of this action, or nil when any of them is missing. +// +// A plugin that lets a user supply their own labels for the Pod templates it builds has to +// re-apply these last, the way it already re-applies the managed label: a user label named +// run, action or attempt would otherwise overwrite the identity, and a Pod whose identity +// has been overwritten is one AttemptPodSelector cannot find. The selector is built from +// the same key set, so the two can never disagree about which labels have to survive. +func AttemptIdentityLabels(taskCtx pluginsCore.TaskExecutionMetadata) map[string]string { + if taskCtx == nil { + return nil + } + + podLabels := taskCtx.GetLabels() + identity := make(map[string]string, len(attemptPodLabelKeys)) + for _, key := range attemptPodLabelKeys { + value, ok := podLabels[key] + if !ok || value == "" { + return nil + } + identity[key] = value + } + + return identity +} + +// PreservedPodLabels are the labels a plugin must apply last to every Pod template it +// builds, after any labels the user supplied, so that nothing a user writes can take the +// Pod out of the executor's reach. That is the managed label, which the Pod cache selects +// on, together with the attempt identity the framework looks a Pod up by. +// +// The managed label is always present. The identity labels are present whenever the task +// carries a complete one, which is the same condition under which AttemptPodSelector will +// go looking for them. +func PreservedPodLabels(taskCtx pluginsCore.TaskExecutionMetadata) map[string]string { + preserved := map[string]string{ManagedLabelKey: ManagedLabelValue} + for key, value := range AttemptIdentityLabels(taskCtx) { + preserved[key] = value + } + return preserved +} + +// AttemptPodSelector selects the Pods belonging to this attempt of this action, including +// the ones an operator expanded from a Pod template a plugin built. It matches only on the +// labels the framework stamps, so it holds whichever operator created the Pod. +// +// It returns nil when any of those labels is missing from the task's execution metadata. +// The run and action labels go through sanitization and are dropped when nothing survives +// it, and a selector missing one of them would match another action's Pods, so refusing to +// select anything is the only safe answer. +func AttemptPodSelector(taskCtx pluginsCore.TaskExecutionMetadata) labels.Selector { + identity := AttemptIdentityLabels(taskCtx) + if identity == nil { + return nil + } + + return labels.SelectorFromValidatedSet(labels.Set(identity)) +} + var migPartitionRegexp = regexp.MustCompile(`^(\d+)g\.\d+gb$`) // parseMigSlices extracts the compute slice count from a MIG partition size string diff --git a/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go b/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go index 01d64848af6..025ba8420e8 100644 --- a/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go +++ b/flyteplugins/go/tasks/pluginmachinery/k8s/plugin.go @@ -4,6 +4,7 @@ import ( "context" "time" + "k8s.io/apimachinery/pkg/labels" "sigs.k8s.io/controller-runtime/pkg/client" pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" @@ -170,6 +171,31 @@ type GarbageCollectable interface { // cleanup, this interface will need to be extended to return those resources. } +// ChildPodDiscovery is an optional interface a Plugin whose tracked resource is not a Pod +// can implement so the framework can find the Pods an operator expanded from that resource. +// +// The framework tracks one object per task. For a plugin that tracks a Pod that object is +// also where the kubelet and the node's daemons record what happened, but for a plugin that +// tracks a CRD the interesting records land on the worker Pods the operator derived from +// the templates the plugin built, which the framework otherwise has no way to name. A +// plugin that does not implement this is treated as having no child Pods, which is what the +// framework assumed before the interface existed. +type ChildPodDiscovery interface { + // ChildPods returns the label selector that selects the Pods belonging to this attempt + // of this resource. They live in the resource's own namespace, which the caller already + // holds, so only the selector is the plugin's to decide. + // + // A nil selector means the plugin cannot name its child Pods and the caller must then + // look at no Pods at all rather than widen the search. Implementations should combine + // flytek8s.AttemptPodSelector, which is exact to one attempt of one action whichever + // operator created the Pod, with a label their own operator applies. + ChildPods( + ctx context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, + ) (labels.Selector, error) +} + // An optional interface a Plugin can implement to override its default OnAbort finalizer (deletion of the underlying resource). type PluginAbortOverride interface { OnAbort(ctx context.Context, tCtx pluginsCore.TaskExecutionContext, resource client.Object) (behavior AbortBehavior, err error) diff --git a/flyteplugins/go/tasks/plugins/k8s/clustered/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/clustered/child_pods_test.go new file mode 100644 index 00000000000..dd776421848 --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/clustered/child_pods_test.go @@ -0,0 +1,139 @@ +package clustered + +import ( + "context" + "testing" + + "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/labels" + jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + coreMocks "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s" + clusteredpb "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/plugins" +) + +// attemptExecutionLabels mirrors what NewTaskExecutionMetadata stamps on every task and +// what build.go then merges onto every child pod template. +func attemptExecutionLabels() map[string]string { + return map[string]string{ + "execution-id": "my-exec", + "node-id": "n1", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", + } +} + +func attemptMetadata(executionLabels map[string]string) pluginsCore.TaskExecutionMetadata { + meta := &coreMocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(executionLabels) + return meta +} + +func TestClusteredChildPods(t *testing.T) { + jobSet := &jobsetv1alpha2.JobSet{ + ObjectMeta: metav1.ObjectMeta{Namespace: testNS, Name: testJobName}, + } + + t.Run("selects on the attempt and the JobSet", func(t *testing.T) { + selector, err := clusteredResourceHandler{}.ChildPods( + context.Background(), attemptMetadata(attemptExecutionLabels()), jobSet) + + require.NoError(t, err) + require.NotNil(t, selector) + + podLabels := attemptExecutionLabels() + podLabels[jobsetv1alpha2.JobSetNameKey] = testJobName + assert.True(t, selector.Matches(labels.Set(podLabels))) + + // Another JobSet in the same namespace is not this one. + podLabels[jobsetv1alpha2.JobSetNameKey] = "some-other-jobset" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("does not select another attempt of the same action", func(t *testing.T) { + selector, err := clusteredResourceHandler{}.ChildPods( + context.Background(), attemptMetadata(attemptExecutionLabels()), jobSet) + require.NoError(t, err) + require.NotNil(t, selector) + + podLabels := attemptExecutionLabels() + podLabels[jobsetv1alpha2.JobSetNameKey] = testJobName + podLabels[flytek8s.AttemptLabel] = "2" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("declines when the attempt cannot be identified", func(t *testing.T) { + executionLabels := attemptExecutionLabels() + delete(executionLabels, flytek8s.AttemptLabel) + + selector, err := clusteredResourceHandler{}.ChildPods( + context.Background(), attemptMetadata(executionLabels), jobSet) + + require.NoError(t, err) + assert.Nil(t, selector) + }) + + t.Run("rejects a resource that is not a JobSet", func(t *testing.T) { + _, err := clusteredResourceHandler{}.ChildPods( + context.Background(), attemptMetadata(attemptExecutionLabels()), &corev1.Pod{}) + + assert.Error(t, err) + }) +} + +// TestClusteredChildPodsMatchTheTemplatesTheyCameFrom is the conformance check between the +// two halves of child pod discovery: the labels this plugin puts on the pod templates the +// JobSet controller expands, and the selector it hands the framework to find the resulting +// pods. Label drift on either side would leave a GPU fault on a worker silently +// unclassified, which is the failure mode this whole path exists to prevent. +func TestClusteredChildPodsMatchTheTemplatesTheyCameFrom(t *testing.T) { + spec := &clusteredpb.ClusteredTaskSpec{ + Replicas: 4, + NprocPerNode: 8, + Runtime: &clusteredpb.Runtime{ + Kind: &clusteredpb.Runtime_Torchrun{ + Torchrun: &clusteredpb.TorchRuntime{ + RdzvBackend: clusteredpb.RdzvBackend_STATIC, + }, + }, + }, + FailurePolicy: &clusteredpb.ClusterFailurePolicy{MaxRestarts: 3}, + } + executionLabels := attemptExecutionLabels() + taskCtx := dummyTaskCtxWithLabels(buildTaskTemplate(spec), testJobName, executionLabels) + + obj, err := clusteredResourceHandler{}.BuildResource(context.Background(), taskCtx) + require.NoError(t, err) + jobSet, ok := obj.(*jobsetv1alpha2.JobSet) + require.True(t, ok) + + selector, err := clusteredResourceHandler{}.ChildPods(context.Background(), attemptMetadata(executionLabels), jobSet) + require.NoError(t, err) + require.NotNil(t, selector) + + require.NotEmpty(t, jobSet.Spec.ReplicatedJobs) + for _, replicatedJob := range jobSet.Spec.ReplicatedJobs { + t.Run(replicatedJob.Name, func(t *testing.T) { + templateLabels := replicatedJob.Template.Spec.Template.GetLabels() + require.NotEmpty(t, templateLabels) + + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + // The JobSet controller stamps its own name on the pods it creates, so the + // template does not carry it and the fixture adds what the operator would. + podLabels[jobsetv1alpha2.JobSetNameKey] = jobSet.Name + + assert.True(t, selector.Matches(labels.Set(podLabels)), + "the %s pod template's labels %v do not satisfy %s", replicatedJob.Name, podLabels, selector) + }) + } +} diff --git a/flyteplugins/go/tasks/plugins/k8s/clustered/clustered_test.go b/flyteplugins/go/tasks/plugins/k8s/clustered/clustered_test.go index b8643434af5..66c83298158 100644 --- a/flyteplugins/go/tasks/plugins/k8s/clustered/clustered_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/clustered/clustered_test.go @@ -69,6 +69,12 @@ func dummyTaskCtx(taskTemplate *core.TaskTemplate) *coreMocks.TaskExecutionConte // dummyTaskCtxWithGeneratedName is dummyTaskCtx with a caller-supplied generated name, used to // exercise the long composed/nested-name truncation path. func dummyTaskCtxWithGeneratedName(taskTemplate *core.TaskTemplate, generatedName string) *coreMocks.TaskExecutionContext { + return dummyTaskCtxWithLabels(taskTemplate, generatedName, map[string]string{"execution-id": "my-exec", "node-id": "n1"}) +} + +// dummyTaskCtxWithLabels is dummyTaskCtxWithGeneratedName with caller-supplied execution +// labels, used to exercise the attempt identity the framework finds child pods by. +func dummyTaskCtxWithLabels(taskTemplate *core.TaskTemplate, generatedName string, executionLabels map[string]string) *coreMocks.TaskExecutionContext { taskCtx := &coreMocks.TaskExecutionContext{} inputReader := &pluginIOMocks.InputReader{} @@ -121,7 +127,7 @@ func dummyTaskCtxWithGeneratedName(taskTemplate *core.TaskTemplate, generatedNam meta.EXPECT().GetTaskExecutionID().Return(tID) meta.EXPECT().GetNamespace().Return(testNS) meta.EXPECT().GetAnnotations().Return(map[string]string{"flyte.org/test-annotation": "av"}) - meta.EXPECT().GetLabels().Return(map[string]string{"execution-id": "my-exec", "node-id": "n1"}) + meta.EXPECT().GetLabels().Return(executionLabels) meta.EXPECT().GetOwnerReference().Return(metav1.OwnerReference{Kind: "node", Name: "n1"}) meta.EXPECT().IsInterruptible().Return(false) meta.EXPECT().GetOverrides().Return(overrides) diff --git a/flyteplugins/go/tasks/plugins/k8s/clustered/plugin.go b/flyteplugins/go/tasks/plugins/k8s/clustered/plugin.go index 3f9293bf2d4..631a0d8b02a 100644 --- a/flyteplugins/go/tasks/plugins/k8s/clustered/plugin.go +++ b/flyteplugins/go/tasks/plugins/k8s/clustered/plugin.go @@ -6,12 +6,15 @@ import ( "time" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/selection" "k8s.io/client-go/kubernetes/scheme" "sigs.k8s.io/controller-runtime/pkg/client" jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2" "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery" pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s" "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/k8s" ) @@ -21,6 +24,10 @@ type clusteredResourceHandler struct{} var _ k8s.Plugin = clusteredResourceHandler{} +// The JobSet's child pods are where a node daemon records what the hardware did, so the +// framework has to be able to find them from the JobSet this plugin tracks. +var _ k8s.ChildPodDiscovery = clusteredResourceHandler{} + func (clusteredResourceHandler) GetProperties() k8s.PluginProperties { // The plugin manager consumes this to stamp the JobSet name via // GetGeneratedNameWith(0, GeneratedNameMaxLength), bounding it so derived child @@ -43,6 +50,36 @@ func (clusteredResourceHandler) IsTerminal(_ context.Context, resource client.Ob return false, nil } +// ChildPods implements k8s.ChildPodDiscovery. The pods that run the task are the ones the +// JobSet controller expands from the templates build.go put in the JobSet, which the +// framework tracks nothing of, since it tracks the JobSet. +// +// The selector is the attempt's own labels, which build.go merges onto every child pod +// template, narrowed by the JobSet the pods belong to. The JobSet's name is its own, so +// the selector is never partial. +func (clusteredResourceHandler) ChildPods( + _ context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, +) (labels.Selector, error) { + jobSet, ok := resource.(*jobsetv1alpha2.JobSet) + if !ok { + return nil, fmt.Errorf("unexpected resource type %T", resource) + } + + selector := flytek8s.AttemptPodSelector(taskCtx) + if selector == nil { + return nil, nil + } + + requirement, err := labels.NewRequirement(jobsetv1alpha2.JobSetNameKey, selection.Equals, []string{jobSet.Name}) + if err != nil { + return nil, err + } + + return selector.Add(*requirement), nil +} + func (clusteredResourceHandler) GetCompletionTime(resource client.Object) (time.Time, error) { jobSet, ok := resource.(*jobsetv1alpha2.JobSet) if !ok { diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/child_pods_test.go new file mode 100644 index 00000000000..cd155da3d38 --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/child_pods_test.go @@ -0,0 +1,83 @@ +package common + +import ( + "context" + "testing" + + kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s" +) + +// AttemptExecutionLabels mirrors what NewTaskExecutionMetadata stamps on every task and +// what ToReplicaSpec then merges onto every replica pod template. +func AttemptExecutionLabels() map[string]string { + return map[string]string{ + "label-key": "label-value", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", + } +} + +func attemptMetadata(executionLabels map[string]string) pluginsCore.TaskExecutionMetadata { + meta := &mocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(executionLabels) + return meta +} + +func TestChildPods(t *testing.T) { + job := &kubeflowv1.PyTorchJob{ + ObjectMeta: metav1.ObjectMeta{Namespace: "test-namespace", Name: "job3"}, + } + + t.Run("selects on the attempt and the job", func(t *testing.T) { + selector, err := ChildPods(context.TODO(), attemptMetadata(AttemptExecutionLabels()), job) + + require.NoError(t, err) + require.NotNil(t, selector) + + podLabels := AttemptExecutionLabels() + podLabels[kubeflowv1.JobNameLabel] = "job3" + podLabels[kubeflowv1.ReplicaTypeLabel] = "worker" + podLabels[kubeflowv1.ReplicaIndexLabel] = "0" + assert.True(t, selector.Matches(labels.Set(podLabels))) + + // Another job in the same namespace is not this one. + podLabels[kubeflowv1.JobNameLabel] = "job4" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("does not select another attempt of the same action", func(t *testing.T) { + selector, err := ChildPods(context.TODO(), attemptMetadata(AttemptExecutionLabels()), job) + require.NoError(t, err) + require.NotNil(t, selector) + + podLabels := AttemptExecutionLabels() + podLabels[kubeflowv1.JobNameLabel] = "job3" + podLabels[flytek8s.AttemptLabel] = "2" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("declines when the attempt cannot be identified", func(t *testing.T) { + executionLabels := AttemptExecutionLabels() + delete(executionLabels, flytek8s.ActionLabel) + + selector, err := ChildPods(context.TODO(), attemptMetadata(executionLabels), job) + + require.NoError(t, err) + assert.Nil(t, selector) + }) + + t.Run("rejects a missing resource", func(t *testing.T) { + _, err := ChildPods(context.TODO(), attemptMetadata(AttemptExecutionLabels()), nil) + assert.Error(t, err) + }) +} diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go index 5f991e82e49..e202ecd66e4 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go @@ -9,6 +9,9 @@ import ( kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" v1 "k8s.io/api/core/v1" meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/selection" + "sigs.k8s.io/controller-runtime/pkg/client" flyteerr "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/errors" "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/logs" @@ -29,6 +32,37 @@ const ( PytorchTaskType = "pytorch" ) +// ChildPods names the replica pods the training operator expands from the templates a +// kubeflow plugin built. It is the shared implementation of k8s.ChildPodDiscovery for +// PyTorchJob, MPIJob and TFJob, which all get their replica pods from the same operator +// and so all carry the same job-name label. +// +// The selector is the attempt's own labels, which every replica template carries through +// ToReplicaSpec, narrowed by the job the pods belong to. Unlike Ray's cluster name the job +// name is the CR's own name, so it is known from the moment the resource exists and the +// selector is never partial. +func ChildPods( + _ context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, +) (labels.Selector, error) { + if resource == nil { + return nil, fmt.Errorf("expected a kubeflow job, got nothing") + } + + selector := flytek8s.AttemptPodSelector(taskCtx) + if selector == nil { + return nil, nil + } + + requirement, err := labels.NewRequirement(kubeflowv1.JobNameLabel, selection.Equals, []string{resource.GetName()}) + if err != nil { + return nil, err + } + + return selector.Add(*requirement), nil +} + // ExtractCurrentCondition will return the first job condition for tensorflow/pytorch func ExtractCurrentCondition(jobConditions []kubeflowv1.JobCondition) (kubeflowv1.JobCondition, error) { if jobConditions != nil { diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/child_pods_test.go new file mode 100644 index 00000000000..4f1da2e8ba6 --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/child_pods_test.go @@ -0,0 +1,58 @@ +package mpi + +import ( + "context" + "testing" + + kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "k8s.io/apimachinery/pkg/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" +) + +func attemptMetadata() pluginsCore.TaskExecutionMetadata { + meta := &mocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(dummyLabels) + return meta +} + +// TestMPIChildPodsMatchTheTemplatesTheyCameFrom is the conformance check between the two +// halves of child pod discovery: the labels this plugin puts on the replica pod templates +// the training operator expands, and the selector it hands the framework to find the +// resulting pods. Label drift on either side would otherwise leave a GPU fault on a worker +// silently unclassified. +func TestMPIChildPodsMatchTheTemplatesTheyCameFrom(t *testing.T) { + taskTemplate := dummyMPITaskTemplate("job3", dummyMPICustomObj(2, 1, 1)) + resource, err := mpiOperatorResourceHandler{}.BuildResource(context.TODO(), dummyMPITaskContext(taskTemplate, resourceRequirements, nil)) + require.NoError(t, err) + + job, ok := resource.(*kubeflowv1.MPIJob) + require.True(t, ok) + // The plugin manager names the resource after it is built, and the operator labels the + // replica pods with that name. + job.Name = "job-name" + + selector, err := mpiOperatorResourceHandler{}.ChildPods(context.TODO(), attemptMetadata(), job) + require.NoError(t, err) + require.NotNil(t, selector) + + require.NotEmpty(t, job.Spec.MPIReplicaSpecs) + for replicaType, replicaSpec := range job.Spec.MPIReplicaSpecs { + t.Run(string(replicaType), func(t *testing.T) { + templateLabels := replicaSpec.Template.GetLabels() + require.NotEmpty(t, templateLabels) + + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + podLabels[kubeflowv1.JobNameLabel] = job.Name + + assert.True(t, selector.Matches(labels.Set(podLabels)), + "the %s replica template's labels %v do not satisfy %s", replicaType, podLabels, selector) + }) + } +} diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go index 833db86bd8e..eb54c0ed4f3 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go @@ -8,6 +8,7 @@ import ( kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" "k8s.io/client-go/kubernetes/scheme" "sigs.k8s.io/controller-runtime/pkg/client" @@ -29,6 +30,10 @@ type mpiOperatorResourceHandler struct { // Sanity test that the plugin implements method of k8s.Plugin var _ k8s.Plugin = mpiOperatorResourceHandler{} +// The job's replica pods are where a node daemon records what the hardware did, so the +// framework has to be able to find them from the CR this plugin tracks. +var _ k8s.ChildPodDiscovery = mpiOperatorResourceHandler{} + func (mpiOperatorResourceHandler) GetProperties() k8s.PluginProperties { return k8s.PluginProperties{} } @@ -216,6 +221,20 @@ func (mpiOperatorResourceHandler) IsTerminal(_ context.Context, resource client. return false, nil } +// ChildPods implements k8s.ChildPodDiscovery. The pods that run the task are the launcher +// and worker pods the training operator expands from the templates this plugin built, +// which the framework tracks nothing of, since it tracks the MPIJob. +func (mpiOperatorResourceHandler) ChildPods( + ctx context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, +) (labels.Selector, error) { + if _, ok := resource.(*kubeflowv1.MPIJob); !ok { + return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "expected an MPIJob, got %T", resource) + } + return common.ChildPods(ctx, taskCtx, resource) +} + // GetCompletionTime returns the completion time of the MPIJob func (mpiOperatorResourceHandler) GetCompletionTime(resource client.Object) (time.Time, error) { job, ok := resource.(*kubeflowv1.MPIJob) diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go index 1db7cdc0e4b..6eb3861f865 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go @@ -48,8 +48,15 @@ var ( dummyAnnotations = map[string]string{ "annotation-key": "annotation-value", } + // The attempt labels are what NewTaskExecutionMetadata stamps on every task, and what + // the framework selects on to find the replica pods the operator derives from these + // templates, so the fixture carries them alongside a plain user label. dummyLabels = map[string]string{ - "label-key": "label-value", + "label-key": "label-value", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", } resourceRequirements = &corev1.ResourceRequirements{ diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/child_pods_test.go new file mode 100644 index 00000000000..90936b20e1c --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/child_pods_test.go @@ -0,0 +1,58 @@ +package pytorch + +import ( + "context" + "testing" + + kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "k8s.io/apimachinery/pkg/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" +) + +func attemptMetadata() pluginsCore.TaskExecutionMetadata { + meta := &mocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(dummyLabels) + return meta +} + +// TestPytorchChildPodsMatchTheTemplatesTheyCameFrom is the conformance check between the two +// halves of child pod discovery: the labels this plugin puts on the replica pod templates +// the training operator expands, and the selector it hands the framework to find the +// resulting pods. Label drift on either side would otherwise leave a GPU fault on a worker +// silently unclassified. +func TestPytorchChildPodsMatchTheTemplatesTheyCameFrom(t *testing.T) { + taskTemplate := dummyPytorchTaskTemplate("job3", dummyPytorchCustomObj(100)) + resource, err := pytorchOperatorResourceHandler{}.BuildResource(context.TODO(), dummyPytorchTaskContext(taskTemplate, resourceRequirements, nil, "")) + require.NoError(t, err) + + job, ok := resource.(*kubeflowv1.PyTorchJob) + require.True(t, ok) + // The plugin manager names the resource after it is built, and the operator labels the + // replica pods with that name. + job.Name = "job-name" + + selector, err := pytorchOperatorResourceHandler{}.ChildPods(context.TODO(), attemptMetadata(), job) + require.NoError(t, err) + require.NotNil(t, selector) + + require.NotEmpty(t, job.Spec.PyTorchReplicaSpecs) + for replicaType, replicaSpec := range job.Spec.PyTorchReplicaSpecs { + t.Run(string(replicaType), func(t *testing.T) { + templateLabels := replicaSpec.Template.GetLabels() + require.NotEmpty(t, templateLabels) + + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + podLabels[kubeflowv1.JobNameLabel] = job.Name + + assert.True(t, selector.Matches(labels.Set(podLabels)), + "the %s replica template's labels %v do not satisfy %s", replicaType, podLabels, selector) + }) + } +} diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go index 0cf67082e42..fc91040fe56 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go @@ -10,6 +10,7 @@ import ( "github.com/samber/lo" v1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" "k8s.io/client-go/kubernetes/scheme" "sigs.k8s.io/controller-runtime/pkg/client" @@ -32,6 +33,10 @@ type pytorchOperatorResourceHandler struct { // Sanity test that the plugin implements method of k8s.Plugin var _ k8s.Plugin = pytorchOperatorResourceHandler{} +// The job's replica pods are where a node daemon records what the hardware did, so the +// framework has to be able to find them from the CR this plugin tracks. +var _ k8s.ChildPodDiscovery = pytorchOperatorResourceHandler{} + func (pytorchOperatorResourceHandler) GetProperties() k8s.PluginProperties { return k8s.PluginProperties{} } @@ -287,6 +292,20 @@ func (pytorchOperatorResourceHandler) IsTerminal(_ context.Context, resource cli return false, nil } +// ChildPods implements k8s.ChildPodDiscovery. The pods that run the task are the replica +// pods the training operator expands from the templates this plugin built, which the +// framework tracks nothing of, since it tracks the PyTorchJob. +func (pytorchOperatorResourceHandler) ChildPods( + ctx context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, +) (labels.Selector, error) { + if _, ok := resource.(*kubeflowv1.PyTorchJob); !ok { + return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "expected a PyTorchJob, got %T", resource) + } + return common.ChildPods(ctx, taskCtx, resource) +} + // GetCompletionTime returns the completion time of the PyTorchJob func (pytorchOperatorResourceHandler) GetCompletionTime(resource client.Object) (time.Time, error) { job, ok := resource.(*kubeflowv1.PyTorchJob) diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go index ac3439156c5..1cab9825252 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go @@ -49,8 +49,15 @@ var ( dummyAnnotations = map[string]string{ "annotation-key": "annotation-value", } + // The attempt labels are what NewTaskExecutionMetadata stamps on every task, and what + // the framework selects on to find the replica pods the operator derives from these + // templates, so the fixture carries them alongside a plain user label. dummyLabels = map[string]string{ - "label-key": "label-value", + "label-key": "label-value", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", } resourceRequirements = &corev1.ResourceRequirements{ diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/child_pods_test.go new file mode 100644 index 00000000000..210de6394a5 --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/child_pods_test.go @@ -0,0 +1,58 @@ +package tensorflow + +import ( + "context" + "testing" + + kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "k8s.io/apimachinery/pkg/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" +) + +func attemptMetadata() pluginsCore.TaskExecutionMetadata { + meta := &mocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(dummyLabels) + return meta +} + +// TestTensorFlowChildPodsMatchTheTemplatesTheyCameFrom is the conformance check between the two +// halves of child pod discovery: the labels this plugin puts on the replica pod templates +// the training operator expands, and the selector it hands the framework to find the +// resulting pods. Label drift on either side would otherwise leave a GPU fault on a worker +// silently unclassified. +func TestTensorFlowChildPodsMatchTheTemplatesTheyCameFrom(t *testing.T) { + taskTemplate := dummyTensorFlowTaskTemplate("job3", dummyTensorFlowCustomObj(2, 1, 1, 1)) + resource, err := tensorflowOperatorResourceHandler{}.BuildResource(context.TODO(), dummyTensorFlowTaskContext(taskTemplate, resourceRequirements, nil)) + require.NoError(t, err) + + job, ok := resource.(*kubeflowv1.TFJob) + require.True(t, ok) + // The plugin manager names the resource after it is built, and the operator labels the + // replica pods with that name. + job.Name = "job-name" + + selector, err := tensorflowOperatorResourceHandler{}.ChildPods(context.TODO(), attemptMetadata(), job) + require.NoError(t, err) + require.NotNil(t, selector) + + require.NotEmpty(t, job.Spec.TFReplicaSpecs) + for replicaType, replicaSpec := range job.Spec.TFReplicaSpecs { + t.Run(string(replicaType), func(t *testing.T) { + templateLabels := replicaSpec.Template.GetLabels() + require.NotEmpty(t, templateLabels) + + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + podLabels[kubeflowv1.JobNameLabel] = job.Name + + assert.True(t, selector.Matches(labels.Set(podLabels)), + "the %s replica template's labels %v do not satisfy %s", replicaType, podLabels, selector) + }) + } +} diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go index 9ab4907c53a..50ea433b87b 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go @@ -7,6 +7,7 @@ import ( kubeflowv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" "k8s.io/client-go/kubernetes/scheme" "sigs.k8s.io/controller-runtime/pkg/client" @@ -26,6 +27,10 @@ type tensorflowOperatorResourceHandler struct { // Sanity test that the plugin implements method of k8s.Plugin var _ k8s.Plugin = tensorflowOperatorResourceHandler{} +// The job's replica pods are where a node daemon records what the hardware did, so the +// framework has to be able to find them from the CR this plugin tracks. +var _ k8s.ChildPodDiscovery = tensorflowOperatorResourceHandler{} + func (tensorflowOperatorResourceHandler) GetProperties() k8s.PluginProperties { return k8s.PluginProperties{} } @@ -214,6 +219,20 @@ func (tensorflowOperatorResourceHandler) IsTerminal(_ context.Context, resource return false, nil } +// ChildPods implements k8s.ChildPodDiscovery. The pods that run the task are the replica +// pods the training operator expands from the templates this plugin built, which the +// framework tracks nothing of, since it tracks the TFJob. +func (tensorflowOperatorResourceHandler) ChildPods( + ctx context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + resource client.Object, +) (labels.Selector, error) { + if _, ok := resource.(*kubeflowv1.TFJob); !ok { + return nil, flyteerr.Errorf(flyteerr.BadTaskSpecification, "expected a TFJob, got %T", resource) + } + return common.ChildPods(ctx, taskCtx, resource) +} + // GetCompletionTime returns the completion time of the TFJob func (tensorflowOperatorResourceHandler) GetCompletionTime(resource client.Object) (time.Time, error) { job, ok := resource.(*kubeflowv1.TFJob) diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go index 068526581e9..5008098ce6a 100644 --- a/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go @@ -47,8 +47,15 @@ var ( dummyAnnotations = map[string]string{ "annotation-key": "annotation-value", } + // The attempt labels are what NewTaskExecutionMetadata stamps on every task, and what + // the framework selects on to find the replica pods the operator derives from these + // templates, so the fixture carries them alongside a plain user label. dummyLabels = map[string]string{ - "label-key": "label-value", + "label-key": "label-value", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", } resourceRequirements = &corev1.ResourceRequirements{ diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/child_pods_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/child_pods_test.go new file mode 100644 index 00000000000..786230e5f15 --- /dev/null +++ b/flyteplugins/go/tasks/plugins/k8s/ray/child_pods_test.go @@ -0,0 +1,197 @@ +package ray + +import ( + "context" + "testing" + + rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" + "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/labels" + + pluginsCore "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s" + "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/flytek8s/config" + "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core" +) + +// attemptExecutionLabels mirrors what NewTaskExecutionMetadata stamps on every task. +func attemptExecutionLabels() map[string]string { + return map[string]string{ + "label-1": "val1", + flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue, + flytek8s.RunLabel: "run-abc", + flytek8s.ActionLabel: "a0", + flytek8s.AttemptLabel: "1", + } +} + +func attemptMetadata(executionLabels map[string]string) pluginsCore.TaskExecutionMetadata { + meta := &mocks.TaskExecutionMetadata{} + meta.EXPECT().GetLabels().Return(executionLabels) + return meta +} + +func rayJobWithCluster(namespace, clusterName string) *rayv1.RayJob { + return &rayv1.RayJob{ + ObjectMeta: metav1.ObjectMeta{Namespace: namespace, Name: "job"}, + Status: rayv1.RayJobStatus{RayClusterName: clusterName}, + } +} + +func TestRayChildPods(t *testing.T) { + t.Run("selects on the attempt and the cluster", func(t *testing.T) { + selector, err := rayJobResourceHandler{}.ChildPods( + context.TODO(), attemptMetadata(attemptExecutionLabels()), rayJobWithCluster("test-namespace", "job-abcde")) + + require.NoError(t, err) + require.NotNil(t, selector) + + podLabels := attemptExecutionLabels() + podLabels[rayClusterLabelKey] = "job-abcde" + assert.True(t, selector.Matches(labels.Set(podLabels))) + + // Another cluster in the same namespace is not this attempt's. + podLabels[rayClusterLabelKey] = "other-fghij" + assert.False(t, selector.Matches(labels.Set(podLabels))) + }) + + t.Run("falls back to the attempt alone before the cluster is named", func(t *testing.T) { + // KubeRay appends a random suffix to the cluster name, so it is only knowable + // from the RayJob's status. Until it is reported the attempt labels stand alone, + // which is still exact to one attempt of one action. + selector, err := rayJobResourceHandler{}.ChildPods( + context.TODO(), attemptMetadata(attemptExecutionLabels()), rayJobWithCluster("test-namespace", "")) + + require.NoError(t, err) + require.NotNil(t, selector) + assert.True(t, selector.Matches(labels.Set(attemptExecutionLabels()))) + + other := attemptExecutionLabels() + other[flytek8s.ActionLabel] = "a1" + assert.False(t, selector.Matches(labels.Set(other))) + }) + + t.Run("declines when the attempt cannot be identified", func(t *testing.T) { + executionLabels := attemptExecutionLabels() + delete(executionLabels, flytek8s.RunLabel) + + selector, err := rayJobResourceHandler{}.ChildPods( + context.TODO(), attemptMetadata(executionLabels), rayJobWithCluster("test-namespace", "job-abcde")) + + require.NoError(t, err) + assert.Nil(t, selector) + }) + + t.Run("rejects a resource that is not a RayJob", func(t *testing.T) { + _, err := rayJobResourceHandler{}.ChildPods( + context.TODO(), attemptMetadata(attemptExecutionLabels()), &corev1.Pod{}) + + assert.Error(t, err) + }) +} + +// TestRayChildPodsIdentityNotOverridable verifies that a task cannot take its own pods out +// of the framework's reach by setting the identity labels in its k8s_pod metadata. A pod +// whose run, action or attempt label has been overwritten is one the selector cannot find, +// and a GPU fault on it would be silently lost rather than reaching the failure. +func TestRayChildPodsIdentityNotOverridable(t *testing.T) { + require.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{})) + + rayJobObj := dummyRayCustomObj() + overrides := &core.K8SPod{ + Metadata: &core.K8SObjectMetadata{ + Labels: map[string]string{ + flytek8s.RunLabel: "a-run-the-user-made-up", + flytek8s.ActionLabel: "not-this-action", + flytek8s.AttemptLabel: "99", + }, + }, + } + rayJobObj.RayCluster.HeadGroupSpec.K8SPod = overrides + rayJobObj.RayCluster.WorkerGroupSpec[0].K8SPod = overrides + + executionLabels := attemptExecutionLabels() + taskTemplate := dummyRayTaskTemplate("ray-id", rayJobObj) + taskCtx := dummyRayTaskContextWithLabels(taskTemplate, resourceRequirements, nil, "", serviceAccount, true, executionLabels) + + resource, err := rayJobResourceHandler{}.BuildResource(context.TODO(), taskCtx) + require.NoError(t, err) + rayJob, ok := resource.(*rayv1.RayJob) + require.True(t, ok) + + rayJob.Status.RayClusterName = "job-abcde" + selector, err := rayJobResourceHandler{}.ChildPods(context.TODO(), attemptMetadata(executionLabels), rayJob) + require.NoError(t, err) + require.NotNil(t, selector) + + for name, templateLabels := range map[string]map[string]string{ + "head": rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.GetLabels(), + "worker": rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.GetLabels(), + } { + t.Run(name, func(t *testing.T) { + assert.Equal(t, executionLabels[flytek8s.RunLabel], templateLabels[flytek8s.RunLabel]) + assert.Equal(t, executionLabels[flytek8s.ActionLabel], templateLabels[flytek8s.ActionLabel]) + assert.Equal(t, executionLabels[flytek8s.AttemptLabel], templateLabels[flytek8s.AttemptLabel]) + + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + podLabels[rayClusterLabelKey] = rayJob.Status.RayClusterName + assert.True(t, selector.Matches(labels.Set(podLabels))) + }) + } +} + +// TestRayChildPodsMatchTheTemplatesTheyCameFrom is the conformance check between the two +// halves of child pod discovery: the labels this plugin puts on the pod templates KubeRay +// expands, and the selector it hands the framework to find the resulting pods. Label drift +// on either side would otherwise leave a GPU fault on a Ray worker silently unclassified. +func TestRayChildPodsMatchTheTemplatesTheyCameFrom(t *testing.T) { + executionLabels := attemptExecutionLabels() + taskTemplate := dummyRayTaskTemplate("ray-id", dummyRayCustomObj()) + taskCtx := dummyRayTaskContextWithLabels(taskTemplate, resourceRequirements, nil, "", serviceAccount, true, executionLabels) + + resource, err := rayJobResourceHandler{}.BuildResource(context.TODO(), taskCtx) + require.NoError(t, err) + rayJob, ok := resource.(*rayv1.RayJob) + require.True(t, ok) + + // KubeRay stamps the cluster label on the pods themselves, so the templates do not + // carry it and the fixture adds what the operator would. + rayJob.Status.RayClusterName = "job-abcde" + selector, err := rayJobResourceHandler{}.ChildPods(context.TODO(), attemptMetadata(executionLabels), rayJob) + require.NoError(t, err) + require.NotNil(t, selector) + + templates := map[string]map[string]string{ + "head": rayJob.Spec.RayClusterSpec.HeadGroupSpec.Template.GetLabels(), + "worker": rayJob.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.GetLabels(), + "submitter": rayJob.Spec.SubmitterPodTemplate.GetLabels(), + } + for name, templateLabels := range templates { + t.Run(name, func(t *testing.T) { + require.NotEmpty(t, templateLabels) + podLabels := make(map[string]string, len(templateLabels)+1) + for k, v := range templateLabels { + podLabels[k] = v + } + // The submitter is a plain Job pod and carries no cluster label, so it is + // checked against the attempt half only, which is what excludes it in + // production. The head and worker pods get the label from KubeRay. + if name != "submitter" { + podLabels[rayClusterLabelKey] = rayJob.Status.RayClusterName + assert.True(t, selector.Matches(labels.Set(podLabels)), + "the %s pod template's labels %v do not satisfy %s", name, podLabels, selector) + return + } + attemptOnly := flytek8s.AttemptPodSelector(attemptMetadata(executionLabels)) + require.NotNil(t, attemptOnly) + assert.True(t, attemptOnly.Matches(labels.Set(podLabels))) + }) + } +} diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go index 3ee0cd148ea..30c31dfd6ae 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go @@ -17,6 +17,8 @@ import ( v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/selection" "k8s.io/client-go/kubernetes/scheme" "sigs.k8s.io/controller-runtime/pkg/client" @@ -47,6 +49,12 @@ const ( DisableUsageStatsStartParameter = "disable-usage-stats" DisableUsageStatsStartParameterVal = "true" RayHeadContainerName = "ray-head" + // rayClusterLabelKey is KubeRay's RayClusterLabelKey. KubeRay puts it on every head + // and worker pod it creates, naming the RayCluster the pod belongs to, and overwrites + // whatever the pod template had in that key. Copied from KubeRay rather than imported + // so that the utils package it lives in, which pulls in an HTTP client, stays out of + // go.mod's build graph here. + rayClusterLabelKey = "ray.io/cluster" ) var logTemplateRegexes = struct { @@ -72,6 +80,10 @@ var submitterDefaultResourceRequirements = v1.ResourceRequirements{ type rayJobResourceHandler struct{} +// The cluster's head and worker pods are where a node daemon records what the hardware did, +// so the framework has to be able to find them from the RayJob this plugin tracks. +var _ k8s.ChildPodDiscovery = rayJobResourceHandler{} + func (rayJobResourceHandler) GetProperties() k8s.PluginProperties { maxLength := 47 return k8s.PluginProperties{GeneratedNameMaxLength: &maxLength} @@ -547,7 +559,10 @@ func buildHeadPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodSpe } cfg := config.GetK8sPluginConfig() podTemplateSpec.SetLabels(utils.UnionMaps(cfg.DefaultLabels, podTemplateSpec.GetLabels(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels()), spec.GetK8SPod().GetMetadata().GetLabels(), - map[string]string{flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue})) + // Applied last: a user label named run, action or attempt would otherwise + // overwrite the identity the framework finds this pod by, and a pod it cannot + // find is one whose GPU faults never reach the failure. + flytek8s.PreservedPodLabels(taskCtx.TaskExecutionMetadata()))) podTemplateSpec.SetAnnotations(utils.UnionMaps(cfg.DefaultAnnotations, podTemplateSpec.GetAnnotations(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations()), spec.GetK8SPod().GetMetadata().GetAnnotations())) return podTemplateSpec, nil @@ -581,7 +596,11 @@ func buildSubmitterPodTemplate(rayClusterSpec *rayv1.RayClusterSpec, taskCtx plu }, } k8sCfg := config.GetK8sPluginConfig() - podTemplateSpec.SetLabels(utils.UnionMaps(k8sCfg.DefaultLabels, utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels()))) + podTemplateSpec.SetLabels(utils.UnionMaps( + k8sCfg.DefaultLabels, + utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels()), + flytek8s.PreservedPodLabels(taskCtx.TaskExecutionMetadata()), + )) podTemplateSpec.SetAnnotations(utils.UnionMaps(k8sCfg.DefaultAnnotations, utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations()))) return podTemplateSpec } @@ -712,7 +731,10 @@ func buildWorkerPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodS } cfg := config.GetK8sPluginConfig() podTemplateSpec.SetLabels(utils.UnionMaps(cfg.DefaultLabels, podTemplateSpec.GetLabels(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetLabels()), spec.GetK8SPod().GetMetadata().GetLabels(), - map[string]string{flytek8s.ManagedLabelKey: flytek8s.ManagedLabelValue})) + // Applied last: a user label named run, action or attempt would otherwise + // overwrite the identity the framework finds this pod by, and a pod it cannot + // find is one whose GPU faults never reach the failure. + flytek8s.PreservedPodLabels(taskCtx.TaskExecutionMetadata()))) podTemplateSpec.SetAnnotations(utils.UnionMaps(cfg.DefaultAnnotations, podTemplateSpec.GetAnnotations(), utils.CopyMap(taskCtx.TaskExecutionMetadata().GetAnnotations()), spec.GetK8SPod().GetMetadata().GetAnnotations())) return podTemplateSpec, nil } @@ -752,6 +774,43 @@ func (rayJobResourceHandler) BuildIdentityResource(ctx context.Context, taskCtx }, nil } +// ChildPods implements k8s.ChildPodDiscovery. The pods that run the task are the head and +// worker pods KubeRay expands from the templates this plugin built, which the framework +// tracks nothing of, since it tracks the RayJob. +// +// The selector is the attempt's own labels, which the head and worker templates carry +// (see buildHeadPodTemplate and buildWorkerPodTemplate), narrowed by the RayCluster the +// pods belong to. The cluster name is only knowable from the RayJob's status, because +// KubeRay appends a random suffix to it; until the operator reports it the attempt labels +// stand alone. That is not a wider search in any sense that matters: run, action and +// attempt already pin the selector to this one attempt of this one action, so the worst it +// can add is the job submitter pod, which has no GPU and so no faults to contribute. +func (rayJobResourceHandler) ChildPods( + _ context.Context, + taskCtx pluginsCore.TaskExecutionMetadata, + obj client.Object, +) (labels.Selector, error) { + rayJob, ok := obj.(*rayv1.RayJob) + if !ok { + return nil, fmt.Errorf("expected a RayJob, got %T", obj) + } + + selector := flytek8s.AttemptPodSelector(taskCtx) + if selector == nil { + return nil, nil + } + + if clusterName := rayJob.Status.RayClusterName; clusterName != "" { + requirement, err := labels.NewRequirement(rayClusterLabelKey, selection.Equals, []string{clusterName}) + if err != nil { + return nil, err + } + selector = selector.Add(*requirement) + } + + return selector, nil +} + func getEventInfoForRayJob(ctx context.Context, logConfig logs.LogConfig, pluginContext k8s.PluginContext, rayJob *rayv1.RayJob) (*pluginsCore.TaskInfo, error) { taskTemplate, err := pluginContext.TaskReader().Read(ctx) if err != nil { diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go index 44592f3e7d2..c57c068e939 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go @@ -119,6 +119,10 @@ func dummyRayTaskContext(taskTemplate *core.TaskTemplate, resources *corev1.Reso } func dummyRayTaskContextInterruptible(taskTemplate *core.TaskTemplate, resources *corev1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage, serviceAccount string, interruptible bool) pluginsCore.TaskExecutionContext { + return dummyRayTaskContextWithLabels(taskTemplate, resources, extendedResources, containerImage, serviceAccount, interruptible, map[string]string{"label-1": "val1"}) +} + +func dummyRayTaskContextWithLabels(taskTemplate *core.TaskTemplate, resources *corev1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage, serviceAccount string, interruptible bool, executionLabels map[string]string) pluginsCore.TaskExecutionContext { taskCtx := &mocks.TaskExecutionContext{} inputReader := &pluginIOMocks.InputReader{} inputReader.EXPECT().GetInputPrefixPath().Return("/input/prefix") @@ -160,7 +164,7 @@ func dummyRayTaskContextInterruptible(taskTemplate *core.TaskTemplate, resources taskExecutionMetadata.EXPECT().GetTaskExecutionID().Return(tID) taskExecutionMetadata.EXPECT().GetNamespace().Return("test-namespace") taskExecutionMetadata.EXPECT().GetAnnotations().Return(map[string]string{"annotation-1": "val1"}) - taskExecutionMetadata.EXPECT().GetLabels().Return(map[string]string{"label-1": "val1"}) + taskExecutionMetadata.EXPECT().GetLabels().Return(executionLabels) taskExecutionMetadata.EXPECT().GetOwnerReference().Return(metav1.OwnerReference{ Kind: "node", Name: "blah",