Merge pull request #304 from mgxd/fix/queue

FIX: SLURM queue support

Author: mgxd

heudiconv/cli/run.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import os
 import os.path as op
 from argparse import ArgumentParser
@@ -215,12 +217,11 @@ def get_parser():
     parser.add_argument('--dcmconfig', default=None,
                         help='JSON file for additional dcm2niix configuration')
     submission = parser.add_argument_group('Conversion submission options')
-    submission.add_argument('-q', '--queue', default=None,
-                            help='select batch system to submit jobs to instead'
-                                 ' of running the conversion serially')
-    submission.add_argument('--sbargs', dest='sbatch_args', default=None,
-                            help='Additional sbatch arguments if running with '
-                                 'queue arg')
+    submission.add_argument('-q', '--queue', choices=("SLURM", None),
+                            default=None,
+                            help='batch system to submit jobs in parallel')
+    submission.add_argument('--queue-args', dest='queue_args', default=None,
+                            help='Additional queue arguments')
     return parser
 
 
@@ -281,27 +282,28 @@ def process_args(args):
             continue
 
         if args.queue:
-            if seqinfo and not dicoms:
-                # flatten them all and provide into batching, which again
-                # would group them... heh
-                dicoms = sum(seqinfo.values(), [])
-                raise NotImplementedError(
-                    "we already grouped them so need to add a switch to avoid "
-                    "any grouping, so no outdir prefix doubled etc")
-
-            progname = op.abspath(inspect.getfile(inspect.currentframe()))
-
-            queue_conversion(progname,
+            # if seqinfo and not dicoms:
+            #     # flatten them all and provide into batching, which again
+            #     # would group them... heh
+            #     dicoms = sum(seqinfo.values(), [])
+            #     raise NotImplementedError(
+            #         "we already grouped them so need to add a switch to avoid "
+            #         "any grouping, so no outdir prefix doubled etc")
+
+            pyscript = op.abspath(inspect.getfile(inspect.currentframe()))
+
+            studyid = sid
+            if session:
+                studyid += "-%s" % session
+            if locator:
+                studyid += "-%s" % locator
+            # remove any separators
+            studyid = studyid.replace(op.sep, '_')
+
+            queue_conversion(pyscript,
                              args.queue,
-                             study_outdir,
-                             heuristic.filename,
-                             dicoms,
-                             sid,
-                             args.anon_cmd,
-                             args.converter,
-                             session,
-                             args.with_prov,
-                             args.bids)
+                             studyid,
+                             args.queue_args)
             continue
 
         anon_sid = anonymize_sid(sid, args.anon_cmd) if args.anon_cmd else None

heudiconv/dicoms.py

@@ -4,7 +4,6 @@
 import logging
 from collections import OrderedDict
 import tarfile
-from nibabel.nicom import csareader
 from heudiconv.external.pydicom import dcm
 
 from .utils import SeqInfo, load_json, set_readonly

heudiconv/parser.py

@@ -186,7 +186,7 @@ def get_study_sessions(dicom_dir_template, files_opt, heuristic, outdir,
                                           "`infotoids` to heuristic file or "
                                           "provide `--subjects` option")
             lgr.warn("Heuristic is missing an `infotoids` method, assigning "
-                     "empty method and using provided subject id %s."
+                     "empty method and using provided subject id %s. "
                      "Provide `session` and `locator` fields for best results."
                      , sid)
             def infotoids(seqinfos, outdir):

heudiconv/queue.py

@@ -1,35 +1,62 @@
+import subprocess
+import sys
 import os
-import os.path as op
 
 import logging
 
 lgr = logging.getLogger(__name__)
 
-# start with SLURM but extend past that #TODO
-def queue_conversion(progname, queue, outdir, heuristic, dicoms, sid,
-                     anon_cmd, converter, session,with_prov, bids):
-
-        # Rework this...
-        convertcmd = ' '.join(['python', progname,
-                               '-o', outdir,
-                               '-f', heuristic,
-                               '-s', sid,
-                               '--anon-cmd', anon_cmd,
-                               '-c', converter])
-        if session:
-            convertcmd += " --ses '%s'" % session
-        if with_prov:
-            convertcmd += " --with-prov"
-        if bids:
-            convertcmd += " --bids"
-        if dicoms:
-            convertcmd += " --files"
-            convertcmd += [" '%s'" % f for f in dicoms]
-
-        script_file = 'dicom-%s.sh' % sid
-        with open(script_file, 'wt') as fp:
-            fp.writelines(['#!/bin/bash\n', convertcmd])
-        outcmd = 'sbatch -J dicom-%s -p %s -N1 -c2 --mem=20G %s' \
-                 % (sid, queue, script_file)
-
-        os.system(outcmd)
+def queue_conversion(pyscript, queue, studyid, queue_args=None):
+        """
+        Write out conversion arguments to file and submit to a job scheduler.
+        Parses `sys.argv` for heudiconv arguments.
+
+        Parameters
+        ----------
+        pyscript: file
+            path to `heudiconv` script
+        queue: string
+            batch scheduler to use
+        studyid: string
+            identifier for conversion
+        queue_args: string (optional)
+            additional queue arguments for job submission
+
+        Returns
+        -------
+        proc: int
+            Queue submission exit code
+        """
+
+        SUPPORTED_QUEUES = {'SLURM': 'sbatch'}
+        if queue not in SUPPORTED_QUEUES:
+            raise NotImplementedError("Queuing with %s is not supported", queue)
+
+        args = sys.argv[1:]
+        # search args for queue flag
+        for i, arg in enumerate(args):
+            if arg in ["-q", "--queue"]:
+                break
+        if i == len(args) - 1:
+            raise RuntimeError(
+                "Queue flag not found (must be provided as a command-line arg)"
+            )
+        # remove queue flag and value
+        del args[i:i+2]
+
+        # make arguments executable again
+        args.insert(0, pyscript)
+        pypath = sys.executable or "python"
+        args.insert(0, pypath)
+        convertcmd = " ".join(args)
+
+        # will overwrite across subjects
+        queue_file = os.path.abspath('heudiconv-%s.sh' % queue)
+        with open(queue_file, 'wt') as fp:
+            fp.writelines(['#!/bin/bash\n', convertcmd, '\n'])
+
+        cmd = [SUPPORTED_QUEUES[queue], queue_file]
+        if queue_args:
+            cmd.insert(1, queue_args)
+        proc = subprocess.call(cmd)
+        return proc

heudiconv/tests/test_queue.py

@@ -0,0 +1,46 @@
+import os
+import sys
+import subprocess
+
+from heudiconv.cli.run import main as runner
+from .utils import TESTS_DATA_PATH
+import pytest
+from nipype.utils.filemanip import which
+
+@pytest.mark.skipif(which("sbatch"), reason="skip a real slurm call")
+@pytest.mark.parametrize(
+    'invocation', [
+        "--files %s/01-fmap_acq-3mm" % TESTS_DATA_PATH,    # our new way with automated groupping
+        "-d %s/{subject}/* -s 01-fmap_acq-3mm" % TESTS_DATA_PATH # "old" way specifying subject
+    ])
+def test_queue_no_slurm(tmpdir, invocation):
+    tmpdir.chdir()
+    hargs = invocation.split(" ")
+    hargs.extend(["-f", "reproin", "-b", "--minmeta", "--queue", "SLURM"])
+
+    # simulate command-line call
+    _sys_args = sys.argv
+    sys.argv = ['heudiconv'] + hargs
+
+    try:
+        with pytest.raises(OSError):
+            runner(hargs)
+        # should have generated a slurm submission script
+        slurm_cmd_file = (tmpdir / 'heudiconv-SLURM.sh').strpath
+        assert slurm_cmd_file
+        # check contents and ensure args match
+        with open(slurm_cmd_file) as fp:
+            lines = fp.readlines()
+        assert lines[0] == "#!/bin/bash\n"
+        cmd = lines[1]
+
+        # check that all flags we gave still being called
+        for arg in hargs:
+            # except --queue <queue>
+            if arg in ['--queue', 'SLURM']:
+                assert arg not in cmd
+            else:
+                assert arg in cmd
+    finally:
+        # revert before breaking something
+        sys.argv = _sys_args