diff --git a/src/glicid_spawner/micromamba.py b/src/glicid_spawner/micromamba.py new file mode 100644 index 0000000..e72d328 --- /dev/null +++ b/src/glicid_spawner/micromamba.py @@ -0,0 +1,75 @@ +"""Micromamba environments module.""" + +from dataclasses import dataclass +from operator import itemgetter +from pathlib import Path + +MICROMAMBA_ROOT = Path('/micromamba') # default micromamba root location + +GLOBAL_USER = 'operator' +GLOBAL_EXCLUDED = 'glicid-jupyterhub' + + +@dataclass +class MicromambaEnv: + """Generic micromamba environment.""" + + scope: str + name: str + path: str + + +def _envs(folder, excluded=None) -> list: + """List micromamba environments.""" + # Convert excluded as as list + if excluded is None: + excluded = [] + elif isinstance(excluded, str): + excluded = [excluded] + + # Get micromamba envs as pathlib.Path + envs = MICROMAMBA_ROOT / folder / 'envs' + + if not envs.exists(): + return [] + + return sorted( + [ + (env.name, str(env)) + for env in envs.iterdir() + if env.is_dir() and env.name not in excluded + ], + key=itemgetter(0), + ) + + +def _envs_user(username: str) -> list: + """Micromamba environment(s) available to the user.""" + return [MicromambaEnv('USER', name, path) for name, path in _envs(username)] + + +def _envs_team(username: str) -> list: + """Micromamba environment(s) available to the user's team. + + Warning + ------- + At the moment micromamba team environments is not available on GLiCID. + + """ + teams = [] # FIXME: pull user teams list from groups + return [ + MicromambaEnv('TEAM', name, path) for team in teams for name, path in _envs(f'teams/{team}') + ] + + +def _envs_global() -> list: + """Micromamba environment(s) available globally.""" + return [ + MicromambaEnv('GLOBAL', name, path) + for name, path in _envs(GLOBAL_USER, excluded=GLOBAL_EXCLUDED) + ] + + +def get_envs(username: str) -> list: + """List of all the micromamba environment available to a user.""" + return _envs_user(username) + _envs_team(username) + _envs_global() diff --git a/src/glicid_spawner/pyenv.py b/src/glicid_spawner/pyenv.py deleted file mode 100644 index eb2e072..0000000 --- a/src/glicid_spawner/pyenv.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Python environment module.""" - -from dataclasses import dataclass -from operator import itemgetter -from pathlib import Path - - -@dataclass -class PyEnv: - """Python environment generic class.""" - - name: str - path: str - - def __eq__(self, item): - return self.path == str(item) - - -class UserPyEnv(PyEnv): - """User python environment.""" - - scope = 'USER' - - -class TeamPyEnv(PyEnv): - """Team python environment.""" - - scope = 'TEAM' - - -class GlobalPyEnv(PyEnv): - """Global python environment.""" - - scope = 'GLOBAL' - - -def _micromamba_path(path: str) -> Path: - """Micromamba path environments locations.""" - return Path(f'/micromamba/{path}/envs') - - -def _micromamba_envs(path, excluded=None) -> list: - """List micromamba environment list.""" - if excluded is None: - excluded = [] - elif isinstance(excluded, str): - excluded = [excluded] - - envs = _micromamba_path(path) - return sorted( - [(f.name, f.absolute()) for f in envs.iterdir() if f.is_dir() and f.name not in excluded], - key=itemgetter(0), - ) - - -def get_pyenv_user(username: str) -> list: - """List of python environment available to the user.""" - return [UserPyEnv(*env) for env in _micromamba_envs(username)] - - -def get_pyenv_team(username: str) -> list: - """List of python environment available to the user's team.""" - teams = [] # FIXME: pull user teams list from groups - return [TeamPyEnv(*env) for team in teams for env in _micromamba_envs(f'teams/{team}')] - - -def get_pyenv_global() -> list: - """List of python environment available globally.""" - return [GlobalPyEnv(*env) for env in _micromamba_envs('operator', excluded='glicid-jupyterhub')] - - -def get_pyenv(username: str) -> list: - """List of all the python environment available to a user.""" - return get_pyenv_user(username) + get_pyenv_team(username) + get_pyenv_global() diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index 8906780..5f42c8e 100644 --- a/src/glicid_spawner/spawner.py +++ b/src/glicid_spawner/spawner.py @@ -3,7 +3,7 @@ from batchspawner import SlurmSpawner from jinja2 import Environment, PackageLoader, select_autoescape -from .pyenv import get_pyenv +from .micromamba import get_envs from .resources import CPU, GPU, RAM @@ -20,16 +20,14 @@ class GlicidSpawner(SlurmSpawner): return template.render( username=self.user.name, - python_envs=get_pyenv(self.user.name), + python_envs=get_envs(self.user.name), cpu_available=CPU, ram_available=RAM, gpu_available=GPU, ) def options_from_form(self, formdata) -> dict: - options = {} - options['pyenv'] = formdata['python-env'][0] - + """Export options from form.""" # Index of user resources choices i_cpu = int(formdata['cpu'][0]) i_ram = int(formdata['ram'][0]) @@ -42,8 +40,12 @@ class GlicidSpawner(SlurmSpawner): ) # Export options + options = {} + + options['pyenv'] = formdata['python-env'][0] + options['cpus-per-task'] = CPU[i_cpu].description - options['mem'] = RAM[i_ram].description + 'GB' + options['mem'] = RAM[i_ram].description.replace(' ', '') if i_gpu: options['gres'] = 'gpu:' + GPU[i_gpu].description diff --git a/tests/data/micromamba/global/envs/baz/.gitkeep b/tests/data/micromamba/global/envs/baz/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/micromamba/global/envs/qux/.gitkeep b/tests/data/micromamba/global/envs/qux/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/micromamba/john-doe/envs/bar/.gitkeep b/tests/data/micromamba/john-doe/envs/bar/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/micromamba/john-doe/envs/foo/.gitkeep b/tests/data/micromamba/john-doe/envs/foo/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_micromamba.py b/tests/test_micromamba.py new file mode 100644 index 0000000..2229abd --- /dev/null +++ b/tests/test_micromamba.py @@ -0,0 +1,36 @@ +"""Test python environment module.""" + +from pathlib import Path + +from glicid_spawner import micromamba + +MICROMAMBA_ROOT = Path(__file__).parent / 'data' / 'micromamba' + + +def test_micromamba_envs_getter(monkeypatch): + """Test micromamba envs getter.""" + monkeypatch.setattr(micromamba, 'MICROMAMBA_ROOT', MICROMAMBA_ROOT) + monkeypatch.setattr(micromamba, 'GLOBAL_USER', 'global') + monkeypatch.setattr(micromamba, 'GLOBAL_EXCLUDED', 'qux') + + # User with micromamba envs + envs = micromamba.get_envs('john-doe') + + assert len(envs) == 3 + assert [env.scope for env in envs] == ['USER', 'USER', 'GLOBAL'] + assert [env.name for env in envs] == ['bar', 'foo', 'baz'] + assert [env.path for env in envs] == [ + str(MICROMAMBA_ROOT / 'john-doe' / 'envs' / 'bar'), + str(MICROMAMBA_ROOT / 'john-doe' / 'envs' / 'foo'), + str(MICROMAMBA_ROOT / 'global' / 'envs' / 'baz'), + ] + + # User without micromamba envs, only global envs listed + envs = micromamba.get_envs('jane-smith') + + assert len(envs) == 1 + assert [env.scope for env in envs] == ['GLOBAL'] + assert [env.name for env in envs] == ['baz'] + assert [env.path for env in envs] == [ + str(MICROMAMBA_ROOT / 'global' / 'envs' / 'baz'), + ]