add on failure parameter to imperative workflow

flyteorg · Nov 6, 2024 · 0e5fba8 · 0e5fba8
1 parent 2fbdc63
commit 0e5fba8
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/flytekit/core/workflow.py b/flytekit/core/workflow.py
@@ -420,10 +420,11 @@ class ImperativeWorkflow(WorkflowBase):
     """
 
     def __init__(
-        self,
-        name: str,
-        failure_policy: Optional[WorkflowFailurePolicy] = None,
-        interruptible: bool = False,
+    self,
+    name: str,
+    failure_policy: Optional[WorkflowFailurePolicy] = None,
+    interruptible: bool = False,
+    on_failure: Optional[Union[WorkflowBase, Task]] = None,
     ):
         metadata = WorkflowMetadata(on_failure=failure_policy or WorkflowFailurePolicy.FAIL_IMMEDIATELY)
         workflow_metadata_defaults = WorkflowMetadataDefaults(interruptible)