m-labs · Apr 4, 2016 · Apr 4, 2016 · Apr 4, 2016 · Apr 4, 2016
Showing with 74 additions and 16 deletions.

+14 −0 artiq/protocols/pc_rpc.py

+20 −1 artiq/protocols/remote_exec.py

+13 −12 examples/master/repository/remote_exec_demo.py

+25 −3 examples/master/repository/remote_exec_processing.py

+2 −0 examples/remote_exec_controller.py
diff --git a/artiq/protocols/pc_rpc.py b/artiq/protocols/pc_rpc.py
@@ -132,6 +132,10 @@ def get_rpc_id(self):
         identification information of the server."""
         return (self.__target_names, self.__description)
 
+    def get_local_host(self):
+        """Returns the address of the local end of the connection."""
+        return self.__socket.getsockname()[0]
+
     def close_rpc(self):
         """Closes the connection to the RPC server.
 
@@ -224,6 +228,10 @@ def get_selected_target(self):
         selected yet."""
         return self.__selected_target
 
+    def get_local_host(self):
+        """Returns the address of the local end of the connection."""
+        return self.__writer.get_extra_info("socket").getsockname()[0]
+
     def get_rpc_id(self):
         """Returns a tuple (target_names, description) containing the
         identification information of the server."""
@@ -396,6 +404,12 @@ def proxy(*args, **kwargs):
             return self.__do_rpc(name, args, kwargs)
         return proxy
 
+    def get_selected_target(self):
+        raise NotImplementedError
+
+    def get_local_host(self):
+        raise NotImplementedError
+
 
 def _format_arguments(arguments):
     fmtargs = []

diff --git a/artiq/protocols/remote_exec.py b/artiq/protocols/remote_exec.py
@@ -4,12 +4,14 @@
 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:
     def __init__(self, initial_namespace):
         self.namespace = dict(initial_namespace)
+        # The module actually has to exist, otherwise it breaks e.g. Numba
+        self.namespace["__name__"] = "artiq.protocols.remote_exec"
 
     def add_code(self, code):
         exec(code, self.namespace)
@@ -21,8 +23,25 @@ def call(self, function, *args, **kwargs):
 def simple_rexec_server_loop(target_name, target, host, port,
                              description=None):
     initial_namespace = {"controller_driver": target}
+    initial_namespace["controller_initial_namespace"] = initial_namespace
     targets = {
         target_name: target,
         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)
diff --git a/examples/master/repository/remote_exec_demo.py b/examples/master/repository/remote_exec_demo.py
@@ -2,6 +2,7 @@
 import inspect
 
 from artiq.experiment import *
+from artiq.protocols.remote_exec import connect_global_rpc
 
 import remote_exec_processing
 
@@ -18,19 +19,16 @@ def build(self):
 
     def prepare(self):
         if self.remote_exec:
+            connect_global_rpc(self.camera_sim_rexec)
             self.camera_sim_rexec.add_code(
                 inspect.getsource(remote_exec_processing))
 
-    def set_dataset(self, name, *args, **kwargs):
-        EnvExperiment.set_dataset(self, "rexec_demo." + name,
-                                  *args, **kwargs)
-
     def transfer_parameters(self, parameters):
         w, h, cx, cy = parameters
-        self.set_dataset("gaussian_w", w, save=False, broadcast=True)
-        self.set_dataset("gaussian_h", h, save=False, broadcast=True)
-        self.set_dataset("gaussian_cx", cx, save=False, broadcast=True)
-        self.set_dataset("gaussian_cy", cy, save=False, broadcast=True)
+        self.set_dataset("rexec_demo.gaussian_w", w, save=False, broadcast=True)
+        self.set_dataset("rexec_demo.gaussian_h", h, save=False, broadcast=True)
+        self.set_dataset("rexec_demo.gaussian_cx", cx, save=False, broadcast=True)
+        self.set_dataset("rexec_demo.gaussian_cy", cy, save=False, broadcast=True)
 
     def fps_meter(self):
         t = time.monotonic()
@@ -39,7 +37,7 @@ def fps_meter(self):
             dt = t - self.last_pt_update
             if dt >= 5:
                 pt = dt/self.iter_count
-                self.set_dataset("picture_pt", pt, save=False, broadcast=True)
+                self.set_dataset("rexec_demo.picture_pt", pt, save=False, broadcast=True)
                 self.last_pt_update = t
                 self.iter_count = 0
         else:
@@ -51,15 +49,18 @@ def run_local(self):
             self.fps_meter()
             data = self.camera_sim.get_picture()
             if self.show_picture:
-                self.set_dataset("picture", data, save=False, broadcast=True)
+                self.set_dataset("rexec_demo.picture", data,
+                                 save=False, broadcast=True)
             if self.enable_fit:
-                self.transfer_parameters(remote_exec_processing.fit(data))
+                p = remote_exec_processing.fit(data, self.get_dataset)
+                self.transfer_parameters(p)
             self.scheduler.pause()
 
     def run_remote(self):
         while True:
             self.fps_meter()
-            self.transfer_parameters(self.camera_sim_rexec.call("get_and_fit"))
+            p = self.camera_sim_rexec.call("get_and_fit")
+            self.transfer_parameters(p)
             self.scheduler.pause()
 
     def run(self):

diff --git a/examples/master/repository/remote_exec_processing.py b/examples/master/repository/remote_exec_processing.py
@@ -1,6 +1,10 @@
 import numpy as np
 from numba import jit
 from scipy.optimize import least_squares
+import logging
+
+
+logger = logging.getLogger(__name__)
 
 
 @jit(nopython=True)
@@ -14,17 +18,35 @@ def compute_gaussian(r, img_w, img_h,
             r[x, y] = np.exp(-ds/2)
 
 
-def fit(data):
+def fit(data, get_dataset):
     img_w, img_h = data.shape
     def err(parameters):
         r = np.empty((img_w, img_h))
         compute_gaussian(r, img_w, img_h, *parameters)
         r -= data
         return r.ravel()
-    guess = [12, 15, img_w/2, img_h/2]
+    guess = [
+        get_dataset("rexec_demo.gaussian_w", 12),
+        get_dataset("rexec_demo.gaussian_h", 15),
+        get_dataset("rexec_demo.gaussian_cx", img_w/2),
+        get_dataset("rexec_demo.gaussian_cy", img_h/2)
+    ]
     res = least_squares(err, guess)
     return res.x
 
 
 def get_and_fit():
-    return fit(controller_driver.get_picture())
+    if "dataset_db" in globals():
+        logger.info("using dataset DB for Gaussian fit guess")
+        def get_dataset(name, default):
+            from artiq.protocols import pc_rpc
+            try:
+                return dataset_db.get(name)
+            except (KeyError, pc_rpc.RemoteError):  # TODO: serializable exceptions
+                return default
+    else:
+        logger.info("using defaults for Gaussian fit guess")
+        def get_dataset(name, default):
+            return default
+        get_dataset = lambda name, default: default 
+    return fit(controller_driver.get_picture(), get_dataset)
diff --git a/examples/remote_exec_controller.py b/examples/remote_exec_controller.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 from numba import jit
+import logging
 
 from artiq.protocols.remote_exec import simple_rexec_server_loop
 
@@ -51,6 +52,7 @@ def get_picture(self):
 
 
 def main():
+    logging.basicConfig(level=logging.INFO)
     simple_rexec_server_loop("camera_sim", CameraSimulation(),
                              "::1", 6283)