Force the user to specify a PodTemplateSpec for PS as well

kubeflow · Jan 25, 2018 · 7864dc8 · 7864dc8
1 parent deb4fc2
commit 7864dc8
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 46 deletions.
diff --git a/pkg/apis/tensorflow/v1alpha1/defaults.go b/pkg/apis/tensorflow/v1alpha1/defaults.go
@@ -2,7 +2,6 @@ package v1alpha1
 
 import (
 	"github.com/golang/protobuf/proto"
-	"k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 )
 
@@ -32,11 +31,6 @@ func SetDefaults_TFJob(obj *TFJob) {
 		if r.Replicas == nil {
 			r.Replicas = proto.Int32(Replicas)
 		}
-
-		//Set the default configuration for a PS server if the user didn't specify a PodTemplateSpec
-		if r.Template == nil && r.TFReplicaType == PS {
-			setDefault_PSPodTemplateSpec(r, c.TFImage)
-		}
 	}
 	if c.TerminationPolicy == nil {
 		c.TerminationPolicy = &TerminationPolicySpec{
@@ -48,23 +42,3 @@ func SetDefaults_TFJob(obj *TFJob) {
 	}
 
 }
-
-func setDefault_PSPodTemplateSpec(r *TFReplicaSpec, tfImage string) {
-	r.Template = &v1.PodTemplateSpec{
-		Spec: v1.PodSpec{
-			Containers: []v1.Container{
-				v1.Container{
-					Image: tfImage,
-					Name:  "tensorflow",
-					VolumeMounts: []v1.VolumeMount{
-						v1.VolumeMount{
-							Name:      "ps-config-volume",
-							MountPath: "/ps-server",
-						},
-					},
-				},
-			},
-			RestartPolicy: v1.RestartPolicyOnFailure,
-		},
-	}
-}
diff --git a/pkg/apis/tensorflow/v1alpha1/defaults_test.go b/pkg/apis/tensorflow/v1alpha1/defaults_test.go
@@ -78,25 +78,8 @@ func TestSetDefaults_TFJob(t *testing.T) {
 				Spec: TFJobSpec{
 					ReplicaSpecs: []*TFReplicaSpec{
 						{
-							Replicas: proto.Int32(1),
-							TFPort:   proto.Int32(2222),
-							Template: &v1.PodTemplateSpec{
-								Spec: v1.PodSpec{
-									Containers: []v1.Container{
-										v1.Container{
-											Image: "tensorflow/tensorflow:1.3.0",
-											Name:  "tensorflow",
-											VolumeMounts: []v1.VolumeMount{
-												v1.VolumeMount{
-													Name:      "ps-config-volume",
-													MountPath: "/ps-server",
-												},
-											},
-										},
-									},
-									RestartPolicy: v1.RestartPolicyOnFailure,
-								},
-							},
+							Replicas:      proto.Int32(1),
+							TFPort:        proto.Int32(2222),
 							TFReplicaType: PS,
 						},
 					},

diff --git a/pkg/apis/tensorflow/validation/validation.go b/pkg/apis/tensorflow/validation/validation.go
@@ -19,7 +19,7 @@ func ValidateTFJobSpec(c *tfv1.TFJobSpec) error {
 	// Check that each replica has a TensorFlow container and a chief.
 	for _, r := range c.ReplicaSpecs {
 		found := false
-		if r.Template == nil && r.TFReplicaType != tfv1.PS {
+		if r.Template == nil {
 			return fmt.Errorf("Replica is missing Template; %v", util.Pformat(r))
 		}