Merge branch 'master' of https://github.com/ChadCumba/nipype

Chad Cumba · Chad Cumba · commit 917b5aa09b72 · 2013-08-07T10:48:44.000-05:00
diff --git a/doc/users/plugins.rst b/doc/users/plugins.rst
@@ -174,20 +174,59 @@ Workflow execution with HTCondor DAGMan is done by calling::
        workflow.run(plugin='CondorDAGMan')
 
 Job execution behavior can be tweaked with the following optional plug-in
-arguments::
-
-  template : submit spec template to use for job submission. The template
-             all generated submit specs are appended to this template. This
-             can be a str or a filename.
-  submit_specs : additional submit specs that are appended to the generated
-                 submit specs to allow for overriding or extending the defaults.
-                 This can be a str or a filename.
-  dagman_args : arguments to be prepended to the job execution script in the
-                dagman call
+arguments. The value of most arguments can be a literal string or a filename,
+where in the latter case the content of the file will be used as the argument
+value::
+
+    submit_template : submit spec template for individual jobs in a DAG (see
+                 CondorDAGManPlugin.default_submit_template for the default.
+    initial_specs : additional submit specs that are prepended to any job's
+                 submit file
+    override_specs : additional submit specs that are appended to any job's
+                 submit file
+    wrapper_cmd : path to an exectuable that will be started instead of a node
+                 script. This is useful for wrapper script that execute certain
+                 functionality prior or after a node runs. If this option is
+                 given the wrapper command is called with the respective Python
+                 exectuable and the path to the node script as final arguments
+    wrapper_args : optional additional arguments to a wrapper command
+    dagman_args : arguments to be prepended to the job execution script in the
+                  dagman call
+    block : if True the plugin call will block until Condor has finished
+            prcoessing the entire workflow (default: False)
 
 Please see the `HTCondor documentation`_ for details on possible configuration
 options and command line arguments.
 
+Using the ``wrapper_cmd`` argument it is possible to combine Nipype workflow
+execution with checkpoint/migration functionality offered by, for example,
+DMTCP_. This is especially useful in the case of workflows with long running
+nodes, such as Freesurfer's recon-all pipeline, where Condor's job
+prioritization algorithm could lead to jobs being evicted from compute
+nodes in order to maximize overall troughput. With checkpoint/migration enabled
+such a job would be checkpointed prior eviction and resume work from the
+checkpointed state after being rescheduled -- instead of restarting from
+scratch.
+
+On a Debian system, executing a workflow with support for checkpoint/migration
+for all nodes could look like this::
+
+  # define common parameters
+  dmtcp_hdr = """
+  should_transfer_files = YES
+  when_to_transfer_output = ON_EXIT_OR_EVICT
+  kill_sig = 2
+  environment = DMTCP_TMPDIR=./;JALIB_STDERR_PATH=/dev/null;DMTCP_PREFIX_ID=$(CLUSTER)_$(PROCESS)
+  """
+  shim_args = "--log %(basename)s.shimlog --stdout %(basename)s.shimout --stderr %(basename)s.shimerr"
+  # run workflow
+  workflow.run(
+        plugin='CondorDAGMan',
+        plugin_args=dict(initial_specs=dmtcp_hdr,
+                         wrapper_cmd='/usr/lib/condor/shim_dmtcp',
+                         wrapper_args=shim_args)
+        )
+
 ``qsub`` emulation
 ~~~~~~~~~~~~~~~~~~
 
@@ -230,3 +269,4 @@ Optional arguments::
 .. _HTCondor: http://www.cs.wisc.edu/htcondor/
 .. _DAGMan: http://research.cs.wisc.edu/htcondor/dagman/dagman.html
 .. _HTCondor documentation: http://research.cs.wisc.edu/htcondor/manual
+.. _DMTCP: http://dmtcp.sourceforge.net
diff --git a/nipype/interfaces/ants/registration.py b/nipype/interfaces/ants/registration.py
@@ -437,7 +437,9 @@ def _formatMetric(self, index):
         if isinstance(name_input, list):
             items = stage_inputs.items()
             indexes = range(0, len(name_input))
-            specs = [{k: v[i] for k, v in items} for i in indexes]
+            # dict-comprehension only works with python 2.7 and up
+            #specs = [{k: v[i] for k, v in items} for i in indexes]
+            specs = [dict([(k, v[i]) for k,v in items]) for i in indexes]
         else:
             specs = [stage_inputs]
         
diff --git a/nipype/interfaces/dipy/tensors.py b/nipype/interfaces/dipy/tensors.py
@@ -19,7 +19,7 @@
 iflogger = logging.getLogger('interface')
 
 try:
-    package_check('dipy', version='0.7.0')
+    package_check('dipy', version='0.6.0')
     import dipy.reconst.dti as dti
     from dipy.core.gradients import GradientTable
 except Exception, e:
diff --git a/nipype/pipeline/plugins/dagman.py b/nipype/pipeline/plugins/dagman.py
@@ -4,6 +4,8 @@
 import os
 import sys
 import uuid
+import time
+from warnings import warn
 
 from .base import (GraphPluginBase, logger)
 
@@ -14,19 +16,46 @@ class CondorDAGManPlugin(GraphPluginBase):
     """Execute using Condor DAGMan
 
     The plugin_args input to run can be used to control the DAGMan execution.
+    The value of most arguments can be a literal string or a filename, where in
+    the latter case the content of the file will be used as the argument value.
+
     Currently supported options are:
 
-    - template : submit spec template for individual jobs in a DAG. All
-                 generated submit spec components (e.g. executable name and
-                 arguments) are appended to this template. This can be a str or
-                 a filename. In the latter case the file content is used as a
-                 template.
-    - submit_specs : additional submit specs that are appended to the generated
-                 submit specs to allow for overriding or extending the defaults.
-                 This can be a str or a filename.
+    - submit_template : submit spec template for individual jobs in a DAG (see
+                 CondorDAGManPlugin.default_submit_template for the default.
+    - initial_specs : additional submit specs that are prepended to any job's
+                 submit file
+    - override_specs : additional submit specs that are appended to any job's
+                 submit file
+    - wrapper_cmd : path to an exectuable that will be started instead of a node
+                 script. This is useful for wrapper script that execute certain
+                 functionality prior or after a node runs. If this option is
+                 given the wrapper command is called with the respective Python
+                 exectuable and the path to the node script as final arguments
+    - wrapper_args : optional additional arguments to a wrapper command
     - dagman_args : arguments to be prepended to the job execution script in the
                   dagman call
+    - block : if True the plugin call will block until Condor has finished
+                 prcoessing the entire workflow (default: False)
     """
+
+    default_submit_template = """
+universe = vanilla
+notification = Never
+executable = %(executable)s
+arguments = %(nodescript)s
+output = %(basename)s.out
+error = %(basename)s.err
+log = %(basename)s.log
+getenv = True
+"""
+    def _get_str_or_file(self, arg):
+        if os.path.isfile(arg):
+            content = open(arg).read()
+        else:
+            content = arg
+        return content
+
     # XXX feature wishlist
     # - infer data file dependencies from jobs
     # - infer CPU requirements from jobs
@@ -35,21 +64,34 @@ class CondorDAGManPlugin(GraphPluginBase):
     #   actually have to run. would be good to be able to decide whether they
     #   actually have to be scheduled (i.e. output already exist).
     def __init__(self, **kwargs):
-        self._template = "universe = vanilla\nnotification = Never"
-        self._submit_specs = ""
-        self._dagman_args = ""
-        if 'plugin_args' in kwargs and not kwargs['plugin_args'] is None:
+        for var, id_, val in \
+            (('_template', 'submit_template', self.default_submit_template),
+             ('_initial_specs', 'template', ''),
+             ('_initial_specs', 'initial_specs', ''),
+             ('_override_specs', 'submit_specs', ''),
+             ('_override_specs', 'override_specs', ''),
+             ('_wrapper_cmd', 'wrapper_cmd', None),
+             ('_wrapper_args', 'wrapper_args', ''),
+             ('_block', 'block', False),
+             ('_dagman_args', 'dagman_args', '')):
+            if 'plugin_args' in kwargs \
+                and not kwargs['plugin_args'] is None \
+                and id_ in kwargs['plugin_args']:
+                    if id_ == 'wrapper_cmd':
+                        val = os.path.abspath(kwargs['plugin_args'][id_])
+                    elif id_ == 'block':
+                        val = kwargs['plugin_args'][id_]
+                    else:
+                        val = self._get_str_or_file(kwargs['plugin_args'][id_])
+            setattr(self, var, val)
+        # TODO remove after some time
+        if 'plugin_args' in kwargs \
+            and not kwargs['plugin_args'] is None:
             plugin_args = kwargs['plugin_args']
             if 'template' in plugin_args:
-                self._template = plugin_args['template']
-                if os.path.isfile(self._template):
-                    self._template = open(self._template).read()
+                warn("the 'template' argument is deprecated, use 'initial_specs' instead")
             if 'submit_specs' in plugin_args:
-                self._submit_specs = plugin_args['submit_specs']
-                if os.path.isfile(self._submit_specs):
-                    self._submit_specs = open(self._submit_specs).read()
-            if 'dagman_args' in plugin_args:
-                self._dagman_args = plugin_args['dagman_args']
+                warn("the 'submit_specs' argument is deprecated, use 'override_specs' instead")
         super(CondorDAGManPlugin, self).__init__(**kwargs)
 
     def _submit_graph(self, pyfiles, dependencies, nodes):
@@ -62,26 +104,35 @@ def _submit_graph(self, pyfiles, dependencies, nodes):
             # as jobs in the DAG
             for idx, pyscript in enumerate(pyfiles):
                 node = nodes[idx]
-                template, submit_specs = self._get_args(
-                    node, ["template", "submit_specs"])
                 # XXX redundant with previous value? or could it change between
                 # scripts?
+                template, initial_specs, override_specs, wrapper_cmd, wrapper_args = \
+                    self._get_args(node,
+                                   ["template", "initial_specs",
+                                    "override_specs", "wrapper_cmd",
+                                    "wrapper_args"])
+                # add required slots to the template
+                template = '%s\n%s\n%s\nqueue\n' % (
+                                '%(initial_specs)s',
+                                template,
+                                '%(override_specs)s')
                 batch_dir, name = os.path.split(pyscript)
                 name = '.'.join(name.split('.')[:-1])
-                submitspec = '\n'.join(
-                    (template,
-                     'executable = %s' % sys.executable,
-                     'arguments = %s' % pyscript,
-                     'output = %s' % os.path.join(batch_dir,
-                                                  '%s.out' % name),
-                     'error = %s' % os.path.join(batch_dir,
-                                                 '%s.err' % name),
-                     'log = %s' % os.path.join(batch_dir,
-                                               '%s.log' % name),
-                     'getenv = True',
-                     submit_specs,
-                     'queue'
-                     ))
+                specs = dict(
+                    # TODO make parameter for this,
+                    initial_specs=initial_specs,
+                    executable=sys.executable,
+                    nodescript=pyscript,
+                    basename=os.path.join(batch_dir, name),
+                    override_specs=override_specs
+                    )
+                if not wrapper_cmd is None:
+                    specs['executable'] = wrapper_cmd
+                    specs['nodescript'] = \
+                            '%s %s %s' % (wrapper_args % specs, # give access to variables
+                                          sys.executable,
+                                          pyscript)
+                submitspec = template % specs
                 # write submit spec for this job
                 submitfile = os.path.join(batch_dir,
                                           '%s.submit' % name)
@@ -105,3 +156,13 @@ def _submit_graph(self, pyfiles, dependencies, nodes):
                                                     self._dagman_args)
         cmd.run()
         logger.info('submitted all jobs to Condor DAGMan')
+        if self._block:
+            # wait for DAGMan to settle down, no time wasted it is already running
+            time.sleep(10)
+            if not os.path.exists('%s.condor.sub' % dagfilename):
+                raise EnvironmentError("DAGMan did not create its submit file, please check the logs")
+            # wait for completion
+            logger.info('waiting for DAGMan to finish')
+            lockfilename = '%s.lock' % dagfilename
+            while os.path.exists(lockfilename):
+                time.sleep(5)
