Add global mamba prefixes

This commit is contained in:
Benoît Seignovert 2024-03-08 18:00:46 +01:00
parent e7ab73de0b
commit d091a6fd99
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
3 changed files with 38 additions and 20 deletions

View file

@ -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:

View file

@ -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):

View file

@ -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