Merge pull request #166 from VOD555/slurm

Support Slurm in gromacs.qsub
Becksteinlab · Feb 26, 2019 · cda0874 · cda0874
2 parents d96c34f + 2dc8dda
commit cda0874
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 16 deletions.
diff --git a/gromacs/qsub.py b/gromacs/qsub.py
@@ -293,6 +293,7 @@ def __repr__(self):
     QueuingSystem('Sun Gridengine', 'sge', '#$', array_variable='SGE_TASK_ID', array_option='-t %d-%d'),
     QueuingSystem('PBS', 'pbs', '#PBS', array_variable='PBS_ARRAY_INDEX', array_option='-J %d-%d'),
     QueuingSystem('LoadLeveler', 'll', '#@'),  # no idea how to do arrays in LL
+    QueuingSystem('Slurm', 'slu', '#SBATCH'), # will add array settings
     ]
 
 def detect_queuing_system(scriptfile):
@@ -365,19 +366,35 @@ def write_script(template):
         submitscript = os.path.join(dirname, prefix + os.path.basename(template))
         logger.info("Setting up queuing system script {submitscript!r}...".format(**vars()))
         # These substitution rules are documented for the user in the module doc string
-        cbook.edit_txt(template,
-                       [('^ *DEFFNM=','(?<==)(.*)', deffnm),
-                        ('^#.*(-N|job_name)', '((?<=-N\s)|(?<=job_name\s))\s*\w+', jobname),
-                        ('^#.*(-A|account_no)', '((?<=-A\s)|(?<=account_no\s))\s*\w+', budget),
-                        ('^#.*(-l walltime|wall_clock_limit)', '(?<==)(\d+:\d+:\d+)', walltime),
-                        ('^ *WALL_HOURS=', '(?<==)(.*)', wall_hours),
-                        ('^ *STARTDIR=', '(?<==)(.*)', startdir),
-                        ('^ *NPME=', '(?<==)(.*)', npme),
-                        ('^ *MDRUN_OPTS=', '(?<==)("")', mdrun_opts),  # only replace literal ""
-                        ('^# JOB_ARRAY_PLACEHOLDER', '^.*$', jobarray_string),
-                    ],
-                       newname=submitscript)
-        ext = os.path.splitext(submitscript)[1]
+        qsystem = detect_queuing_system(template)
+        if qsystem is not None and (qsystem.name == 'Slurm'):
+            cbook.edit_txt(template,
+                           [('^ *DEFFNM=','(?<==)(.*)', deffnm),
+                            ('^#.*(-J)', '((?<=-J\s))\s*\w+', jobname),
+                            ('^#.*(-A|account_no)', '((?<=-A\s)|(?<=account_no\s))\s*\w+', budget),
+                            ('^#.*(-t)', '(?<=-t\s)(\d+:\d+:\d+)', walltime),
+                            ('^ *WALL_HOURS=', '(?<==)(.*)', wall_hours),
+                            ('^ *STARTDIR=', '(?<==)(.*)', startdir),
+                            ('^ *NPME=', '(?<==)(.*)', npme),
+                            ('^ *MDRUN_OPTS=', '(?<==)("")', mdrun_opts),  # only replace literal ""
+                            ('^# JOB_ARRAY_PLACEHOLDER', '^.*$', jobarray_string),
+                            ],
+                           newname=submitscript)
+            ext = os.path.splitext(submitscript)[1]
+        else:
+            cbook.edit_txt(template,
+                           [('^ *DEFFNM=','(?<==)(.*)', deffnm),
+                            ('^#.*(-N|job_name)', '((?<=-N\s)|(?<=job_name\s))\s*\w+', jobname),
+                            ('^#.*(-A|account_no)', '((?<=-A\s)|(?<=account_no\s))\s*\w+', budget),
+                            ('^#.*(-l walltime|wall_clock_limit)', '(?<==)(\d+:\d+:\d+)', walltime),
+                            ('^ *WALL_HOURS=', '(?<==)(.*)', wall_hours),
+                            ('^ *STARTDIR=', '(?<==)(.*)', startdir),
+                            ('^ *NPME=', '(?<==)(.*)', npme),
+                            ('^ *MDRUN_OPTS=', '(?<==)("")', mdrun_opts),  # only replace literal ""
+                            ('^# JOB_ARRAY_PLACEHOLDER', '^.*$', jobarray_string),
+                           ],
+                           newname=submitscript)
+            ext = os.path.splitext(submitscript)[1]
         if ext in ('.sh', '.csh', '.bash'):
             os.chmod(submitscript, 0o755)
         return submitscript
@@ -430,4 +447,3 @@ def write_script(template):
 
     # must use config.get_templates() because we need to access the file for detecting
     return [write_script(template) for template in config.get_templates(templates)]
-
diff --git a/gromacs/tests/test_qsub.py b/gromacs/tests/test_qsub.py
@@ -3,7 +3,7 @@
 
 import gromacs.qsub
 
-def test_queuing_systems(known=("Sun Gridengine", "PBS", "LoadLeveler")):
+def test_queuing_systems(known=("Sun Gridengine", "PBS", "LoadLeveler", 'Slurm')):
     assert len(gromacs.qsub.queuing_systems) == len(known)
     for qs in gromacs.qsub.queuing_systems:
         assert qs.name in known
@@ -18,7 +18,8 @@ def test_queuing_systems(known=("Sun Gridengine", "PBS", "LoadLeveler")):
 @pytest.mark.parametrize("scriptfile,name", [
     ("foo.sge", "Sun Gridengine"),
     ("foo.pbs", "PBS"),
-    ("foo.ll", "LoadLeveler")])
+    ("foo.ll", "LoadLeveler"),
+    ("foo.slu", "Slurm")])
 def test_detect_queuing_system(scriptfile, name):
     qs = gromacs.qsub.detect_queuing_system(scriptfile)
     assert qs.name == name