From c165cd4214ecbcf7e369c7702da600f8233895cf Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Tue, 25 Aug 2026 18:12:52 +0530 Subject: [PATCH 1/2] feat(executor): classify GPU faults recorded on a CRD task's child pods Classification so far covered tasks whose tracked resource is the Pod itself. A RayJob, PyTorchJob, MPIJob or TFJob records the Xid on a worker pod while the executor tracks the CRD, so the fault the node saw never reached the failure the user read, and a GPU that fell off the bus during distributed training still looked like the workload's own crash. Nothing new has to be watched to fix that. The event informer is not scoped to the tracked resource, so a worker pod's events are already cached, and the Pod informer is already running and already selects on the managed label that every plugin merges into the pod templates it builds. What was missing was a way to name the pods: an operator's child pod names carry a random suffix, so the executor cannot derive the keys the events are cached under. ChildPodDiscovery is the optional interface that closes that, in the shape PluginAbortOverride already established. A plugin returns the label selector that picks out this attempt's pods; the manager lists them from the resource's own namespace out of the cache it already holds, and reads each one's faults. Because the pods arrive as objects, their UIDs are known, so the identity check that already guarded the Pod path applies unchanged and never has to fall back to matching on a name. AttemptPodSelector builds the half of that selector which holds whichever operator created the pod. The run, action and attempt labels move to flytek8s next to the managed label, for the reason that one already lives there: the plugins that build pod templates have to keep them intact. It refuses to build a selector when any of them is missing, since run and action are dropped when sanitization empties them and selecting on what is left would reach another action's pods. PreservedPodLabels exposes the same key set for plugins to re-apply after any labels a user supplied, so that the two halves cannot disagree about which labels have to survive. Which pods are consulted is narrower than everything the selector matches. A worker that finished its work cannot be why the job failed, so a succeeded pod contributes nothing even if it logged a fault mid-run; crediting that to a launcher's user-code failure would turn a user error into a hardware one and stop charging the retry. A pod still running is kept, because a worker wedged on a GPU that fell off the bus is exactly the case this path exists for. Relevance is anchored per pod on that pod's own container termination, because every CRD plugin stamps the failure with the time of the reconcile that noticed it rather than with anything the kubelet recorded, and anchoring on the reconcile would age out a fault that killed a worker minutes earlier. Init containers are not eligible anchors: they finish before the workload starts, and a native sidecar declared among them is reaped after everything else. The anchor is never later than the attempt's own failure, so a pod torn down long afterwards cannot let faults recorded in the meantime explain it. Which fault wins needs no new rule: ClassifyFailure already takes the worst severity and the earliest within it, so the faults gathered across pods are only sorted back into order before being handed over. That order is on when each fault was first recorded rather than when it was last seen, because 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. The last observation still decides relevance, which is all it was ever for. The one pod whose fault settled the verdict is named as an additional reason, because the fault sentence carries the Xid, the GPU and the node but not the pod, which on a job with sixty-four workers is not enough to act on. Only that pod is named: pointing at every pod that saw any fault would put a misleading reason on a failure the faults did not explain, and would emit one cluster event per pod when a whole node's worth of GPUs faults at once. The single pod path names nothing, since the failure is already reported against that pod. A gang-scheduled restart creates fresh pods under the same attempt and its faults count, while an earlier Flyte attempt's pods are excluded by the attempt label. A plugin that does not implement the interface contributes no pods and so no faults, which is exactly what it did before. A pod already torn down with its job is not recovered. The events are cached under a name nothing can reconstruct, so discovery logs at debug and classifies nothing rather than guessing. Co-Authored-By: Claude Fable 5 Signed-off-by: Samhita Alla --- .../plugin/k8s/child_pod_gpu_fault_test.go | 524 ++++++++++++++++++ executor/pkg/plugin/k8s/plugin_manager.go | 260 ++++++++- .../pkg/plugin/k8s/plugin_manager_test.go | 24 +- executor/pkg/plugin/task_exec_metadata.go | 24 +- .../flytek8s/attempt_pod_selector_test.go | 169 ++++++ .../pluginmachinery/flytek8s/pod_helper.go | 85 +++ .../go/tasks/pluginmachinery/k8s/plugin.go | 26 + 7 files changed, 1062 insertions(+), 50 deletions(-) create mode 100644 executor/pkg/plugin/k8s/child_pod_gpu_fault_test.go create mode 100644 flyteplugins/go/tasks/pluginmachinery/flytek8s/attempt_pod_selector_test.go 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) From 23a0ecea9139bee3c1fd1341109d9eb20b13a0a7 Mon Sep 17 00:00:00 2001 From: Samhita Alla Date: Tue, 25 Aug 2026 18:13:09 +0530 Subject: [PATCH 2/2] feat(ray,kfoperators,clustered): name the child pods of a distributed task The Ray, Kubeflow and clustered plugins track a CRD, so they implement ChildPodDiscovery to tell the framework where the pods that actually ran the task are. All three build the same selector: the attempt's own labels, which the pod templates they hand the operator already carry, narrowed by one label the operator applies to the pods themselves. For Kubeflow that is training.kubeflow.org/job-name, which is the CR's own name and so is known from the moment the resource exists; one implementation in common covers PyTorchJob, MPIJob and TFJob, because all three get their replica pods from the same operator. For the clustered plugin it is jobset.sigs.k8s.io/jobset-name, exported by the vendored JobSet API the plugin already imports. For Ray it is ray.io/cluster, which KubeRay stamps on head and worker pods and overwrites on the template, and whose value is only knowable from the RayJob's status because KubeRay appends a random suffix to the cluster name. Until that status is reported Ray falls back to the attempt labels alone. That is not a wider search in any sense that matters, since run, action and attempt already pin the selector to one attempt of one action; the most it can add is the job submitter pod, which has no GPU and so no faults to contribute. The alternative, declining until the operator catches up, would lose the fault on exactly the runs where the cluster never came up healthy. Ray also has to stop a task from taking its own pods out of reach. The head and worker templates let a task supply k8s_pod metadata, and that metadata was unioned after the execution labels with only the managed label re-forced afterwards, so a user label named run, action or attempt would overwrite the identity the selector looks a pod up by. All three templates now re-apply flytek8s.PreservedPodLabels last, which is the same key set the selector is built from. Dask and Spark are left for a follow-up. The conformance tests are the point of the test additions. Each asserts that the labels the plugin puts on the pod templates satisfy the selector the same plugin hands the framework, for every replica type it builds. The two halves live far apart and drift on either one would leave a GPU fault on a worker silently unclassified, which is the failure mode this whole path exists to prevent. Co-Authored-By: Claude Fable 5 Signed-off-by: Samhita Alla --- .../plugins/k8s/clustered/child_pods_test.go | 139 ++++++++++++ .../plugins/k8s/clustered/clustered_test.go | 8 +- .../go/tasks/plugins/k8s/clustered/plugin.go | 37 ++++ .../k8s/kfoperators/common/child_pods_test.go | 83 ++++++++ .../k8s/kfoperators/common/common_operator.go | 34 +++ .../k8s/kfoperators/mpi/child_pods_test.go | 58 ++++++ .../tasks/plugins/k8s/kfoperators/mpi/mpi.go | 19 ++ .../plugins/k8s/kfoperators/mpi/mpi_test.go | 9 +- .../kfoperators/pytorch/child_pods_test.go | 58 ++++++ .../k8s/kfoperators/pytorch/pytorch.go | 19 ++ .../k8s/kfoperators/pytorch/pytorch_test.go | 9 +- .../kfoperators/tensorflow/child_pods_test.go | 58 ++++++ .../k8s/kfoperators/tensorflow/tensorflow.go | 19 ++ .../kfoperators/tensorflow/tensorflow_test.go | 9 +- .../tasks/plugins/k8s/ray/child_pods_test.go | 197 ++++++++++++++++++ flyteplugins/go/tasks/plugins/k8s/ray/ray.go | 65 +++++- .../go/tasks/plugins/k8s/ray/ray_test.go | 6 +- 17 files changed, 819 insertions(+), 8 deletions(-) create mode 100644 flyteplugins/go/tasks/plugins/k8s/clustered/child_pods_test.go create mode 100644 flyteplugins/go/tasks/plugins/k8s/kfoperators/common/child_pods_test.go create mode 100644 flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/child_pods_test.go create mode 100644 flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/child_pods_test.go create mode 100644 flyteplugins/go/tasks/plugins/k8s/kfoperators/tensorflow/child_pods_test.go create mode 100644 flyteplugins/go/tasks/plugins/k8s/ray/child_pods_test.go 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",