send {{workflow.failures}} to exit_handler

.gitlab-ci.yml

@@ -1,6 +1,7 @@
 include:
   - project: 'analytics/artificial-intelligence/ai-platform/aip-infrastructure/ci-templates/ci-cd-template'
-    ref: &include_ref 'v4'
+#    ref: &include_ref 'v4'
+    ref: &include_ref '4e25c7200369aec17fcf8de723db360a7726151a'
     file: 'environments/devex.yml'
   - project: 'analytics/artificial-intelligence/ai-platform/aip-infrastructure/ci-templates/ci-cd-template'
     ref: *include_ref

metaflow/plugins/aip/aip.py

@@ -1811,6 +1811,7 @@ def _get_user_defined_exit_handler_op(
                 "  --status {{workflow.status}}"
                 f" --metaflow_configs_json {json.dumps(json.dumps(metaflow_configs))}"
                 " --retries {{retries}}"
+                " --failures {{workflow.failures}}"
             ),
         ]
 

metaflow/plugins/aip/aip_cli.py

@@ -55,6 +55,7 @@ def kubeflow_pipelines(obj):
 @click.option("--flow_parameters_json")
 @click.option("--metaflow_configs_json")
 @click.option("--retries")
+@click.option("--failures")
 @click.pass_obj
 def user_defined_exit_handler(
     obj,
@@ -65,6 +66,7 @@ def user_defined_exit_handler(
     flow_parameters_json: str,
     metaflow_configs_json: str,
     retries: int,
+    failures: str,
 ):
     # call user defined exit handler
     invoke_user_defined_exit_handler(
@@ -76,6 +78,7 @@ def user_defined_exit_handler(
         flow_parameters_json,
         metaflow_configs_json,
         retries,
+        failures,
     )
 
 

metaflow/plugins/aip/aip_udf_exit_handler.py

@@ -17,6 +17,7 @@ def invoke_user_defined_exit_handler(
     flow_parameters_json: str,
     metaflow_configs_json: str,
     retries: int,
+    failures: str,
 ):
     """
     The environment variables that this depends on:
@@ -35,6 +36,8 @@ def invoke_user_defined_exit_handler(
     import json
     import os
 
+    print(f"{failures=}")
+
     env_variables: Dict[str, str] = json.loads(env_variables_json)
 
     def get_env(name, default=None) -> str:
@@ -74,4 +77,5 @@ def get_env(name, default=None) -> str:
         metaflow_run_id=run_id,
         argo_ui_url=argo_ui_url,
         retries=int(retries),
+        failures=failures,
     )

metaflow/tutorials/00-helloworld/helloworld.py

@@ -1,6 +1,26 @@
-from metaflow import FlowSpec, step
+from metaflow import FlowSpec, step, resources, exit_handler, Parameter, JSONType
 
+from metaflow.plugins.aip import exit_handler_resources, exit_handler_retry
 
+
+@exit_handler_retry(times=1)
+# @exit_handler_resources(cpu="600m", memory="600Mi")
+def my_exit_handler(
+    status: str,
+    flow_parameters: dict,
+    argo_workflow_run_name: str,
+    metaflow_run_id: str,
+    argo_ui_url: str,
+    retries: int,
+    failures: str,
+) -> None:
+    print(f"{failures=}")
+    print(f"{flow_parameters=}")
+    if retries == 0:
+        raise Exception("oopsie")
+
+
+@exit_handler(func=my_exit_handler, on_success=True)
 class HelloFlow(FlowSpec):
     """
     A flow where Metaflow prints 'Hi'.
@@ -9,23 +29,23 @@ class HelloFlow(FlowSpec):
 
     """
 
+    s3_keys = Parameter("s3_keys", type=JSONType)
+    hi = Parameter("hi", default=None, help="The word to print.")
+    hi2 = Parameter("hi2")
+
     @step
     def start(self):
         """
         This is the 'start' step. All flows must have a step named 'start' that
         is the first step in the flow.
 
         """
+        self.x = 42
+        self.y = {"foo": "bar", "baz": "qux", "quux": "corge"}
+        # raise Exception("foo")
         print("HelloFlow is starting.")
-        self.next(self.hello)
-
-    @step
-    def hello(self):
-        """
-        A step for metaflow to introduce itself.
-
-        """
-        print("Metaflow says: Hi!")
+        # import time
+        # time.sleep(3)
         self.next(self.end)
 
     @step