diff --git a/src/glicid_spawner/micromamba.py b/src/glicid_spawner/micromamba.py index 0f62bb4..1d42063 100644 --- a/src/glicid_spawner/micromamba.py +++ b/src/glicid_spawner/micromamba.py @@ -5,9 +5,11 @@ from operator import itemgetter from pathlib import Path MICROMAMBA_ROOT = Path('/micromamba') # default micromamba root location - GLOBAL_USER = 'operator' +MAMBA_ROOT_PREFIX = f'{MICROMAMBA_ROOT}/{GLOBAL_USER}' +MAMBA_EXE = f'{MAMBA_ROOT_PREFIX}/bin/micromamba' + @dataclass class MicromambaEnv: diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index 5152f33..93265a4 100644 --- a/src/glicid_spawner/spawner.py +++ b/src/glicid_spawner/spawner.py @@ -7,12 +7,38 @@ 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 .progress import ElapseTime, get_progress class GlicidSpawner(SlurmSpawner): """Glicid SLURM Spawner.""" + batchspawner_singleuser_cmd = Unicode( + 'glicid-spawner-singleuser', + help='Spawner singleuser command.', + ).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', + ).tag(config=True) + + req_job_name = Unicode( + 'jupyterhub_glicid', + help='SLURM job name', + ).tag(config=True) + + req_qos = Unicode( + 'short', + help='QoS name to submit job to resource manager', + ).tag(config=True) + batch_script = Unicode( """#!/bin/bash #SBATCH --job-name={{job_name}} @@ -47,8 +73,8 @@ export JUPYTER_JOB_LOG=${JUPYTER_LOG_DIR}/$(date "+%Y-%m-%d")_{{job_name}}_${SLU scontrol write batch_script ${SLURM_JOB_ID} -; {# Micromamba config #} - export MAMBA_ROOT_PREFIX=/micromamba/operator; - export MAMBA_EXE=$MAMBA_ROOT_PREFIX/bin/micromamba; + export MAMBA_EXE={{mamba_exe}}; + export MAMBA_ROOT_PREFIX={{mamba_root_prefix}}; source $MAMBA_ROOT_PREFIX/etc/profile.d/micromamba.sh; {# Activate micromamba env requested by the user #} @@ -87,21 +113,6 @@ export JUPYTER_JOB_LOG=${JUPYTER_LOG_DIR}/$(date "+%Y-%m-%d")_{{job_name}}_${SLU """The URL the single-user server should start in.""" return '/lab/tree' + self.user_options.get('workdir', '/home/{username}') + '?reset' - batchspawner_singleuser_cmd = Unicode( - 'glicid-spawner-singleuser', - help='Spawner singleuser command.', - ).tag(config=True) - - req_job_name = Unicode( - 'jupyterhub_glicid', - help='SLURM job name', - ).tag(config=True) - - req_qos = Unicode( - 'short', - help='QoS name to submit job to resource manager', - ).tag(config=True) - slurm_job_id_re = Unicode(r'(\d+)(?:;(\w+))?').tag(config=True) def parse_job_id(self, output): diff --git a/tests/test_spawner.py b/tests/test_spawner.py index a6e4f01..2eb88e0 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -22,11 +22,14 @@ def test_spawner_config(): assert spawner.batchspawner_singleuser_cmd == 'glicid-spawner-singleuser' + assert spawner.req_mamba_root_prefix == '/micromamba/operator' + assert spawner.req_mamba_exe == '/micromamba/operator/bin/micromamba' + assert spawner.req_job_name == 'jupyterhub_glicid' + assert spawner.req_qos == 'short' + assert spawner.disable_user_config assert spawner.notebook_dir == '/' assert spawner.default_url == '/lab/tree/home/{username}?reset' - assert spawner.req_job_name == 'jupyterhub_glicid' - assert spawner.req_qos == 'short' assert spawner.progress_rate == 10 @@ -56,6 +59,8 @@ def test_spawner_batch_script(monkeypatch): assert 'export JUPYTER_LOG_DIR="/home/john-doe/.jupyter/spawner/logs"' in script + assert 'export MAMBA_EXE=/micromamba/operator/bin/micromamba;' in script + assert 'export MAMBA_ROOT_PREFIX=/micromamba/operator;' in script assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script