mirror of https://github.com/m-labs/artiq.git
153 lines
4.2 KiB
Python
153 lines
4.2 KiB
Python
import sys
|
|
import time
|
|
|
|
from artiq.protocols import pyon
|
|
from artiq.tools import file_import
|
|
from artiq.master.worker_db import DeviceManager, ResultDB, get_hdf5_output
|
|
from artiq.language.environment import is_experiment
|
|
from artiq.language.core import set_watchdog_factory
|
|
|
|
|
|
def get_object():
|
|
line = sys.__stdin__.readline()
|
|
return pyon.decode(line)
|
|
|
|
|
|
def put_object(obj):
|
|
ds = pyon.encode(obj)
|
|
sys.__stdout__.write(ds)
|
|
sys.__stdout__.write("\n")
|
|
sys.__stdout__.flush()
|
|
|
|
|
|
class ParentActionError(Exception):
|
|
pass
|
|
|
|
|
|
def make_parent_action(action, argnames, exception=ParentActionError):
|
|
argnames = argnames.split()
|
|
def parent_action(*args):
|
|
request = {"action": action}
|
|
for argname, arg in zip(argnames, args):
|
|
request[argname] = arg
|
|
put_object(request)
|
|
reply = get_object()
|
|
if "action" in reply:
|
|
if reply["action"] == "terminate":
|
|
sys.exit()
|
|
else:
|
|
raise ValueError
|
|
if reply["status"] == "ok":
|
|
return reply["data"]
|
|
else:
|
|
raise exception(reply["message"])
|
|
return parent_action
|
|
|
|
|
|
class ParentDDB:
|
|
get = make_parent_action("get_device", "name", KeyError)
|
|
|
|
|
|
class ParentPDB:
|
|
get = make_parent_action("get_parameter", "name", KeyError)
|
|
set = make_parent_action("set_parameter", "name value")
|
|
|
|
|
|
update_rt_results = make_parent_action("update_rt_results", "mod")
|
|
|
|
|
|
class Watchdog:
|
|
_create = make_parent_action("create_watchdog", "t")
|
|
_delete = make_parent_action("delete_watchdog", "wid")
|
|
|
|
def __init__(self, t):
|
|
self.t = t
|
|
|
|
def __enter__(self):
|
|
self.wid = Watchdog._create(self.t)
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
Watchdog._delete(self.wid)
|
|
|
|
|
|
set_watchdog_factory(Watchdog)
|
|
|
|
|
|
class Scheduler:
|
|
pause = staticmethod(make_parent_action("pause", ""))
|
|
|
|
submit = staticmethod(make_parent_action("scheduler_submit",
|
|
"pipeline_name expid priority due_date flush"))
|
|
cancel = staticmethod(make_parent_action("scheduler_cancel", "rid"))
|
|
|
|
def set_run_info(self, pipeline_name, expid, priority):
|
|
self.pipeline_name = pipeline_name
|
|
self.expid = expid
|
|
self.priority = priority
|
|
|
|
|
|
def get_exp(file, exp):
|
|
module = file_import(file)
|
|
if exp is None:
|
|
exps = [v for k, v in module.__dict__.items()
|
|
if is_experiment(v)]
|
|
if len(exps) != 1:
|
|
raise ValueError("Found {} experiments in module"
|
|
.format(len(exps)))
|
|
return exps[0]
|
|
else:
|
|
return getattr(module, exp)
|
|
|
|
|
|
def main():
|
|
sys.stdout = sys.stderr
|
|
|
|
start_time = None
|
|
rid = None
|
|
expid = None
|
|
exp = None
|
|
exp_inst = None
|
|
|
|
dmgr = DeviceManager(ParentDDB,
|
|
virtual_devices={"scheduler": Scheduler()})
|
|
rdb = ResultDB()
|
|
rdb.rt.publish = update_rt_results
|
|
|
|
try:
|
|
while True:
|
|
obj = get_object()
|
|
action = obj["action"]
|
|
if action == "build":
|
|
start_time = time.localtime()
|
|
rid = obj["rid"]
|
|
expid = obj["expid"]
|
|
exp = get_exp(expid["file"], expid["experiment"])
|
|
dmgr.virtual_devices["scheduler"].set_run_info(
|
|
obj["pipeline_name"], expid, obj["priority"])
|
|
exp_inst = exp(dmgr, ParentPDB, rdb,
|
|
**expid["arguments"])
|
|
put_object({"action": "completed"})
|
|
elif action == "prepare":
|
|
exp_inst.prepare()
|
|
put_object({"action": "completed"})
|
|
elif action == "run":
|
|
exp_inst.run()
|
|
put_object({"action": "completed"})
|
|
elif action == "analyze":
|
|
exp_inst.analyze()
|
|
put_object({"action": "completed"})
|
|
elif action == "write_results":
|
|
f = get_hdf5_output(start_time, rid, exp.__name__)
|
|
try:
|
|
rdb.write_hdf5(f)
|
|
finally:
|
|
f.close()
|
|
put_object({"action": "completed"})
|
|
elif action == "terminate":
|
|
break
|
|
finally:
|
|
dmgr.close_devices()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|