Merge pull request #1582 from satra/fix/slurm_variant

fix: slurm job submission and result checking

Author: satra

doc/users/plugins.rst

@@ -178,6 +178,7 @@ Optional arguments::
 
   template: custom template file to use
   sbatch_args: any other command line args to be passed to bsub.
+  jobid_re: regular expression for custom job submission id search
 
 
 SLURMGraph

nipype/pipeline/engine/nodes.py

@@ -57,7 +57,8 @@
 from ...interfaces.base import (traits, InputMultiPath, CommandLine,
                                 Undefined, TraitedSpec, DynamicTraitedSpec,
                                 Bunch, InterfaceResult, md5, Interface,
-                                TraitDictObject, TraitListObject, isdefined)
+                                TraitDictObject, TraitListObject, isdefined,
+                                runtime_profile)
 from ...utils.misc import (getsource, create_function_from_source,
                            flatten, unflatten)
 from ...utils.filemanip import (save_json, FileNotFoundError,
@@ -743,11 +744,12 @@ def write_report(self, report_type=None, cwd=None):
             rst_dict = {'hostname' : self.result.runtime.hostname,
                         'duration' : self.result.runtime.duration}
             # Try and insert memory/threads usage if available
-            try:
-                rst_dict['runtime_memory_gb'] = self.result.runtime.runtime_memory_gb
-                rst_dict['runtime_threads'] = self.result.runtime.runtime_threads
-            except AttributeError:
-                logger.info('Runtime memory and threads stats unavailable')
+            if runtime_profile:
+                try:
+                    rst_dict['runtime_memory_gb'] = self.result.runtime.runtime_memory_gb
+                    rst_dict['runtime_threads'] = self.result.runtime.runtime_threads
+                except AttributeError:
+                    logger.info('Runtime memory and threads stats unavailable')
             if hasattr(self.result.runtime, 'cmdline'):
                 rst_dict['command'] = self.result.runtime.cmdline
                 fp.writelines(write_rst_dict(rst_dict))

nipype/pipeline/plugins/slurm.py

@@ -13,7 +13,7 @@
 
 from .base import (SGELikeBatchManagerBase, logger, iflogger, logging)
 
-from nipype.interfaces.base import CommandLine
+from ...interfaces.base import CommandLine
 
 
 class SLURMPlugin(SGELikeBatchManagerBase):
@@ -38,12 +38,15 @@ def __init__(self, **kwargs):
         self._max_tries = 2
         self._template = template
         self._sbatch_args = None
+        self._jobid_re = "Submitted batch job ([0-9]*)"
 
         if 'plugin_args' in kwargs and kwargs['plugin_args']:
             if 'retry_timeout' in kwargs['plugin_args']:
                 self._retry_timeout = kwargs['plugin_args']['retry_timeout']
             if 'max_tries' in kwargs['plugin_args']:
                 self._max_tries = kwargs['plugin_args']['max_tries']
+            if 'jobid_re' in kwargs['plugin_args']:
+                self._jobid_re = kwargs['plugin_args']['jobid_re']
             if 'template' in kwargs['plugin_args']:
                 self._template = kwargs['plugin_args']['template']
                 if os.path.isfile(self._template):
@@ -55,17 +58,16 @@ def __init__(self, **kwargs):
 
     def _is_pending(self, taskid):
         #  subprocess.Popen requires taskid to be a string
-        proc = subprocess.Popen(["squeue", '-j', '%s' % taskid],
-                                stdout=subprocess.PIPE,
-                                stderr=subprocess.PIPE)
-        o, _ = proc.communicate()
-
-        return o.find(str(taskid)) > -1
+        res = CommandLine('squeue',
+                          args=' '.join(['-j', '%s' % taskid]),
+                          terminal_output='allatonce').run()
+        return res.runtime.stdout.find(str(taskid)) > -1
 
     def _submit_batchtask(self, scriptfile, node):
         """
-        This is more or less the _submit_batchtask from sge.py with flipped variable
-        names, different command line switches, and different output formatting/processing
+        This is more or less the _submit_batchtask from sge.py with flipped
+        variable names, different command line switches, and different output
+        formatting/processing
         """
         cmd = CommandLine('sbatch', environ=dict(os.environ),
                           terminal_output='allatonce')
@@ -118,7 +120,7 @@ def _submit_batchtask(self, scriptfile, node):
         iflogger.setLevel(oldlevel)
         # retrieve taskid
         lines = [line for line in result.runtime.stdout.split('\n') if line]
-        taskid = int(re.match("Submitted batch job ([0-9]*)",
+        taskid = int(re.match(self._jobid_re,
                               lines[-1]).groups()[0])
         self._pending[taskid] = node.output_dir()
         logger.debug('submitted sbatch task: %d for node %s' % (taskid, node._id))