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
@@ -138,8 +134,9 @@ func (r *TrainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 	}
 
 	setSuspendedCondition(&trainJob)
-	if terminalCondErr := setTerminalCondition(ctx, runtime, &trainJob); terminalCondErr != nil {
-		err = errors.Join(err, terminalCondErr)
+
+	if statusErr := setTrainJobStatus(ctx, runtime, &trainJob); statusErr != nil {
+		err = errors.Join(err, statusErr)
 	}
 
 	if !equality.Semantic.DeepEqual(&trainJob.Status, originStatus) {
@@ -256,22 +253,17 @@ func removeFailedCondition(trainJob *trainer.TrainJob) {
 	meta.RemoveStatusCondition(&trainJob.Status.Conditions, trainer.TrainJobFailed)
 }
 
-func setTerminalCondition(ctx context.Context, runtime jobruntimes.Runtime, trainJob *trainer.TrainJob) error {
-	terminalCond, err := runtime.TerminalCondition(ctx, trainJob)
+func setTrainJobStatus(ctx context.Context, runtime jobruntimes.Runtime, trainJob *trainer.TrainJob) error {
+	status, err := runtime.TrainJobStatus(ctx, trainJob)
 	if err != nil {
 		return err
 	}
-	if terminalCond != nil {
-		meta.SetStatusCondition(&trainJob.Status.Conditions, *terminalCond)
+	if status != nil {
+		trainJob.Status = *status
 	}
 	return nil
 }
 
-func isTrainJobFinished(trainJob *trainer.TrainJob) bool {
-	return meta.IsStatusConditionTrue(trainJob.Status.Conditions, trainer.TrainJobComplete) ||
-		meta.IsStatusConditionTrue(trainJob.Status.Conditions, trainer.TrainJobFailed)
-}
-
 func (r *TrainJobReconciler) SetupWithManager(mgr ctrl.Manager, options controller.Options) error {
 	b := builder.TypedControllerManagedBy[reconcile.Request](mgr).
 		Named("trainjob_controller").

pkg/runtime/core/clustertrainingruntime.go

@@ -21,7 +21,6 @@ import (
 	"errors"
 	"fmt"
 
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -71,8 +70,8 @@ func (r *ClusterTrainingRuntime) RuntimeInfo(
 	return r.TrainingRuntime.RuntimeInfo(trainJob, runtimeTemplateSpec, mlPolicy, podGroupPolicy)
 }
 
-func (r *ClusterTrainingRuntime) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {
-	return r.TrainingRuntime.TerminalCondition(ctx, trainJob)
+func (r *ClusterTrainingRuntime) TrainJobStatus(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error) {
+	return r.TrainingRuntime.TrainJobStatus(ctx, trainJob)
 }
 
 func (r *ClusterTrainingRuntime) EventHandlerRegistrars() []runtime.ReconcilerBuilder {

pkg/runtime/core/trainingruntime.go

@@ -246,8 +246,8 @@ func syncPodSets(info *runtime.Info) {
 	}
 }
 
-func (r *TrainingRuntime) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {
-	return r.framework.RunTerminalConditionPlugins(ctx, trainJob)
+func (r *TrainingRuntime) TrainJobStatus(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error) {
+	return r.framework.RunTrainJobStatusPlugin(ctx, trainJob)
 }
 
 func (r *TrainingRuntime) EventHandlerRegistrars() []runtime.ReconcilerBuilder {

pkg/runtime/framework/core/framework.go

@@ -20,7 +20,6 @@ import (
 	"context"
 	"errors"
 
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -32,7 +31,7 @@ import (
 	index "github.com/kubeflow/trainer/v2/pkg/runtime/indexer"
 )
 
-var errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
+var errorTooManyTrainJobStatusPlugin = errors.New("too many TrainJobStatus plugins are registered")
 
 type Framework struct {
 	registry                     fwkplugins.Registry
@@ -43,7 +42,7 @@ type Framework struct {
 	watchExtensionPlugins        []framework.WatchExtensionPlugin
 	podNetworkPlugins            []framework.PodNetworkPlugin
 	componentBuilderPlugins      []framework.ComponentBuilderPlugin
-	terminalConditionPlugins     []framework.TerminalConditionPlugin
+	trainJobStatusPlugin         framework.TrainJobStatusPlugin
 }
 
 func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer client.FieldIndexer) (*Framework, error) {
@@ -79,8 +78,11 @@ func New(ctx context.Context, c client.Client, r fwkplugins.Registry, indexer cl
 		if p, ok := plugin.(framework.ComponentBuilderPlugin); ok {
 			f.componentBuilderPlugins = append(f.componentBuilderPlugins, p)
 		}
-		if p, ok := plugin.(framework.TerminalConditionPlugin); ok {
-			f.terminalConditionPlugins = append(f.terminalConditionPlugins, p)
+		if p, ok := plugin.(framework.TrainJobStatusPlugin); ok {
+			if f.trainJobStatusPlugin != nil {
+				return nil, errorTooManyTrainJobStatusPlugin
+			}
+			f.trainJobStatusPlugin = p
 		}
 	}
 	f.plugins = plugins
@@ -141,13 +143,9 @@ func (f *Framework) RunComponentBuilderPlugins(ctx context.Context, info *runtim
 	return objs, nil
 }
 
-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 len(f.terminalConditionPlugins) != 0 {
-		return f.terminalConditionPlugins[0].TerminalCondition(ctx, trainJob)
+func (f *Framework) RunTrainJobStatusPlugin(ctx context.Context, trainJob *trainer.TrainJob) (*trainer.TrainJobStatus, error) {
+	if f.trainJobStatusPlugin != nil {
+		return f.trainJobStatusPlugin.Status(ctx, trainJob)
 	}
 	return nil, nil
 }