AIP-7511 user defined @exit_handler

zillow · Nov 30, 2023 · ee87619 · ee87619
1 parent e75b38f
commit ee87619
Show file tree

Hide file tree

Showing 9 changed files with 506 additions and 81 deletions.
diff --git a/metaflow/plugins/__init__.py b/metaflow/plugins/__init__.py
@@ -159,12 +159,14 @@ def get_plugin_cli():
 from .aws.step_functions.schedule_decorator import ScheduleDecorator
 from .project_decorator import ProjectDecorator
 from .aip.s3_sensor_decorator import S3SensorDecorator
+from .aip.exit_handler_decorator import ExitHandlerDecorator
 
 FLOW_DECORATORS = [
     CondaFlowDecorator,
     ScheduleDecorator,
     ProjectDecorator,
     S3SensorDecorator,
+    ExitHandlerDecorator,
 ]
 _merge_lists(FLOW_DECORATORS, _ext_plugins["FLOW_DECORATORS"], "name")
 

diff --git a/metaflow/plugins/aip/__init__.py b/metaflow/plugins/aip/__init__.py
@@ -10,3 +10,8 @@
     delete_argo_workflow,
     to_metaflow_run_id,
 )
+
+from .exit_handler_decorator import (
+    exit_handler_resources,
+    exit_handler_retry,
+)