From d2c0631956012725ee35be19d470835aef4a5258 Mon Sep 17 00:00:00 2001
From: Benoit Seignovert <benoit.seignovert@univ-nantes.fr>
Date: Tue, 26 Mar 2024 08:05:00 +0100
Subject: [PATCH] Add explicit PYTHONUSERBASE path to micromamba endpoint

---
 src/glicid_spawner/micromamba.py                |  1 +
 src/glicid_spawner/spawner.py                   | 13 +++++++++----
 src/glicid_spawner/templates/slurm_script.jinja |  1 +
 tests/test_spawner.py                           |  7 ++++++-
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/glicid_spawner/micromamba.py b/src/glicid_spawner/micromamba.py
index 1d42063..0fe1641 100644
--- a/src/glicid_spawner/micromamba.py
+++ b/src/glicid_spawner/micromamba.py
@@ -9,6 +9,7 @@ GLOBAL_USER = 'operator'
 
 MAMBA_ROOT_PREFIX = f'{MICROMAMBA_ROOT}/{GLOBAL_USER}'
 MAMBA_EXE = f'{MAMBA_ROOT_PREFIX}/bin/micromamba'
+MAMBA_USER_BASE = f'{MICROMAMBA_ROOT}/$USER'
 
 
 @dataclass
diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py
index cdd7631..b65af99 100644
--- a/src/glicid_spawner/spawner.py
+++ b/src/glicid_spawner/spawner.py
@@ -9,7 +9,7 @@ from batchspawner import JobStatus, SlurmSpawner
 from traitlets import Bool, Integer, Unicode, default
 
 from .form import options_form, options_from_form
-from .micromamba import MAMBA_EXE, MAMBA_ROOT_PREFIX
+from .micromamba import MAMBA_EXE, MAMBA_ROOT_PREFIX, MAMBA_USER_BASE
 from .progress import ElapseTime, get_progress
 from .templates import get_template_src
 
@@ -24,14 +24,19 @@ class GlicidSpawner(SlurmSpawner):
         help='Spawner singleuser command.',
     ).tag(config=True)
 
+    req_mamba_exe = Unicode(
+        MAMBA_EXE,
+        help='Micromamba global exe',
+    ).tag(config=True)
+
     req_mamba_root_prefix = Unicode(
         MAMBA_ROOT_PREFIX,
         help='Micromamba global root prefix',
     ).tag(config=True)
 
-    req_mamba_exe = Unicode(
-        MAMBA_EXE,
-        help='Micromamba global exe',
+    req_mamba_user_base = Unicode(
+        MAMBA_USER_BASE,
+        help='Micromamba user base prefix',
     ).tag(config=True)
 
     req_job_name = Unicode(
diff --git a/src/glicid_spawner/templates/slurm_script.jinja b/src/glicid_spawner/templates/slurm_script.jinja
index 3c01679..c3e4515 100644
--- a/src/glicid_spawner/templates/slurm_script.jinja
+++ b/src/glicid_spawner/templates/slurm_script.jinja
@@ -33,6 +33,7 @@ echo "The {{job_name}} logs are located in: ${JUPYTER_LOG_DIR}"
     {# Micromamba config -#}
     export MAMBA_EXE={{mamba_exe}};
     export MAMBA_ROOT_PREFIX={{mamba_root_prefix}};
+    export PYTHONUSERBASE={{ mamba_user_base }};
     source $MAMBA_ROOT_PREFIX/etc/profile.d/micromamba.sh;
 
     {# Activate micromamba env requested by the user -#}
diff --git a/tests/test_spawner.py b/tests/test_spawner.py
index 8ac208b..7990185 100644
--- a/tests/test_spawner.py
+++ b/tests/test_spawner.py
@@ -27,8 +27,9 @@ def test_spawner_config():
     assert 'bin/glicid-spawner-singleuser' in cmd
     assert 'bin/jupyterhub-singleuser' in cmd
 
-    assert spawner.req_mamba_root_prefix == '/micromamba/operator'
     assert spawner.req_mamba_exe == '/micromamba/operator/bin/micromamba'
+    assert spawner.req_mamba_root_prefix == '/micromamba/operator'
+    assert spawner.req_mamba_user_base == '/micromamba/$USER'
     assert spawner.req_job_name == 'jupyterhub_glicid'
     assert spawner.req_qos == 'short'
 
@@ -70,8 +71,12 @@ def test_spawner_batch_script(monkeypatch):
 
     assert 'export MAMBA_EXE=/micromamba/operator/bin/micromamba;' in script
     assert 'export MAMBA_ROOT_PREFIX=/micromamba/operator;' in script
+    assert 'export PYTHONUSERBASE=/micromamba/$USER;' in script
     assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script
 
+    assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script
+    assert 'export JUPYTER_PATH=/micromamba/john-doe/envs/foo/share/jupyter;' in script
+
     assert re.search(r'.*/bin/glicid-spawner-singleuser .*/bin/jupyterhub-singleuser', script)