Merge pull request #170 from Hoeze/master

add 'gres' option to SlurmSpawner
jupyterhub · Jul 23, 2020 · ec5e972 · ec5e972
2 parents 9e7d557 + 50da68f
commit ec5e972
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -24,6 +24,7 @@ Added (user)
 * Add `req_prologue` and `req_epilogue` options to scripts which are inserted before/after the main jupyterhub-singleuser command, which allow for generic setup/cleanup without overriding the entire script.  #96
 * SlurmSpawner: add the `req_reservation` option.  #91
 * Add basic support for JupyterHub progress updates, but this is not used much yet.  #86
+# PR #170: SlurmSpawner: add `req_gres` to specify `-go-res`.
 
 Added (developer)
 

diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py
@@ -598,6 +598,7 @@ class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
 {% if partition  %}#SBATCH --partition={{partition}}
 {% endif %}{% if runtime    %}#SBATCH --time={{runtime}}
 {% endif %}{% if memory     %}#SBATCH --mem={{memory}}
+{% endif %}{% if gres       %}#SBATCH --gres={{gres}}
 {% endif %}{% if nprocs     %}#SBATCH --cpus-per-task={{nprocs}}
 {% endif %}{% if reservation%}#SBATCH --reservation={{reservation}}
 {% endif %}{% if options    %}#SBATCH {{options}}{% endif %}
@@ -631,6 +632,10 @@ class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
         help="Reservation name to submit to resource manager"
         ).tag(config=True)
 
+    req_gres = Unicode('', \
+        help="Additional resources (e.g. GPUs) requested"
+    ).tag(config=True)
+
     # outputs line like "Submitted batch job 209"
     batch_submit_cmd = Unicode('sbatch --parsable').tag(config=True)
     # outputs status and exec node like "RUNNING hostname"

diff --git a/batchspawner/tests/test_spawners.py b/batchspawner/tests/test_spawners.py
@@ -375,13 +375,15 @@ def test_slurm(db, io_loop):
         'req_prologue': 'PROLOGUE',
         'req_epilogue': 'EPILOGUE',
         'req_reservation': 'RES123',
+        'req_gres': 'GRES123',
         }
     batch_script_re_list = [
         re.compile(r'PROLOGUE.*srun batchspawner-singleuser singleuser_command.*EPILOGUE', re.S),
         re.compile(r'^#SBATCH \s+ --cpus-per-task=5', re.X|re.M),
         re.compile(r'^#SBATCH \s+ --time=3-05:10:10', re.X|re.M),
         re.compile(r'^#SBATCH \s+ some_option_asdf', re.X|re.M),
         re.compile(r'^#SBATCH \s+ --reservation=RES123', re.X|re.M),
+        re.compile(r'^#SBATCH \s+ --gres=GRES123', re.X|re.M),
         ]
     from .. import SlurmSpawner
     run_spawner_script(db, io_loop, SlurmSpawner, normal_slurm_script,