fix(databricks): Handle FAILED state as retryable error (#5277)

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Apr 29, 2024 · f32132b · f32132b
1 parent 95e9ac8
commit f32132b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/flyteplugins/go/tasks/plugins/webapi/databricks/plugin.go b/flyteplugins/go/tasks/plugins/webapi/databricks/plugin.go
@@ -282,12 +282,14 @@ func (p Plugin) Status(ctx context.Context, taskCtx webapi.StatusContext) (phase
 				return core.PhaseInfoFailure(string(rune(http.StatusInternalServerError)), "failed to write output", taskInfo), nil
 			}
 			return core.PhaseInfoSuccess(taskInfo), nil
+		} else if resultState == "FAILED" {
+			return core.PhaseInfoRetryableFailure("job failed", message, taskInfo), nil
 		}
 		return core.PhaseInfoFailure(pluginErrors.TaskFailedWithError, message, taskInfo), nil
 	case "SKIPPED":
 		return core.PhaseInfoFailure(string(rune(http.StatusConflict)), message, taskInfo), nil
 	case "INTERNAL_ERROR":
-		return core.PhaseInfoFailure(string(rune(http.StatusInternalServerError)), message, taskInfo), nil
+		return core.PhaseInfoRetryableFailure(string(rune(http.StatusInternalServerError)), message, taskInfo), nil
 	}
 	return core.PhaseInfoUndefined, pluginErrors.Errorf(pluginsCore.SystemErrorCode, "unknown execution phase [%v].", lifeCycleState)
 }