fix: minimize processing before batch submission

Author: mgxd

heudiconv/cli/run.py

@@ -246,6 +246,13 @@ def process_args(args):
     if not args.heuristic:
         raise RuntimeError("No heuristic specified - add to arguments and rerun")
 
+    if args.queue:
+        lgr.info("Queuing %s conversion", args.queue)
+        iterarg, iterables = ("files", len(args.files)) if args.files else \
+                             ("subjects", len(args.subjs))
+        queue_conversion(args.queue, iterarg, iterables, args.queue_args)
+        sys.exit(0)
+
     heuristic = load_heuristic(args.heuristic)
 
     study_sessions = get_study_sessions(args.dicom_dir_template, args.files,
@@ -281,21 +288,6 @@ def process_args(args):
             lgr.warning("Skipping unknown locator dataset")
             continue
 
-        if args.queue:
-
-            studyid = sid
-            if session:
-                studyid += "-%s" % session
-            if locator:
-                studyid += "-%s" % locator
-            # remove any separators
-            studyid = studyid.replace(op.sep, '_')
-
-            queue_conversion(args.queue,
-                             studyid,
-                             args.queue_args)
-            continue
-
         anon_sid = anonymize_sid(sid, args.anon_cmd) if args.anon_cmd else None
         if args.anon_cmd:
             lgr.info('Anonymized {} to {}'.format(sid, anon_sid))

heudiconv/queue.py

@@ -7,64 +7,107 @@
 
 lgr = logging.getLogger(__name__)
 
-def queue_conversion(queue, studyid, queue_args=None):
+def queue_conversion(queue, iterarg, iterables, queue_args=None):
     """
     Write out conversion arguments to file and submit to a job scheduler.
     Parses `sys.argv` for heudiconv arguments.
 
     Parameters
     ----------
     queue: string
-        batch scheduler to use
-    studyid: string
-        identifier for conversion
+        Batch scheduler to use
+    iterarg: str
+        Multi-argument to index (`subjects` OR `files`)
+    iterables: int
+        Number of `iterarg` arguments
     queue_args: string (optional)
-        additional queue arguments for job submission
+        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 = clean_args(sys.argv[1:])
-    # make arguments executable
-    heudiconv_exec = which("heudiconv") or "heudiconv"
-    args.insert(0, heudiconv_exec)
-    convertcmd = " ".join(args)
-
-    # will overwrite across subjects
-    queue_file = os.path.abspath('heudiconv-%s.sh' % queue)
-    with open(queue_file, 'wt') as fp:
-        fp.write("#!/bin/bash\n")
-        if queue_args:
-            for qarg in queue_args.split():
-                fp.write("#SBATCH %s\n" % qarg)
-        fp.write(convertcmd + "\n")
-
-    cmd = [SUPPORTED_QUEUES[queue], queue_file]
-    proc = subprocess.call(cmd)
-    return proc
-
-def clean_args(hargs, keys=['-q', '--queue', '--queue-args']):
+    for i in range(iterables):
+        args = clean_args(sys.argv[1:], iterarg, i)
+        # make arguments executable
+        heudiconv_exec = which("heudiconv") or "heudiconv"
+        args.insert(0, heudiconv_exec)
+        convertcmd = " ".join(args)
+
+        print(convertcmd)
+        # will overwrite across subjects
+        queue_file = os.path.abspath('heudiconv-%s.sh' % queue)
+        with open(queue_file, 'wt') as fp:
+            fp.write("#!/bin/bash\n")
+            if queue_args:
+                for qarg in queue_args.split():
+                    fp.write("#SBATCH %s\n" % qarg)
+            fp.write(convertcmd + "\n")
+
+        cmd = [SUPPORTED_QUEUES[queue], queue_file]
+        proc = subprocess.call(cmd)
+    lgr.info("Submitted %d jobs", iterables)
+
+def clean_args(hargs, iterarg, iteridx):
     """
-    Filters out unwanted arguments
+    Filters arguments for batch submission.
+
+    Parameters
+    ----------
+    hargs: list
+        Command-line arguments
+    iterarg: str
+        Multi-argument to index (`subjects` OR `files`)
+    iteridx: int
+        `iterarg` index to submit
+
+    Returns
+    -------
+    cmdargs : list
+        Filtered arguments for batch submission
 
-    :param hargs: Arguments passed
-    :type hargs: Iterable
-    :param keys: Unwanted arguments
-    :type keys: Iterable
-    :return: Filtered arguments 
+    Example
+    --------
+    >>> from heudiconv.queue import clean_args
+    >>> cmd = ['heudiconv', '-d', '/some/{subject}/path',
+    ...                     '-q', 'SLURM',
+    ...                     '-s', 'sub-1', 'sub-2', 'sub-3', 'sub-4']
+    >>> clean_args(cmd, 'subjects', 0)
+    ['heudiconv', '-d', '/some/{subject}/path', '-s', 'sub-1']
     """
+
+    if iterarg == "subjects":
+        iterarg = ['-s', '--subjects']
+    elif iterarg == "files":
+        iterarg = ['--files']
+    else:
+        raise ValueError("Cannot index %s" % iterarg)
+
+    # remove these or cause an infinite loop
+    queue_args = ['-q', '--queue', '--queue-args']
+
+    # control variables for multi-argument parsing
+    is_iterarg = False
+    itercount = 0
+
     indicies = []
+    cmdargs = hargs[:]
+
     for i, arg in enumerate(hargs):
-        if arg in keys:
+        if arg.startswith('-') and is_iterarg:
+            # moving on to another argument
+            is_iterarg = False
+        if is_iterarg:
+            if iteridx != itercount:
+                indicies.append(i)
+            itercount += 1
+        if arg in iterarg:
+            is_iterarg = True
+        if arg in queue_args:
             indicies.extend([i, i+1])
-    for j in sorted(indicies, reverse=True):
-        del hargs[j]
-    return hargs
 
+    for j in sorted(indicies, reverse=True):
+        del cmdargs[j]
+    return cmdargs

heudiconv/tests/test_queue.py

@@ -23,7 +23,7 @@ def test_queue_no_slurm(tmpdir, invocation):
     sys.argv = ['heudiconv'] + hargs
 
     try:
-        with pytest.raises(OSError):
+        with pytest.raises(OSError):  # SLURM should not be installed
             runner(hargs)
         # should have generated a slurm submission script
         slurm_cmd_file = (tmpdir / 'heudiconv-SLURM.sh').strpath
@@ -46,17 +46,48 @@ def test_queue_no_slurm(tmpdir, invocation):
         sys.argv = _sys_args
 
 def test_argument_filtering(tmpdir):
-    cmdargs = [
-        'heudiconv', 
-        '--files', 
-        '/fake/path/to/files', 
-        '-f', 
-        'convertall', 
-        '-q', 
-        'SLURM', 
-        '--queue-args', 
+    cmd_files = [
+        'heudiconv',
+        '--files',
+        '/fake/path/to/files',
+        '/another/fake/path',
+        '-f',
+        'convertall',
+        '-q',
+        'SLURM',
+        '--queue-args',
         '--cpus-per-task=4 --contiguous --time=10'
     ]
-    filtered = cmdargs[:-4]
+    filtered = [
+        'heudiconv',
+        '--files',
+        '/another/fake/path',
+        '-f',
+        'convertall',
+    ]
+    assert clean_args(cmd_files, 'files', 1) == filtered
 
-    assert(clean_args(cmdargs) == filtered)
+    cmd_subjects = [
+        'heudiconv',
+        '-d',
+        '/some/{subject}/path',
+        '--queue',
+        'SLURM',
+        '--subjects',
+        'sub1',
+        'sub2',
+        'sub3',
+        'sub4',
+        '-f',
+        'convertall'
+    ]
+    filtered = [
+        'heudiconv',
+        '-d',
+        '/some/{subject}/path',
+        '--subjects',
+        'sub3',
+        '-f',
+        'convertall'
+    ]
+    assert clean_args(cmd_subjects, 'subjects', 2) == filtered