Skip to content

GitLab

"...composable_kernel.git" did not exist on "3276a5e9b94dc9ffccacf469e57b1f8502dabccb"
Commit 0b3de18a authored by Jonas Kaufmann's avatar Jonas Kaufmann Committed by Antoine Kaufmann
Browse files

experiments/run.py: improve readabilty 

There are now no more inline functions and a proper main function
instead of top-level code. Functionality-wise, nothing was changed. Code
was only rearranged and some comments added.
parent d3118b9b
Hide whitespace changes
Inline Side-by-side
Showing with 270 additions and 254 deletions
experiments/run.py
View file @ 0b3de18a
...@@ -19,6 +19,8 @@ ...@@ -19,6 +19,8 @@
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, # CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""This is the top-level module of the SimBricks orchestration framework that
users interact with."""
import argparse import argparse
import asyncio import asyncio
...@@ -32,11 +34,13 @@ import sys ...@@ -32,11 +34,13 @@ import sys
import typing as tp import typing as tp
from signal import SIGINT, signal from signal import SIGINT, signal
from simbricks.orchestration import exectools
from simbricks.orchestration.exectools import LocalExecutor, RemoteExecutor from simbricks.orchestration.exectools import LocalExecutor, RemoteExecutor
from simbricks.orchestration.experiment.experiment_environment import ExpEnv from simbricks.orchestration.experiment.experiment_environment import ExpEnv
from simbricks.orchestration.experiments import ( from simbricks.orchestration.experiments import (
DistributedExperiment, Experiment DistributedExperiment, Experiment
) )
from simbricks.orchestration.runtime import common as rt_common
from simbricks.orchestration.runtime.common import Run from simbricks.orchestration.runtime.common import Run
from simbricks.orchestration.runtime.distributed import ( from simbricks.orchestration.runtime.distributed import (
DistributedSimpleRuntime, auto_dist DistributedSimpleRuntime, auto_dist
...@@ -47,186 +51,186 @@ from simbricks.orchestration.runtime.local import ( ...@@ -47,186 +51,186 @@ from simbricks.orchestration.runtime.local import (
from simbricks.orchestration.runtime.slurm import SlurmRuntime from simbricks.orchestration.runtime.slurm import SlurmRuntime
# pylint: disable=redefined-outer-name def parse_args() -> argparse.Namespace:
def mkdir_if_not_exists(path): parser = argparse.ArgumentParser()
if not os.path.exists(path): # general arguments for experiments
os.mkdir(path) parser.add_argument(
'experiments',
metavar='EXP',
parser = argparse.ArgumentParser() type=str,
parser.add_argument( nargs='+',
'experiments', help='Python modules to load the experiments from'
metavar='EXP', )
type=str, parser.add_argument(
nargs='+', '--list',
help='Python modules to load the experiments from' action='store_const',
) const=True,
parser.add_argument( default=False,
'--list', help='List available experiment names'
action='store_const', )
const=True, parser.add_argument(
default=False, '--filter',
help='List available experiment names' metavar='PATTERN',
) type=str,
parser.add_argument( nargs='+',
'--filter', help='Only run experiments matching the given Unix shell style patterns'
metavar='PATTERN', )
type=str, parser.add_argument(
nargs='+', '--pickled',
help='Only run experiments matching the given Unix shell style patterns' action='store_const',
) const=True,
parser.add_argument( default=False,
'--pickled', help='Interpret experiment modules as pickled runs instead of .py files'
action='store_const', )
const=True, parser.add_argument(
default=False, '--runs',
help='Interpret experiment modules as pickled runs instead of .py files' metavar='N',
) type=int,
parser.add_argument( default=1,
'--runs', help='Number of repetition of each experiment'
metavar='N', )
type=int, parser.add_argument(
default=1, '--firstrun', metavar='N', type=int, default=1, help='ID for first run'
help='Number of repetition of each experiment' )
) parser.add_argument(
parser.add_argument( '--force',
'--firstrun', metavar='N', type=int, default=1, help='ID for first run' action='store_const',
) const=True,
parser.add_argument( default=False,
'--force', help='Run experiments even if output already exists (overwrites output)'
action='store_const', )
const=True, parser.add_argument(
default=False, '--verbose',
help='Run experiments even if output already exists (overwrites output)' action='store_const',
) const=True,
parser.add_argument( default=False,
'--verbose', help='Verbose output, for example, print component simulators\' output'
action='store_const', )
const=True, parser.add_argument(
default=False, '--pcap',
help='Verbose output, for example, print component simulators\' output' action='store_const',
) const=True,
parser.add_argument( default=False,
'--pcap', help='Dump pcap file (if supported by component simulator)'
action='store_const', )
const=True,
default=False,
help='Dump pcap file (if supported by component simulator)'
)
g_env = parser.add_argument_group('Environment') # arguments for the experiment environment
g_env.add_argument( g_env = parser.add_argument_group('Environment')
'--repo', g_env.add_argument(
metavar='DIR', '--repo',
type=str, metavar='DIR',
default='..', type=str,
help='SimBricks repository directory' default='..',
) help='SimBricks repository directory'
g_env.add_argument( )
'--workdir', g_env.add_argument(
metavar='DIR', '--workdir',
type=str, metavar='DIR',
default='./out/', type=str,
help='Work directory base' default='./out/',
) help='Work directory base'
g_env.add_argument( )
'--outdir', g_env.add_argument(
metavar='DIR', '--outdir',
type=str, metavar='DIR',
default='./out/', type=str,
help='Output directory base' default='./out/',
) help='Output directory base'
g_env.add_argument( )
'--cpdir', g_env.add_argument(
metavar='DIR', '--cpdir',
type=str, metavar='DIR',
default='./out/', type=str,
help='Checkpoint directory base' default='./out/',
) help='Checkpoint directory base'
g_env.add_argument( )
'--hosts', g_env.add_argument(
metavar='JSON_FILE', '--hosts',
type=str, metavar='JSON_FILE',
default=None, type=str,
help='List of hosts to use (json)' default=None,
) help='List of hosts to use (json)'
g_env.add_argument( )
'--shmdir', g_env.add_argument(
metavar='DIR', '--shmdir',
type=str, metavar='DIR',
default=None, type=str,
help='Shared memory directory base (workdir if not set)' default=None,
) help='Shared memory directory base (workdir if not set)'
)
g_par = parser.add_argument_group('Parallel Runtime') # arguments for the parallel runtime
g_par.add_argument( g_par = parser.add_argument_group('Parallel Runtime')
'--parallel', g_par.add_argument(
dest='runtime', '--parallel',
action='store_const', dest='runtime',
const='parallel', action='store_const',
default='sequential', const='parallel',
help='Use parallel instead of sequential runtime' default='sequential',
) help='Use parallel instead of sequential runtime'
g_par.add_argument( )
'--cores', g_par.add_argument(
metavar='N', '--cores',
type=int, metavar='N',
default=len(os.sched_getaffinity(0)), type=int,
help='Number of cores to use for parallel runs' default=len(os.sched_getaffinity(0)),
) help='Number of cores to use for parallel runs'
g_par.add_argument( )
'--mem', g_par.add_argument(
metavar='N', '--mem',
type=int, metavar='N',
default=None, type=int,
help='Memory limit for parallel runs (in MB)' default=None,
) help='Memory limit for parallel runs (in MB)'
)
g_slurm = parser.add_argument_group('Slurm Runtime') # arguments for the slurm runtime
g_slurm.add_argument( g_slurm = parser.add_argument_group('Slurm Runtime')
'--slurm', g_slurm.add_argument(
dest='runtime', '--slurm',
action='store_const', dest='runtime',
const='slurm', action='store_const',
default='sequential', const='slurm',
help='Use slurm instead of sequential runtime' default='sequential',
) help='Use slurm instead of sequential runtime'
g_slurm.add_argument( )
'--slurmdir', g_slurm.add_argument(
metavar='DIR', '--slurmdir',
type=str, metavar='DIR',
default='./slurm/', type=str,
help='Slurm communication directory' default='./slurm/',
) help='Slurm communication directory'
)
g_dist = parser.add_argument_group('Distributed Runtime') # arguments for the distributed runtime
g_dist.add_argument( g_dist = parser.add_argument_group('Distributed Runtime')
'--dist', g_dist.add_argument(
dest='runtime', '--dist',
action='store_const', dest='runtime',
const='dist', action='store_const',
default='sequential', const='dist',
help='Use sequential distributed runtime instead of local' default='sequential',
) help='Use sequential distributed runtime instead of local'
g_dist.add_argument( )
'--auto-dist', g_dist.add_argument(
action='store_const', '--auto-dist',
const=True, action='store_const',
default=False, const=True,
help='Automatically distribute non-distributed experiments' default=False,
) help='Automatically distribute non-distributed experiments'
g_dist.add_argument( )
'--proxy-type', g_dist.add_argument(
metavar='TYPE', '--proxy-type',
type=str, metavar='TYPE',
default='sockets', type=str,
help='Proxy type to use (sockets,rdma) for auto distribution' default='sockets',
) help='Proxy type to use (sockets,rdma) for auto distribution'
args = parser.parse_args() )
return parser.parse_args()
# pylint: disable=redefined-outer-name
def load_executors(path): def load_executors(path: str) -> tp.List[exectools.Executor]:
"""Load hosts list from json file and return list of executors.""" """Load hosts list from json file and return list of executors."""
with open(path, 'r', encoding='utf-8') as f: with open(path, 'r', encoding='utf-8') as f:
hosts = json.load(f) hosts = json.load(f)
...@@ -248,13 +252,7 @@ def load_executors(path): ...@@ -248,13 +252,7 @@ def load_executors(path):
return exs return exs
if args.hosts is None: def warn_multi_exec(executors: tp.List[exectools.Executor]):
executors = [LocalExecutor()]
else:
executors = load_executors(args.hosts)
def warn_multi_exec():
if len(executors) > 1: if len(executors) > 1:
print( print(
'Warning: multiple hosts specified, only using first one for now', 'Warning: multiple hosts specified, only using first one for now',
...@@ -262,32 +260,16 @@ def warn_multi_exec(): ...@@ -262,32 +260,16 @@ def warn_multi_exec():
) )
# initialize runtime
if args.runtime == 'parallel':
warn_multi_exec()
rt = LocalParallelRuntime(
cores=args.cores,
mem=args.mem,
verbose=args.verbose,
executor=executors[0]
)
elif args.runtime == 'slurm':
rt = SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
elif args.runtime == 'dist':
rt = DistributedSimpleRuntime(executors, verbose=args.verbose)
else:
warn_multi_exec()
rt = LocalSimpleRuntime(verbose=args.verbose, executor=executors[0])
# pylint: disable=redefined-outer-name # pylint: disable=redefined-outer-name
def add_exp( def add_exp(
e: Experiment, e: Experiment,
rt: rt_common.Runtime,
run: int, run: int,
prereq: tp.Optional[Run], prereq: tp.Optional[Run],
create_cp: bool, create_cp: bool,
restore_cp: bool, restore_cp: bool,
no_simbricks: bool no_simbricks: bool,
args: argparse.Namespace
): ):
outpath = f'{args.outdir}/{e.name}-{run}.json' outpath = f'{args.outdir}/{e.name}-{run}.json'
if os.path.exists(outpath) and not args.force: if os.path.exists(outpath) and not args.force:
...@@ -314,60 +296,94 @@ def add_exp( ...@@ -314,60 +296,94 @@ def add_exp(
return run return run
# load experiments def main():
if not args.pickled: args = parse_args()
# default: load python modules with experiments if args.hosts is None:
experiments = [] executors = [exectools.LocalExecutor()]
for path in args.experiments: else:
modname, _ = os.path.splitext(os.path.basename(path)) executors = load_executors(args.hosts)
# initialize runtime
if args.runtime == 'parallel':
warn_multi_exec(executors)
rt = LocalParallelRuntime(
cores=args.cores,
mem=args.mem,
verbose=args.verbose,
executor=executors[0]
)
elif args.runtime == 'slurm':
rt = SlurmRuntime(args.slurmdir, args, verbose=args.verbose)
elif args.runtime == 'dist':
rt = DistributedSimpleRuntime(executors, verbose=args.verbose)
else:
warn_multi_exec(executors)
rt = LocalSimpleRuntime(verbose=args.verbose, executor=executors[0])
# 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))
class ExperimentModuleLoadError(Exception):
pass
spec = importlib.util.spec_from_file_location(modname, path)
if spec is None:
raise ExperimentModuleLoadError('spec is None')
mod = importlib.util.module_from_spec(spec)
if spec.loader is None:
raise ExperimentModuleLoadError('spec.loader is None')
spec.loader.exec_module(mod)
experiments += mod.experiments
if args.list:
for e in experiments:
print(e.name)
sys.exit(0)
class ExperimentModuleLoadError(Exception): for e in experiments:
pass if args.auto_dist and not isinstance(e, DistributedExperiment):
e = auto_dist(e, executors, args.proxy_type)
# apply filter if any specified
if (args.filter) and (len(args.filter) > 0):
match = False
for f in args.filter:
match = fnmatch.fnmatch(e.name, f)
if match:
break
if not match:
continue
# if this is an experiment with a checkpoint we might have to create
# it
no_simbricks = e.no_simbricks
if e.checkpoint:
prereq = add_exp(
e, rt, 0, None, True, False, no_simbricks, args
)
else:
prereq = None
spec = importlib.util.spec_from_file_location(modname, path) for run in range(args.firstrun, args.firstrun + args.runs):
if spec is None: add_exp(
raise ExperimentModuleLoadError('spec is None') e, rt, run, prereq, False, e.checkpoint, no_simbricks, args
mod = importlib.util.module_from_spec(spec) )
if spec.loader is None: else:
raise ExperimentModuleLoadError('spec.loader is None') # otherwise load pickled run object
spec.loader.exec_module(mod) for path in args.experiments:
experiments += mod.experiments with open(path, 'rb') as f:
rt.add_run(pickle.load(f))
if args.list: # register interrupt handler
for e in experiments: signal(SIGINT, lambda *_: rt.interrupt())
print(e.name)
sys.exit(0) # invoke runtime to run experiments
asyncio.run(rt.start())
for e in experiments:
if args.auto_dist and not isinstance(e, DistributedExperiment):
e = auto_dist(e, executors, args.proxy_type) if __name__ == '__main__':
# apply filter if any specified main()
if (args.filter) and (len(args.filter) > 0):
match = False
for f in args.filter:
match = fnmatch.fnmatch(e.name, f)
if match:
break
if not match:
continue
# if this is an experiment with a checkpoint we might have to create it
no_simbricks = e.no_simbricks
if e.checkpoint:
prereq = add_exp(e, 0, None, True, False, no_simbricks)
else:
prereq = None
for run in range(args.firstrun, args.firstrun + args.runs):
add_exp(e, run, prereq, False, e.checkpoint, no_simbricks)
else:
# otherwise load pickled run object
for path in args.experiments:
with open(path, 'rb') as f:
rt.add_run(pickle.load(f))
# register interrupt handler
signal(SIGINT, lambda *_: rt.interrupt())
asyncio.run(rt.start())
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment