diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index f3f6853..5152f33 100644 --- a/src/glicid_spawner/spawner.py +++ b/src/glicid_spawner/spawner.py @@ -13,6 +13,65 @@ from .progress import ElapseTime, get_progress class GlicidSpawner(SlurmSpawner): """Glicid SLURM Spawner.""" + batch_script = Unicode( + """#!/bin/bash +#SBATCH --job-name={{job_name}} +#SBATCH --output={{homedir}}/.{{job_name}}.log +#SBATCH --chdir={{workdir}} +#SBATCH --export={{keepvars}} + +{% if cluster -%}#SBATCH --cluster={{cluster}}{%- endif %} +{% if partition -%}#SBATCH --partition={{partition}}{%- endif %} +{% if node -%}#SBATCH --nodelist={{node}}{%- endif %} +{% if qos -%}#SBATCH --qos={{qos}}{%- endif %} + +{% if runtime -%}#SBATCH --time={{runtime}}{%- endif %} +{% if nprocs -%}#SBATCH --cpus-per-task={{nprocs}}{%- endif %} +{% if memory -%}#SBATCH --mem={{memory}}{%- endif %} +{% if gres -%}#SBATCH --gres={{gres}}{%- endif %} + +# Redirect logs +export JUPYTER_LOG_DIR="{{homedir}}/.jupyter/spawner/logs" +mkdir -p ${JUPYTER_LOG_DIR} + +echo "The {{job_name}} logs are located in: ${JUPYTER_LOG_DIR}" + +export JUPYTER_JOB_LOG=${JUPYTER_LOG_DIR}/$(date "+%Y-%m-%d")_{{job_name}}_${SLURM_JOB_ID}.log + +{ PS4='[$(date "+%Y-%m-%d %T")]\011 '; + + set -xeo pipefail; + trap 'echo SIGTERM received' TERM; + + {# SLURM config #} + scontrol write batch_script ${SLURM_JOB_ID} -; + + {# Micromamba config #} + export MAMBA_ROOT_PREFIX=/micromamba/operator; + export MAMBA_EXE=$MAMBA_ROOT_PREFIX/bin/micromamba; + source $MAMBA_ROOT_PREFIX/etc/profile.d/micromamba.sh; + + {# Activate micromamba env requested by the user #} + micromamba activate {{ pyenv }}; + export JUPYTER_PATH={{ pyenv }}/share/jupyter; + + {# Prologue #} + {%- if prologue -%} + {{prologue}}; + {%- endif -%} + + {# Start Jupyter single-user command #} + {{cmd}}; + + {# Epilogue #} + {%- if epilogue -%} + {{epilogue}}; + {%- endif -%} +} > ${JUPYTER_JOB_LOG} 2>&1 +""", + help='Template for SLURM job submission batch script.', + ).tag(config=True) + disable_user_config = Bool( True, help='Disable per-user configuration of single-user servers.', diff --git a/tests/test_spawner.py b/tests/test_spawner.py index fb84036..a6e4f01 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -5,11 +5,12 @@ from itertools import repeat import glicid_spawner.spawner import pytest -from batchspawner import SlurmSpawner +from batchspawner import SlurmSpawner, format_template from glicid_spawner import GlicidSpawner from glicid_spawner.spawner import asyncio -User = namedtuple('User', 'name') +User = namedtuple('User', 'name,url') +Hub = namedtuple('Hub', 'public_host,api_url,base_url') def test_spawner_config(): @@ -29,6 +30,35 @@ def test_spawner_config(): assert spawner.progress_rate == 10 +def test_spawner_batch_script(monkeypatch): + """Test spawner SLURM batch script.""" + monkeypatch.setattr(GlicidSpawner, 'get_env', lambda _: {}) + + spawner = GlicidSpawner( + user=User('john-doe', '/john-doe/'), + req_homedir='/home/john-doe', + hub=Hub('public.hostname', 'api.url', 'base.url'), + user_options={ + 'cluster': 'nautilus', + 'workdir': '/home/john-doe', + 'pyenv': '/micromamba/john-doe/envs/foo', + }, + ) + + script = format_template( + spawner.batch_script, **(spawner.get_req_subvars() | spawner.user_options) + ) + + assert '#SBATCH --job-name=jupyterhub_glicid' in script + assert '#SBATCH --chdir=/home/john-doe' in script + assert '#SBATCH --cluster=nautilus' in script + assert '#SBATCH --qos=short' in script + + assert 'export JUPYTER_LOG_DIR="/home/john-doe/.jupyter/spawner/logs"' in script + + assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script + + def test_spawner_parse_job_id(): """Test spawner job id parser.""" spawner = GlicidSpawner() @@ -50,7 +80,7 @@ def test_spawner_options_form(monkeypatch): lambda formdata: f'options_from_form({formdata})', ) - spawner = GlicidSpawner(user=User('john-doe')) + spawner = GlicidSpawner(user=User('john-doe', '/john-doe/')) assert spawner.user.name == 'john-doe' assert spawner.options_form == "options_form('john-doe')"