From 5daca01c14b81b1cccdf48e54d1e917cfd2f44b3 Mon Sep 17 00:00:00 2001
From: Taleb Zeghmi <talebz@zillowgroup.com>
Date: Tue, 23 Jan 2024 07:24:38 -0800
Subject: [PATCH 1/2] initial commit

---
 metaflow/plugins/aip/aip.py                   |  1 +
 metaflow/plugins/aip/aip_cli.py               |  3 ++
 metaflow/plugins/aip/aip_udf_exit_handler.py  |  4 ++
 .../tutorials/00-helloworld/helloworld.py     | 40 ++++++++++++++-----
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/metaflow/plugins/aip/aip.py b/metaflow/plugins/aip/aip.py
index 4d5e2c5c044..60028ae7336 100644
--- a/metaflow/plugins/aip/aip.py
+++ b/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}}"
             ),
         ]
 
diff --git a/metaflow/plugins/aip/aip_cli.py b/metaflow/plugins/aip/aip_cli.py
index d9a88ea51e9..d4ff9e84519 100644
--- a/metaflow/plugins/aip/aip_cli.py
+++ b/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,
     )
 
 
diff --git a/metaflow/plugins/aip/aip_udf_exit_handler.py b/metaflow/plugins/aip/aip_udf_exit_handler.py
index d21ce2e36f4..7f3049497bf 100644
--- a/metaflow/plugins/aip/aip_udf_exit_handler.py
+++ b/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,
     )
diff --git a/metaflow/tutorials/00-helloworld/helloworld.py b/metaflow/tutorials/00-helloworld/helloworld.py
index 2e073b5131c..32a5fb7e779 100644
--- a/metaflow/tutorials/00-helloworld/helloworld.py
+++ b/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,6 +29,10 @@ 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):
         """
@@ -16,16 +40,12 @@ def start(self):
         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

From 9768cbf253867cf73d4d5852ba71ecfb504edf8f Mon Sep 17 00:00:00 2001
From: Taleb Zeghmi <talebz@zillowgroup.com>
Date: Tue, 23 Jan 2024 07:30:27 -0800
Subject: [PATCH 2/2] ref CICD hash commit before WORKING_DIR

---
 .gitlab-ci.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 9ecad46198c..402d8fa2570 100644
--- a/.gitlab-ci.yml
+++ b/.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