Add support for setting Kubernetes affinity in pipeline configuration (#276)

* Add support for setting Kubernetes affinity in pipeline configuration

* chore(ci): update upload/download-artifact actions to v4

Author: rdziadosz

.github/workflows/test.yml

@@ -44,7 +44,7 @@ jobs:
         tox -v
 
     - name: Store coverage reports
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: coverage-${{ matrix.python-version }}
         path: coverage.xml
@@ -59,7 +59,7 @@ jobs:
       with:
         fetch-depth: 0
 
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: coverage-3.9
         path: .

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+-   Add support for setting Kubernetes affinity in pipeline configuration
+
 ## [0.8.0] - 2024-09-04
 
 -   Added support for python 3.11. Python 3.12 is still blocked by kedro-mlflow

docs/source/02_installation/02_configuration.md

@@ -104,6 +104,25 @@ run_config:
       value: "gpu_workload"
       effect: "NoSchedule"
 
+  # Optional section allowing adjustment of the affinity for the nodes
+  affinity:
+    __default__:
+      node_affinity:
+        required_during_scheduling_ignored_during_execution:
+          nodeSelectorTerms:
+            - matchExpressions:
+                - key: "node-type"
+                  operator: "In"
+                  values: ["general"]
+    node_a:
+      node_affinity:
+        required_during_scheduling_ignored_during_execution:
+          nodeSelectorTerms:
+            - matchExpressions:
+                - key: "gpu"
+                  operator: "In"
+                  values: ["true"]
+
   # Optional section to allow mounting additional volumes (such as EmptyDir)
   # to specific nodes
   extra_volumes:

kedro_kubeflow/config.py

@@ -259,6 +259,25 @@ def volume_validator(cls, value):
         return value
 
 
+class AffinityConfig(defaultdict):
+    def __init__(self, default_factory=None, *args, **kwargs):
+        super().__init__(default_factory, *args, **kwargs)
+
+    def is_set_for(self, node_name):
+        return bool(self.get_for(node_name))
+
+    def get_for(self, node_name):
+        node_value = self._get_or_default(node_name, None)
+        if node_value is not None:
+            return node_value
+        return self._get_or_default("__default__", None)
+
+    def _get_or_default(self, node_name, default):
+        if node_name in self:
+            return self[node_name]
+        return self.get("__default__", default)
+
+
 class RunConfig(BaseModel):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -293,6 +312,10 @@ def _validate_tolerations(cls, value):
     def _validate_extra_volumes(cls, value):
         return RunConfig._create_default_dict_with(value, [], defaultdict)
 
+    @validator("affinity", always=True)
+    def _validate_affinity(cls, value):
+        return RunConfig._create_default_dict_with(value, None, AffinityConfig)
+
     image: str
     image_pull_policy: str = "IfNotPresent"
     root: Optional[str]
@@ -305,6 +328,7 @@ def _validate_extra_volumes(cls, value):
     retry_policy: Optional[Dict[str, Optional[RetryPolicyConfig]]]
     volume: Optional[VolumeConfig] = None
     extra_volumes: Optional[Dict[str, List[ExtraVolumeConfig]]] = None
+    affinity: Optional[Dict[str, Any]] = None
     wait_for_completion: bool = False
     store_kedro_outputs_as_kfp_artifacts: bool = True
     max_cache_staleness: Optional[str] = None

kedro_kubeflow/generators/utils.py

@@ -145,6 +145,29 @@ def create_pipeline_exit_handler(pipeline, image, image_pull_policy, run_config,
     return dsl.ExitHandler(customize_op(exit_container_op, image_pull_policy, run_config))
 
 
+def dict_to_v1affinity(affinity_dict):
+    def convert_node_affinity(na):
+        if not na:
+            return None
+        return k8s.V1NodeAffinity(**na)
+
+    def convert_pod_affinity(pa):
+        if not pa:
+            return None
+        return k8s.V1PodAffinity(**pa)
+
+    def convert_pod_anti_affinity(paa):
+        if not paa:
+            return None
+        return k8s.V1PodAntiAffinity(**paa)
+
+    return k8s.V1Affinity(
+        node_affinity=convert_node_affinity(affinity_dict.get("node_affinity")),
+        pod_affinity=convert_pod_affinity(affinity_dict.get("pod_affinity")),
+        pod_anti_affinity=convert_pod_anti_affinity(affinity_dict.get("pod_anti_affinity")),
+    )
+
+
 def customize_op(op, image_pull_policy, run_config: RunConfig):
     op.container.set_image_pull_policy(image_pull_policy)
     if run_config.volume and run_config.volume.owner is not None:
@@ -162,6 +185,12 @@ def customize_op(op, image_pull_policy, run_config: RunConfig):
     for toleration in run_config.tolerations[op.name]:
         op.add_toleration(k8s.V1Toleration(**toleration.dict()))
 
+    if run_config.affinity.is_set_for(op.name):
+        affinity_dict = run_config.affinity.get_for(op.name)
+        if affinity_dict:
+            affinity_obj = dict_to_v1affinity(affinity_dict)
+            op.add_affinity(affinity_obj)
+
     if extra_volumes := run_config.extra_volumes[op.name]:
         op.add_pvolumes({ev.mount_path: dsl.PipelineVolume(volume=ev.as_v1volume()) for ev in extra_volumes})
     return op

tests/test_config.py

@@ -22,6 +22,23 @@
     size: 3Gi
     access_modes: [ReadWriteOnce]
     keep: True
+  affinity:
+    __default__:
+      node_affinity:
+        required_during_scheduling_ignored_during_execution:
+          nodeSelectorTerms:
+            - matchExpressions:
+                - key: "default-key"
+                  operator: "In"
+                  values: ["default-value"]
+    node1:
+      node_affinity:
+        required_during_scheduling_ignored_during_execution:
+          nodeSelectorTerms:
+            - matchExpressions:
+                - key: "node1-key"
+                  operator: "In"
+                  values: ["node1-value"]
 """
 
 
@@ -264,3 +281,48 @@ def test_retry_policy_no_default(self):
         )
 
         self.assertIsNone(cfg.run_config.retry_policy["node2"])
+
+    def test_affinity_node_specific(self):
+        cfg = PluginConfig(**yaml.safe_load(CONFIG_YAML))
+        affinity = cfg.run_config.affinity.get_for("node1")
+        assert affinity is not None
+        assert "node_affinity" in affinity
+        assert (
+            affinity["node_affinity"]["required_during_scheduling_ignored_during_execution"]["nodeSelectorTerms"][0][
+                "matchExpressions"
+            ][0]["key"]
+            == "node1-key"
+        )
+
+    def test_affinity_default(self):
+        cfg = PluginConfig(**yaml.safe_load(CONFIG_YAML))
+        affinity = cfg.run_config.affinity.get_for("node2")
+        assert affinity is not None
+        assert "node_affinity" in affinity
+        assert (
+            affinity["node_affinity"]["required_during_scheduling_ignored_during_execution"]["nodeSelectorTerms"][0][
+                "matchExpressions"
+            ][0]["key"]
+            == "default-key"
+        )
+
+    def test_affinity_is_set_for(self):
+        cfg = PluginConfig(**yaml.safe_load(CONFIG_YAML))
+        assert cfg.run_config.affinity.is_set_for("node1")
+        assert cfg.run_config.affinity.is_set_for("node2")  # falls back to default
+        # Should be False if neither node nor default is set
+        empty_affinity = PluginConfig(
+            **yaml.safe_load(
+                """
+        host: https://example.com
+        run_config:
+          image: "gcr.io/project-image/test"
+          image_pull_policy: "Always"
+          experiment_name: "Test Experiment"
+          run_name: "test run"
+          scheduled_run_name: "scheduled run"
+          affinity: {}
+        """
+            )
+        )
+        assert not empty_affinity.run_config.affinity.is_set_for("unknown_node")