Fix initial dask job state (#357)

* Fix initial dask state Signed-off-by: Bernhard Stadlbauer <11799671+bstadlbauer@users.noreply.github.com> * Fix linting Signed-off-by: Bernhard Stadlbauer <11799671+bstadlbauer@users.noreply.github.com> --------- Signed-off-by: Bernhard Stadlbauer <11799671+bstadlbauer@users.noreply.github.com>
flyteorg · Jun 8, 2023 · 2a7a6f9 · 2a7a6f9
1 parent aa1ed67
commit 2a7a6f9
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/go/tasks/plugins/k8s/dask/dask.go b/go/tasks/plugins/k8s/dask/dask.go
@@ -320,7 +320,10 @@ func (p daskResourceHandler) GetTaskPhase(ctx context.Context, pluginContext k8s
 		OccurredAt: &occurredAt,
 	}
 
-	isQueued := status == daskAPI.DaskJobCreated ||
+	// There is a short period between the `DaskJob` resource being created and `Status.JobStatus` being set by the `dask-operator`.
+	// In that period, the `JobStatus` will be an empty string. We're treating this as Initializing/Queuing.
+	isQueued := status == "" ||
+		status == daskAPI.DaskJobCreated ||
 		status == daskAPI.DaskJobClusterCreated
 
 	if !isQueued {
@@ -337,6 +340,8 @@ func (p daskResourceHandler) GetTaskPhase(ctx context.Context, pluginContext k8s
 	}
 
 	switch status {
+	case "":
+		return pluginsCore.PhaseInfoInitializing(occurredAt, pluginsCore.DefaultPhaseVersion, "unknown", &info), nil
 	case daskAPI.DaskJobCreated:
 		return pluginsCore.PhaseInfoInitializing(occurredAt, pluginsCore.DefaultPhaseVersion, "job created", &info), nil
 	case daskAPI.DaskJobClusterCreated:

diff --git a/go/tasks/plugins/k8s/dask/dask_test.go b/go/tasks/plugins/k8s/dask/dask_test.go
@@ -491,7 +491,14 @@ func TestGetTaskPhaseDask(t *testing.T) {
 	daskResourceHandler := daskResourceHandler{}
 	ctx := context.TODO()
 
-	taskPhase, err := daskResourceHandler.GetTaskPhase(ctx, nil, dummyDaskJob(daskAPI.DaskJobCreated))
+	taskPhase, err := daskResourceHandler.GetTaskPhase(ctx, nil, dummyDaskJob(""))
+	assert.NoError(t, err)
+	assert.Equal(t, taskPhase.Phase(), pluginsCore.PhaseInitializing)
+	assert.NotNil(t, taskPhase.Info())
+	assert.Nil(t, taskPhase.Info().Logs)
+	assert.Nil(t, err)
+
+	taskPhase, err = daskResourceHandler.GetTaskPhase(ctx, nil, dummyDaskJob(daskAPI.DaskJobCreated))
 	assert.NoError(t, err)
 	assert.Equal(t, taskPhase.Phase(), pluginsCore.PhaseInitializing)
 	assert.NotNil(t, taskPhase.Info())