m-labs · Oct 23, 2014 · Oct 25, 2014 · Oct 25, 2014 · Oct 25, 2014 · Oct 25, 2014
Showing with 121 additions and 21 deletions.

+23 −8 artiq/management/pc_rpc.py

+73 −0 artiq/management/pyon.py

+8 −2 frontend/artiq

+17 −11 frontend/artiqd
diff --git a/artiq/management/pc_rpc.py b/artiq/management/pc_rpc.py
@@ -1,6 +1,8 @@
 import socket
-import json
 import asyncio
+import traceback
+
+from artiq.management import pyon
 
 
 class RemoteError(Exception):
@@ -11,12 +13,12 @@ class Client:
     def __init__(self, host, port):
         self.socket = socket.create_connection((host, port))
 
-    def close(self):
+    def close_rpc(self):
         self.socket.close()
 
     def do_rpc(self, name, args, kwargs):
         obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
-        line = json.dumps(obj) + "\n"
+        line = pyon.encode(obj) + "\n"
         self.socket.sendall(line.encode())
 
         buf = self.socket.recv(4096).decode()
@@ -25,11 +27,11 @@ def do_rpc(self, name, args, kwargs):
             if not more:
                 break
             buf += more.decode()
-        obj = json.loads(buf)
+        obj = pyon.decode(buf)
         if obj["result"] == "ok":
             return obj["ret"]
         elif obj["result"] == "error":
-            raise RemoteError(obj["message"])
+            raise RemoteError(obj["message"] + "\n" + obj["traceback"])
         else:
             raise ValueError
 
@@ -72,7 +74,7 @@ def handle_connection_task(self, reader, writer):
                 line = yield from reader.readline()
                 if not line:
                     break
-                obj = json.loads(line.decode())
+                obj = pyon.decode(line.decode())
                 action = obj["action"]
                 if action == "call":
                     method = getattr(self.target, obj["name"])
@@ -81,8 +83,21 @@ def handle_connection_task(self, reader, writer):
                         obj = {"result": "ok", "ret": ret}
                     except Exception as e:
                         obj = {"result": "error",
-                               "message": type(e).__name__ + ": " + str(e)}
-                    line = json.dumps(obj) + "\n"
+                               "message": type(e).__name__ + ": " + str(e),
+                               "traceback": traceback.format_exc()}
+                    line = pyon.encode(obj) + "\n"
                     writer.write(line.encode())
         finally:
             writer.close()
+
+
+class WaitQuit:
+    def __init__(self):
+        self.terminate_notify = asyncio.Semaphore(0)
+
+    @asyncio.coroutine
+    def wait_quit(self):
+        yield from self.terminate_notify.acquire()
+
+    def quit(self):
+        self.terminate_notify.release()
diff --git a/artiq/management/pyon.py b/artiq/management/pyon.py
@@ -0,0 +1,73 @@
+import base64
+
+import numpy
+
+
+def _encode_none(x):
+    return "None"
+
+
+def _encode_number(x):
+    return str(x)
+
+
+def _encode_str(x):
+    return repr(x)
+
+
+def _encode_tuple(x):
+    if len(x) == 1:
+        return "(" + encode(x[0]) + ", )"
+    else:
+        r = "("
+        r += ", ".join([encode(item) for item in x])
+        r += ")"
+        return r
+
+
+def _encode_list(x):
+    r = "["
+    r += ", ".join([encode(item) for item in x])
+    r += "]"
+    return r
+
+
+def _encode_dict(x):
+    r = "{"
+    r += ", ".join([encode(k) + ": " + encode(v) for k, v in x.items()])
+    r += "}"
+    return r
+
+
+def _encode_nparray(x):
+    r = "nparray("
+    r += encode(x.shape) + ", "
+    r += encode(str(x.dtype)) + ", "
+    r += encode(base64.b64encode(x).decode())
+    r += ")"
+    return r
+
+
+_encode_map = {
+    type(None): _encode_none,
+    int: _encode_number,
+    float: _encode_number,
+    str: _encode_str,
+    tuple: _encode_tuple,
+    list: _encode_list,
+    dict: _encode_dict,
+    numpy.ndarray: _encode_nparray
+}
+
+
+def encode(x):
+    return _encode_map[type(x)](x)
+
+
+def _nparray(shape, dtype, data):
+    a = numpy.frombuffer(base64.b64decode(data), dtype=dtype)
+    return a.reshape(shape)
+
+
+def decode(s):
+    return eval(s, {"__builtins__": None, "nparray": _nparray}, {})
diff --git a/frontend/artiq b/frontend/artiq
@@ -7,6 +7,12 @@ from artiq.management.pc_rpc import Client
 
 def _get_args():
     parser = argparse.ArgumentParser(description="ARTIQ client")
+    parser.add_argument(
+        "-s", "--server", default="::1",
+        help="hostname or IP of the master to connect to")
+    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",
@@ -20,7 +26,7 @@ def _get_args():
 
 def main():
     args = _get_args()
-    remote = Client("::1", 8888)
+    remote = Client(args.server, args.port)
     try:
         for path, name, timeout in args.run_once:
             remote.run_once(
@@ -31,7 +37,7 @@ def main():
         if args.quit_master:
             remote.quit()
     finally:
-        remote.close()
+        remote.close_rpc()
 
 if __name__ == "__main__":
     main()
diff --git a/frontend/artiqd b/frontend/artiqd
@@ -1,36 +1,42 @@
 #!/usr/bin/env python3
 
 import asyncio
+import argparse
 
-from artiq.management.pc_rpc import Server
+from artiq.management.pc_rpc import Server, WaitQuit
 from artiq.management.scheduler import Scheduler
 
 
-class Master:
+class Master(WaitQuit):
     def __init__(self, scheduler):
+        WaitQuit.__init__(self)
         self.scheduler = scheduler
-        self.terminate_notify = asyncio.Semaphore(0)
-
-    @asyncio.coroutine
-    def wait_quit(self):
-        yield from self.terminate_notify.acquire()
-
-    def quit(self):
-        self.terminate_notify.release()
 
     def run_once(self, run_params, timeout):
         self.scheduler.run_once(run_params, timeout)
 
 
+def _get_args():
+    parser = argparse.ArgumentParser(description="PDQ2 controller")
+    parser.add_argument(
+        "--bind", default="::1",
+        help="hostname or IP address to bind to")
+    parser.add_argument(
+        "--port", default=8888, type=int,
+        help="TCP port to listen to")
+    return parser.parse_args()
+
+
 def main():
+    args = _get_args()
     loop = asyncio.get_event_loop()
     try:
         scheduler = Scheduler()
         loop.run_until_complete(scheduler.start())
         try:
             master = Master(scheduler)
             server = Server(master)
-            loop.run_until_complete(server.start("::1", 8888))
+            loop.run_until_complete(server.start(args.bind, args.port))
             try:
                 loop.run_until_complete(master.wait_quit())
             finally: