*: Add cleanpod policy for v1alpha2 (#691)

* *: Support cleanpod policy Signed-off-by: Ce Gao <[email protected]> * Gopkg: Keep code-generator Signed-off-by: Ce Gao <[email protected]> * vendor: Update Signed-off-by: Ce Gao <[email protected]> * test: Fix test cases Signed-off-by: Ce Gao <[email protected]> * test: Fix Signed-off-by: Ce Gao <[email protected]> * test: Fix style Signed-off-by: Ce Gao <[email protected]> * defaults: Update Signed-off-by: Ce Gao <[email protected]>
kubeflow · Jun 21, 2018 · 38b886a · 38b886a
1 parent bfff36c
commit 38b886a
Show file tree

Hide file tree

Showing 599 changed files with 9,945 additions and 38,509 deletions.
diff --git a/Gopkg.toml b/Gopkg.toml
@@ -85,3 +85,8 @@ required = [
 [prune]
   go-tests = true
   unused-packages = true
+  non-go = true
+  [[prune.project]]
+    name = "k8s.io/code-generator"
+    unused-packages = false
+    non-go = false
diff --git a/pkg/apis/tensorflow/v1alpha1/zz_generated.deepcopy.go b/pkg/apis/tensorflow/v1alpha1/zz_generated.deepcopy.go
diff --git a/pkg/apis/tensorflow/v1alpha2/defaults.go b/pkg/apis/tensorflow/v1alpha2/defaults.go
@@ -88,9 +88,19 @@ func setTypeNameToCamelCase(tfJob *TFJob, typ TFReplicaType) {
 
 // SetDefaults_TFJob sets any unspecified values to defaults.
 func SetDefaults_TFJob(tfjob *TFJob) {
+	// Set default cleanpod policy to All.
+	if tfjob.Spec.CleanPodPolicy == nil {
+		running := CleanPodPolicyRunning
+		tfjob.Spec.CleanPodPolicy = &running
+	}
+
+	// Update the key of TFReplicaSpecs to camel case.
 	setTypeNamesToCamelCase(tfjob)
+
 	for _, spec := range tfjob.Spec.TFReplicaSpecs {
+		// Set default replicas to 1.
 		setDefaultReplicas(spec)
+		// Set default port to tensorFlow container.
 		setDefaultPort(&spec.Template.Spec)
 	}
 }
diff --git a/pkg/apis/tensorflow/v1alpha2/defaults_test.go b/pkg/apis/tensorflow/v1alpha2/defaults_test.go
@@ -27,7 +27,7 @@ const (
 	testImage = "test-image:latest"
 )
 
-func expectedTFJob(restartPolicy RestartPolicy, portName string, port int32) *TFJob {
+func expectedTFJob(cleanPodPolicy CleanPodPolicy, restartPolicy RestartPolicy, portName string, port int32) *TFJob {
 	ports := []v1.ContainerPort{}
 
 	// port not set
@@ -52,6 +52,7 @@ func expectedTFJob(restartPolicy RestartPolicy, portName string, port int32) *TF
 
 	return &TFJob{
 		Spec: TFJobSpec{
+			CleanPodPolicy: &cleanPodPolicy,
 			TFReplicaSpecs: map[TFReplicaType]*TFReplicaSpec{
 				TFReplicaTypeWorker: &TFReplicaSpec{
 					Replicas:      Int32(1),
@@ -146,7 +147,7 @@ func TestSetDefaultTFJob(t *testing.T) {
 					},
 				},
 			},
-			expected: expectedTFJob(customRestartPolicy, DefaultPortName, DefaultPort),
+			expected: expectedTFJob(CleanPodPolicyRunning, customRestartPolicy, DefaultPortName, DefaultPort),
 		},
 		"set replicas with default restartpolicy": {
 			original: &TFJob{
@@ -173,7 +174,7 @@ func TestSetDefaultTFJob(t *testing.T) {
 					},
 				},
 			},
-			expected: expectedTFJob(DefaultRestartPolicy, DefaultPortName, DefaultPort),
+			expected: expectedTFJob(CleanPodPolicyRunning, DefaultRestartPolicy, DefaultPortName, DefaultPort),
 		},
 		"set replicas with default port": {
 			original: &TFJob{
@@ -196,7 +197,7 @@ func TestSetDefaultTFJob(t *testing.T) {
 					},
 				},
 			},
-			expected: expectedTFJob(customRestartPolicy, "", 0),
+			expected: expectedTFJob(CleanPodPolicyRunning, customRestartPolicy, "", 0),
 		},
 		"set replicas adding default port": {
 			original: &TFJob{
@@ -225,7 +226,37 @@ func TestSetDefaultTFJob(t *testing.T) {
 					},
 				},
 			},
-			expected: expectedTFJob(customRestartPolicy, customPortName, customPort),
+			expected: expectedTFJob(CleanPodPolicyRunning, customRestartPolicy, customPortName, customPort),
+		},
+		"set custom cleanpod policy": {
+			original: &TFJob{
+				Spec: TFJobSpec{
+					CleanPodPolicy: cleanPodPolicyPointer(CleanPodPolicyAll),
+					TFReplicaSpecs: map[TFReplicaType]*TFReplicaSpec{
+						TFReplicaTypeWorker: &TFReplicaSpec{
+							Replicas:      Int32(1),
+							RestartPolicy: customRestartPolicy,
+							Template: v1.PodTemplateSpec{
+								Spec: v1.PodSpec{
+									Containers: []v1.Container{
+										v1.Container{
+											Name:  DefaultContainerName,
+											Image: testImage,
+											Ports: []v1.ContainerPort{
+												v1.ContainerPort{
+													Name:          customPortName,
+													ContainerPort: customPort,
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: expectedTFJob(CleanPodPolicyAll, customRestartPolicy, customPortName, customPort),
 		},
 	}
 
@@ -236,3 +267,8 @@ func TestSetDefaultTFJob(t *testing.T) {
 		}
 	}
 }
+
+func cleanPodPolicyPointer(cleanPodPolicy CleanPodPolicy) *CleanPodPolicy {
+	c := cleanPodPolicy
+	return &c
+}