run.py 9.76 KB
Newer Older
Antoine Kaufmann's avatar
Antoine Kaufmann committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Copyright 2021 Max Planck Institute for Software Systems, and
# National University of Singapore
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

23
import argparse
24
import asyncio
Jonas Kaufmann's avatar
Jonas Kaufmann committed
25
import fnmatch
26
import importlib
Hejing Li's avatar
Hejing Li committed
27
import importlib.util
28
import json
Jonas Kaufmann's avatar
Jonas Kaufmann committed
29
import os
30
import pickle
Jonas Kaufmann's avatar
Jonas Kaufmann committed
31
import sys
32
import typing as tp
33
from signal import SIGINT, signal
34

35
36
37
from simbricks.exectools import LocalExecutor, RemoteExecutor
from simbricks.experiment.experiment_environment import ExpEnv
from simbricks.experiments import DistributedExperiment, Experiment
38
39
40
41
from simbricks.runtime.common import Run
from simbricks.runtime.distributed import DistributedSimpleRuntime, auto_dist
from simbricks.runtime.local import LocalParallelRuntime, LocalSimpleRuntime
from simbricks.runtime.slurm import SlurmRuntime
Jonas Kaufmann's avatar
Jonas Kaufmann committed
42
43


44
# pylint: disable=redefined-outer-name
45
46
47
48
def mkdir_if_not_exists(path):
    if not os.path.exists(path):
        os.mkdir(path)

49

50
parser = argparse.ArgumentParser()
Jonas Kaufmann's avatar
Jonas Kaufmann committed
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
101
102
103
104
105
106
107
108
109
parser.add_argument(
    'experiments',
    metavar='EXP',
    type=str,
    nargs='+',
    help='An experiment file to run'
)
parser.add_argument(
    '--list',
    action='store_const',
    const=True,
    default=False,
    help='Only list available experiment names'
)
parser.add_argument(
    '--filter',
    metavar='PATTERN',
    type=str,
    nargs='+',
    help='Pattern to match experiment names against'
)
parser.add_argument(
    '--pickled',
    action='store_const',
    const=True,
    default=False,
    help='Read exp files as pickled runs instead of exp.py files'
)
parser.add_argument(
    '--runs',
    metavar='N',
    type=int,
    default=1,
    help='Number of repetition for each experiment'
)
parser.add_argument(
    '--firstrun', metavar='N', type=int, default=1, help='ID for first run'
)
parser.add_argument(
    '--force',
    action='store_const',
    const=True,
    default=False,
    help='Run experiments even if output already exists'
)
parser.add_argument(
    '--verbose',
    action='store_const',
    const=True,
    default=False,
    help='Verbose output'
)
parser.add_argument(
    '--pcap',
    action='store_const',
    const=True,
    default=False,
    help='Dump pcap file (if supported by simulator)'
)
110
111

g_env = parser.add_argument_group('Environment')
Jonas Kaufmann's avatar
Jonas Kaufmann committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
g_env.add_argument(
    '--repo', metavar='DIR', type=str, default='..', help='Repo directory'
)
g_env.add_argument(
    '--workdir',
    metavar='DIR',
    type=str,
    default='./out/',
    help='Work directory base'
)
g_env.add_argument(
    '--outdir',
    metavar='DIR',
    type=str,
    default='./out/',
    help='Output directory base'
)
g_env.add_argument(
    '--cpdir',
    metavar='DIR',
    type=str,
    default='./out/',
    help='Checkpoint directory base'
)
g_env.add_argument(
    '--hosts',
    metavar='JSON_FILE',
    type=str,
    default=None,
    help='List of hosts to use (json)'
)
g_env.add_argument(
    '--shmdir',
    metavar='DIR',
    type=str,
    default=None,
    help='Shared memory directory base (workdir if not set)'
)
150

151
g_par = parser.add_argument_group('Parallel Runtime')
Jonas Kaufmann's avatar
Jonas Kaufmann committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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)'
)
174

175
g_slurm = parser.add_argument_group('Slurm Runtime')
Jonas Kaufmann's avatar
Jonas Kaufmann committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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'
)
191

192
g_dist = parser.add_argument_group('Distributed Runtime')
Jonas Kaufmann's avatar
Jonas Kaufmann committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
g_dist.add_argument(
    '--dist',
    dest='runtime',
    action='store_const',
    const='dist',
    default='sequential',
    help='Use sequential distributed runtime instead of local'
)
g_dist.add_argument(
    '--auto-dist',
    action='store_const',
    const=True,
    default=False,
    help='Automatically distribute non-distributed experiments'
)
g_dist.add_argument(
    '--proxy-type',
    metavar='TYPE',
    type=str,
    default='sockets',
    help='Proxy type to use (sockets,rdma) for auto distribution'
)
215
args = parser.parse_args()
216

217

218
# pylint: disable=redefined-outer-name
219
def load_executors(path):
Jonas Kaufmann's avatar
Jonas Kaufmann committed
220
    """Load hosts list from json file and return list of executors."""
221
    with open(path, 'r', encoding='utf-8') as f:
222
223
224
225
226
        hosts = json.load(f)

        exs = []
        for h in hosts:
            if h['type'] == 'local':
227
                ex = LocalExecutor()
228
            elif h['type'] == 'remote':
229
                ex = RemoteExecutor(h['host'], h['workdir'])
230
231
232
233
                if 'ssh_args' in h:
                    ex.ssh_extra_args += h['ssh_args']
                if 'scp_args' in h:
                    ex.scp_extra_args += h['scp_args']
234
235
            else:
                raise RuntimeError('invalid host type "' + h['type'] + '"')
236
237
            ex.ip = h['ip']
            exs.append(ex)
238
239
    return exs

Jonas Kaufmann's avatar
Jonas Kaufmann committed
240

241
if args.hosts is None:
242
    executors = [LocalExecutor()]
243
244
245
else:
    executors = load_executors(args.hosts)

Jonas Kaufmann's avatar
Jonas Kaufmann committed
246

247
248
def warn_multi_exec():
    if len(executors) > 1:
Jonas Kaufmann's avatar
Jonas Kaufmann committed
249
250
251
252
253
        print(
            'Warning: multiple hosts specified, only using first one for now',
            file=sys.stderr
        )

254

255
# initialize runtime
256
if args.runtime == 'parallel':
257
    warn_multi_exec()
Jonas Kaufmann's avatar
Jonas Kaufmann committed
258
    rt = LocalParallelRuntime(
259
260
261
262
        cores=args.cores,
        mem=args.mem,
        verbose=args.verbose,
        executor=executors[0]
Jonas Kaufmann's avatar
Jonas Kaufmann committed
263
    )
264
elif args.runtime == 'slurm':
Jonas Kaufmann's avatar
Jonas Kaufmann committed
265
    rt = SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
266
elif args.runtime == 'dist':
Jonas Kaufmann's avatar
Jonas Kaufmann committed
267
    rt = DistributedSimpleRuntime(executors, verbose=args.verbose)
268
else:
269
    warn_multi_exec()
270
    rt = LocalSimpleRuntime(verbose=args.verbose, executor=executors[0])
271

272

273
# pylint: disable=redefined-outer-name
274
def add_exp(
275
    e: Experiment,
Jonas Kaufmann's avatar
Jonas Kaufmann committed
276
277
278
279
280
    run: int,
    prereq: tp.Optional[Run],
    create_cp: bool,
    restore_cp: bool,
    no_simbricks: bool
281
):
282
    outpath = f'{args.outdir}/{e.name}-{run}.json'
283
    if os.path.exists(outpath) and not args.force:
284
        print(f'skip {e.name} run {run}')
285
286
        return None

287
288
    workdir = f'{args.workdir}/{e.name}/{run}'
    cpdir = f'{args.cpdir}/{e.name}/0'
289
    if args.shmdir is not None:
290
        shmdir = f'{args.shmdir}/{e.name}/{run}'
291

292
    env = ExpEnv(args.repo, workdir, cpdir)
293
294
    env.create_cp = create_cp
    env.restore_cp = restore_cp
Jonas Kaufmann's avatar
Jonas Kaufmann committed
295
    env.no_simbricks = no_simbricks
296
297
    env.pcap_file = ''
    if args.pcap:
Jonas Kaufmann's avatar
Jonas Kaufmann committed
298
        env.pcap_file = workdir + '/pcap'
299
300
    if args.shmdir is not None:
        env.shm_base = os.path.abspath(shmdir)
301

Jonas Kaufmann's avatar
Jonas Kaufmann committed
302
    run = Run(e, run, env, outpath, prereq)
303
304
305
    rt.add_run(run)
    return run

Jonas Kaufmann's avatar
Jonas Kaufmann committed
306

307
308
309
310
311
312
313
# load experiments
if not args.pickled:
    # default: load python modules with experiments
    experiments = []
    for path in args.experiments:
        modname, _ = os.path.splitext(os.path.basename(path))

314
315
316
        class ExperimentModuleLoadError(Exception):
            pass

317
        spec = importlib.util.spec_from_file_location(modname, path)
318
319
        if spec is None:
            raise ExperimentModuleLoadError('spec is None')
320
        mod = importlib.util.module_from_spec(spec)
321
322
        if spec.loader is None:
            raise ExperimentModuleLoadError('spec.loader is None')
323
324
        spec.loader.exec_module(mod)
        experiments += mod.experiments
325

326
327
328
329
330
    if args.list:
        for e in experiments:
            print(e.name)
        sys.exit(0)

331
    for e in experiments:
332
        if args.auto_dist and not isinstance(e, DistributedExperiment):
Jonas Kaufmann's avatar
Jonas Kaufmann committed
333
            e = auto_dist(e, executors, args.proxy_type)
334
        # apply filter if any specified
335
        if (args.filter) and (len(args.filter) > 0):
336
337
            match = False
            for f in args.filter:
338
339
                match = fnmatch.fnmatch(e.name, f)
                if match:
340
                    break
341

342
343
344
            if not match:
                continue

345
        # if this is an experiment with a checkpoint we might have to create it
346
        no_simbricks = e.no_simbricks
347
        if e.checkpoint:
348
            prereq = add_exp(e, 0, None, True, False, no_simbricks)
349
350
        else:
            prereq = None
351

352
        for run in range(args.firstrun, args.firstrun + args.runs):
353
            add_exp(e, run, prereq, False, e.checkpoint, no_simbricks)
354
355
356
357
358
else:
    # otherwise load pickled run object
    for path in args.experiments:
        with open(path, 'rb') as f:
            rt.add_run(pickle.load(f))
359

360
361
362
363
364
365
366
367
368
369

# register interrupt handler
def handle_interrupt(signalnum: int, handler: int):
    # pylint: disable=unused-argument
    rt.interrupt()


signal(SIGINT, handle_interrupt)

asyncio.run(rt.start())