From 99276b9b6959f40847b2ca0b5450e6c94449ab95 Mon Sep 17 00:00:00 2001
From: Robert Nishihara <robertnishihara@gmail.com>
Date: Thu, 31 Mar 2016 15:43:04 -0700
Subject: [PATCH] properly serialize/deserialize calls

---
 lib/orchpy/orchpy/serialization.py | 40 +++++++++++++++++++-----------
 lib/orchpy/orchpy/worker.py        |  4 +--
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/lib/orchpy/orchpy/serialization.py b/lib/orchpy/orchpy/serialization.py
index dc5c0952c..3cf7aac3a 100644
--- a/lib/orchpy/orchpy/serialization.py
+++ b/lib/orchpy/orchpy/serialization.py
@@ -2,25 +2,37 @@ import importlib
 
 import orchpy
 
-def serialize(obj):
+def to_primitive(obj):
   if hasattr(obj, "serialize"):
     primitive_obj = ((type(obj).__module__, type(obj).__name__), obj.serialize())
   else:
-    # TODO(rkn): Right now we don't handle arbitrary python objects, but later
-    # we can unpack the fields of a python object into a list and call
-    # orchpy.lib.serialize_object.
     primitive_obj = ("primitive", obj)
+  return primitive_obj
+
+def from_primitive(primitive_obj):
+  if primitive_obj[0] == "primitive":
+    obj = primitive_obj[1]
+  else:
+    # This code assumes that the type module.__dict__[type_name] knows how to deserialize itself
+    type_module, type_name = primitive_obj[0]
+    module = importlib.import_module(type_module)
+    obj = module.__dict__[type_name]()
+    obj.deserialize(primitive_obj[1])
+  return obj
+
+def serialize(obj):
+  primitive_obj = to_primitive(obj)
   return orchpy.lib.serialize_object(primitive_obj)
 
 def deserialize(capsule):
   primitive_obj = orchpy.lib.deserialize_object(capsule)
-  if primitive_obj[0] == "primitive":
-    return primitive_obj[1]
-  else:
-    # assert primitive_obj[0] must be a tuple of module and class name
-    type_module, type_name = primitive_obj[0]
-    module = importlib.import_module(type_module)
-    if hasattr(module.__dict__[type_name], "deserialize"):
-      obj = module.__dict__[type_name]()
-      obj.deserialize(primitive_obj[1])
-      return obj
+  return from_primitive(primitive_obj)
+
+def serialize_call(func_name, args):
+  primitive_args = [(arg if isinstance(arg, orchpy.lib.ObjRef) else to_primitive(arg)) for arg in args]
+  return orchpy.lib.serialize_call(func_name, primitive_args)
+
+def deserialize_call(call):
+  func_name, primitive_args, return_objrefs = orchpy.lib.deserialize_call(call)
+  args = [(arg if isinstance(arg, orchpy.lib.ObjRef) else from_primitive(arg)) for arg in primitive_args]
+  return func_name, args, return_objrefs
diff --git a/lib/orchpy/orchpy/worker.py b/lib/orchpy/orchpy/worker.py
index c7dc77680..e20a6db7d 100644
--- a/lib/orchpy/orchpy/worker.py
+++ b/lib/orchpy/orchpy/worker.py
@@ -29,7 +29,7 @@ class Worker(object):
 
   def remote_call(self, func_name, args):
     """Tell the scheduler to schedule the execution of the function with name `func_name` with arguments `args`. Retrieve object references for the outputs of the function from the scheduler and immediately return them."""
-    call_capsule = orchpy.lib.serialize_call(func_name, args)
+    call_capsule = serialization.serialize_call(func_name, args)
     objrefs = orchpy.lib.remote_call(self.handle, call_capsule)
     return objrefs
 
@@ -66,7 +66,7 @@ def main_loop(worker=global_worker):
   orchpy.lib.start_worker_service(worker.handle)
   while True:
     call = orchpy.lib.wait_for_next_task(worker.handle)
-    func_name, args, return_objrefs = orchpy.lib.deserialize_call(call)
+    func_name, args, return_objrefs = serialization.deserialize_call(call)
     arguments = get_arguments_for_execution(worker.functions[func_name], args, worker) # get args from objstore
     outputs = worker.functions[func_name].executor(arguments) # execute the function
     store_outputs_in_objstore(return_objrefs, outputs, worker) # store output in local object store