/*
Copyright 2024 The Kubeflow Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package core

import (
	"context"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/klog/v2/ktesting"
	"k8s.io/utils/ptr"
	jobsetv1alpha2 "sigs.k8s.io/jobset/api/jobset/v1alpha2"
	schedulerpluginsv1alpha1 "sigs.k8s.io/scheduler-plugins/apis/scheduling/v1alpha1"

	trainer "github.com/kubeflow/trainer/pkg/apis/trainer/v1alpha1"
	"github.com/kubeflow/trainer/pkg/constants"
	jobsetplgconsts "github.com/kubeflow/trainer/pkg/runtime/framework/plugins/jobset/constants"
	testingutil "github.com/kubeflow/trainer/pkg/util/testing"
)

func TestTrainingRuntimeNewObjects(t *testing.T) {
	resRequests := corev1.ResourceList{
		corev1.ResourceCPU: resource.MustParse("1"),
	}

	// TODO (andreyvelich): Add more test cases.
	cases := map[string]struct {
		trainingRuntime *trainer.TrainingRuntime
		trainJob        *trainer.TrainJob
		ObjCmpOpts      []cmp.Option
		wantObjs        []runtime.Object
		wantError       error
	}{
		// Test cases for the PlainML MLPolicy.
		"succeeded to build PodGroup and JobSet with NumNodes from the TrainJob and container from the Runtime.": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
				Label("conflictLabel", "overridden").
				Annotation("conflictAnnotation", "overridden").
				RuntimeSpec(
					testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
						WithMLPolicy(
							testingutil.MakeMLPolicyWrapper().
								WithNumNodes(100).
								Obj(),
						).
						PodGroupPolicyCoschedulingSchedulingTimeout(120).
						Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
						Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
						Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
						Obj(),
				).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				Suspend(true).
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				SpecLabel("conflictLabel", "override").
				SpecAnnotation("conflictAnnotation", "override").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						NumNodes(30).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Suspend(true).
					Label("conflictLabel", "override").
					Annotation("conflictAnnotation", "override").
					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
					NumNodes(30).
					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Obj(),
				testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					MinMember(32). // 32 replicas = 30 Trainer nodes + 2 Initializer.
					MinResources(corev1.ResourceList{
						// Trainer node has 30 CPUs + 2 CPUs from 2 initializer containers.
						corev1.ResourceCPU: resource.MustParse("32"),
					}).
					SchedulingTimeout(120).
					Obj(),
			},
		},
		"succeeded to build JobSet with NumNodes from the Runtime and container from the TrainJob.": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
					WithMLPolicy(
						testingutil.MakeMLPolicyWrapper().
							WithNumNodes(100).
							Obj(),
					).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  "TRAIN_JOB",
								Value: "original",
							},
							{
								Name:  "RUNTIME",
								Value: "test:runtime",
							},
						}...,
					).
					Obj(),
			).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
						Env(
							[]corev1.EnvVar{
								{
									Name:  "TRAIN_JOB",
									Value: "override",
								},
								{
									Name:  "TRAIN_JOB_CUSTOM",
									Value: "test:trainjob",
								},
							}...,
						).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
					NumNodes(100).
					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  "TRAIN_JOB",
								Value: "override",
							},
							{
								Name:  "TRAIN_JOB_CUSTOM",
								Value: "test:trainjob",
							},
							{
								Name:  "RUNTIME",
								Value: "test:runtime",
							},
						}...,
					).
					Obj(),
			},
		},
		"succeeded to build JobSet with dataset and model initializer from the TrainJob.": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
					WithMLPolicy(
						testingutil.MakeMLPolicyWrapper().
							WithNumNodes(100).
							Obj(),
					).
					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Obj(),
			).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				Initializer(
					testingutil.MakeTrainJobInitializerWrapper().
						DatasetInitializer(
							testingutil.MakeTrainJobDatasetInitializerWrapper().
								StorageUri("hf://trainjob-dataset").
								Env(
									[]corev1.EnvVar{
										{
											Name:  "TRAIN_JOB",
											Value: "test:trainjob:dataset",
										},
									}...,
								).
								SecretRef(corev1.LocalObjectReference{Name: "trainjob-secret-dataset"}).
								Obj(),
						).
						ModelInitializer(
							testingutil.MakeTrainJobModelInitializerWrapper().
								StorageUri("hf://trainjob-model").
								Env(
									[]corev1.EnvVar{
										{
											Name:  "TRAIN_JOB",
											Value: "test:trainjob:model",
										},
									}...,
								).
								SecretRef(corev1.LocalObjectReference{Name: "trainjob-secret-model"}).
								Obj(),
						).
						Obj(),
				).
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
					NumNodes(100).
					Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.ModelInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Env(constants.DatasetInitializer, constants.DatasetInitializer,
						[]corev1.EnvVar{
							{
								Name:  jobsetplgconsts.InitializerEnvStorageUri,
								Value: "hf://trainjob-dataset",
							},
							{
								Name:  "TRAIN_JOB",
								Value: "test:trainjob:dataset",
							},
						}...,
					).
					EnvFrom(constants.DatasetInitializer, constants.DatasetInitializer,
						[]corev1.EnvFromSource{
							{
								SecretRef: &corev1.SecretEnvSource{
									LocalObjectReference: corev1.LocalObjectReference{
										Name: "trainjob-secret-dataset",
									},
								},
							},
						}...,
					).
					Env(constants.ModelInitializer, constants.ModelInitializer,
						[]corev1.EnvVar{
							{
								Name:  jobsetplgconsts.InitializerEnvStorageUri,
								Value: "hf://trainjob-model",
							},
							{
								Name:  "TRAIN_JOB",
								Value: "test:trainjob:model",
							},
						}...,
					).
					EnvFrom(constants.ModelInitializer, constants.ModelInitializer,
						[]corev1.EnvFromSource{
							{
								SecretRef: &corev1.SecretEnvSource{
									LocalObjectReference: corev1.LocalObjectReference{
										Name: "trainjob-secret-model",
									},
								},
							},
						}...,
					).
					Obj(),
			},
		},
		// Test cases for the Torch MLPolicy.
		"succeeded to build JobSet with Torch values from the TrainJob": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
					WithMLPolicy(
						testingutil.MakeMLPolicyWrapper().
							WithNumNodes(100).
							WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
								TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
								Obj(),
							).
							Obj(),
					).
					JobSetSpec(
						testingutil.MakeJobSetWrapper("", "").
							DependsOn(constants.Node,
								[]jobsetv1alpha2.DependsOn{
									{
										Name:   constants.DatasetInitializer,
										Status: jobsetv1alpha2.DependencyComplete,
									},
									{
										Name:   constants.ModelInitializer,
										Status: jobsetv1alpha2.DependencyComplete,
									},
								}...,
							).
							Obj().
							Spec,
					).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Obj(),
			).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						NumNodes(30).
						NumProcPerNode(intstr.FromInt32(3)).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
					NumNodes(30).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  constants.TorchEnvNumNodes,
								Value: "30",
							},
							{
								Name:  constants.TorchEnvNumProcPerNode,
								Value: "3",
							},
							{
								Name: constants.TorchEnvNodeRank,
								ValueFrom: &corev1.EnvVarSource{
									FieldRef: &corev1.ObjectFieldSelector{
										FieldPath: constants.JobCompletionIndexFieldPath,
									},
								},
							},
							{
								Name:  constants.TorchEnvMasterAddr,
								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
							},
							{
								Name:  constants.TorchEnvMasterPort,
								Value: fmt.Sprintf("%d", constants.ContainerTrainerPort),
							},
						}...,
					).
					DependsOn(constants.Node,
						[]jobsetv1alpha2.DependsOn{
							{
								Name:   constants.DatasetInitializer,
								Status: jobsetv1alpha2.DependencyComplete,
							},
							{
								Name:   constants.ModelInitializer,
								Status: jobsetv1alpha2.DependencyComplete,
							},
						}...,
					).
					Obj(),
			},
		},
		"succeeded to build JobSet with Torch values from the Runtime and envs.": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").RuntimeSpec(
				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
					WithMLPolicy(
						testingutil.MakeMLPolicyWrapper().
							WithNumNodes(100).
							WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
								TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
								Obj(),
							).
							Obj(),
					).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  "TRAIN_JOB",
								Value: "original",
							},
							{
								Name:  "RUNTIME",
								Value: "test:runtime",
							},
						}...,
					).
					Obj(),
			).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
						Env(
							[]corev1.EnvVar{
								{
									Name:  "TRAIN_JOB",
									Value: "override",
								},
								{
									Name:  "TRAIN_JOB_CUSTOM",
									Value: "test:trainjob",
								},
							}...,
						).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
					NumNodes(100).
					Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
					ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  "TRAIN_JOB",
								Value: "override",
							},
							{
								Name:  "TRAIN_JOB_CUSTOM",
								Value: "test:trainjob",
							},
							{
								Name:  constants.TorchEnvNumNodes,
								Value: "100",
							},
							{
								Name:  constants.TorchEnvNumProcPerNode,
								Value: "1",
							},
							{
								Name: constants.TorchEnvNodeRank,
								ValueFrom: &corev1.EnvVarSource{
									FieldRef: &corev1.ObjectFieldSelector{
										FieldPath: constants.JobCompletionIndexFieldPath,
									},
								},
							},
							{
								Name:  constants.TorchEnvMasterAddr,
								Value: fmt.Sprintf("test-job-%s-0-0.test-job", constants.Node),
							},
							{
								Name:  constants.TorchEnvMasterPort,
								Value: fmt.Sprintf("%d", constants.ContainerTrainerPort),
							},
							{
								Name:  "RUNTIME",
								Value: "test:runtime",
							},
						}...,
					).
					Obj(),
			},
		},
		"succeeded to build JobSet with TorchTune values from the TrainJob": {
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "torchtune-llama3.3-70b").RuntimeSpec(
				testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "torchtune-llama3.3-70b").Spec).
					WithMLPolicy(
						testingutil.MakeMLPolicyWrapper().
							WithNumNodes(100).
							WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
								TorchPolicy(ptr.To(intstr.FromString("auto")), nil).
								Obj(),
							).
							Obj(),
					).
					JobSetSpec(
						testingutil.MakeJobSetWrapper("", "").
							DependsOn(constants.Node,
								[]jobsetv1alpha2.DependsOn{
									{
										Name:   constants.DatasetInitializer,
										Status: jobsetv1alpha2.DependencyComplete,
									},
									{
										Name:   constants.ModelInitializer,
										Status: jobsetv1alpha2.DependencyComplete,
									},
								}...,
							).
							Obj().
							Spec,
					).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Obj(),
			).Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "torchtune-llama3.3-70b").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						Container(
							"test:trainjob",
							[]string{"tune", "run"},
							[]string{
								"dtype=fp16",
								"batch_size=32",
								"epochs=10",
								"loss=torchtune.modules.loss.CEWithChunkedOutputLoss",
							},
							resRequests).
						NumNodes(30).
						NumProcPerNode(intstr.FromInt32(3)).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
					NumNodes(30).
					Container(
						constants.Node,
						constants.Node,
						"test:trainjob",
						[]string{
							"tune",
							"run",
							fmt.Sprintf("%s %s", constants.TorchTuneArgRdzvEndpoint, "test-job-node-0-0.test-job:29500"),
							constants.TorchTuneFullFinetuneDistributed,
							"--config llama3_3/70B_full_multinode.yaml",
						},
						[]string{
							"dtype=fp16",
							"batch_size=32",
							"epochs=10",
							"loss=torchtune.modules.loss.CEWithChunkedOutputLoss",
						},
						resRequests,
					).
					ContainerTrainerPorts([]corev1.ContainerPort{{ContainerPort: constants.ContainerTrainerPort}}).
					Env(constants.Node, constants.Node,
						[]corev1.EnvVar{
							{
								Name:  constants.TorchEnvNumNodes,
								Value: "30",
							},
							{
								Name:  constants.TorchEnvNumProcPerNode,
								Value: "3",
							},
							{
								Name: constants.TorchEnvNodeRank,
								ValueFrom: &corev1.EnvVarSource{
									FieldRef: &corev1.ObjectFieldSelector{
										FieldPath: constants.JobCompletionIndexFieldPath,
									},
								},
							},
						}...,
					).
					DependsOn(constants.Node,
						[]jobsetv1alpha2.DependsOn{
							{
								Name:   constants.DatasetInitializer,
								Status: jobsetv1alpha2.DependencyComplete,
							},
							{
								Name:   constants.ModelInitializer,
								Status: jobsetv1alpha2.DependencyComplete,
							},
						}...,
					).
					Obj(),
			},
		},
		"succeeded to build JobSet with OpenMPI values from the TrainJob": {
			ObjCmpOpts: cmp.Options{
				cmp.Comparer(testingutil.MPISecretDataComparer),
			},
			trainingRuntime: testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
				RuntimeSpec(
					testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").Spec).
						WithMLPolicy(
							testingutil.MakeMLPolicyWrapper().
								WithNumNodes(1).
								WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
									MPIPolicy(ptr.To[int32](1), ptr.To(trainer.MPIImplementationOpenMPI), ptr.To("/root/.ssh"), ptr.To(false)).
									Obj(),
								).
								Obj(),
						).
						LauncherReplica().
						Replicas(1, constants.Launcher).
						Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
						Obj(),
				).
				Obj(),
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						NumNodes(2).
						NumProcPerNode(intstr.FromInt32(8)).
						Obj(),
				).
				Obj(),
			wantObjs: []runtime.Object{
				testingutil.MakeConfigMapWrapper(fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix), metav1.NamespaceDefault).
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					WithData(map[string]string{
						constants.MPIHostfileName: `test-job-node-0-0.test-job slots=8
test-job-node-0-1.test-job slots=8
`,
					}).
					Obj(),
				testingutil.MakeSecretWrapper(fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix), metav1.NamespaceDefault).
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					WithImmutable(true).
					WithData(map[string][]byte{
						corev1.SSHAuthPrivateKey:  []byte("EXIST"),
						constants.MPISSHPublicKey: []byte("EXIST"),
					}).
					WithType(corev1.SecretTypeSSHAuth).
					Obj(),
				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
					ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), "test-job", "uid").
					LauncherReplica().
					Replicas(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Node, constants.Launcher).
					Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
					Completions(1, constants.DatasetInitializer, constants.ModelInitializer, constants.Launcher).
					NumNodes(2).
					Volumes(constants.Launcher,
						corev1.Volume{
							Name: constants.MPISSHAuthVolumeName,
							VolumeSource: corev1.VolumeSource{
								Secret: &corev1.SecretVolumeSource{
									SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
									Items: []corev1.KeyToPath{
										{
											Key:  corev1.SSHAuthPrivateKey,
											Path: constants.MPISSHPrivateKeyFile,
										},
										{
											Key:  constants.MPISSHPublicKey,
											Path: constants.MPISSHPublicKeyFile,
										},
										{
											Key:  constants.MPISSHPublicKey,
											Path: constants.MPISSHAuthorizedKeys,
										},
									},
								},
							},
						},
						corev1.Volume{
							Name: constants.MPIHostfileVolumeName,
							VolumeSource: corev1.VolumeSource{
								ConfigMap: &corev1.ConfigMapVolumeSource{
									LocalObjectReference: corev1.LocalObjectReference{
										Name: fmt.Sprintf("test-job%s", constants.MPIHostfileConfigMapSuffix),
									},
									Items: []corev1.KeyToPath{{
										Key:  constants.MPIHostfileName,
										Path: constants.MPIHostfileName,
										Mode: ptr.To[int32](0444),
									}},
								},
							},
						},
					).
					Volumes(constants.Node,
						corev1.Volume{
							Name: constants.MPISSHAuthVolumeName,
							VolumeSource: corev1.VolumeSource{
								Secret: &corev1.SecretVolumeSource{
									SecretName: fmt.Sprintf("test-job%s", constants.MPISSHAuthSecretSuffix),
									Items: []corev1.KeyToPath{
										{
											Key:  corev1.SSHAuthPrivateKey,
											Path: constants.MPISSHPrivateKeyFile,
										},
										{
											Key:  constants.MPISSHPublicKey,
											Path: constants.MPISSHPublicKeyFile,
										},
										{
											Key:  constants.MPISSHPublicKey,
											Path: constants.MPISSHAuthorizedKeys,
										},
									},
								},
							},
						},
					).
					VolumeMounts(constants.Launcher, constants.Node,
						corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
						corev1.VolumeMount{Name: constants.MPIHostfileVolumeName, MountPath: constants.MPIHostfileDir},
					).
					VolumeMounts(constants.Node, constants.Node,
						corev1.VolumeMount{Name: constants.MPISSHAuthVolumeName, MountPath: "/root/.ssh"},
					).
					Env(constants.Launcher, constants.Node,
						corev1.EnvVar{
							Name:  constants.OpenMPIEnvHostFileLocation,
							Value: fmt.Sprintf("%s/%s", constants.MPIHostfileDir, constants.MPIHostfileName),
						},
						corev1.EnvVar{
							Name:  constants.OpenMPIEnvKeepFQDNHostNames,
							Value: "true",
						},
						corev1.EnvVar{
							Name:  constants.OpenMPIEnvDefaultSlots,
							Value: "8",
						},
						corev1.EnvVar{
							Name:  constants.OpenMPIEnvKeyRSHArgs,
							Value: constants.OpenMPIEnvDefaultValueRSHArgs,
						},
					).
					Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
					Obj(),
			},
		},
		// Failed test cases.
		"missing trainingRuntime resource": {
			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job-3").
				UID("uid").
				RuntimeRef(trainer.SchemeGroupVersion.WithKind(trainer.TrainingRuntimeKind), "test-runtime-3").
				Trainer(
					testingutil.MakeTrainJobTrainerWrapper().
						Obj(),
				).
				Obj(),
			wantError: errorNotFoundSpecifiedTrainingRuntime,
		},
	}
	cmpOpts := []cmp.Option{
		cmpopts.SortSlices(func(a, b runtime.Object) bool {
			return a.GetObjectKind().GroupVersionKind().String() < b.GetObjectKind().GroupVersionKind().String()
		}),
		cmpopts.SortSlices(func(a, b corev1.EnvVar) bool {
			return a.Name < b.Name
		}),
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			_, ctx := ktesting.NewTestContext(t)
			var cancel func()
			ctx, cancel = context.WithCancel(ctx)
			t.Cleanup(cancel)
			clientBuilder := testingutil.NewClientBuilder()
			if tc.trainingRuntime != nil {
				clientBuilder.WithObjects(tc.trainingRuntime)
			}
			c := clientBuilder.Build()

			trainingRuntime, err := NewTrainingRuntime(ctx, c, testingutil.AsIndex(clientBuilder))
			if err != nil {
				t.Fatal(err)
			}

			objs, err := trainingRuntime.NewObjects(ctx, tc.trainJob)
			if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
				t.Errorf("Unexpected error (-want,+got):\n%s", diff)
			}

			resultObjs, err := testingutil.ToObject(c.Scheme(), objs...)
			if err != nil {
				t.Errorf("Pipeline built unrecognizable objects: %v", err)
			}

			if diff := cmp.Diff(tc.wantObjs, resultObjs, append(cmpOpts, tc.ObjCmpOpts...)...); len(diff) != 0 {
				t.Errorf("Unexpected objects (-want,+got):\n%s", diff)
			}
		})
	}
}