|
1 | 1 | import subprocess |
2 | 2 | import sys |
3 | 3 | import os |
4 | | - |
5 | 4 | import logging |
6 | 5 |
|
| 6 | +from .utils import which |
| 7 | + |
7 | 8 | lgr = logging.getLogger(__name__) |
8 | 9 |
|
9 | | -def queue_conversion(pyscript, queue, studyid, queue_args=None): |
10 | | - """ |
11 | | - Write out conversion arguments to file and submit to a job scheduler. |
12 | | - Parses `sys.argv` for heudiconv arguments. |
13 | | -
|
14 | | - Parameters |
15 | | - ---------- |
16 | | - pyscript: file |
17 | | - path to `heudiconv` script |
18 | | - queue: string |
19 | | - batch scheduler to use |
20 | | - studyid: string |
21 | | - identifier for conversion |
22 | | - queue_args: string (optional) |
23 | | - additional queue arguments for job submission |
24 | | -
|
25 | | - Returns |
26 | | - ------- |
27 | | - proc: int |
28 | | - Queue submission exit code |
29 | | - """ |
30 | | - |
31 | | - SUPPORTED_QUEUES = {'SLURM': 'sbatch'} |
32 | | - if queue not in SUPPORTED_QUEUES: |
33 | | - raise NotImplementedError("Queuing with %s is not supported", queue) |
34 | | - |
35 | | - args = sys.argv[1:] |
36 | | - # search args for queue flag |
37 | | - for i, arg in enumerate(args): |
38 | | - if arg in ["-q", "--queue"]: |
39 | | - break |
40 | | - if i == len(args) - 1: |
41 | | - raise RuntimeError( |
42 | | - "Queue flag not found (must be provided as a command-line arg)" |
43 | | - ) |
44 | | - # remove queue flag and value |
45 | | - del args[i:i+2] |
46 | | - |
47 | | - # make arguments executable again |
48 | | - args.insert(0, pyscript) |
49 | | - pypath = sys.executable or "python" |
50 | | - args.insert(0, pypath) |
| 10 | +def queue_conversion(queue, iterarg, iterables, queue_args=None): |
| 11 | + """ |
| 12 | + Write out conversion arguments to file and submit to a job scheduler. |
| 13 | + Parses `sys.argv` for heudiconv arguments. |
| 14 | +
|
| 15 | + Parameters |
| 16 | + ---------- |
| 17 | + queue: string |
| 18 | + Batch scheduler to use |
| 19 | + iterarg: str |
| 20 | + Multi-argument to index (`subjects` OR `files`) |
| 21 | + iterables: int |
| 22 | + Number of `iterarg` arguments |
| 23 | + queue_args: string (optional) |
| 24 | + Additional queue arguments for job submission |
| 25 | +
|
| 26 | + """ |
| 27 | + |
| 28 | + SUPPORTED_QUEUES = {'SLURM': 'sbatch'} |
| 29 | + if queue not in SUPPORTED_QUEUES: |
| 30 | + raise NotImplementedError("Queuing with %s is not supported", queue) |
| 31 | + |
| 32 | + for i in range(iterables): |
| 33 | + args = clean_args(sys.argv[1:], iterarg, i) |
| 34 | + # make arguments executable |
| 35 | + heudiconv_exec = which("heudiconv") or "heudiconv" |
| 36 | + args.insert(0, heudiconv_exec) |
51 | 37 | convertcmd = " ".join(args) |
52 | 38 |
|
53 | 39 | # will overwrite across subjects |
54 | 40 | queue_file = os.path.abspath('heudiconv-%s.sh' % queue) |
55 | 41 | with open(queue_file, 'wt') as fp: |
56 | | - fp.writelines(['#!/bin/bash\n', convertcmd, '\n']) |
| 42 | + fp.write("#!/bin/bash\n") |
| 43 | + if queue_args: |
| 44 | + for qarg in queue_args.split(): |
| 45 | + fp.write("#SBATCH %s\n" % qarg) |
| 46 | + fp.write(convertcmd + "\n") |
57 | 47 |
|
58 | 48 | cmd = [SUPPORTED_QUEUES[queue], queue_file] |
59 | | - if queue_args: |
60 | | - cmd.insert(1, queue_args) |
61 | 49 | proc = subprocess.call(cmd) |
62 | | - return proc |
| 50 | + lgr.info("Submitted %d jobs", iterables) |
| 51 | + |
| 52 | +def clean_args(hargs, iterarg, iteridx): |
| 53 | + """ |
| 54 | + Filters arguments for batch submission. |
| 55 | +
|
| 56 | + Parameters |
| 57 | + ---------- |
| 58 | + hargs: list |
| 59 | + Command-line arguments |
| 60 | + iterarg: str |
| 61 | + Multi-argument to index (`subjects` OR `files`) |
| 62 | + iteridx: int |
| 63 | + `iterarg` index to submit |
| 64 | +
|
| 65 | + Returns |
| 66 | + ------- |
| 67 | + cmdargs : list |
| 68 | + Filtered arguments for batch submission |
| 69 | +
|
| 70 | + Example |
| 71 | + -------- |
| 72 | + >>> from heudiconv.queue import clean_args |
| 73 | + >>> cmd = ['heudiconv', '-d', '/some/{subject}/path', |
| 74 | + ... '-q', 'SLURM', |
| 75 | + ... '-s', 'sub-1', 'sub-2', 'sub-3', 'sub-4'] |
| 76 | + >>> clean_args(cmd, 'subjects', 0) |
| 77 | + ['heudiconv', '-d', '/some/{subject}/path', '-s', 'sub-1'] |
| 78 | + """ |
| 79 | + |
| 80 | + if iterarg == "subjects": |
| 81 | + iterarg = ['-s', '--subjects'] |
| 82 | + elif iterarg == "files": |
| 83 | + iterarg = ['--files'] |
| 84 | + else: |
| 85 | + raise ValueError("Cannot index %s" % iterarg) |
| 86 | + |
| 87 | + # remove these or cause an infinite loop |
| 88 | + queue_args = ['-q', '--queue', '--queue-args'] |
| 89 | + |
| 90 | + # control variables for multi-argument parsing |
| 91 | + is_iterarg = False |
| 92 | + itercount = 0 |
| 93 | + |
| 94 | + indicies = [] |
| 95 | + cmdargs = hargs[:] |
| 96 | + |
| 97 | + for i, arg in enumerate(hargs): |
| 98 | + if arg.startswith('-') and is_iterarg: |
| 99 | + # moving on to another argument |
| 100 | + is_iterarg = False |
| 101 | + if is_iterarg: |
| 102 | + if iteridx != itercount: |
| 103 | + indicies.append(i) |
| 104 | + itercount += 1 |
| 105 | + if arg in iterarg: |
| 106 | + is_iterarg = True |
| 107 | + if arg in queue_args: |
| 108 | + indicies.extend([i, i+1]) |
| 109 | + |
| 110 | + for j in sorted(indicies, reverse=True): |
| 111 | + del cmdargs[j] |
| 112 | + return cmdargs |
0 commit comments