From 245a3aa075e78563d59fc961c0850220f94d23db Mon Sep 17 00:00:00 2001
From: Taleb Zeghmi <talebz@zillowgroup.com>
Date: Mon, 4 Dec 2023 09:24:15 -0800
Subject: [PATCH] PR Feedback

---
 metaflow/plugins/aip/aip.py                   | 10 -------
 .../aip/tests/flows/raise_error_flow.py       |  3 +-
 .../aip/tests/run_integration_tests.py        | 28 ++++++++++++++++++-
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/metaflow/plugins/aip/aip.py b/metaflow/plugins/aip/aip.py
index 8d306586ed1..eebe5ee5405 100644
--- a/metaflow/plugins/aip/aip.py
+++ b/metaflow/plugins/aip/aip.py
@@ -1720,21 +1720,11 @@ def _create_user_defined_exit_handler_op(
         env_variables: dict = {
             key: from_conf(key)
             for key in [
-                "METAFLOW_NOTIFY_EMAIL_FROM",
-                "METAFLOW_NOTIFY_EMAIL_SMTP_HOST",
-                "METAFLOW_NOTIFY_EMAIL_SMTP_PORT",
-                "METAFLOW_NOTIFY_EMAIL_BODY",
                 "ARGO_RUN_URL_PREFIX",
             ]
             if from_conf(key)
         }
 
-        if self.notify_on_error:
-            env_variables["METAFLOW_NOTIFY_ON_ERROR"] = self.notify_on_error
-
-        if self.notify_on_success:
-            env_variables["METAFLOW_NOTIFY_ON_SUCCESS"] = self.notify_on_success
-
         return self._get_user_defined_exit_handler_op(
             udf_handler,
             flow_parameters,
diff --git a/metaflow/plugins/aip/tests/flows/raise_error_flow.py b/metaflow/plugins/aip/tests/flows/raise_error_flow.py
index 9cb96bc84b8..c7d15fb039f 100644
--- a/metaflow/plugins/aip/tests/flows/raise_error_flow.py
+++ b/metaflow/plugins/aip/tests/flows/raise_error_flow.py
@@ -1,7 +1,8 @@
 from metaflow import FlowSpec, step, exit_handler
-from metaflow.plugins.aip import exit_handler_retry
+from metaflow.plugins.aip import exit_handler_retry, exit_handler_resources
 
 
+@exit_handler_resources(cpu="501m", memory="601Mi")
 @exit_handler_retry(times=1, minutes_between_retries=0)
 def my_exit_handler(
     status: str,
diff --git a/metaflow/plugins/aip/tests/run_integration_tests.py b/metaflow/plugins/aip/tests/run_integration_tests.py
index 85910830278..d7d6346174b 100644
--- a/metaflow/plugins/aip/tests/run_integration_tests.py
+++ b/metaflow/plugins/aip/tests/run_integration_tests.py
@@ -360,9 +360,35 @@ def test_user_defined_exit_handler(pytestconfig) -> None:
         for task in exit_handler_template["dag"]["tasks"]
         if task["name"] == "user-defined-exit-handler"
     )
-
     assert user_defined_exit_handler
 
+    user_defined_exit_handler_template = next(
+        template
+        for template in flow_yaml["spec"]["templates"]
+        if template["name"] == "user-defined-exit-handler"
+    )
+
+    # check that the exit-handler spec template has the correct
+    # resource requirements
+    assert (
+        user_defined_exit_handler_template["container"]["resources"]["requests"]["cpu"]
+        == "501m"
+    )
+    assert (
+        user_defined_exit_handler_template["container"]["resources"]["limits"]["cpu"]
+        == "501m"
+    )
+    assert (
+        user_defined_exit_handler_template["container"]["resources"]["requests"][
+            "memory"
+        ]
+        == "601Mi"
+    )
+    assert (
+        user_defined_exit_handler_template["container"]["resources"]["limits"]["memory"]
+        == "601Mi"
+    )
+
 
 def test_kubernetes_service_account_compile_only(pytestconfig) -> None:
     service_account = "test-service-account"