small code cleanup

Author: talebzeghmi

metaflow/plugins/aip/aip.py

@@ -382,15 +382,17 @@ def _create_workflow_yaml(
                     "template": "user-defined-exit-handler",
                 }
 
-                if udf_handler.attributes.get(
-                    "on_success", True
-                ) and udf_handler.attributes.get("on_failure", True):
-                    # always run
+                on_success = udf_handler.attributes.get("on_success", True)
+                on_failure = udf_handler.attributes.get("on_failure", True)
+                if on_success and on_failure:
+                    # always run, no condition
                     pass
-                elif udf_handler.attributes.get("on_success", True):
+                elif on_success:
                     udf_task["when"] = "{{workflow.status}} == 'Succeeded'"
-                else:
+                elif on_failure:
                     udf_task["when"] = "{{workflow.status}} != 'Succeeded'"
+                else:
+                    raise AIPException("on_success and on_failure cannot both be False")
 
                 exit_handler_template["dag"]["tasks"].append(udf_task)
 

metaflow/plugins/aip/exit_handler_decorator.py

@@ -96,7 +96,7 @@ class ExitHandlerDecorator(FlowDecorator):
     >>> @exit_handler_resources(memory="2G")
     >>> def my_exit_handler(
     >>>     status: str,
-    >>>     flow_parameters: Dict[str, str],
+    >>>     flow_parameters: dict,
     >>>     argo_workflow_run_name: str,
     >>>     metaflow_run_id: str,
     >>>     argo_ui_url: str,