status: Fail the TFJob if PS is failed (kubeflow#690)

* status: Fail the TFJob if PS is failed Signed-off-by: Ce Gao <[email protected]> * status: Fix Signed-off-by: Ce Gao <[email protected]>
jetmuffin · Jun 21, 2018 · f9b9551 · f9b9551
1 parent 5e3bd2b
commit f9b9551
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 42 deletions.
diff --git a/pkg/controller.v2/controller_status.go b/pkg/controller.v2/controller_status.go
@@ -77,7 +77,7 @@ func updateStatusSingle(tfjob *tfv1alpha2.TFJob, rtype tfv1alpha2.TFReplicaType,
 	failed := int(tfjob.Status.TFReplicaStatuses[rtype].Failed)
 
 	// All workers are running, set StartTime.
-	if running == replicas {
+	if running == replicas && tfjob.Status.StartTime == nil {
 		now := metav1.Now()
 		tfjob.Status.StartTime = &now
 	}
@@ -102,25 +102,6 @@ func updateStatusSingle(tfjob *tfv1alpha2.TFJob, rtype tfv1alpha2.TFReplicaType,
 					return err
 				}
 			}
-
-			if failed > 0 {
-				if restart {
-					msg := fmt.Sprintf("TFJob %s is restarting.", tfjob.Name)
-					err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobRestarting, tfJobRestartingReason, msg)
-					if err != nil {
-						loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
-						return err
-					}
-				} else {
-					msg := fmt.Sprintf("TFJob %s is failed.", tfjob.Name)
-					err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobFailed, tfJobFailedReason, msg)
-					if err != nil {
-						loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
-						return err
-					}
-				}
-			}
-			return nil
 		}
 	} else {
 		if rtype == tfv1alpha2.TFReplicaTypeWorker {
@@ -145,26 +126,24 @@ func updateStatusSingle(tfjob *tfv1alpha2.TFJob, rtype tfv1alpha2.TFReplicaType,
 					return err
 				}
 			}
+		}
+	}
 
-			// Some workers or pss are failed , leave a failed condition.
-			if failed > 0 {
-				if restart {
-					msg := fmt.Sprintf("TFJob %s is restarting.", tfjob.Name)
-					err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobRestarting, tfJobRestartingReason, msg)
-					if err != nil {
-						loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
-						return err
-					}
-				} else {
-					msg := fmt.Sprintf("TFJob %s is failed.", tfjob.Name)
-					err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobFailed, tfJobFailedReason, msg)
-					if err != nil {
-						loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
-						return err
-					}
-				}
+	if failed > 0 {
+		if restart {
+			msg := fmt.Sprintf("TFJob %s is restarting.", tfjob.Name)
+			err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobRestarting, tfJobRestartingReason, msg)
+			if err != nil {
+				loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
+				return err
+			}
+		} else {
+			msg := fmt.Sprintf("TFJob %s is failed.", tfjob.Name)
+			err := updateTFJobConditions(tfjob, tfv1alpha2.TFJobFailed, tfJobFailedReason, msg)
+			if err != nil {
+				loggerForTFJob(tfjob).Infof("Append tfjob condition error: %v", err)
+				return err
 			}
-			return nil
 		}
 	}
 	return nil

diff --git a/pkg/controller.v2/controller_status_test.go b/pkg/controller.v2/controller_status_test.go
@@ -237,6 +237,21 @@ func TestStatus(t *testing.T) {
 			restart:                 false,
 			expectedType:            tfv1alpha2.TFJobRunning,
 		},
+		testCase{
+			description:             "Chief is running, a PS is failed",
+			tfJob:                   testutil.NewTFJobWithChief(4, 2),
+			expectedFailedPS:        1,
+			expectedSucceededPS:     0,
+			expectedActivePS:        1,
+			expectedFailedWorker:    0,
+			expectedSucceededWorker: 4,
+			expectedActiveWorker:    0,
+			expectedFailedChief:     0,
+			expectedSucceededChief:  0,
+			expectedActiveChief:     1,
+			restart:                 false,
+			expectedType:            tfv1alpha2.TFJobFailed,
+		},
 		testCase{
 			description:             "Chief is failed, workers are succeeded",
 			tfJob:                   testutil.NewTFJobWithChief(4, 2),
@@ -298,11 +313,34 @@ func TestStatus(t *testing.T) {
 			if err != nil {
 				t.Errorf("%s: Expected error %v to be nil", c.description, err)
 			}
+			if c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker] != nil {
+				replicas := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Replicas
+				err := updateStatusSingle(c.tfJob, tfv1alpha2.TFReplicaTypeWorker, int(*replicas), c.restart)
+				if err != nil {
+					t.Errorf("%s: Expected error %v to be nil", c.description, err)
+				}
+			}
+			if c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS] != nil {
+				replicas := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS].Replicas
+				err := updateStatusSingle(c.tfJob, tfv1alpha2.TFReplicaTypePS, int(*replicas), c.restart)
+				if err != nil {
+					t.Errorf("%s: Expected error %v to be nil", c.description, err)
+				}
+			}
 		} else {
-			replicas := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Replicas
-			err := updateStatusSingle(c.tfJob, tfv1alpha2.TFReplicaTypeWorker, int(*replicas), c.restart)
-			if err != nil {
-				t.Errorf("%s: Expected error %v to be nil", c.description, err)
+			if c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker] != nil {
+				replicas := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Replicas
+				err := updateStatusSingle(c.tfJob, tfv1alpha2.TFReplicaTypeWorker, int(*replicas), c.restart)
+				if err != nil {
+					t.Errorf("%s: Expected error %v to be nil", c.description, err)
+				}
+			}
+			if c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS] != nil {
+				replicas := c.tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS].Replicas
+				err := updateStatusSingle(c.tfJob, tfv1alpha2.TFReplicaTypePS, int(*replicas), c.restart)
+				if err != nil {
+					t.Errorf("%s: Expected error %v to be nil", c.description, err)
+				}
 			}
 		}
 		found := false