Validation test for V1beta1 apis

pkg/apis/pytorch/validation/validation.go

@@ -22,6 +22,7 @@ import (
 
 	torchv1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1alpha1"
 	torchv2 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1alpha2"
+	torchv1beta1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1beta1"
 	"github.com/kubeflow/pytorch-operator/pkg/util"
 )
 
@@ -137,3 +138,61 @@ func ValidateAlphaTwoPyTorchJobSpec(c *torchv2.PyTorchJobSpec) error {
 	}
 	return nil
 }
+
+func ValidateBetaOnePyTorchJobSpec(c *torchv1beta1.PyTorchJobSpec) error {
+	if c.PyTorchReplicaSpecs == nil {
+		return fmt.Errorf("PyTorchJobSpec is not valid")
+	}
+	masterExists := false
+	for rType, value := range c.PyTorchReplicaSpecs {
+		if value == nil || len(value.Template.Spec.Containers) == 0 {
+			return fmt.Errorf("PyTorchJobSpec is not valid")
+		}
+		// Make sure the replica type is valid.
+		validReplicaTypes := []torchv1beta1.PyTorchReplicaType{torchv1beta1.PyTorchReplicaTypeMaster, torchv1beta1.PyTorchReplicaTypeWorker}
+
+		isValidReplicaType := false
+		for _, t := range validReplicaTypes {
+			if t == rType {
+				isValidReplicaType = true
+				break
+			}
+		}
+
+		if !isValidReplicaType {
+			return fmt.Errorf("PyTorchReplicaType is %v but must be one of %v", rType, validReplicaTypes)
+		}
+
+		//Make sure the image is defined in the container
+		defaultContainerPresent := false
+		for _, container := range value.Template.Spec.Containers {
+			if container.Image == "" {
+				log.Warn("Image is undefined in the container")
+				return fmt.Errorf("PyTorchJobSpec is not valid")
+			}
+			if container.Name == torchv1beta1.DefaultContainerName {
+				defaultContainerPresent = true
+			}
+		}
+		//Make sure there has at least one container named "pytorch"
+		if !defaultContainerPresent {
+			log.Warnf("There is no container named pytorch in %v", rType)
+			return fmt.Errorf("PyTorchJobSpec is not valid")
+		}
+		if rType == torchv1beta1.PyTorchReplicaTypeMaster {
+			masterExists = true
+			if value.Replicas != nil && int(*value.Replicas) != 1 {
+				log.Warnf("There must be only 1 master replica")
+				return fmt.Errorf("PyTorchJobSpec is not valid")
+			}
+		}
+
+	}
+
+	if !masterExists {
+		log.Warnf("Master ReplicaSpec must be present")
+		return fmt.Errorf("PyTorchJobSpec is not valid")
+	}
+	return nil
+
+}

pkg/apis/pytorch/validation/validation_test.go

@@ -19,11 +19,103 @@ import (
 
 	torchv1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1alpha1"
 	torchv2 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1alpha2"
+	torchv1beta1 "github.com/kubeflow/pytorch-operator/pkg/apis/pytorch/v1beta1"
+	common "github.com/kubeflow/tf-operator/pkg/apis/common/v1beta1"
 
 	"github.com/gogo/protobuf/proto"
 	"k8s.io/api/core/v1"
 )
 
+func TestValidateBetaOnePyTorchJobSpec(t *testing.T) {
+	testCases := []torchv1beta1.PyTorchJobSpec{
+		{
+			PyTorchReplicaSpecs: nil,
+		},
+		{
+			PyTorchReplicaSpecs: map[torchv1beta1.PyTorchReplicaType]*common.ReplicaSpec{
+				torchv1beta1.PyTorchReplicaTypeWorker: &common.ReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{},
+						},
+					},
+				},
+			},
+		},
+		{
+			PyTorchReplicaSpecs: map[torchv1beta1.PyTorchReplicaType]*common.ReplicaSpec{
+				torchv1beta1.PyTorchReplicaTypeWorker: &common.ReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Image: "",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			PyTorchReplicaSpecs: map[torchv1beta1.PyTorchReplicaType]*common.ReplicaSpec{
+				torchv1beta1.PyTorchReplicaTypeWorker: &common.ReplicaSpec{
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name:  "",
+									Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			PyTorchReplicaSpecs: map[torchv1beta1.PyTorchReplicaType]*common.ReplicaSpec{
+				torchv1beta1.PyTorchReplicaTypeMaster: &common.ReplicaSpec{
+					Replicas: torchv1beta1.Int32(2),
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name:  "pytorch",
+									Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			PyTorchReplicaSpecs: map[torchv1beta1.PyTorchReplicaType]*common.ReplicaSpec{
+				torchv1beta1.PyTorchReplicaTypeWorker: &common.ReplicaSpec{
+					Replicas: torchv1beta1.Int32(1),
+					Template: v1.PodTemplateSpec{
+						Spec: v1.PodSpec{
+							Containers: []v1.Container{
+								v1.Container{
+									Name:  "pytorch",
+									Image: "gcr.io/kubeflow-ci/pytorch-dist-mnist_test:1.0",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	for _, c := range testCases {
+		err := ValidateBetaOnePyTorchJobSpec(&c)
+		if err.Error() != "PyTorchJobSpec is not valid" {
+			t.Error("Failed validate the v1beta1.PyTorchJobSpec")
+		}
+	}
+}
+
 func TestValidateAlphaTwoPyTorchJobSpec(t *testing.T) {
 	testCases := []torchv2.PyTorchJobSpec{
 		{