mirror of https://github.com/m-labs/artiq.git
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
||
|
import argparse
|
||
|
import importlib
|
||
|
import sys
|
||
|
from inspect import isclass
|
||
|
from operator import itemgetter
|
||
|
|
||
|
from artiq.language.context import *
|
||
|
from artiq.management import pyon
|
||
|
from artiq.management.dpdb import DeviceParamDB
|
||
|
|
||
|
|
||
|
class ELFRunner(AutoContext):
|
||
|
comm = Device("comm")
|
||
|
implicit_core = False
|
||
|
|
||
|
def run(self, filename, function):
|
||
|
with open(filename, "rb") as f:
|
||
|
binary = f.read()
|
||
|
comm.load(binary)
|
||
|
comm.run(function)
|
||
|
comm.serve(dict(), dict())
|
||
|
|
||
|
|
||
|
def _get_args():
|
||
|
parser = argparse.ArgumentParser(description="Experiment running tool")
|
||
|
|
||
|
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("-f", "--function", default="run",
|
||
|
help="function to run")
|
||
|
parser.add_argument("-u", "--unit", default=None,
|
||
|
help="unit to run")
|
||
|
parser.add_argument("module",
|
||
|
help="module containing the unit to run")
|
||
|
|
||
|
return parser.parse_args()
|
||
|
|
||
|
|
||
|
def main():
|
||
|
args = _get_args()
|
||
|
|
||
|
devices = pyon.load_file(args.ddb)
|
||
|
parameters = pyon.load_file(args.pdb)
|
||
|
dpdb = DeviceParamDB(devices, parameters)
|
||
|
try:
|
||
|
if args.elf:
|
||
|
unit_inst = ELFRunner(dpdb)
|
||
|
unit_inst.run(args.file, args.function)
|
||
|
else:
|
||
|
sys.path.append(".")
|
||
|
module = importlib.import_module(args.module)
|
||
|
if args.unit is None:
|
||
|
units = [(k, v) for k, v in module.__dict__.items()
|
||
|
if k[0] != "_" and isclass(v) and issubclass(v, AutoContext) and v is not AutoContext]
|
||
|
l = len(units)
|
||
|
if l == 0:
|
||
|
print("No units found in module")
|
||
|
sys.exit(1)
|
||
|
elif l > 1:
|
||
|
print("More than one unit found in module:")
|
||
|
for k, v in sorted(units, key=itemgetter(0)):
|
||
|
print(" " + k)
|
||
|
print("Use -u to specify which unit to use.")
|
||
|
sys.exit(1)
|
||
|
else:
|
||
|
unit = units[0][1]
|
||
|
else:
|
||
|
unit = getattr(module, args.unit)
|
||
|
unit_inst = unit(dpdb)
|
||
|
f = getattr(unit_inst, args.function)
|
||
|
f()
|
||
|
finally:
|
||
|
dpdb.close()
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
main()
|