artiq/frontend/artiq_run.py

84 lines
2.6 KiB
Python
Executable File

#!/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()