diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index b033c6d638f..159dc32866e 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/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() diff --git a/metaflow/plugins/argo/exit_hooks.py b/metaflow/plugins/argo/exit_hooks.py index 59d17664a91..6ae433ea4e5 100644 --- a/metaflow/plugins/argo/exit_hooks.py +++ b/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: diff --git a/test/ux/core/test_argo_compilation.py b/test/ux/core/test_argo_compilation.py index 54f46653116..c296292d8e9 100644 --- a/test/ux/core/test_argo_compilation.py +++ b/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}"