Repository: couler-proj/couler

Issue:
Support node selector for scheduling pods
An example in Argo Workflows: https://github.com/argoproj/argo/blob/master/examples/node-selector.yaml

cc @inohmonton99 

Comments:

Patch:
diff --git a/couler/core/run_templates.py b/couler/core/run_templates.py
index b94d5b6..35a08b4 100644
--- a/couler/core/run_templates.py
+++ b/couler/core/run_templates.py
@@ -99,6 +99,7 @@ def run_container(
     daemon=False,
     volume_mounts=None,
     working_dir=None,
+    node_selector=None,
 ):
     """
     Generate an Argo container template.  For example, the template whalesay
@@ -176,6 +177,7 @@ def run_container(
             daemon=daemon,
             volume_mounts=volume_mounts,
             working_dir=working_dir,
+            node_selector=node_selector,
         )
         states.workflow.add_template(template)
 
diff --git a/couler/core/templates/container.py b/couler/core/templates/container.py
index e3569f0..c73a489 100644
--- a/couler/core/templates/container.py
+++ b/couler/core/templates/container.py
@@ -42,6 +42,7 @@ class Container(Template):
         daemon=False,
         volume_mounts=None,
         working_dir=None,
+        node_selector=None,
     ):
         Template.__init__(
             self,
@@ -63,6 +64,7 @@ class Container(Template):
         self.image_pull_policy = image_pull_policy
         self.volume_mounts = volume_mounts
         self.working_dir = working_dir
+        self.node_selector = node_selector
 
     def to_dict(self):
         template = Template.to_dict(self)
@@ -108,6 +110,11 @@ class Container(Template):
 
                 template["inputs"]["artifacts"] = _input_list
 
+        # Node selector
+        if self.node_selector is not None:
+            # TODO: Support inferring node selector values from Argo parameters
+            template["nodeSelector"] = self.node_selector
+
         # Container
         if not utils.gpu_requested(self.resources):
             if self.env is None:


Test patch:
diff --git a/couler/tests/argo_test.py b/couler/tests/argo_test.py
index 8e31ce1..8b4d925 100644
--- a/couler/tests/argo_test.py
+++ b/couler/tests/argo_test.py
@@ -90,6 +90,22 @@ class ArgoTest(unittest.TestCase):
         )
         couler._cleanup()
 
+    def test_run_container_with_node_selector(self):
+        couler.run_container(
+            image="docker/whalesay:latest",
+            args=["echo -n hello world"],
+            command=["bash", "-c"],
+            step_name="A",
+            node_selector={"beta.kubernetes.io/arch": "amd64"},
+        )
+
+        wf = couler.workflow_yaml()
+        self.assertEqual(
+            wf["spec"]["templates"][1]["nodeSelector"],
+            {"beta.kubernetes.io/arch": "amd64"},
+        )
+        couler._cleanup()
+
     def test_run_container_with_workflow_volume(self):
         pvc = VolumeClaimTemplate("workdir")
         volume_mount = VolumeMount("workdir", "/mnt/vol")
@@ -161,8 +177,8 @@ class ArgoTest(unittest.TestCase):
         self.assertTrue(
             params["value"]
             in [
-                '"{{workflow.outputs.parameters.output-id-117}}"',
-                '"{{workflow.outputs.parameters.output-id-118}}"',
+                '"{{workflow.outputs.parameters.output-id-133}}"',
+                '"{{workflow.outputs.parameters.output-id-134}}"',
             ]
         )
         # Check input parameters for step B