Move SLURM script to jinja template

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,
     )
 

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:

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)
 

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)

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

tests/test_form.py

@@ -1,4 +1,4 @@
-"""Test form templates module."""
+"""Test template form module."""
 
 import re
 from pathlib import Path

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')