feat(chart): add session tolerations, affinity, nodeSelector (#806)

SwissDataScienceCenter · Nov 15, 2021 · 49a2d54 · 49a2d54
1 parent 30df71b
commit 49a2d54
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 10 deletions.
diff --git a/helm-chart/renku-notebooks/templates/deployment.yaml b/helm-chart/renku-notebooks/templates/deployment.yaml
@@ -99,7 +99,21 @@ spec:
             - name: SENTRY_ENV
               value: {{ .Values.sentry.env | quote }}
             {{ end }}
-
+            {{- with .Values.sessionNodeSelector }}
+            - name: SESSION_NODE_SELECTOR
+              value: |
+                {{- toYaml . | nindent 16 }}
+            {{- end }}
+            {{- with .Values.sessionAffinity }}
+            - name: SESSION_AFFINITY
+              value: |
+                {{- toYaml . | nindent 16 }}
+            {{- end }}
+            {{- with .Values.sessionTolerations }}
+            - name: SESSION_TOLERATIONS
+              value: |
+                {{- toYaml . | nindent 16 }}
+            {{- end }}
             - name: KUBERNETES_NAMESPACE
               valueFrom:
                 fieldRef:

diff --git a/helm-chart/renku-notebooks/templates/test.yaml b/helm-chart/renku-notebooks/templates/test.yaml
@@ -61,6 +61,21 @@ spec:
             fieldRef:
               apiVersion: v1
               fieldPath: metadata.namespace
+        {{- with $.Values.sessionNodeSelector }}
+        - name: SESSION_NODE_SELECTOR
+          value: |
+            {{- toYaml . | nindent 12 }}
+        {{- end }}
+        {{- with $.Values.sessionAffinity }}
+        - name: SESSION_AFFINITY
+          value: |
+            {{- toYaml . | nindent 12 }}
+        {{- end }}
+        {{- with $.Values.sessionTolerations }}
+        - name: SESSION_TOLERATIONS
+          value: |
+            {{- toYaml . | nindent 12 }}
+        {{- end }}
       command:
         - pipenv
         - run

diff --git a/helm-chart/renku-notebooks/values.yaml b/helm-chart/renku-notebooks/values.yaml
@@ -79,6 +79,12 @@ sessionIngress:
     nginx.ingress.kubernetes.io/proxy-buffer-size: 8k
     nginx.ingress.kubernetes.io/proxy-request-buffering: "off"
 
+## Setup node selector, tolerations and node affinities for user sessions that are launched
+## by the notebook service.
+sessionNodeSelector: {}
+sessionTolerations: []
+sessionAffinity: {}
+
 ## Default server options - these will be provided to the user in the UI
 ## Note that requests are also limits, i.e. a user's jupyter kernel
 ## that exceeds the requested memory limit will be terminated/restarted.

diff --git a/renku_notebooks/api/classes/server.py b/renku_notebooks/api/classes/server.py
@@ -306,14 +306,6 @@ def _get_session_manifest(self):
             + "projectName": self.gl_project.path.lower(),
             f"{current_app.config['RENKU_ANNOTATION_PREFIX']}requested-image": self.image,
         }
-        tolerations = [
-            {
-                "key": f"{current_app.config['RENKU_ANNOTATION_PREFIX']}dedicated",
-                "operator": "Equal",
-                "value": "user",
-                "effect": "NoSchedule",
-            }
-        ]
         # Add image pull secret if image is private
         if self.is_image_private:
             image_pull_secret_name = self.server_name + "-image-secret"
@@ -565,7 +557,31 @@ def _get_session_manifest(self):
                     {
                         "op": "add",
                         "path": "/statefulset/spec/template/spec/tolerations",
-                        "value": tolerations,
+                        "value": current_app.config["SESSION_TOLERATIONS"],
+                    }
+                ],
+            }
+        )
+        patches.append(
+            {
+                "type": "application/json-patch+json",
+                "patch": [
+                    {
+                        "op": "add",
+                        "path": "/statefulset/spec/template/spec/affinity",
+                        "value": current_app.config["SESSION_AFFINITY"],
+                    }
+                ],
+            }
+        )
+        patches.append(
+            {
+                "type": "application/json-patch+json",
+                "patch": [
+                    {
+                        "op": "add",
+                        "path": "/statefulset/spec/template/spec/nodeSelector",
+                        "value": current_app.config["SESSION_NODE_SELECTOR"],
                     }
                 ],
             }

diff --git a/renku_notebooks/config.py b/renku_notebooks/config.py
@@ -114,3 +114,7 @@
 CULLING_ANONYMOUS_IDLE_SESSIONS_THRESHOLD_SECONDS = int(
     os.getenv("CULLING_ANONYMOUS_IDLE_SESSIONS_THRESHOLD_SECONDS", 43200)
 )
+
+SESSION_NODE_SELECTOR = safe_load(os.environ.get("SESSION_NODE_SELECTOR", "{}"))
+SESSION_AFFINITY = safe_load(os.environ.get("SESSION_AFFINITY", "{}"))
+SESSION_TOLERATIONS = safe_load(os.environ.get("SESSION_TOLERATIONS", "[]"))