fix: properly patch tolerations and affinities (#1695)

SwissDataScienceCenter · Nov 6, 2023 · a45cee3 · a45cee3
1 parent 8410b4c
commit a45cee3
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 68 deletions.
diff --git a/README.md b/README.md
@@ -123,3 +123,4 @@ service in a full Renku deployment.
   [chartpress]: https://github.com/jupyterhub/chartpress
   [k8s python client]: https://github.com/kubernetes-client/python
   [renkulab]: https://renkulab.io
+
diff --git a/renku_notebooks/api/amalthea_patches/general.py b/renku_notebooks/api/amalthea_patches/general.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import Any, Dict, List, TYPE_CHECKING
 
 from ...config import config
 from ..classes.user import RegisteredUser
@@ -10,29 +10,28 @@
 def session_tolerations(server: "UserServer"):
     """Patch for node taint tolerations, the static tolerations from the configuration are ignored
     if the tolerations are set in the server options (coming from CRC)."""
-    if not server.server_options.tolerations:
-        key = f"{config.session_get_endpoint_annotations.renku_annotation_prefix}dedicated"
-        tolerations = [
-            {
-                "key": key,
-                "operator": "Equal",
-                "value": "user",
-                "effect": "NoSchedule",
-            },
-        ] + config.sessions.tolerations
-        return [
-            {
-                "type": "application/json-patch+json",
-                "patch": [
-                    {
-                        "op": "add",
-                        "path": "/statefulset/spec/template/spec/tolerations",
-                        "value": tolerations,
-                    }
-                ],
-            }
-        ]
-    return [i.json_patch() for i in server.server_options.tolerations]
+    key = f"{config.session_get_endpoint_annotations.renku_annotation_prefix}dedicated"
+    default_tolerations: List[Dict[str, str]] = [
+        {
+            "key": key,
+            "operator": "Equal",
+            "value": "user",
+            "effect": "NoSchedule",
+        },
+    ] + config.sessions.tolerations
+    return [
+        {
+            "type": "application/json-patch+json",
+            "patch": [
+                {
+                    "op": "add",
+                    "path": "/statefulset/spec/template/spec/tolerations",
+                    "value": default_tolerations
+                    + [i.json_match_expression() for i in server.server_options.tolerations],
+                }
+            ],
+        }
+    ]
 
 
 def session_affinity(server: "UserServer"):
@@ -51,7 +50,54 @@ def session_affinity(server: "UserServer"):
                 ],
             }
         ]
-    return [i.json_patch() for i in server.server_options.node_affinities]
+    default_preferred_selector_terms: List[Dict[str, Any]] = config.sessions.affinity.get(
+        "nodeAffinity", {}
+    ).get("preferredDuringSchedulingIgnoredDuringExecution", [])
+    default_required_selector_terms: List[Dict[str, Any]] = (
+        config.sessions.affinity.get("nodeAffinity", {})
+        .get("requiredDuringSchedulingIgnoredDuringExecution", {})
+        .get("nodeSelectorTerms", [])
+    )
+    preferred_match_expressions: List[Dict[str, str]] = []
+    required_match_expressions: List[Dict[str, str]] = []
+    for affintiy in server.server_options.node_affinities:
+        if affintiy.required_during_scheduling:
+            required_match_expressions.append(affintiy.json_match_expression())
+        else:
+            preferred_match_expressions.append(affintiy.json_match_expression())
+    return [
+        {
+            "type": "application/json-patch+json",
+            "patch": [
+                {
+                    "op": "add",
+                    "path": "/statefulset/spec/template/spec/affinity",
+                    "value": {
+                        "nodeAffinity": {
+                            "preferredDuringScheduling"
+                            "IgnoredDuringExecution": default_preferred_selector_terms
+                            + [
+                                {
+                                    "weight": 1,
+                                    "preference": {
+                                        "matchExpressions": preferred_match_expressions,
+                                    },
+                                }
+                            ],
+                            "requiredDuringSchedulingIgnoredDuringExecution": {
+                                "nodeSelectorTerms": default_required_selector_terms
+                                + [
+                                    {
+                                        "matchExpressions": required_match_expressions,
+                                    }
+                                ],
+                            },
+                        },
+                    },
+                }
+            ],
+        }
+    ]
 
 
 def session_node_selector(server: "UserServer"):

diff --git a/renku_notebooks/api/schemas/server_options.py b/renku_notebooks/api/schemas/server_options.py
@@ -15,38 +15,10 @@ class NodeAffinity:
     key: str
     required_during_scheduling: bool = False
 
-    def json_patch(self) -> Dict[str, Any]:
-        match_expressions = {
-            "matchExpressions": {
-                "key": self.key,
-                "operator": "Exists",
-            },
-        }
-        if self.required_during_scheduling:
-            return {
-                "type": "application/json-patch+json",
-                "patch": [
-                    {
-                        "op": "add",
-                        "path": "/statefulset/spec/template/spec/affinity/nodeAffinity"
-                        "/requiredDuringSchedulingIgnoredDuringExecution/nodeSelectorTerms/-",
-                        "value": match_expressions,
-                    }
-                ],
-            }
+    def json_match_expression(self) -> Dict[str, str]:
         return {
-            "type": "application/json-patch+json",
-            "patch": [
-                {
-                    "op": "add",
-                    "path": "/statefulset/spec/template/spec/affinity/nodeAffinity"
-                    "/preferredDuringSchedulingIgnoredDuringExecution/-",
-                    "value": {
-                        "weight": 1,
-                        "preference": match_expressions,
-                    },
-                }
-            ],
+            "key": self.key,
+            "operator": "Exists",
         }
 
 
@@ -56,19 +28,10 @@ class Toleration:
 
     key: str
 
-    def json_patch(self) -> Dict[str, Any]:
+    def json_match_expression(self) -> Dict[str, Any]:
         return {
-            "type": "application/json-patch+json",
-            "patch": [
-                {
-                    "op": "add",
-                    "path": "/statefulset/spec/template/spec/tolerations/-",
-                    "value": {
-                        "key": self.key,
-                        "operator": "Exists",
-                    },
-                }
-            ],
+            "key": self.key,
+            "operator": "Exists",
         }