From 8217d63c41f6ef627cdc152dd13b18f16ba5ac61 Mon Sep 17 00:00:00 2001 From: Benoit Seignovert Date: Wed, 13 Mar 2024 18:38:29 +0100 Subject: [PATCH] Move SLURM script to jinja template --- render/__main__.py | 5 +- src/glicid_spawner/form.py | 13 +---- src/glicid_spawner/spawner.py | 57 +------------------ src/glicid_spawner/templates.py | 25 ++++++++ .../templates/slurm_script.jinja | 54 ++++++++++++++++++ tests/test_form.py | 2 +- tests/test_templates.py | 20 +++++++ 7 files changed, 108 insertions(+), 68 deletions(-) create mode 100644 src/glicid_spawner/templates.py create mode 100644 src/glicid_spawner/templates/slurm_script.jinja create mode 100644 tests/test_templates.py diff --git a/render/__main__.py b/render/__main__.py index 6f84032..f554c5a 100644 --- a/render/__main__.py +++ b/render/__main__.py @@ -9,10 +9,11 @@ from pathlib import Path from traceback import format_exc from flask import Flask, render_template, request -from glicid_spawner.form import TEMPLATES, options_from_form +from glicid_spawner.form import options_from_form from glicid_spawner.micromamba import MicromambaEnv from glicid_spawner.resources import CPU, MEMORY, get_folders, gpu_max_duration from glicid_spawner.slurm import gres, sinfo_from_file +from glicid_spawner.templates import render_template as glicid_template from livereload import Server # Dummy username and python environments @@ -58,7 +59,7 @@ def home(): """Form spawner home page.""" return render_template( 'form.html', - spawner_options_form=TEMPLATES.get_template('spawner_form.jinja').render(**OPTIONS), + spawner_options_form=glicid_template('spawner_form.jinja', **OPTIONS), options=OPTIONS, ) diff --git a/src/glicid_spawner/form.py b/src/glicid_spawner/form.py index 56e6db4..312d010 100644 --- a/src/glicid_spawner/form.py +++ b/src/glicid_spawner/form.py @@ -1,15 +1,9 @@ -"""GLiCID form templates module.""" - -from jinja2 import Environment, PackageLoader, select_autoescape +"""GLiCID template form module.""" from .micromamba import get_envs from .resources import CPU, GPU_DEFAULTS, MEMORY, get_folders, gpu_max_duration from .slurm import gres, sinfo - -TEMPLATES = Environment( - loader=PackageLoader('glicid_spawner'), - autoescape=select_autoescape(), -) +from .templates import render_template GRES_CLUSTERS_LOWERCASE = [ 'waves', @@ -36,8 +30,7 @@ def options_attrs(username: str) -> dict: def options_form(username: str) -> str: """Render default spawner form.""" - template = TEMPLATES.get_template('spawner_form.jinja') - return template.render(**options_attrs(username)) + return render_template('spawner_form.jinja', **options_attrs(username)) def options_from_form(formdata) -> dict: diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index 93265a4..65ccbf7 100644 --- a/src/glicid_spawner/spawner.py +++ b/src/glicid_spawner/spawner.py @@ -9,6 +9,7 @@ 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 +from .templates import get_template_src class GlicidSpawner(SlurmSpawner): @@ -40,61 +41,7 @@ class GlicidSpawner(SlurmSpawner): ).tag(config=True) 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_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 #} - 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 -""", + get_template_src('slurm_script.jinja'), help='Template for SLURM job submission batch script.', ).tag(config=True) diff --git a/src/glicid_spawner/templates.py b/src/glicid_spawner/templates.py new file mode 100644 index 0000000..adc5de4 --- /dev/null +++ b/src/glicid_spawner/templates.py @@ -0,0 +1,25 @@ +"""GLiCID Jinja templates module.""" +from pathlib import Path + +from jinja2 import Environment, PackageLoader, Template, select_autoescape + +TEMPLATES = Environment( + loader=PackageLoader('glicid_spawner'), + autoescape=select_autoescape(), +) + + +def get_template(name: str) -> Template: + """Get Jinja2 template.""" + return TEMPLATES.get_template(name) + + +def get_template_src(name: str) -> str: + """Get Jinja2 template source.""" + template = get_template(name) + return Path(template.filename).read_text() + + +def render_template(name: str, **kwargs) -> str: + """Render Jinja2 template.""" + return get_template(name).render(**kwargs) diff --git a/src/glicid_spawner/templates/slurm_script.jinja b/src/glicid_spawner/templates/slurm_script.jinja new file mode 100644 index 0000000..ad403d4 --- /dev/null +++ b/src/glicid_spawner/templates/slurm_script.jinja @@ -0,0 +1,54 @@ +#!/bin/bash +{#- SLURM batch script for GLiCID spawner #} + +#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 %} +{# Export logs to $JUPYTER_LOG_DIR folder -#} +export JUPYTER_LOG_DIR="{{homedir}}/.jupyter/spawner/logs" +export JUPYTER_JOB_LOG=${JUPYTER_LOG_DIR}/$(date "+%Y-%m-%d")_{{job_name}}_${SLURM_JOB_ID}.log + +mkdir -p ${JUPYTER_LOG_DIR} +echo "The {{job_name}} logs are located in: ${JUPYTER_LOG_DIR}" + +{# Redirect all Jupyter logs into $JUPYTER_LOG_DIR -#} +{ 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_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 -#} + 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 diff --git a/tests/test_form.py b/tests/test_form.py index 9fd17fb..99b8895 100644 --- a/tests/test_form.py +++ b/tests/test_form.py @@ -1,4 +1,4 @@ -"""Test form templates module.""" +"""Test template form module.""" import re from pathlib import Path diff --git a/tests/test_templates.py b/tests/test_templates.py new file mode 100644 index 0000000..ad71595 --- /dev/null +++ b/tests/test_templates.py @@ -0,0 +1,20 @@ +"""Test Jinja templates module.""" + +from glicid_spawner.templates import get_template_src, render_template + + +def test_template_src(): + """Test jinja template source.""" + src = get_template_src('slurm_script.jinja') + + assert src.startswith('#!/bin/bash\n{#- SLURM batch script for GLiCID spawner #}') + assert src.endswith('} > ${JUPYTER_JOB_LOG} 2>&1\n') + + +def test_template_render(): + """Test jinja template render.""" + src = render_template('slurm_script.jinja', job_name='foo', qos='short', gres='bar', cmd='cmd') + + assert src.startswith('#!/bin/bash\n\n#SBATCH --job-name=foo\n') + assert '#SBATCH --export=\n#SBATCH --qos=short\n#SBATCH --gres=bar\n' in src + assert src.endswith('cmd;} > ${JUPYTER_JOB_LOG} 2>&1')