fix: better support for queue args

Author: mgxd

heudiconv/cli/run.py

@@ -221,7 +221,9 @@ def get_parser():
                             default=None,
                             help='batch system to submit jobs in parallel')
     submission.add_argument('--queue-args', dest='queue_args', default=None,
-                            help='Additional queue arguments')
+                            help='Additional queue arguments passed as '
+                            'single string of Argument=Value pairs space '
+                            'separated.')
     return parser
 
 

heudiconv/queue.py

@@ -7,56 +7,66 @@
 lgr = logging.getLogger(__name__)
 
 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]
+    """
+    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 = clean_args(sys.argv[1:])
+    # make arguments executable
+    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.write("#!/bin/bash\n")
         if queue_args:
-            cmd.insert(1, queue_args)
-        proc = subprocess.call(cmd)
-        return proc
+            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']):
+    """
+    Filters out unwanted arguments
+
+    :param hargs: Arguments passed
+    :type hargs: Iterable
+    :param keys: Unwanted arguments
+    :type keys: Iterable
+    :return: Filtered arguments 
+    """
+    indicies = []
+    for i, arg in enumerate(hargs):
+        if arg in keys:
+            indicies.extend([i, i+1])
+    for j in sorted(indicies, reverse=True):
+        del hargs[j]
+    return hargs
+

heudiconv/tests/test_queue.py

@@ -3,6 +3,7 @@
 import subprocess
 
 from heudiconv.cli.run import main as runner
+from heudiconv.queue import clean_args
 from .utils import TESTS_DATA_PATH
 import pytest
 from nipype.utils.filemanip import which
@@ -44,3 +45,19 @@ def test_queue_no_slurm(tmpdir, invocation):
     finally:
         # revert before breaking something
         sys.argv = _sys_args
+
+def test_argument_filtering(tmpdir):
+    cmdargs = [
+        'heudiconv', 
+        '--files', 
+        '/fake/path/to/files', 
+        '-f', 
+        'convertall', 
+        '-q', 
+        'SLURM', 
+        '--queue-args', 
+        '--cpus-per-task=4 --contiguous --time=10'
+    ]
+    filtered = cmdargs[:-4]
+
+    assert(clean_args(cmdargs) == filtered)