devito/benchmarks/user/make-pbs.py at b5f2534efb41f01ae0488fc6dad1aa10e3c343a6 · devitocodes/devito · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import os
import argparse
import devito


def generate(args):
    join = lambda l: ' '.join('%d' % i for i in l)
    args = vars(args)
    args['shape'] = join(args['shape'])
    args['space_order'] = join(args['space_order'])

    args['home'] = os.path.dirname(os.path.dirname(devito.__file__))

    args['load'] = '\n'.join('module load %s' % i for i in args['load'])
    args['export'] = '\n'.join('export %s' % i for i in args['export'])

    template_header = """\
#!/bin/bash

#PBS -lselect=%(nn)s:ncpus=%(ncpus)s:mem=120gb:mpiprocs=%(np)s:ompthreads=%(nt)s
#PBS -lwalltime=02:00:00

lscpu

%(load)s

cd %(home)s

source activate devito

export DEVITO_HOME=%(home)s
export DEVITO_ARCH=intel
export DEVITO_LANGUAGE=openmp
export DEVITO_LOGGING=DEBUG

%(export)s

cd benchmarks/user
"""  # noqa
    template_cmd = """\
DEVITO_MPI=%(mpi)s mpiexec python benchmark.py bench -P %(problem)s -bm O2 -d %(shape)s -so %(space_order)s --tn %(tn)s -x 1 --arch %(arch)s -r %(resultsdir)s\
"""  # noqa

    # Generate one PBS file for each `np` value
    for nn in args['nn']:
        args['nn'] = nn

        cmds = []
        for i in args['mpi']:
            args['mpi'] = i
            cmds.append(template_cmd % args)
        cmds = ' \n'.join(cmds)

        body = ' \n'.join([template_header % args, cmds])

        with open('pbs_nn%d.gen.sh' % int(nn), 'w') as f:
            f.write(body)


def cleanup(args):
    for f in os.listdir():
        if f.endswith('.gen.sh'):
            os.remove(f)

def create_parser():
    parser = argparse.ArgumentParser(description='Generate PBS scripts')
    subparsers = parser.add_subparsers(dest='command', required=True)

    gen = subparsers.add_parser('generate')
    gen.add_argument('-P', '--problem', required=True)
    gen.add_argument('-d', '--shape', type=int, nargs='+', default=[50, 50, 50])
    gen.add_argument('-so', '--space-order', type=int, nargs='+', default=[2])
    gen.add_argument('--tn', type=float, default=250)
    gen.add_argument('-nn', type=int, nargs='+', default=[1])
    gen.add_argument('-ncpus', type=int, default=1)
    gen.add_argument('-mem', type=int, default=120)
    gen.add_argument('-np', type=int, default=1)
    gen.add_argument('-nt', type=int, default=1)
    gen.add_argument('--mpi', nargs='+', default=['basic'])
    gen.add_argument('--arch', default='unknown')
    gen.add_argument('-r', '--resultsdir', default='results')
    gen.add_argument('--load', nargs='*', default=[])
    gen.add_argument('--export', nargs='*', default=[])

    subparsers.add_parser('cleanup')
    return parser


def main():
    parser = create_parser()
    args = parser.parse_args()

    if args.command == 'generate':
        generate(args)
    elif args.command == 'cleanup':
        cleanup(args)


if __name__ == "__main__":
    main()