[release-2.0] feat: Add a public function to create runtime info objects (#2846)

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Co-authored-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

Author: google-oss-robot

pkg/runtime/core/clustertrainingruntime.go

@@ -57,7 +57,18 @@ func (r *ClusterTrainingRuntime) NewObjects(ctx context.Context, trainJob *train
 	if err := r.client.Get(ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &clTrainingRuntime); err != nil {
 		return nil, fmt.Errorf("%w: %w", errorNotFoundSpecifiedClusterTrainingRuntime, err)
 	}
-	return r.buildObjects(ctx, trainJob, clTrainingRuntime.Spec.Template, clTrainingRuntime.Spec.MLPolicy, clTrainingRuntime.Spec.PodGroupPolicy)
+
+	info, err := r.RuntimeInfo(trainJob, clTrainingRuntime.Spec.Template, clTrainingRuntime.Spec.MLPolicy, clTrainingRuntime.Spec.PodGroupPolicy)
+	if err != nil {
+		return nil, err
+	}
+	return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
+}
+
+func (r *ClusterTrainingRuntime) RuntimeInfo(
+	trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
+) (*runtime.Info, error) {
+	return r.TrainingRuntime.RuntimeInfo(trainJob, runtimeTemplateSpec, mlPolicy, podGroupPolicy)
 }
 
 func (r *ClusterTrainingRuntime) TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error) {

pkg/runtime/core/trainingruntime.go

@@ -85,13 +85,21 @@ func (r *TrainingRuntime) NewObjects(ctx context.Context, trainJob *trainer.Trai
 	if err != nil {
 		return nil, fmt.Errorf("%w: %w", errorNotFoundSpecifiedTrainingRuntime, err)
 	}
-	return r.buildObjects(ctx, trainJob, trainingRuntime.Spec.Template, trainingRuntime.Spec.MLPolicy, trainingRuntime.Spec.PodGroupPolicy)
+	info, err := r.RuntimeInfo(trainJob, trainingRuntime.Spec.Template, trainingRuntime.Spec.MLPolicy, trainingRuntime.Spec.PodGroupPolicy)
+	if err != nil {
+		return nil, err
+	}
+	return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
 }
 
-func (r *TrainingRuntime) buildObjects(
-	ctx context.Context, trainJob *trainer.TrainJob, jobSetTemplateSpec trainer.JobSetTemplateSpec, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
-) ([]any, error) {
+func (r *TrainingRuntime) RuntimeInfo(
+	trainJob *trainer.TrainJob, runtimeTemplateSpec any, mlPolicy *trainer.MLPolicy, podGroupPolicy *trainer.PodGroupPolicy,
+) (*runtime.Info, error) {
 
+	jobSetTemplateSpec, ok := runtimeTemplateSpec.(trainer.JobSetTemplateSpec)
+	if !ok {
+		return nil, fmt.Errorf("unsupported runtimeTemplateSpec")
+	}
 	info, err := r.newRuntimeInfo(trainJob, jobSetTemplateSpec, mlPolicy, podGroupPolicy)
 	if err != nil {
 		return nil, err
@@ -108,7 +116,7 @@ func (r *TrainingRuntime) buildObjects(
 		return nil, err
 	}
 
-	return r.framework.RunComponentBuilderPlugins(ctx, info, trainJob)
+	return info, nil
 }
 
 func (r *TrainingRuntime) newRuntimeInfo(

pkg/runtime/interface.go

@@ -37,6 +37,7 @@ type Runtime interface {
 	// https://github.com/kubernetes/kubernetes/pull/129313 becomes available
 
 	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)
 	EventHandlerRegistrars() []ReconcilerBuilder
 	ValidateObjects(ctx context.Context, old, new *trainer.TrainJob) (admission.Warnings, field.ErrorList)