Add ability to pass user groups to jupyterhub context

docs/user-guide.md

@@ -4,10 +4,50 @@
 
 ### Adding and updating Conda environments
 
-Managing conda environments is done via
-[Conda-Store](https://conda-store.readthedocs.io/en/latest/). You can
-visit Conda-Store within your QHub-HPC deployment at
-`https://<hpc_master>/conda-store/`.
+#### What is `conda-store`?
+
+[`conda-store`][conda-store-docs] is a Python package that serves _identical_
+`conda` environments by controlling the environment lifecycle.
+It ensures that the management, building, and serving of environments is as
+identical as possible and seamless for the end users.
+
+All environments in Nebari are served through `conda-store`.
+
+Using `conda-store`, Nebari admins can track specific files or directories for
+changes in environment specifications. They can manage environments using the
+web interface, REST API, or the command-line utility (CLI).
+
+#### Exploring the conda-store Interface
+
+Access conda-store through your domain `<nebari-slurm-domain/conda-store>`, log
+in to authenticate, and navigate the dashboard to view account details and permissions.
+
+- Key sections include User, Namespaces, and Permissions, which dictate access
+  levels and capabilities.
+
+![conda-store default main page, before authentication, login button highlighted](https://conda.store/assets/images/login-1346a06ed408f74937da23b0a1c6fda3.png)
+
+#### Creating a New Environment
+
+Environments are created in conda-store using a YAML file. Post-creation, the
+environment can be managed through the conda-store UI, allowing for edits and
+build status monitoring.
+
+More details on creating environments can be found in the [conda-store documentation](https://conda.store/conda-store-ui/tutorials/create-envs).
+
+Package installation should be done via the conda-store web interface to avoid
+inconsistencies and limitations associated with command line installations.
+
+#### **Note on Shared Namespaces**
+
+Access to shared namespaces in conda-store depends on user assignment to groups
+in Keycloak and the corresponding permissions within those groups.
+
+By default, NebarSlurm is deployed with the following groups: `admin`, `developer`,
+and `analyst` (in roughly descending order of permissions and scope). Note that
+such group names will differ on a per-instance basis. Check
+[Conda-store authorization model](https://conda-store.readthedocs.io/en/latest/contributing.html#authorization-model)
+for more details on conda-store authorization.
 
 ## ContainDS Dashboards
 
@@ -43,3 +83,4 @@ Dask support is based on:
 - [dask distributed](https://distributed.dask.org/en/latest/)
 - [dask gateway](https://gateway.dask.org/)
 
+[conda-store-docs]: https://conda-store.readthedocs.io/

roles/jupyterhub/templates/jupyterhub_config.py

@@ -17,21 +17,6 @@
 from jupyterhub.traitlets import Callable
 from tornado import gen
 
-# Find all conda environments that have dask jupyterlab, batchspawner, and jupyterhub installed
-jupyterlab_packages = ['jupyterlab', 'batchspawner', 'jupyterhub']
-def conda_envs_w_packages(packages, names_only=False):
-    _environments = []
-    output = subprocess.check_output(['conda', 'env', 'list', '--json'])
-    environments = json.loads(output)['envs']
-    for environment in environments:
-        output = subprocess.check_output(['conda', 'list', '-p', environment, '--json'])
-        if set(packages)  <= {_['name'] for _ in json.loads(output)}:
-            _environments.append((os.path.basename(environment), environment))
-    if names_only:
-        return [env_name for env_name, path in _environments]
-    return _environments
-
-
 # Allow gathering of jupyterhub prometheus metrics
 c.JupyterHub.authenticate_prometheus = False
 
@@ -148,13 +133,33 @@ async def _get_batch_script(self, **subvars):
 class QHubHPCSpawner(QHubHPCSpawnerBase):
     pass
 
+# batchspawner does not support auth_state correctly right now, using Keycloak client to retrieve user-info
+@gen.coroutine
+def _get_user_groups_state(spawner):
+    user_name = spawner.user.name
+    keycloak_admin = keycloak.KeycloakAdmin(
+            server_url="https://{{ traefik_domain | default(hostvars[groups['hpc_master'][0]].ansible_ssh_host) }}/auth/",
+            username="{{ keycloak_admin_username }}",
+            password="{{ keycloak_admin_password }}",
+            realm_name="{{ keycloak_realm }}",
+            user_realm_name="master",
+            verify=False)
+    username_uid = keycloak_admin.get_user_id(user_name)
+    _user_groups = keycloak_admin.get_user_groups(username_uid)
+    spawner.environment.update(
+        {
+            "USER_GROUPS": json.dumps([group['name'] for group in _user_groups])
+        }
+    )
+
 c.JupyterHub.allow_named_servers = True
 c.JupyterHub.default_url = '/hub/home'
 
 c.JupyterHub.template_paths = []
 c.JupyterHub.extra_handlers = []
 
 c.JupyterHub.spawner_class = 'wrapspawner.ProfilesSpawner'
+c.Spawner.pre_spawn_hook = _get_user_groups_state
 
 c.SlurmSpawner.start_timeout = {{ jupyterhub_config.spawner.start_timeout }}
 c.QHubHPCSpawner.default_url = '/lab'
@@ -176,11 +181,13 @@ class QHubHPCSpawner(QHubHPCSpawnerBase):
 export PATH={{ miniforge_home }}/condabin:$PATH
 '''
 
-def populate_condarc(username):
+
+def populate_condarc(username, groups):
     """Generate condarc configuration string for the given username."""
+    # # only run if conda-store is enabled and the jupyterhub service token is available
     condarc = json.dumps({
         "envs_dirs": [
-            f"/opt/conda-store/conda-store/{dir_name}/envs" for dir_name in [username, "filesystem"]
+            f"/opt/conda-store/conda-store/{dir_name}/envs" for dir_name in [username, "filesystem", *groups]
         ]
     })
     return f"printf '{condarc}' > /home/{username}/.condarc\n"
@@ -189,11 +196,7 @@ def populate_condarc(username):
 def generate_batch_script(spawner):
     """Generate a batch script for SLURM and JupyterHub based on spawner settings."""
     username = spawner.user.name
-
-    auth_state = yield spawner.user.get_auth_state()
-    if auth_state:
-        print(f"auth_state: {auth_state}")
-        print("#######################")
+    _groups = json.loads(spawner.environment.get("USER_GROUPS", "[]"))
 
     print(f"Generating batch script for {username}")
 
@@ -244,7 +247,7 @@ def generate_batch_script(spawner):
 
     return "".join([
         sbatch_headers,
-        populate_condarc(username),
+        populate_condarc(username, _groups),
         conda_store_headers,
         srun_jupyterhub_single_user
     ])