Discard fields when replica number equals zero to avoid api client error

go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go

@@ -86,31 +86,27 @@ func (tensorflowOperatorResourceHandler) BuildResource(ctx context.Context, task
 	chiefReplicas := tensorflowTaskExtraArgs.GetChiefReplicas()
 
 	jobSpec := kubeflowv1.TFJobSpec{
-		TFReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{
-			kubeflowv1.TFJobReplicaTypePS: {
-				Replicas: &psReplicas,
-				Template: v1.PodTemplateSpec{
-					ObjectMeta: objectMeta,
-					Spec:       *podSpec,
-				},
-				RestartPolicy: commonOp.RestartPolicyNever,
-			},
-			kubeflowv1.TFJobReplicaTypeChief: {
-				Replicas: &chiefReplicas,
+		TFReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{},
+	}
+
+	for _, t := range []struct {
+		replicaNum  *int32
+		replicaType commonOp.ReplicaType
+	}{
+		{&workers, kubeflowv1.TFJobReplicaTypeWorker},
+		{&psReplicas, kubeflowv1.TFJobReplicaTypePS},
+		{&chiefReplicas, kubeflowv1.TFJobReplicaTypeChief},
+	} {
+		if *t.replicaNum > 0 {
+			jobSpec.TFReplicaSpecs[t.replicaType] = &commonOp.ReplicaSpec{
+				Replicas: t.replicaNum,
 				Template: v1.PodTemplateSpec{
 					ObjectMeta: objectMeta,
 					Spec:       *podSpec,
 				},
 				RestartPolicy: commonOp.RestartPolicyNever,
-			},
-			kubeflowv1.TFJobReplicaTypeWorker: {
-				Replicas: &workers,
-				Template: v1.PodTemplateSpec{
-					Spec: *podSpec,
-				},
-				RestartPolicy: commonOp.RestartPolicyNever,
-			},
-		},
+			}
+		}
 	}
 
 	job := &kubeflowv1.TFJob{

go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go

@@ -70,7 +70,7 @@ func dummyTensorFlowCustomObj(workers int32, psReplicas int32, chiefReplicas int
 	}
 }
 
-func dummySparkTaskTemplate(id string, tensorflowCustomObj *plugins.DistributedTensorflowTrainingTask) *core.TaskTemplate {
+func dummyTensorFlowTaskTemplate(id string, tensorflowCustomObj *plugins.DistributedTensorflowTrainingTask) *core.TaskTemplate {
 
 	tfObjJSON, err := utils.MarshalToString(tensorflowCustomObj)
 	if err != nil {
@@ -251,7 +251,7 @@ func dummyTensorFlowJobResource(tensorflowResourceHandler tensorflowOperatorReso
 	}
 
 	tfObj := dummyTensorFlowCustomObj(workers, psReplicas, chiefReplicas)
-	taskTemplate := dummySparkTaskTemplate("the job", tfObj)
+	taskTemplate := dummyTensorFlowTaskTemplate("the job", tfObj)
 	resource, err := tensorflowResourceHandler.BuildResource(context.TODO(), dummyTensorFlowTaskContext(taskTemplate))
 	if err != nil {
 		panic(err)
@@ -277,7 +277,7 @@ func TestBuildResourceTensorFlow(t *testing.T) {
 	tensorflowResourceHandler := tensorflowOperatorResourceHandler{}
 
 	tfObj := dummyTensorFlowCustomObj(100, 50, 1)
-	taskTemplate := dummySparkTaskTemplate("the job", tfObj)
+	taskTemplate := dummyTensorFlowTaskTemplate("the job", tfObj)
 
 	resource, err := tensorflowResourceHandler.BuildResource(context.TODO(), dummyTensorFlowTaskContext(taskTemplate))
 	assert.NoError(t, err)
@@ -371,3 +371,22 @@ func TestGetProperties(t *testing.T) {
 	expected := k8s.PluginProperties{}
 	assert.Equal(t, expected, tensorflowResourceHandler.GetProperties())
 }
+
+func TestZeroReplicas(t *testing.T) {
+	// if the number of replicas is zero, the field should not be created because the api client might complain.
+
+	tensorflowResourceHandler := tensorflowOperatorResourceHandler{}
+
+	tfObj := dummyTensorFlowCustomObj(10, 0, 0)
+	taskTemplate := dummyTensorFlowTaskTemplate("the job", tfObj)
+
+	resource, err := tensorflowResourceHandler.BuildResource(context.TODO(), dummyTensorFlowTaskContext(taskTemplate))
+	assert.NoError(t, err)
+	assert.NotNil(t, resource)
+
+	tensorflowJob, ok := resource.(*kubeflowv1.TFJob)
+	assert.True(t, ok)
+
+	assert.NotContains(t, tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypePS)
+	assert.NotContains(t, tensorflowJob.Spec.TFReplicaSpecs, kubeflowv1.TFJobReplicaTypeChief)
+}