PR Feedback

zillow · Dec 4, 2023 · 245a3aa · 245a3aa
1 parent 9207a96
commit 245a3aa
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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"