From 4794fdebc20341854a00ae90a042cd61bf72f487 Mon Sep 17 00:00:00 2001
From: viniciusdc <vinivdc2009@hotmail.com>
Date: Fri, 22 Mar 2024 18:33:36 -0300
Subject: [PATCH] Switch user selector template to wrapspawner.ProfilesSpawner
 and add default resource profile

---
 roles/jupyterhub/defaults/main.yml            |  7 ++
 .../templates/environments/jupyterhub.yaml    |  1 +
 .../jupyterhub/templates/jupyterhub_config.py | 83 +++++--------------
 3 files changed, 28 insertions(+), 63 deletions(-)

diff --git a/roles/jupyterhub/defaults/main.yml b/roles/jupyterhub/defaults/main.yml
index d21dee55..c32cf873 100644
--- a/roles/jupyterhub/defaults/main.yml
+++ b/roles/jupyterhub/defaults/main.yml
@@ -24,6 +24,13 @@ jupyterhub_config:
   spawner:
     start_timeout: 180 # seconds
 
+jupyterhub_profiles:
+  - default:
+      display_name: default
+      options:
+        req_memory: "2"
+        req_nprocs: "1"
+
 jupyterhub_services:
   dask_gateway: CStgn1NN8DogQR1KajuoQfye1qNRqx6zsh
 
diff --git a/roles/jupyterhub/templates/environments/jupyterhub.yaml b/roles/jupyterhub/templates/environments/jupyterhub.yaml
index 35e13704..ac0668a9 100644
--- a/roles/jupyterhub/templates/environments/jupyterhub.yaml
+++ b/roles/jupyterhub/templates/environments/jupyterhub.yaml
@@ -16,4 +16,5 @@ dependencies:
       - jupyterhub-traefik-proxy==1.1.0
       # jupyterhub-ssh has not made a release yet
       - git+https://github.com/yuvipanda/jupyterhub-ssh.git
+      - git+https://github.com/jupyterhub/wrapspawner.git
       - jhub-apps==2024.2.1rc1
diff --git a/roles/jupyterhub/templates/jupyterhub_config.py b/roles/jupyterhub/templates/jupyterhub_config.py
index 9001c6ea..3cf195f8 100644
--- a/roles/jupyterhub/templates/jupyterhub_config.py
+++ b/roles/jupyterhub/templates/jupyterhub_config.py
@@ -144,68 +144,6 @@ async def _get_batch_script(self, **subvars):
         # Could be overridden by subclasses, but mainly useful for testing
         return format_template(self.batch_script, **subvars)
 
-
-
-
-{% if jupyterhub_qhub_options_form %}
-    # data from form submission is {key: [value]}
-    # we need to convert the formdata to a key value dict
-    def options_from_form(self, data):
-        return {key: value[0] for key, value in data.items()}
-
-    main_options_form = f'''
-    <div class="form-group row">
-    <label for="memory" class="col-2 col-form-label">JupyterLab Memory (GB)</label>
-    <div class="col-10">
-        <input class="form-control" type="number" value="1" id="memory" name="memory">
-    </div>
-    </div>
-    <div class="form-group row">
-    <label for="nprocs" class="col-2 col-form-label">JupyterLab CPUs</label>
-    <div class="col-10">
-        <input class="form-control" type="number" value="1" id="nprocs" name="nprocs">
-    </div>
-    </div>
-    <div class="form-group row">
-    <label for="partition" class="col-2 col-form-label">Slurm Partition</label>
-    <div class="col-10">
-        <select class="form-select" aria-label="Slurm Queue" id="partition" name="partition">
-        <option value="general">general</option>
-    {% for item in groups %}
-    {% if item.startswith('partition-')%}
-        <option value="{{ item[10:] }}">{{ item[10:] }}</option>
-    {% endif %}
-    {% endfor %}
-        </select>
-    </div>
-    </div>
-    '''
-
-    conda_options_form = f'''
-    {% raw %}
-    <div class="form-group row">
-    <label for="conda_environment_prefix" class="col-2 col-form-label">Conda Environment</label>
-    <div class="col-10">
-        <select class="form-select" aria-label="Conda Environment" id="conda_environment_prefix" name="conda_environment_prefix">
-    {''.join([f'<option value="{_[1]}">{_[0]}</option>' for _ in conda_envs_w_packages(jupyterlab_packages)])}
-        </select>
-    </div>
-    </div>
-    {% endraw %}
-    '''
-
-    def options_form(self, spawner):
-        ## Not currently working - idea is to omit conda env from spawner form
-        ## for dashboards, since already selected a conda env
-        #if spawner.orm_spawner.user_options and 'presentation_type' in spawner.orm_spawner.user_options:
-        #  self.log.info("In options_form")
-        #  if spawner.user_options['presentation_type']:
-        #    return self.main_options_form # Omit the conda env dropdown since that is chosen per dashboard
-
-        # Display full form including conda env dropdown
-        return ''.join([self.main_options_form, self.conda_options_form])
-{% endif %}
-
 # Assign Qhub Spawner
 class QHubHPCSpawner(QHubHPCSpawnerBase):
     pass
@@ -216,7 +154,7 @@ class QHubHPCSpawner(QHubHPCSpawnerBase):
 c.JupyterHub.template_paths = []
 c.JupyterHub.extra_handlers = []
 
-c.JupyterHub.spawner_class = QHubHPCSpawner
+c.JupyterHub.spawner_class = 'wrapspawner.ProfilesSpawner'
 
 c.SlurmSpawner.start_timeout = {{ jupyterhub_config.spawner.start_timeout }}
 c.QHubHPCSpawner.default_url = '/lab'
@@ -224,6 +162,7 @@ class QHubHPCSpawner(QHubHPCSpawnerBase):
 # default values for batch spawner
 c.QHubHPCSpawner.req_memory = '1' # GB
 c.QHubHPCSpawner.req_nprocs = '1'
+c.QHubHPCSpawner.req_partition = 'general'
 c.QHubHPCSpawner.req_conda_environment_prefix = '{{ miniforge_home }}/envs/{{ jupyterhub_lab_environment | basename | splitext | first }}'
 c.QHubHPCSpawner.req_prologue = '''
 # ensure user has link to the shared directory, if it exists
@@ -312,6 +251,24 @@ def generate_batch_script(spawner):
 
 c.QHubHPCSpawner.batch_script = generate_batch_script
 
+USER_PROFILES = {{ jupyterhub_profiles | tojson }}
+
+def set_profiles(user_profiles: list):
+    """Set the profiles for the ProfilesSpawner."""
+    profiles = []
+    for profile_dict in user_profiles:
+        profile_name, profile_data = next(iter(profile_dict.items()))
+        profiles.append((
+            profile_data['display_name'],
+            profile_name,
+            QHubHPCSpawner,
+            dict(**profile_data['options'])
+        ))
+    return profiles
+
+
+# User profiles
+c.ProfilesSpawner.profiles = set_profiles(USER_PROFILES)
 
 # ===== adding api tokens for external services =======
 c.JupyterHub.services = [