feat: support for custom node selector

we want jobs to be able to select different node types

Signed-off-by: vsoch <[email protected]>

Author: vsoch

README.md

@@ -90,6 +90,8 @@ For each job script section, the following environment variables are provided fo
 - registry: the registry where your artifact will be pushed
   - pull_tag: the pull tag to use (if the workflow is pulling)
   - push_tag: the push tag to use (if the workflow is pushing)
+- properties:
+  - node-selector: key value pair to be added as node selectors for the job (Kubernetes only). E.g., `node.kubernetes.io/instance-type: c7a.4xlarge`
 
 Take a look at the simple example [examples/state-machine.yaml](examples/state-machine.yaml) to see how push/pull is defined between steps. Given that these are found (with a tag) your artifact will be named `<registry>:<jobid>:<tag>` to be moved between steps.
 

python/state_machine_operator/schema.py

@@ -80,6 +80,7 @@
         "config": {"$ref": "#/definitions/config"},
         "script": {"type": "string"},
         "image": {"type": "string"},
+        "properties": {"$ref": "#/definitions/properties"},
         "registry": {"$ref": "#/definitions/registry"},
         "workdir": {"type": "string", "default": "/tmp/out"},
         "additionalProperties": False,
@@ -96,6 +97,10 @@
             },
             "additionalProperties": False,
         },
+        "properties": {
+            "type": "object",
+            "properties": {"type": "object", "additionalProperties": True},
+        },
         "config": {
             "type": "object",
             "properties": {

python/state_machine_operator/tracker/kubernetes/tracker.py

@@ -90,6 +90,12 @@ def delete_configmap(self, name):
             except Exception as e:
                 LOGGER.warning(f"Issue deleting configmap {name}: {e}")
 
+    def get_node_selector(self):
+        """
+        Node selector is in properties -> node-selector
+        """
+        return self.job_desc.get("properties", {}).get("node-selector")
+
     def generate_batch_job(self, step, jobid):
         """
         Generate the job CRD assuming the config map entrypoint.
@@ -186,6 +192,12 @@ def generate_batch_job(self, step, jobid):
             },
         }
 
+        # Add node selectors? E.g.,
+        # node.kubernetes.io/instance-type: c7a.4xlarge
+        node_selector = self.get_node_selector()
+        if node_selector is not None:
+            template["spec"]["nodeSelector"] = node_selector
+
         # Only add walltime if it's > 0 and not None
         if walltime:
             template["spec"]["activeDeadlineSeconds"] = int(walltime)