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 <noreply@anthropic.com>
Signed-off-by: Samhita Alla <aallasamhita@gmail.com>

Author: samhita-alla

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)
+		})
+	}
+}

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)

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 {

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)
+	})
+}

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 {