Implement TrainingRuntime ReplicatedJob validation

Signed-off-by: Yuki Iwai <[email protected]>
kubeflow · Oct 16, 2024 · ba81beb · ba81beb
1 parent 498ff28
commit ba81beb
Show file tree

Hide file tree

Showing 15 changed files with 203 additions and 66 deletions.
diff --git a/pkg/runtime.v2/core/clustertrainingruntime_test.go b/pkg/runtime.v2/core/clustertrainingruntime_test.go
@@ -35,7 +35,7 @@ import (
 )
 
 func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
-	baseRuntime := testingutil.MakeClusterTrainingRuntimeWrapper(t, "test-runtime").
+	baseRuntime := testingutil.MakeClusterTrainingRuntimeWrapper("test-runtime").
 		Clone()
 
 	cases := map[string]struct {
@@ -45,17 +45,17 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 		wantError              error
 	}{
 		"succeeded to build JobSet and PodGroup": {
-			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
 				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.ClusterTrainingRuntimeKind), "test-runtime").
 				Trainer(
-					testingutil.MakeTrainJobTrainerWrapper(t).
+					testingutil.MakeTrainJobTrainerWrapper().
 						ContainerImage("test:trainjob").
 						Obj(),
 				).
 				Obj(),
 			clusterTrainingRuntime: baseRuntime.RuntimeSpec(
-				testingutil.MakeTrainingRuntimeSpecWrapper(t, baseRuntime.Spec).
+				testingutil.MakeTrainingRuntimeSpecWrapper(baseRuntime.Spec).
 					ContainerImage("test:runtime").
 					PodGroupPolicySchedulingTimeout(120).
 					MLPolicyNumNodes(20).
@@ -68,7 +68,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 					Obj(),
 			).Obj(),
 			wantObjs: []client.Object{
-				testingutil.MakeJobSetWrapper(t, metav1.NamespaceDefault, "test-job").
+				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
 					ContainerImage(ptr.To("test:trainjob")).
 					JobCompletionMode(batchv1.IndexedCompletion).
@@ -80,7 +80,7 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 					}).
 					ControllerReference(kubeflowv2.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
 					Obj(),
-				testingutil.MakeSchedulerPluginsPodGroup(t, metav1.NamespaceDefault, "test-job").
+				testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
 					ControllerReference(kubeflowv2.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
 					MinMember(40).
 					SchedulingTimeout(120).
@@ -91,11 +91,11 @@ func TestClusterTrainingRuntimeNewObjects(t *testing.T) {
 			},
 		},
 		"missing trainingRuntime resource": {
-			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
 				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.ClusterTrainingRuntimeKind), "test-runtime").
 				Trainer(
-					testingutil.MakeTrainJobTrainerWrapper(t).
+					testingutil.MakeTrainJobTrainerWrapper().
 						ContainerImage("test:trainjob").
 						Obj(),
 				).

diff --git a/pkg/runtime.v2/core/trainingruntime.go b/pkg/runtime.v2/core/trainingruntime.go
@@ -107,9 +107,6 @@ func (r *TrainingRuntime) buildObjects(
 		runtime.WithPodGroupPolicy(podGroupPolicy),
 	}
 	for idx, rJob := range jobSetTemplateSpec.Spec.ReplicatedJobs {
-		if rJob.Replicas == 0 {
-			jobSetTemplateSpec.Spec.ReplicatedJobs[idx].Replicas = 1
-		}
 		replicas := jobSetTemplateSpec.Spec.ReplicatedJobs[idx].Replicas * ptr.Deref(rJob.Template.Spec.Completions, 1)
 		opts = append(opts, runtime.WithPodSpecReplicas(rJob.Name, replicas, rJob.Template.Spec.Template.Spec))
 	}

diff --git a/pkg/runtime.v2/core/trainingruntime_test.go b/pkg/runtime.v2/core/trainingruntime_test.go
@@ -35,7 +35,7 @@ import (
 )
 
 func TestTrainingRuntimeNewObjects(t *testing.T) {
-	baseRuntime := testingutil.MakeTrainingRuntimeWrapper(t, metav1.NamespaceDefault, "test-runtime").
+	baseRuntime := testingutil.MakeTrainingRuntimeWrapper(metav1.NamespaceDefault, "test-runtime").
 		Clone()
 
 	cases := map[string]struct {
@@ -45,13 +45,13 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 		wantError       error
 	}{
 		"succeeded to build JobSet and PodGroup": {
-			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
 				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.TrainingRuntimeKind), "test-runtime").
 				SpecLabel("conflictLabel", "override").
 				SpecAnnotation("conflictAnnotation", "override").
 				Trainer(
-					testingutil.MakeTrainJobTrainerWrapper(t).
+					testingutil.MakeTrainJobTrainerWrapper().
 						ContainerImage("test:trainjob").
 						Obj(),
 				).
@@ -60,7 +60,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 				Label("conflictLabel", "overridden").
 				Annotation("conflictAnnotation", "overridden").
 				RuntimeSpec(
-					testingutil.MakeTrainingRuntimeSpecWrapper(t, baseRuntime.Spec).
+					testingutil.MakeTrainingRuntimeSpecWrapper(baseRuntime.Spec).
 						ContainerImage("test:runtime").
 						PodGroupPolicySchedulingTimeout(120).
 						MLPolicyNumNodes(20).
@@ -73,7 +73,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 						Obj(),
 				).Obj(),
 			wantObjs: []client.Object{
-				testingutil.MakeJobSetWrapper(t, metav1.NamespaceDefault, "test-job").
+				testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 					Label("conflictLabel", "override").
 					Annotation("conflictAnnotation", "override").
 					PodLabel(schedulerpluginsv1alpha1.PodGroupLabel, "test-job").
@@ -87,7 +87,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 					}).
 					ControllerReference(kubeflowv2.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
 					Obj(),
-				testingutil.MakeSchedulerPluginsPodGroup(t, metav1.NamespaceDefault, "test-job").
+				testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
 					ControllerReference(kubeflowv2.SchemeGroupVersion.WithKind("TrainJob"), "test-job", "uid").
 					MinMember(40).
 					SchedulingTimeout(120).
@@ -98,11 +98,11 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
 			},
 		},
 		"missing trainingRuntime resource": {
-			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
 				TrainingRuntimeRef(kubeflowv2.SchemeGroupVersion.WithKind(kubeflowv2.TrainingRuntimeKind), "test-runtime").
 				Trainer(
-					testingutil.MakeTrainJobTrainerWrapper(t).
+					testingutil.MakeTrainJobTrainerWrapper().
 						ContainerImage("test:trainjob").
 						Obj(),
 				).

diff --git a/pkg/runtime.v2/framework/core/framework.go b/pkg/runtime.v2/framework/core/framework.go
@@ -89,7 +89,7 @@ func (f *Framework) RunEnforcePodGroupPolicyPlugins(trainJob *kubeflowv2.TrainJo
 	return nil
 }
 
-func (f *Framework) RunCustomValidationPlugins(oldObj, newObj client.Object) (admission.Warnings, field.ErrorList) {
+func (f *Framework) RunCustomValidationPlugins(oldObj, newObj *kubeflowv2.TrainJob) (admission.Warnings, field.ErrorList) {
 	var aggregatedWarnings admission.Warnings
 	var aggregatedErrors field.ErrorList
 	for _, plugin := range f.customValidationPlugins {
@@ -101,9 +101,6 @@ func (f *Framework) RunCustomValidationPlugins(oldObj, newObj client.Object) (ad
 			aggregatedErrors = append(aggregatedErrors, errs...)
 		}
 	}
-	if len(aggregatedErrors) == 0 {
-		return aggregatedWarnings, nil
-	}
 	return aggregatedWarnings, aggregatedErrors
 }
 

diff --git a/pkg/runtime.v2/framework/core/framework_test.go b/pkg/runtime.v2/framework/core/framework_test.go
@@ -287,24 +287,21 @@ func TestRunEnforcePodGroupPolicyPlugins(t *testing.T) {
 
 func TestRunCustomValidationPlugins(t *testing.T) {
 	cases := map[string]struct {
-		trainJob     *kubeflowv2.TrainJob
 		registry     fwkplugins.Registry
-		oldObj       client.Object
-		newObj       client.Object
+		oldObj       *kubeflowv2.TrainJob
+		newObj       *kubeflowv2.TrainJob
 		wantWarnings admission.Warnings
 		wantError    field.ErrorList
 	}{
 		// Need to implement more detail testing after we implement custom validator in any plugins.
 		"there are not any custom validations": {
-			trainJob: &kubeflowv2.TrainJob{ObjectMeta: metav1.ObjectMeta{Name: "test-job", Namespace: metav1.NamespaceDefault}},
 			registry: fwkplugins.NewRegistry(),
-			oldObj:   testingutil.MakeTrainingRuntimeWrapper(t, metav1.NamespaceDefault, "test").Obj(),
-			newObj:   testingutil.MakeTrainingRuntimeWrapper(t, metav1.NamespaceDefault, "test").Obj(),
+			oldObj:   testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").Obj(),
+			newObj:   testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").Obj(),
 		},
 		"an empty registry": {
-			trainJob: &kubeflowv2.TrainJob{ObjectMeta: metav1.ObjectMeta{Name: "test-job", Namespace: metav1.NamespaceDefault}},
-			oldObj:   testingutil.MakeTrainingRuntimeWrapper(t, metav1.NamespaceDefault, "test").Obj(),
-			newObj:   testingutil.MakeTrainingRuntimeWrapper(t, metav1.NamespaceDefault, "test").Obj(),
+			oldObj: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").Obj(),
+			newObj: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").Obj(),
 		},
 	}
 	for name, tc := range cases {
@@ -329,7 +326,7 @@ func TestRunCustomValidationPlugins(t *testing.T) {
 }
 
 func TestRunComponentBuilderPlugins(t *testing.T) {
-	jobSetBase := testingutil.MakeJobSetWrapper(t, metav1.NamespaceDefault, "test-job").
+	jobSetBase := testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 		ResourceRequests(0, corev1.ResourceList{
 			corev1.ResourceCPU:    resource.MustParse("2"),
 			corev1.ResourceMemory: resource.MustParse("4Gi"),
@@ -354,10 +351,10 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 		wantObjs        []client.Object
 	}{
 		"coscheduling and jobset are performed": {
-			trainJob: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			trainJob: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				UID("uid").
 				Trainer(
-					testingutil.MakeTrainJobTrainerWrapper(t).
+					testingutil.MakeTrainJobTrainerWrapper().
 						ContainerImage("foo:bar").
 						Obj(),
 				).
@@ -396,7 +393,7 @@ func TestRunComponentBuilderPlugins(t *testing.T) {
 			},
 			registry: fwkplugins.NewRegistry(),
 			wantObjs: []client.Object{
-				testingutil.MakeSchedulerPluginsPodGroup(t, metav1.NamespaceDefault, "test-job").
+				testingutil.MakeSchedulerPluginsPodGroup(metav1.NamespaceDefault, "test-job").
 					SchedulingTimeout(300).
 					MinMember(20).
 					MinResources(corev1.ResourceList{

diff --git a/pkg/runtime.v2/framework/interface.go b/pkg/runtime.v2/framework/interface.go
@@ -48,7 +48,7 @@ type EnforceMLPolicyPlugin interface {
 
 type CustomValidationPlugin interface {
 	Plugin
-	Validate(oldObj, newObj client.Object) (admission.Warnings, field.ErrorList)
+	Validate(oldObj, newObj *kubeflowv2.TrainJob) (admission.Warnings, field.ErrorList)
 }
 
 type ComponentBuilderPlugin interface {

diff --git a/pkg/runtime.v2/framework/plugins/mpi/mpi.go b/pkg/runtime.v2/framework/plugins/mpi/mpi.go
@@ -23,6 +23,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
+	kubeflowv2 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v2alpha1"
 	runtime "github.com/kubeflow/training-operator/pkg/runtime.v2"
 	"github.com/kubeflow/training-operator/pkg/runtime.v2/framework"
 )
@@ -55,6 +56,6 @@ func (m *MPI) EnforceMLPolicy(info *runtime.Info) error {
 }
 
 // TODO: Need to implement validations for MPIJob.
-func (m *MPI) Validate(oldObj, newObj client.Object) (admission.Warnings, field.ErrorList) {
+func (m *MPI) Validate(oldObj, newObj *kubeflowv2.TrainJob) (admission.Warnings, field.ErrorList) {
 	return nil, nil
 }
diff --git a/pkg/runtime.v2/framework/plugins/torch/torch.go b/pkg/runtime.v2/framework/plugins/torch/torch.go
@@ -23,6 +23,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
+	kubeflowv2 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v2alpha1"
 	runtime "github.com/kubeflow/training-operator/pkg/runtime.v2"
 	"github.com/kubeflow/training-operator/pkg/runtime.v2/framework"
 )
@@ -51,6 +52,6 @@ func (t *Torch) EnforceMLPolicy(info *runtime.Info) error {
 }
 
 // TODO: Need to implement validateions for TorchJob.
-func (t *Torch) Validate(oldObj, newObj client.Object) (admission.Warnings, field.ErrorList) {
+func (t *Torch) Validate(oldObj, newObj *kubeflowv2.TrainJob) (admission.Warnings, field.ErrorList) {
 	return nil, nil
 }
diff --git a/pkg/runtime.v2/runtime_test.go b/pkg/runtime.v2/runtime_test.go
@@ -32,7 +32,7 @@ import (
 )
 
 func TestNewInfo(t *testing.T) {
-	jobSetBase := testingutil.MakeJobSetWrapper(t, metav1.NamespaceDefault, "test-job").
+	jobSetBase := testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 		Clone()
 
 	cases := map[string]struct {
@@ -159,7 +159,7 @@ func TestNewInfo(t *testing.T) {
 }
 
 func TestUpdate(t *testing.T) {
-	jobSetBase := testingutil.MakeJobSetWrapper(t, metav1.NamespaceDefault, "test-job").
+	jobSetBase := testingutil.MakeJobSetWrapper(metav1.NamespaceDefault, "test-job").
 		Clone()
 
 	cases := map[string]struct {
@@ -172,7 +172,7 @@ func TestUpdate(t *testing.T) {
 			info: &Info{
 				Obj: jobSetBase.Obj(),
 			},
-			obj: testingutil.MakeTrainJobWrapper(t, metav1.NamespaceDefault, "test-job").
+			obj: testingutil.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
 				Obj(),
 			wantInfo: &Info{
 				Obj: jobSetBase.Obj(),