qsub.py 1.67 KB
Newer Older
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
1
#!/usr/bin/env python3
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
2
"""GPAW-wrapper for qsub."""
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
3
4
5
6
7

import argparse
import subprocess

description = 'Submit a GPAW Python script via qsub.'
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
8
maxcores = 8  # XeonX5550
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


def main():
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('-p', '--processes', type=int, default=4,
                        help='Number of processes.')
    parser.add_argument('-t', '--time', type=int, default=1,
                        help='Max running time in hours.')
    parser.add_argument('-z', '--dry-run', action='store_true',
                        help='Don\'t actually submit script.')
    parser.add_argument('script', help='Python script')
    parser.add_argument('argument', nargs='*',
                        help='Command-line argument for Python script.')
    args = parser.parse_args()
    arguments = ' '.join(args.argument)
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
24
    cmd = f'gpaw python {args.script} {arguments}'
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
25
26
27
28
    nodes, rest = divmod(args.processes, maxcores)
    if nodes > 0:
        ppn = maxcores
        if rest > 0:
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
29
30
            msg = f'Please use a multiple of {maxcores} processes!'
            raise SystemExit(msg)
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
31
32
33
34
35
36
37
38
39
    else:
        nodes = 1
        ppn = args.processes
    lines = [
        '#!/bin/sh',
        '#PBS -q hpc',
        f'#PBS -N {args.script}',
        f'#PBS -l nodes={nodes}:ppn={ppn}',
        f'#PBS -l walltime={args.time}:00:00',
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
40
        # '#PBS -l feature=XeonX5550',
Jens Jørgen Mortensen's avatar
Jens Jørgen Mortensen committed
41
42
43
44
45
46
47
48
49
50
51
52
        'cd $PBS_O_WORKDIR',
        f'OMP_NUM_THREADS=1 mpiexec {cmd}']
    script = '\n'.join(lines) + '\n'
    if args.dry_run:
        print(script)
    else:
        p = subprocess.Popen(['qsub'], stdin=subprocess.PIPE)
        p.communicate(script.encode())


if __name__ == '__main__':
    main()