run.py 3.35 KB
Newer Older
1
2
3
4
import argparse
import sys
import os
import importlib
5
import pickle
6
import modes.experiments as exp
7
import modes.runtime as runtime
8
9
10
11
12

def mkdir_if_not_exists(path):
    if not os.path.exists(path):
        os.mkdir(path)

13

14
15
16
17
parser = argparse.ArgumentParser()
parser.add_argument('experiments', metavar='EXP', type=str, nargs='+',
        help='An experiment file to run')
parser.add_argument('--runs', metavar='N', type=int, default=1,
18
        help='Number of repetition for each experiment')
19
20
parser.add_argument('--firstrun', metavar='N', type=int, default=1,
        help='ID for first run')
21
22
23
parser.add_argument('--verbose', action='store_const', const=True,
        default=False,
        help='Verbose output')
24
25
26

g_env = parser.add_argument_group('Environment')
g_env.add_argument('--repo', metavar='DIR', type=str,
27
        default='..', help='Repo directory')
28
g_env.add_argument('--workdir', metavar='DIR', type=str,
29
        default='./out/', help='Work directory base')
30
g_env.add_argument('--outdir', metavar='DIR',  type=str,
31
32
        default='./out/', help='Output directory base')

33
34
35
36
37
38
39
40
41
42
g_par = parser.add_argument_group('Parallel Runtime')
g_par.add_argument('--parallel', dest='runtime', action='store_const',
        const='parallel', default='sequential',
        help='Use parallel instead of sequential runtime')
g_par.add_argument('--cores', metavar='N', type=int,
        default=len(os.sched_getaffinity(0)),
        help='Number of cores to use for parallel runs')
g_par.add_argument('--mem', metavar='N', type=int, default=None,
        help='Memory limit for parallel runs (in MB)')

43
44
45
46
47
48
49
g_slurm = parser.add_argument_group('Slurm Runtime')
g_slurm.add_argument('--slurm', dest='runtime', action='store_const',
        const='slurm', default='sequential',
        help='Use slurm instead of sequential runtime')
g_slurm.add_argument('--slurmdir', metavar='DIR',  type=str,
        default='./slurm/', help='Slurm communication directory')

50
51
52
53
args = parser.parse_args()

experiments = []
for path in args.experiments:
54
55
56
57
58
59
60
61
62
63
    modname, modext = os.path.splitext(os.path.basename(path))

    if modext == '.py':
        spec = importlib.util.spec_from_file_location(modname, path)
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        experiments += mod.experiments
    else:
        with open(path, 'rb') as f:
            experiments.append(pickle.load(f))
64

65
66
if args.runtime != 'slurm':
    mkdir_if_not_exists(args.workdir)
67
68
mkdir_if_not_exists(args.outdir)

69
if args.runtime == 'parallel':
70
71
    rt = runtime.LocalParallelRuntime(cores=args.cores, mem=args.mem,
            verbose=args.verbose)
72
73
elif args.runtime == 'slurm':
    rt = runtime.SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
74
else:
75
    rt = runtime.LocalSimpleRuntime(verbose=args.verbose)
76

77
78
for e in experiments:
    workdir_base = '%s/%s' % (args.workdir, e.name)
79
80
    if args.runtime != 'slurm':
        mkdir_if_not_exists(workdir_base)
81

82
    for run in range(args.firstrun, args.firstrun + args.runs):
83
84
85
86
87
88
        outpath = '%s/%s-%d.json' % (args.outdir, e.name, run)
        if os.path.exists(outpath):
            print('skip %s run %d' % (e.name, run))
            continue

        workdir = '%s/%d' % (workdir_base, run)
89
90
        if args.runtime != 'slurm':
            mkdir_if_not_exists(workdir)
91
92

        env = exp.ExpEnv(args.repo, workdir)
93
        rt.add_run(runtime.Run(e, run, env, outpath))
94

95
rt.start()