Remove chief!=0 or ps!=0 assertion in tensorflow (flyteorg#327)

Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]>
eapolinario · Mar 6, 2023 · be833e2 · be833e2
1 parent 114645a
commit be833e2
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 4 deletions.
diff --git a/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go b/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow.go
@@ -75,9 +75,6 @@ func (tensorflowOperatorResourceHandler) BuildResource(ctx context.Context, task
 	if workers == 0 {
 		return nil, fmt.Errorf("number of worker should be more then 0")
 	}
-	if psReplicas == 0 && chiefReplicas == 0 {
-		return nil, fmt.Errorf("either number of chief or parameter servers needs to be be more then 0")
-	}
 
 	jobSpec := kubeflowv1.TFJobSpec{
 		TFReplicaSpecs: map[commonOp.ReplicaType]*commonOp.ReplicaSpec{},

diff --git a/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go b/go/tasks/plugins/k8s/kfoperators/tensorflow/tensorflow_test.go
@@ -383,7 +383,6 @@ func TestReplicaCounts(t *testing.T) {
 		notContains        []commonOp.ReplicaType
 	}{
 		{"NoWorkers", 1, 1, 0, true, nil, nil},
-		{"NoChiefOrPS", 0, 0, 1, true, nil, nil},
 		{"SingleChief", 1, 0, 1, false,
 			[]commonOp.ReplicaType{kubeflowv1.TFJobReplicaTypeChief, kubeflowv1.TFJobReplicaTypeWorker},
 			[]commonOp.ReplicaType{kubeflowv1.TFJobReplicaTypePS}},