diff --git a/artiq/management/scheduler.py b/artiq/management/scheduler.py index 9372b563b..775b14c19 100644 --- a/artiq/management/scheduler.py +++ b/artiq/management/scheduler.py @@ -4,8 +4,8 @@ from artiq.management.worker import Worker class Scheduler: - def __init__(self): - self.worker = Worker() + def __init__(self, *args, **kwargs): + self.worker = Worker(*args, **kwargs) self.queue = asyncio.Queue() @asyncio.coroutine diff --git a/artiq/management/worker.py b/artiq/management/worker.py index 7746fecd3..e20dc6f36 100644 --- a/artiq/management/worker.py +++ b/artiq/management/worker.py @@ -11,8 +11,10 @@ class WorkerFailed(Exception): class Worker: - def __init__(self, send_timeout=0.5, start_reply_timeout=1.0, - term_timeout=1.0): + def __init__(self, ddb, pdb, + send_timeout=0.5, start_reply_timeout=1.0, term_timeout=1.0): + self.ddb = ddb + self.pdb = pdb self.send_timeout = send_timeout self.start_reply_timeout = start_reply_timeout self.term_timeout = term_timeout @@ -21,6 +23,7 @@ class Worker: def create_process(self): self.process = yield from asyncio.create_subprocess_exec( sys.executable, "-m", "artiq.management.worker_impl", + self.ddb, self.pdb, stdout=subprocess.PIPE, stdin=subprocess.PIPE) @asyncio.coroutine diff --git a/artiq/management/worker_impl.py b/artiq/management/worker_impl.py index 79492512d..d3682f08d 100644 --- a/artiq/management/worker_impl.py +++ b/artiq/management/worker_impl.py @@ -1,23 +1,28 @@ import sys -import importlib +from inspect import isclass from artiq.management import pyon +from artiq.management.file_import import file_import +from artiq.language.context import AutoContext +from artiq.management.dpdb import DeviceParamDB -def import_in_folder(path, name): - try: - del sys.modules[name] # force path search - except KeyError: - pass - loader = importlib.find_loader(name, [path]) - if loader is None: - raise ImportError("Could not find loader") - return loader.load_module() - - -def run(path, name): - module = import_in_folder(path, name) - module.main() +def run(dpdb, file, unit, function): + module = file_import(file) + if unit is None: + units = [v for k, v in module.__dict__.items() + if k[0] != "_" + and isclass(v) + and issubclass(v, AutoContext) + and v is not AutoContext] + if len(units) != 1: + raise ValueError("Found {} units in module".format(len(units))) + unit = units[0] + else: + unit = getattr(module, unit) + unit_inst = unit(dpdb) + f = getattr(unit_inst, function) + f() def put_object(obj): @@ -30,13 +35,17 @@ def put_object(obj): def main(): sys.stdout = sys.stderr + devices = pyon.load_file(sys.argv[1]) + parameters = pyon.load_file(sys.argv[2]) + dpdb = DeviceParamDB(devices, parameters) + while True: line = sys.__stdin__.readline() obj = pyon.decode(line) put_object("ack") try: - run(**obj) + run(dpdb, **obj) except Exception as e: put_object({"status": "failed", "message": str(e)}) else: diff --git a/frontend/artiq_client.py b/frontend/artiq_client.py index 41ea7e01a..a75f1f260 100755 --- a/frontend/artiq_client.py +++ b/frontend/artiq_client.py @@ -13,10 +13,22 @@ def _get_args(): parser.add_argument( "--port", default=8888, type=int, help="TCP port to use to connect to the master") - parser.add_argument( - "-o", "--run-once", default=[], nargs=3, - action="append", - help="run experiment once. arguments: ") + + subparsers = parser.add_subparsers(dest="action") + + parser_add = subparsers.add_parser("add", help="add an experiment") + parser_add.add_argument( + "-p", "--periodic", default=None, type=float, + help="run the experiment periodically every given number of seconds") + parser_add.add_argument( + "-t", "--timeout", default=None, type=float, + help="specify a timeout for the experiment to complete") + parser_add.add_argument("-f", "--function", default="run", + help="function to run") + parser_add.add_argument("-u", "--unit", default=None, + help="unit to run") + parser_add.add_argument("file", help="file containing the unit to run") + return parser.parse_args() @@ -24,12 +36,16 @@ def main(): args = _get_args() remote = Client(args.server, args.port, "master") try: - for path, name, timeout in args.run_once: - remote.run_once( - { - "path": path, - "name": name - }, int(timeout)) + if args.action == "add": + if args.periodic is None: + remote.run_once( + { + "file": args.file, + "unit": args.unit, + "function": args.function + }, args.timeout) + else: + raise NotImplementedError finally: remote.close_rpc() diff --git a/frontend/artiq_master.py b/frontend/artiq_master.py index 1006fdab9..615c1b052 100755 --- a/frontend/artiq_master.py +++ b/frontend/artiq_master.py @@ -8,7 +8,7 @@ from artiq.management.scheduler import Scheduler def _get_args(): - parser = argparse.ArgumentParser(description="PDQ2 controller") + parser = argparse.ArgumentParser(description="ARTIQ master") parser.add_argument( "--bind", default="::1", help="hostname or IP address to bind to") @@ -22,7 +22,7 @@ def main(): args = _get_args() loop = asyncio.get_event_loop() try: - scheduler = Scheduler() + scheduler = Scheduler("ddb.pyon", "pdb.pyon") loop.run_until_complete(scheduler.start()) try: server = Server(scheduler, "master") diff --git a/frontend/artiq_run.py b/frontend/artiq_run.py index 9b261ad6b..3afd622ee 100755 --- a/frontend/artiq_run.py +++ b/frontend/artiq_run.py @@ -57,7 +57,10 @@ def main(): module = file_import(args.file) 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] + 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")