Add SLURM batch script

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.',

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