From a45cee3523453273cfa7281058c044a02516be74 Mon Sep 17 00:00:00 2001 From: Tasko Olevski Date: Mon, 6 Nov 2023 16:42:49 +0100 Subject: [PATCH] fix: properly patch tolerations and affinities (#1695) --- README.md | 1 + .../api/amalthea_patches/general.py | 96 ++++++++++++++----- renku_notebooks/api/schemas/server_options.py | 49 ++-------- 3 files changed, 78 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 6cad16c6d..8e6e59c20 100644 --- 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 index dbe401994..0744ab841 100644 --- 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 index a78ed896d..fa7f5f47d 100644 --- 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", }