Add validation for evaluator

Add validation for evaluator so that the number of evaluator should no more than one.
kubeflow · Dec 19, 2018 · c10b155 · c10b155
1 parent 451069f
commit c10b155
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 2 deletions.
diff --git a/pkg/apis/tensorflow/v1alpha2/util.go b/pkg/apis/tensorflow/v1alpha2/util.go
@@ -18,3 +18,8 @@ package v1alpha2
 func IsChieforMaster(typ TFReplicaType) bool {
 	return typ == TFReplicaTypeChief || typ == TFReplicaTypeMaster
 }
+
+// IsEvaluator returns true if the type is Evaluator.
+func IsEvaluator(typ TFReplicaType) bool {
+	return typ == TFReplicaTypeEval
+}
diff --git a/pkg/apis/tensorflow/v1beta1/util.go b/pkg/apis/tensorflow/v1beta1/util.go
@@ -19,7 +19,12 @@ func IsChieforMaster(typ TFReplicaType) bool {
 	return typ == TFReplicaTypeChief || typ == TFReplicaTypeMaster
 }
 
-// IsWorker returns true if the type is Worker
+// IsWorker returns true if the type is Worker.
 func IsWorker(typ TFReplicaType) bool {
 	return typ == TFReplicaTypeWorker
 }
+
+// IsEvaluator returns true if the type is Evaluator.
+func IsEvaluator(typ TFReplicaType) bool {
+	return typ == TFReplicaTypeEval
+}
diff --git a/pkg/apis/tensorflow/validation/validation.go b/pkg/apis/tensorflow/validation/validation.go
@@ -34,13 +34,17 @@ func validateBetaOneReplicaSpecs(specs map[tfv1beta1.TFReplicaType]*common.Repli
 		return fmt.Errorf("TFJobSpec is not valid")
 	}
 	foundChief := 0
+	var foundEvaluator int32 = 0
 	for rType, value := range specs {
 		if value == nil || len(value.Template.Spec.Containers) == 0 {
 			return fmt.Errorf("TFJobSpec is not valid")
 		}
 		if tfv1beta1.IsChieforMaster(rType) {
 			foundChief++
 		}
+		if tfv1beta1.IsEvaluator(rType) {
+			foundEvaluator = foundEvaluator + *value.Replicas
+		}
 		// Make sure the image is defined in the container.
 		numNamedTensorflow := 0
 		for _, container := range value.Template.Spec.Containers {
@@ -61,6 +65,9 @@ func validateBetaOneReplicaSpecs(specs map[tfv1beta1.TFReplicaType]*common.Repli
 	if foundChief > 1 {
 		return fmt.Errorf("More than 1 chief/master found")
 	}
+	if foundEvaluator > 1 {
+		return fmt.Errorf("More than 1 evaluator found")
+	}
 	return nil
 }
 
@@ -74,13 +81,17 @@ func validateAlphaTwoReplicaSpecs(specs map[tfv2.TFReplicaType]*tfv2.TFReplicaSp
 		return fmt.Errorf("TFJobSpec is not valid")
 	}
 	foundChief := 0
+	var foundEvaluator int32 = 0
 	for rType, value := range specs {
 		if value == nil || len(value.Template.Spec.Containers) == 0 {
 			return fmt.Errorf("TFJobSpec is not valid")
 		}
 		if tfv2.IsChieforMaster(rType) {
 			foundChief++
 		}
+		if tfv2.IsEvaluator(rType) {
+			foundEvaluator = foundEvaluator + *value.Replicas
+		}
 		// Make sure the image is defined in the container.
 		numNamedTensorflow := 0
 		for _, container := range value.Template.Spec.Containers {
@@ -101,5 +112,8 @@ func validateAlphaTwoReplicaSpecs(specs map[tfv2.TFReplicaType]*tfv2.TFReplicaSp
 	if foundChief > 1 {
 		return fmt.Errorf("More than 1 chief/master found")
 	}
+	if foundEvaluator > 1 {
+		return fmt.Errorf("More than 1 evaluator found")
+	}
 	return nil
 }
diff --git a/pkg/apis/tensorflow/validation/validation_test.go b/pkg/apis/tensorflow/validation/validation_test.go
@@ -17,7 +17,8 @@ package validation
 import (
 	"testing"
 
-        common "github.com/kubeflow/tf-operator/pkg/apis/common/v1beta1"
+	"github.com/golang/protobuf/proto"
+	common "github.com/kubeflow/tf-operator/pkg/apis/common/v1beta1"
 	tfv2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
 	tfv1beta1 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1beta1"
 
@@ -89,6 +90,23 @@ func TestValidateAlphaTwoTFJobSpec(t *testing.T) {
 				},
 			},
 		},
+		{
+			TFReplicaSpecs: map[tfv2.TFReplicaType]*tfv2.TFReplicaSpec{
+				tfv2.TFReplicaTypeEval: &tfv2.TFReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name:  "tensorflow",
+									Image: "kubeflow/tf-dist-mnist-test:1.0",
+								},
+							},
+						},
+					},
+					Replicas: proto.Int32(2),
+				},
+			},
+		},
 	}
 	for _, c := range testCases {
 		err := ValidateAlphaTwoTFJobSpec(&c)
@@ -163,6 +181,23 @@ func TestValidateBetaOneTFJobSpec(t *testing.T) {
 				},
 			},
 		},
+		{
+			TFReplicaSpecs: map[tfv1beta1.TFReplicaType]*common.ReplicaSpec{
+				tfv1beta1.TFReplicaTypeEval: &common.ReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name:  "tensorflow",
+									Image: "kubeflow/tf-dist-mnist-test:1.0",
+								},
+							},
+						},
+					},
+					Replicas: proto.Int32(2),
+				},
+			},
+		},
 	}
 	for _, c := range testCases {
 		err := ValidateBetaOneTFJobSpec(&c)