Transition to Queue if the JobCondition is empty (#387)

Signed-off-by: Kevin Su <[email protected]>

flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator.go

@@ -31,23 +31,6 @@ type ReplicaEntry struct {
 	RestartPolicy commonOp.RestartPolicy
 }
 
-// ExtractMPICurrentCondition will return the first job condition for MPI
-func ExtractMPICurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.JobCondition, error) {
-	if jobConditions != nil {
-		sort.Slice(jobConditions, func(i, j int) bool {
-			return jobConditions[i].LastTransitionTime.Time.After(jobConditions[j].LastTransitionTime.Time)
-		})
-
-		for _, jc := range jobConditions {
-			if jc.Status == v1.ConditionTrue {
-				return jc, nil
-			}
-		}
-	}
-
-	return commonOp.JobCondition{}, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions)
-}
-
 // ExtractCurrentCondition will return the first job condition for tensorflow/pytorch
 func ExtractCurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.JobCondition, error) {
 	if jobConditions != nil {
@@ -60,14 +43,17 @@ func ExtractCurrentCondition(jobConditions []commonOp.JobCondition) (commonOp.Jo
 				return jc, nil
 			}
 		}
+		return commonOp.JobCondition{}, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions)
 	}
-
-	return commonOp.JobCondition{}, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions)
+	return commonOp.JobCondition{}, nil
 }
 
 // GetPhaseInfo will return the phase of kubeflow job
 func GetPhaseInfo(currentCondition commonOp.JobCondition, occurredAt time.Time,
 	taskPhaseInfo pluginsCore.TaskInfo) (pluginsCore.PhaseInfo, error) {
+	if len(currentCondition.Type) == 0 {
+		return pluginsCore.PhaseInfoQueued(occurredAt, pluginsCore.DefaultPhaseVersion, "JobCreated"), nil
+	}
 	switch currentCondition.Type {
 	case commonOp.JobCreated:
 		return pluginsCore.PhaseInfoQueued(occurredAt, pluginsCore.DefaultPhaseVersion, "JobCreated"), nil

flyteplugins/go/tasks/plugins/k8s/kfoperators/common/common_operator_test.go

@@ -18,7 +18,7 @@ import (
 	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func TestExtractMPICurrentCondition(t *testing.T) {
+func TestExtractCurrentCondition(t *testing.T) {
 	jobCreated := commonOp.JobCondition{
 		Type:   commonOp.JobCreated,
 		Status: corev1.ConditionTrue,
@@ -31,46 +31,39 @@ func TestExtractMPICurrentCondition(t *testing.T) {
 		jobCreated,
 		jobRunningActive,
 	}
-	currentCondition, err := ExtractMPICurrentCondition(jobConditions)
+	currentCondition, err := ExtractCurrentCondition(jobConditions)
 	assert.NoError(t, err)
 	assert.Equal(t, currentCondition, jobCreated)
 
 	jobConditions = nil
-	currentCondition, err = ExtractMPICurrentCondition(jobConditions)
-	assert.Error(t, err)
+	currentCondition, err = ExtractCurrentCondition(jobConditions)
+	assert.NoError(t, err)
 	assert.Equal(t, currentCondition, commonOp.JobCondition{})
-	assert.Equal(t, err, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions))
-}
 
-func TestExtractCurrentCondition(t *testing.T) {
-	jobCreated := commonOp.JobCondition{
-		Type:   commonOp.JobCreated,
-		Status: corev1.ConditionTrue,
-	}
-	jobRunningActive := commonOp.JobCondition{
-		Type:   commonOp.JobRunning,
-		Status: corev1.ConditionFalse,
-	}
-	jobConditions := []commonOp.JobCondition{
-		jobCreated,
-		jobRunningActive,
-	}
-	currentCondition, err := ExtractCurrentCondition(jobConditions)
+	currentCondition, err = ExtractCurrentCondition(nil)
 	assert.NoError(t, err)
-	assert.Equal(t, currentCondition, jobCreated)
+	assert.Equal(t, currentCondition, commonOp.JobCondition{})
 
-	jobConditions = nil
+	jobUnknown := commonOp.JobCondition{Type: "unknown"}
+	jobConditions = []commonOp.JobCondition{jobUnknown}
 	currentCondition, err = ExtractCurrentCondition(jobConditions)
 	assert.Error(t, err)
 	assert.Equal(t, currentCondition, commonOp.JobCondition{})
 	assert.Equal(t, err, fmt.Errorf("found no current condition. Conditions: %+v", jobConditions))
 }
 
 func TestGetPhaseInfo(t *testing.T) {
+	jobCreating := commonOp.JobCondition{}
+	taskPhase, err := GetPhaseInfo(jobCreating, time.Now(), pluginsCore.TaskInfo{})
+	assert.NoError(t, err)
+	assert.Equal(t, pluginsCore.PhaseQueued, taskPhase.Phase())
+	assert.NotNil(t, taskPhase.Info())
+	assert.Nil(t, err)
+
 	jobCreated := commonOp.JobCondition{
 		Type: commonOp.JobCreated,
 	}
-	taskPhase, err := GetPhaseInfo(jobCreated, time.Now(), pluginsCore.TaskInfo{})
+	taskPhase, err = GetPhaseInfo(jobCreated, time.Now(), pluginsCore.TaskInfo{})
 	assert.NoError(t, err)
 	assert.Equal(t, pluginsCore.PhaseQueued, taskPhase.Phase())
 	assert.NotNil(t, taskPhase.Info())

flyteplugins/go/tasks/plugins/k8s/kfoperators/common/config.go

@@ -0,0 +1,32 @@
+package common
+
+import (
+	"time"
+
+	pluginsConfig "github.com/flyteorg/flyteplugins/go/tasks/config"
+	"github.com/flyteorg/flytestdlib/config"
+)
+
+//go:generate pflags Config --default-var=defaultConfig
+
+var (
+	defaultConfig = Config{
+		Timeout: config.Duration{Duration: 1 * time.Minute},
+	}
+
+	configSection = pluginsConfig.MustRegisterSubSection("kf-operator", &defaultConfig)
+)
+
+// Config is config for 'pytorch' plugin
+type Config struct {
+	// If kubeflow operator doesn't update the status of the task after this timeout, the task will be considered failed.
+	Timeout config.Duration `json:"timeout,omitempty"`
+}
+
+func GetConfig() *Config {
+	return configSection.GetConfig().(*Config)
+}
+
+func SetConfig(cfg *Config) error {
+	return configSection.SetConfig(cfg)
+}

flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi.go

@@ -209,7 +209,10 @@ func (mpiOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err
 	}
-	currentCondition, err := common.ExtractMPICurrentCondition(app.Status.Conditions)
+	if app.Status.StartTime == nil && app.CreationTimestamp.Add(common.GetConfig().Timeout.Duration).Before(time.Now()) {
+		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("kubeflow operator hasn't updated the mpi custom resource since creation time %v", app.CreationTimestamp)
+	}
+	currentCondition, err := common.ExtractCurrentCondition(app.Status.Conditions)
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err
 	}
@@ -223,7 +226,6 @@ func (mpiOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginContext
 	}
 
 	return common.GetMPIPhaseInfo(currentCondition, occurredAt, taskPhaseInfo)
-
 }
 
 func init() {

flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go

@@ -281,7 +281,7 @@ func dummyMPIJobResource(mpiResourceHandler mpiOperatorResourceHandler,
 		Status: mpiOp.JobStatus{
 			Conditions:        jobConditions,
 			ReplicaStatuses:   nil,
-			StartTime:         nil,
+			StartTime:         &v1.Time{Time: time.Now()},
 			CompletionTime:    nil,
 			LastReconcileTime: nil,
 		},

flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch.go

@@ -231,6 +231,9 @@ func (pytorchOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginCont
 		return pluginsCore.PhaseInfoUndefined, err
 	}
 
+	if app.Status.StartTime == nil && app.CreationTimestamp.Add(common.GetConfig().Timeout.Duration).Before(time.Now()) {
+		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("kubeflow operator hasn't updated the pytorch custom resource since creation time %v", app.CreationTimestamp)
+	}
 	currentCondition, err := common.ExtractCurrentCondition(app.Status.Conditions)
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err

flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go

@@ -293,8 +293,9 @@ func dummyPytorchJobResource(pytorchResourceHandler pytorchOperatorResourceHandl
 
 	return &kubeflowv1.PyTorchJob{
 		ObjectMeta: v1.ObjectMeta{
-			Name:      jobName,
-			Namespace: jobNamespace,
+			CreationTimestamp: v1.Time{Time: time.Now()},
+			Name:              jobName,
+			Namespace:         jobNamespace,
 		},
 		Spec: resource.(*kubeflowv1.PyTorchJob).Spec,
 		Status: commonOp.JobStatus{

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

@@ -209,6 +209,10 @@ func (tensorflowOperatorResourceHandler) GetTaskPhase(_ context.Context, pluginC
 		return pluginsCore.PhaseInfoUndefined, err
 	}
 
+	if app.Status.StartTime == nil && app.CreationTimestamp.Add(common.GetConfig().Timeout.Duration).Before(time.Now()) {
+		return pluginsCore.PhaseInfoUndefined, fmt.Errorf("kubeflow operator hasn't updated the tensorflow custom resource since creation time %v", app.CreationTimestamp)
+	}
+
 	currentCondition, err := common.ExtractCurrentCondition(app.Status.Conditions)
 	if err != nil {
 		return pluginsCore.PhaseInfoUndefined, err

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

@@ -283,7 +283,7 @@ func dummyTensorFlowJobResource(tensorflowResourceHandler tensorflowOperatorReso
 		Status: commonOp.JobStatus{
 			Conditions:        jobConditions,
 			ReplicaStatuses:   nil,
-			StartTime:         nil,
+			StartTime:         &v1.Time{Time: time.Now()},
 			CompletionTime:    nil,
 			LastReconcileTime: nil,
 		},