diff --git a/artiq/protocols/remote_exec.py b/artiq/protocols/remote_exec.py index 7d82dc373..6dd2ad5e9 100644 --- a/artiq/protocols/remote_exec.py +++ b/artiq/protocols/remote_exec.py @@ -4,7 +4,7 @@ import inspect from artiq.protocols.pc_rpc import simple_server_loop -__all__ = ["RemoteExecServer", "RemoteExecClient", "simple_rexec_server_loop"] +__all__ = ["RemoteExecServer", "simple_rexec_server_loop", "connect_global_rpc"] class RemoteExecServer: @@ -29,3 +29,19 @@ def simple_rexec_server_loop(target_name, target, host, port, target_name + "_rexec": lambda: RemoteExecServer(initial_namespace) } simple_server_loop(targets, host, port, description) + + +def connect_global_rpc(controller_rexec, host=None, port=3251, + target="master_dataset_db", name="dataset_db"): + if host is None: + host = controller_rexec.get_local_host() + code = """ +if "{name}" not in controller_initial_namespace: + import atexit + from artiq.protocols.pc_rpc import Client + + {name} = Client("{host}", {port}, "{target}") + atexit.register({name}.close_rpc) + controller_initial_namespace["{name}"] = {name} +""".format(host=host, port=port, target=target, name=name) + controller_rexec.add_code(code)