From adf232116cc834f0bfffa3b1f977edd4693b7f61 Mon Sep 17 00:00:00 2001
From: Xi Chen <xchen27@tesla.com>
Date: Tue, 17 Dec 2024 21:06:51 +0800
Subject: [PATCH] Support extra pod spec for user placeholder pods

Add configuration `scheduling.userPlaceholder.extraPodSpec` for arbitrary extra pod spec for the user placeholder statefulset.

This change arose from a desire to add pod topology spread constraints to the placeholder pods. Our singleuser pods run in different availability zones at runtime, as their corresponding persistent volumes are created in different availability zones. Adding pod topology spread constraints across zones helps improving singleuser pod startup time in different availability zones.

For example,
```
scheduling:
  userPlaceholder:
    extraPodSpec:
      topologySpreadConstraints:
        - maxSkew: 1
          topologyKey: zone
          whenUnsatisfiable: DoNotSchedule
          labelSelector:
            matchLabels:
              foo: bar
```
---
 .../scheduling/user-placeholder/statefulset.yaml          | 3 +++
 jupyterhub/values.schema.yaml                             | 2 ++
 jupyterhub/values.yaml                                    | 1 +
 tools/templates/lint-and-validate-values.yaml             | 8 ++++++++
 4 files changed, 14 insertions(+)

diff --git a/jupyterhub/templates/scheduling/user-placeholder/statefulset.yaml b/jupyterhub/templates/scheduling/user-placeholder/statefulset.yaml
index 7f2c785b99..a68c52a679 100644
--- a/jupyterhub/templates/scheduling/user-placeholder/statefulset.yaml
+++ b/jupyterhub/templates/scheduling/user-placeholder/statefulset.yaml
@@ -77,4 +77,7 @@ spec:
           securityContext:
             {{- . | toYaml | nindent 12 }}
           {{- end }}
+      {{- with .Values.scheduling.userPlaceholder.extraPodSpec }}
+      {{- . | toYaml | nindent 6 }}
+      {{- end }}
 {{- end }}
diff --git a/jupyterhub/values.schema.yaml b/jupyterhub/values.schema.yaml
index 4d735ca20c..25494f1c62 100644
--- a/jupyterhub/values.schema.yaml
+++ b/jupyterhub/values.schema.yaml
@@ -1526,6 +1526,7 @@ properties:
           2. proxy.chp.extraPodSpec
           3. proxy.traefik.extraPodSpec
           4. scheduling.userScheduler.extraPodSpec
+          5. scheduling.userPlaceholder.extraPodSpec
 
           One real-world use of these settings is to enable host networking. For
           example, to configure host networking for the hub pod, add the
@@ -2648,6 +2649,7 @@ properties:
               Unless specified here, the placeholder pods will request the same
               resources specified for the real singleuser pods.
           containerSecurityContext: *containerSecurityContext-spec
+          extraPodSpec: *extraPodSpec-spec
       corePods:
         type: object
         additionalProperties: false
diff --git a/jupyterhub/values.yaml b/jupyterhub/values.yaml
index d8c95bafe1..d23354b3aa 100644
--- a/jupyterhub/values.yaml
+++ b/jupyterhub/values.yaml
@@ -600,6 +600,7 @@ scheduling:
       seccompProfile:
         type: "RuntimeDefault"
     resources: {}
+    extraPodSpec: {}
   corePods:
     tolerations:
       - key: hub.jupyter.org/dedicated
diff --git a/tools/templates/lint-and-validate-values.yaml b/tools/templates/lint-and-validate-values.yaml
index 5af74b65af..ef64f82ee3 100644
--- a/tools/templates/lint-and-validate-values.yaml
+++ b/tools/templates/lint-and-validate-values.yaml
@@ -521,6 +521,14 @@ scheduling:
     revisionHistoryLimit: 1
     replicas: 1
     resources: *resources
+    extraPodSpec:
+      topologySpreadConstraints:
+        - maxSkew: 1
+          topologyKey: zone
+          whenUnsatisfiable: DoNotSchedule
+          labelSelector:
+            matchLabels:
+              foo: bar
   corePods:
     tolerations:
       - key: mock-taint-key-corePods