forked from M-Labs/artiq
185 lines
5.6 KiB
Python
Executable File
185 lines
5.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import sys
|
|
import time
|
|
from operator import itemgetter
|
|
from itertools import chain
|
|
|
|
import h5py
|
|
|
|
from artiq.language.db import *
|
|
from artiq.language.experiment import is_experiment
|
|
from artiq.protocols import pyon
|
|
from artiq.protocols.file_db import FlatFileDB
|
|
from artiq.master.worker_db import DBHub, ResultDB
|
|
from artiq.tools import file_import, verbosity_args, init_logger
|
|
|
|
|
|
class ELFRunner(AutoDB):
|
|
class DBKeys:
|
|
comm = Device()
|
|
|
|
def run(self, filename):
|
|
with open(filename, "rb") as f:
|
|
binary = f.read()
|
|
comm.load(binary)
|
|
comm.run("run")
|
|
comm.serve(dict(), dict())
|
|
|
|
|
|
class SimpleParamLogger:
|
|
def set(self, timestamp, name, value):
|
|
print("Parameter change: {} -> {}".format(name, value))
|
|
|
|
|
|
class DummyWatchdog:
|
|
def __init__(self, t):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
pass
|
|
|
|
|
|
class DummyScheduler:
|
|
def __init__(self):
|
|
self.next_rid = 0
|
|
self.next_trid = 0
|
|
|
|
def run_queued(self, run_params):
|
|
rid = self.next_rid
|
|
self.next_rid += 1
|
|
print("Queuing: {}, RID={}".format(run_params, rid))
|
|
return rid
|
|
|
|
def cancel_queued(self, rid):
|
|
print("Cancelling RID {}".format(rid))
|
|
|
|
def run_timed(self, run_params, next_run):
|
|
trid = self.next_trid
|
|
self.next_trid += 1
|
|
next_run_s = time.strftime("%m/%d %H:%M:%S", time.localtime(next_run))
|
|
print("Timing: {} at {}, TRID={}".format(run_params, next_run_s, trid))
|
|
return trid
|
|
|
|
def cancel_timed(self, trid):
|
|
print("Cancelling TRID {}".format(trid))
|
|
|
|
watchdog = DummyWatchdog
|
|
|
|
|
|
def get_argparser():
|
|
parser = argparse.ArgumentParser(
|
|
description="Local experiment running tool")
|
|
|
|
verbosity_args(parser)
|
|
parser.add_argument("-d", "--ddb", default="ddb.pyon",
|
|
help="device database file")
|
|
parser.add_argument("-p", "--pdb", default="pdb.pyon",
|
|
help="parameter database file")
|
|
|
|
parser.add_argument("-E", "--elf", default=False, action="store_true",
|
|
help="run ELF binary")
|
|
parser.add_argument("-e", "--experiment", default=None,
|
|
help="experiment to run")
|
|
parser.add_argument("-o", "--hdf5", default=None,
|
|
help="write results to specified HDF5 file"
|
|
" (default: print them)")
|
|
parser.add_argument("file",
|
|
help="file containing the experiment to run")
|
|
parser.add_argument("arguments", nargs="*",
|
|
help="run arguments")
|
|
|
|
return parser
|
|
|
|
|
|
def _parse_arguments(arguments):
|
|
d = {}
|
|
for argument in arguments:
|
|
name, value = argument.split("=")
|
|
d[name] = pyon.decode(value)
|
|
return d
|
|
|
|
|
|
def main():
|
|
args = get_argparser().parse_args()
|
|
init_logger(args)
|
|
|
|
ddb = FlatFileDB(args.ddb)
|
|
pdb = FlatFileDB(args.pdb)
|
|
pdb.hooks.append(SimpleParamLogger())
|
|
rdb = ResultDB(lambda description: None, lambda mod: None)
|
|
dbh = DBHub(ddb, pdb, rdb)
|
|
try:
|
|
if args.elf:
|
|
if args.arguments:
|
|
print("Run arguments are not supported in ELF mode")
|
|
sys.exit(1)
|
|
exp_inst = ELFRunner(dbh)
|
|
rdb.build()
|
|
exp_inst.run(args.file)
|
|
else:
|
|
module = file_import(args.file)
|
|
if args.experiment is None:
|
|
exps = [(k, v) for k, v in module.__dict__.items()
|
|
if is_experiment(v)]
|
|
l = len(exps)
|
|
if l == 0:
|
|
print("No experiments found in module")
|
|
sys.exit(1)
|
|
elif l > 1:
|
|
print("More than one experiment found in module:")
|
|
for k, v in sorted(experiments, key=itemgetter(0)):
|
|
if v.__doc__ is None:
|
|
print(" {}".format(k))
|
|
else:
|
|
print(" {} ({})".format(
|
|
k, v.__doc__.splitlines()[0].strip()))
|
|
print("Use -u to specify which experiment to use.")
|
|
sys.exit(1)
|
|
else:
|
|
exp = exps[0][1]
|
|
else:
|
|
exp = getattr(module, args.experiment)
|
|
|
|
try:
|
|
arguments = _parse_arguments(args.arguments)
|
|
except:
|
|
print("Failed to parse run arguments")
|
|
sys.exit(1)
|
|
|
|
run_params = {
|
|
"file": args.file,
|
|
"experiment": args.experiment,
|
|
"arguments": arguments
|
|
}
|
|
exp_inst = exp(dbh,
|
|
scheduler=DummyScheduler(),
|
|
run_params=run_params,
|
|
**run_params["arguments"])
|
|
rdb.build()
|
|
exp_inst.run()
|
|
exp_inst.analyze()
|
|
|
|
if args.hdf5 is not None:
|
|
f = h5py.File(args.hdf5, "w")
|
|
try:
|
|
rdb.write_hdf5(f)
|
|
finally:
|
|
f.close()
|
|
else:
|
|
if rdb.data.read or rdb.realtime_data.read:
|
|
print("Results:")
|
|
for k, v in sorted(chain(rdb.realtime_data.read.items(),
|
|
rdb.data.read.items()),
|
|
key=itemgetter(0)):
|
|
print("{}: {}".format(k, v))
|
|
finally:
|
|
dbh.close_devices()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|