feat(api): Sync TrainJob JobsStatus from JobSet ReplicatedJobsStatus

pkg/controller/trainjob_controller.go

@@ -104,10 +104,6 @@ func (r *TrainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 	log := ctrl.LoggerFrom(ctx).WithValues("trainJob", klog.KObj(&trainJob))
 	ctx = ctrl.LoggerInto(ctx, log)
 	log.V(2).Info("Reconciling TrainJob")
-	if isTrainJobFinished(&trainJob) {
-		log.V(5).Info("TrainJob has already been finished")
-		return ctrl.Result{}, nil
-	}
 
 	var err error
 	// Keep track of the origin TrainJob status
@@ -142,6 +138,10 @@ func (r *TrainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		err = errors.Join(err, terminalCondErr)
 	}
 
+	if jobsStatusErr := setJobsStatus(ctx, runtime, &trainJob); jobsStatusErr != nil {
+		err = errors.Join(err, jobsStatusErr)
+	}
+
 	if !equality.Semantic.DeepEqual(&trainJob.Status, originStatus) {
 		return ctrl.Result{}, errors.Join(err, r.client.Status().Update(ctx, &trainJob))
 	}
@@ -267,9 +267,13 @@ func setTerminalCondition(ctx context.Context, runtime jobruntimes.Runtime, trai
 	return nil
 }
 
-func isTrainJobFinished(trainJob *trainer.TrainJob) bool {
-	return meta.IsStatusConditionTrue(trainJob.Status.Conditions, trainer.TrainJobComplete) ||
-		meta.IsStatusConditionTrue(trainJob.Status.Conditions, trainer.TrainJobFailed)
+func setJobsStatus(ctx context.Context, runtime jobruntimes.Runtime, trainJob *trainer.TrainJob) error {
+	statuses, err := runtime.JobsStatus(ctx, trainJob)
+	if err != nil {
+		return err
+	}
+	trainJob.Status.JobsStatus = statuses
+	return nil
 }
 
 func (r *TrainJobReconciler) SetupWithManager(mgr ctrl.Manager, options controller.Options) error {

pkg/runtime/core/clustertrainingruntime.go

@@ -75,6 +75,10 @@ func (r *ClusterTrainingRuntime) TerminalCondition(ctx context.Context, trainJob
 	return r.TrainingRuntime.TerminalCondition(ctx, trainJob)
 }
 
+func (r *ClusterTrainingRuntime) JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error) {
+	return r.TrainingRuntime.JobsStatus(ctx, trainJob)
+}
+
 func (r *ClusterTrainingRuntime) EventHandlerRegistrars() []runtime.ReconcilerBuilder {
 	return nil
 }

pkg/runtime/core/trainingruntime.go

@@ -250,6 +250,10 @@ func (r *TrainingRuntime) TerminalCondition(ctx context.Context, trainJob *train
 	return r.framework.RunTerminalConditionPlugins(ctx, trainJob)
 }
 
+func (r *TrainingRuntime) JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error) {
+	return r.framework.RunJobsStatusPlugins(ctx, trainJob)
+}
+
 func (r *TrainingRuntime) EventHandlerRegistrars() []runtime.ReconcilerBuilder {
 	var builders []runtime.ReconcilerBuilder
 	for _, ex := range r.framework.WatchExtensionPlugins() {

pkg/runtime/framework/core/framework.go

@@ -32,7 +32,10 @@ import (
 	index "github.com/kubeflow/trainer/v2/pkg/runtime/indexer"
 )
 
-var errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
+var (
+	errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
+	errorTooManyJobsStatusPlugin        = errors.New("too many JobsStatus plugins are registered")
+)
 
 type Framework struct {
 	registry                     fwkplugins.Registry
@@ -43,7 +46,8 @@ type Framework struct {
 	watchExtensionPlugins        []framework.WatchExtensionPlugin
 	podNetworkPlugins            []framework.PodNetworkPlugin
 	componentBuilderPlugins      []framework.ComponentBuilderPlugin
-	terminalConditionPlugins     []framework.TerminalConditionPlugin
+	terminalConditionPlugin      framework.TerminalConditionPlugin
+	jobsStatusPlugin             framework.JobsStatusPlugin
 }
 
 func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer client.FieldIndexer) (*Framework, error) {
@@ -80,7 +84,16 @@ func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer cl
 			f.componentBuilderPlugins = append(f.componentBuilderPlugins, p)
 		}
 		if p, ok := plugin.(framework.TerminalConditionPlugin); ok {
-			f.terminalConditionPlugins = append(f.terminalConditionPlugins, p)
+			if f.terminalConditionPlugin != nil {
+				return nil, errorTooManyTerminalConditionPlugin
+			}
+			f.terminalConditionPlugin = p
+		}
+		if p, ok := plugin.(framework.JobsStatusPlugin); ok {
+			if f.jobsStatusPlugin != nil {
+				return nil, errorTooManyJobsStatusPlugin
+			}
+			f.jobsStatusPlugin = p
 		}
 	}
 	f.plugins = plugins
@@ -142,12 +155,15 @@ func (f *Framework) RunComponentBuilderPlugins(ctx context.Context, info *runtim
 }
 
 func (f *Framework) RunTerminalConditionPlugins(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {
-	// TODO (tenzen-y): Once we provide the Configuration API, we should validate which plugin should have terminalCondition execution points.
-	if len(f.terminalConditionPlugins) > 1 {
-		return nil, errorTooManyTerminalConditionPlugin
+	if f.terminalConditionPlugin != nil {
+		return f.terminalConditionPlugin.TerminalCondition(ctx, trainJob)
 	}
-	if len(f.terminalConditionPlugins) != 0 {
-		return f.terminalConditionPlugins[0].TerminalCondition(ctx, trainJob)
+	return nil, nil
+}
+
+func (f *Framework) RunJobsStatusPlugins(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error) {
+	if f.jobsStatusPlugin != nil {
+		return f.jobsStatusPlugin.JobsStatus(ctx, trainJob)
 	}
 	return nil, nil
 }

pkg/runtime/framework/core/framework_test.go

@@ -111,9 +111,8 @@ func TestNew(t *testing.T) {
 					&jobset.JobSet{},
 					&mpi.MPI{},
 				},
-				terminalConditionPlugins: []framework.TerminalConditionPlugin{
-					&jobset.JobSet{},
-				},
+				terminalConditionPlugin: &jobset.JobSet{},
+				jobsStatusPlugin:        &jobset.JobSet{},
 			},
 		},
 		"indexer key for trainingRuntime and runtimeClass is an empty": {
@@ -136,7 +135,7 @@ func TestNew(t *testing.T) {
 		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, volcano.Volcano{}, mpi.MPI{}, plainml.PlainML{}, torch.Torch{}, jobset.JobSet{}),
 		cmpopts.IgnoreFields(coscheduling.CoScheduling{}, "client"),
 		cmpopts.IgnoreFields(volcano.Volcano{}, "client"),
-		cmpopts.IgnoreFields(jobset.JobSet{}, "client"),
+		cmpopts.IgnoreFields(jobset.JobSet{}, "client", "restMapper", "scheme", "logger"),
 		cmpopts.IgnoreTypes(apiruntime.Scheme{}, meta.DefaultRESTMapper{}, fwkplugins.Registry{}),
 		cmpopts.SortMaps(func(a, b string) bool { return a < b }),
 		cmpopts.SortSlices(func(a, b framework.Plugin) bool { return a.Name() < b.Name() }),
@@ -1534,7 +1533,10 @@ func TestTerminalConditionPlugins(t *testing.T) {
 
 			fwk, err := New(ctx, c, tc.registry, testingutil.AsIndex(clientBuilder))
 			if err != nil {
-				t.Fatal(err)
+				if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
+					t.Errorf("Unexpected error (-want,+got):\n%s", diff)
+				}
+				return
 			}
 
 			gotCond, gotErr := fwk.RunTerminalConditionPlugins(ctx, tc.trainJob)
@@ -1549,6 +1551,163 @@ func TestTerminalConditionPlugins(t *testing.T) {
 	}
 }
 
+type fakeJobsStatusPlugin struct{}
+
+var _ framework.JobsStatusPlugin = (*fakeJobsStatusPlugin)(nil)
+
+func newFakeJobsStatusPlugin(context.Context, client.Client, client.FieldIndexer) (framework.Plugin, error) {
+	return &fakeJobsStatusPlugin{}, nil
+}
+
+const fakeJobsStatusPluginName = "fake-jobs-status"
+
+func (f fakeJobsStatusPlugin) Name() string { return fakeJobsStatusPluginName }
+func (f fakeJobsStatusPlugin) JobsStatus(context.Context, *trainer.TrainJob) ([]trainer.JobStatus, error) {
+	return []trainer.JobStatus{
+		{Name: "fake-job", Ready: 1, Succeeded: 0, Failed: 0, Active: 1, Suspended: 0},
+	}, nil
+}
+
+func TestJobsStatusPlugins(t *testing.T) {
+	cases := map[string]struct {
+		registry     fwkplugins.Registry
+		trainJob     *trainer.TrainJob
+		jobSet       *jobsetv1alpha2.JobSet
+		wantStatuses []trainer.JobStatus
+		wantError    error
+	}{
+		"JobSet with empty replicated jobs status": {
+			registry: fwkplugins.NewRegistry(),
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "testing").
+				Obj(),
+			jobSet: testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "testing").
+				Obj(),
+			wantStatuses: nil,
+		},
+		"succeeded to obtain JobsStatus from JobSet with multiple replicated jobs": {
+			registry: fwkplugins.NewRegistry(),
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "testing").
+				Obj(),
+			jobSet: testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "testing").
+				ReplicatedJobsStatuses([]jobsetv1alpha2.ReplicatedJobStatus{
+					{
+						Name:      constants.DatasetInitializer,
+						Ready:     1,
+						Succeeded: 1,
+						Failed:    0,
+						Active:    0,
+						Suspended: 0,
+					},
+					{
+						Name:      constants.ModelInitializer,
+						Ready:     1,
+						Succeeded: 1,
+						Failed:    0,
+						Active:    0,
+						Suspended: 0,
+					},
+					{
+						Name:      constants.Node,
+						Ready:     2,
+						Succeeded: 0,
+						Failed:    0,
+						Active:    2,
+						Suspended: 0,
+					},
+				}).
+				Obj(),
+			wantStatuses: []trainer.JobStatus{
+				{
+					Name:      constants.DatasetInitializer,
+					Ready:     1,
+					Succeeded: 1,
+					Failed:    0,
+					Active:    0,
+					Suspended: 0,
+				},
+				{
+					Name:      constants.ModelInitializer,
+					Ready:     1,
+					Succeeded: 1,
+					Failed:    0,
+					Active:    0,
+					Suspended: 0,
+				},
+				{
+					Name:      constants.Node,
+					Ready:     2,
+					Succeeded: 0,
+					Failed:    0,
+					Active:    2,
+					Suspended: 0,
+				},
+			},
+		},
+		"succeeded to obtain JobsStatus from JobSet with failed job": {
+			registry: fwkplugins.NewRegistry(),
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "testing").
+				Obj(),
+			jobSet: testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "testing").
+				ReplicatedJobsStatuses([]jobsetv1alpha2.ReplicatedJobStatus{
+					{
+						Name:      constants.Node,
+						Ready:     0,
+						Succeeded: 0,
+						Failed:    1,
+						Active:    0,
+						Suspended: 0,
+					},
+				}).
+				Obj(),
+			wantStatuses: []trainer.JobStatus{
+				{
+					Name:      constants.Node,
+					Ready:     0,
+					Succeeded: 0,
+					Failed:    1,
+					Active:    0,
+					Suspended: 0,
+				},
+			},
+		},
+		"failed to obtain JobsStatus due to multiple JobsStatusPlugins": {
+			registry: fwkplugins.Registry{
+				jobset.Name:              jobset.New,
+				fakeJobsStatusPluginName: newFakeJobsStatusPlugin,
+			},
+			wantError: errorTooManyJobsStatusPlugin,
+		},
+	}
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			t.Cleanup(cancel)
+			clientBuilder := testingutil.NewClientBuilder()
+			if tc.jobSet != nil {
+				clientBuilder = clientBuilder.WithObjects(tc.jobSet)
+			}
+			c := clientBuilder.Build()
+
+			fwk, err := New(ctx, c, tc.registry, testingutil.AsIndex(clientBuilder))
+			if err != nil {
+				if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
+					t.Errorf("Unexpected error (-want,+got):\n%s", diff)
+				}
+				return
+			}
+
+			gotStatuses, gotErr := fwk.RunJobsStatusPlugins(ctx, tc.trainJob)
+			if diff := cmp.Diff(tc.wantError, gotErr, cmpopts.EquateErrors()); len(diff) != 0 {
+				t.Errorf("Unexpected error (-want,+got):\n%s", diff)
+			}
+
+			if diff := cmp.Diff(tc.wantStatuses, gotStatuses); len(diff) != 0 {
+				t.Errorf("Unexpected jobs status (-want,+got):\n%s", diff)
+			}
+		})
+	}
+}
+
 func TestPodNetworkPlugins(t *testing.T) {
 	cases := map[string]struct {
 		registry        fwkplugins.Registry

pkg/runtime/framework/interface.go

@@ -65,3 +65,8 @@ type TerminalConditionPlugin interface {
 	Plugin
 	TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)
 }
+
+type JobsStatusPlugin interface {
+	Plugin
+	JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)
+}

pkg/runtime/framework/plugins/jobset/jobset.go

@@ -63,6 +63,7 @@ var _ framework.WatchExtensionPlugin = (*JobSet)(nil)
 var _ framework.PodNetworkPlugin = (*JobSet)(nil)
 var _ framework.ComponentBuilderPlugin = (*JobSet)(nil)
 var _ framework.TerminalConditionPlugin = (*JobSet)(nil)
+var _ framework.JobsStatusPlugin = (*JobSet)(nil)
 var _ framework.CustomValidationPlugin = (*JobSet)(nil)
 
 const Name = constants.JobSetKind
@@ -304,3 +305,22 @@ func (j *JobSet) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJ
 	}
 	return nil, nil
 }
+
+func (j *JobSet) JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error) {
+	jobSet := &jobsetv1alpha2.JobSet{}
+	if err := j.client.Get(ctx, client.ObjectKeyFromObject(trainJob), jobSet); err != nil {
+		return nil, err
+	}
+	var statuses []trainer.JobStatus
+	for _, status := range jobSet.Status.ReplicatedJobsStatus {
+		statuses = append(statuses, trainer.JobStatus{
+			Name:      status.Name,
+			Ready:     status.Ready,
+			Succeeded: status.Succeeded,
+			Failed:    status.Failed,
+			Active:    status.Active,
+			Suspended: status.Suspended,
+		})
+	}
+	return statuses, nil
+}

pkg/runtime/interface.go

@@ -39,6 +39,7 @@ type Runtime interface {
 	NewObjects(ctx context.Context, trainJob *trainer.TrainJob) ([]any, error)
 	RuntimeInfo(trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy) (*Info, error)
 	TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)
+	JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)
 	EventHandlerRegistrars() []ReconcilerBuilder
 	ValidateObjects(ctx context.Context, old, new *trainer.TrainJob) (admission.Warnings, field.ErrorList)
 }

pkg/util/testing/wrapper.go

@@ -503,6 +503,11 @@ func (j *JobSetWrapper) DependsOn(rJobName string, dependsOn ...jobsetv1alpha2.D
 	return j
 }
 
+func (j *JobSetWrapper) ReplicatedJobsStatuses(statuses []jobsetv1alpha2.ReplicatedJobStatus) *JobSetWrapper {
+	j.Status.ReplicatedJobsStatus = statuses
+	return j
+}
+
 func (j *JobSetWrapper) Obj() *jobsetv1alpha2.JobSet {
 	return &j.JobSet
 }