fix(argo): propagate nodeSelector/tolerations to error and exit hooks

metaflow/plugins/argo/argo_workflows.py

@@ -3163,6 +3163,8 @@ def _container(cmds):
                 name=f"success-{success_fn_name.replace('_', '-')}",
                 container=_container(cmds=_cmd(success_fn_name)),
                 service_account_name=resources["service_account"],
+                node_selectors=resources.get("node_selector"),
+                tolerations=resources.get("tolerations"),
                 on_success=True,
             )
             hooks.append(hook)
@@ -3172,6 +3174,8 @@ def _container(cmds):
                 name=f"error-{error_fn_name.replace('_', '-')}",
                 service_account_name=resources["service_account"],
                 container=_container(cmds=_cmd(error_fn_name)),
+                node_selectors=resources.get("node_selector"),
+                tolerations=resources.get("tolerations"),
                 on_error=True,
             )
             hooks.append(hook)
@@ -3320,7 +3324,9 @@ def _error_msg_capture_hook_templates(self):
                         ),
                     ).to_dict()
                 )
-            ),
+            )
+            .node_selectors(resources.get("node_selector"))
+            .tolerations(resources.get("tolerations")),
             Template("capture-error-hook-fn-preflight").steps(
                 [
                     WorkflowStep()

metaflow/plugins/argo/exit_hooks.py

@@ -53,6 +53,17 @@ def service_account_name(self, service_account_name):
         self.payload["serviceAccountName"] = service_account_name
         return self
 
+    def node_selectors(self, node_selectors):
+        if "nodeSelector" not in self.payload:
+            self.payload["nodeSelector"] = {}
+        if node_selectors:
+            self.payload["nodeSelector"].update(node_selectors)
+        return self
+
+    def tolerations(self, tolerations):
+        self.payload["tolerations"] = tolerations
+        return self
+
 
 class Hook(object):
     """
@@ -175,6 +186,8 @@ def __init__(
         name: str,
         container: Dict,
         service_account_name: str = None,
+        node_selectors: Optional[Dict] = None,
+        tolerations: Optional[List] = None,
         on_success: bool = False,
         on_error: bool = False,
     ):
@@ -185,6 +198,12 @@ def __init__(
 
         self.template.container(container)
 
+        if node_selectors is not None:
+            self.template.node_selectors(node_selectors)
+
+        if tolerations is not None:
+            self.template.tolerations(tolerations)
+
         self.lifecycle_hooks = []
 
         if on_success and on_error:

test/ux/core/test_argo_compilation.py

@@ -138,3 +138,131 @@ def test_late_attached_kubernetes_mutator_is_reflected_in_argo_template(
 
     assert end_resources["requests"]["cpu"] == "1"
     assert end_resources["requests"]["memory"] == "4096M"
+
+
+def test_argo_error_hook_inherits_node_selector_and_tolerations(
+    exec_mode, tag, scheduler_config, monkeypatch
+):
+    if exec_mode != "deployer":
+        pytest.skip("Argo compilation tests require deployer mode")
+    if scheduler_config.scheduler_type != "argo-workflows":
+        pytest.skip("Argo compilation tests require the argo-workflows scheduler")
+
+    from metaflow import Deployer
+
+    from .test_utils import _resolve_flow_path, prepare_runner_deployer_args
+
+    # Set environment variables for nodeSelector and tolerations
+    monkeypatch.setenv(
+        "METAFLOW_KUBERNETES_NODE_SELECTOR", '{"disktype": "ssd", "role": "compute"}'
+    )
+    monkeypatch.setenv(
+        "METAFLOW_KUBERNETES_TOLERATIONS",
+        '[{"key": "dedicated", "operator": "Equal", "value": "ml", "effect": "NoSchedule"}]',
+    )
+
+    deployed_flow = (
+        Deployer(
+            flow_file=_resolve_flow_path("basic/helloworld.py"),
+            show_output=False,
+            **prepare_runner_deployer_args({}),
+        )
+        .argo_workflows()
+        .create(
+            only_json=True,
+            tags=tag + ["test_argo_error_hook_node_selector_tolerations"],
+            **(scheduler_config.deploy_args or {}),
+        )
+    )
+
+    workflow_template = deployed_flow.workflow_template
+    assert workflow_template is not None
+
+    # Find the error-msg-capture-hook template
+    error_hook_template = None
+    start_template = None
+    for template in workflow_template.get("spec", {}).get("templates", []):
+        if template.get("name") == "error-msg-capture-hook":
+            error_hook_template = template
+        annotations = template.get("metadata", {}).get("annotations", {})
+        if annotations.get("metaflow/step_name") == "start":
+            start_template = template
+
+    assert error_hook_template is not None, "error-msg-capture-hook template not found"
+    assert start_template is not None, "start step template not found"
+
+    # Verify nodeSelector is present in both error hook and step template
+    assert "nodeSelector" in error_hook_template
+    assert error_hook_template["nodeSelector"] == {
+        "disktype": "ssd",
+        "role": "compute",
+    }
+
+    assert "nodeSelector" in start_template
+    assert start_template["nodeSelector"] == {"disktype": "ssd", "role": "compute"}
+
+    # Verify tolerations are present in both error hook and step template
+    assert "tolerations" in error_hook_template
+    assert error_hook_template["tolerations"] == [
+        {
+            "key": "dedicated",
+            "operator": "Equal",
+            "value": "ml",
+            "effect": "NoSchedule",
+        }
+    ]
+
+    assert "tolerations" in start_template
+    assert start_template["tolerations"] == [
+        {"key": "dedicated", "operator": "Equal", "value": "ml", "effect": "NoSchedule"}
+    ]
+
+
+def test_argo_error_hook_without_node_selector_and_tolerations(
+    exec_mode, tag, scheduler_config, monkeypatch
+):
+    if exec_mode != "deployer":
+        pytest.skip("Argo compilation tests require deployer mode")
+    if scheduler_config.scheduler_type != "argo-workflows":
+        pytest.skip("Argo compilation tests require the argo-workflows scheduler")
+
+    from metaflow import Deployer
+
+    from .test_utils import _resolve_flow_path, prepare_runner_deployer_args
+
+    # Ensure the env vars are not set
+    monkeypatch.delenv("METAFLOW_KUBERNETES_NODE_SELECTOR", raising=False)
+    monkeypatch.delenv("METAFLOW_KUBERNETES_TOLERATIONS", raising=False)
+
+    deployed_flow = (
+        Deployer(
+            flow_file=_resolve_flow_path("basic/helloworld.py"),
+            show_output=False,
+            **prepare_runner_deployer_args({}),
+        )
+        .argo_workflows()
+        .create(
+            only_json=True,
+            tags=tag + ["test_argo_error_hook_no_node_selector_tolerations"],
+            **(scheduler_config.deploy_args or {}),
+        )
+    )
+
+    workflow_template = deployed_flow.workflow_template
+    assert workflow_template is not None
+
+    # Find the error-msg-capture-hook template
+    error_hook_template = None
+    for template in workflow_template.get("spec", {}).get("templates", []):
+        if template.get("name") == "error-msg-capture-hook":
+            error_hook_template = template
+
+    assert error_hook_template is not None, "error-msg-capture-hook template not found"
+
+    # Verify nodeSelector is either empty or not present
+    node_selector = error_hook_template.get("nodeSelector", {})
+    assert node_selector == {}, f"Expected empty nodeSelector, got {node_selector}"
+
+    # Verify tolerations are not present or None
+    tolerations = error_hook_template.get("tolerations")
+    assert tolerations is None, f"Expected no tolerations, got {tolerations}"