implement reference counting and much more (#43)

lib/orchpy/orchpy/__init__.py

@@ -1,3 +1,3 @@
 import liborchpylib as lib
 import serialization
-from worker import register_module, connect, pull, push, distributed
+from worker import scheduler_info, register_module, connect, disconnect, pull, push, distributed

lib/orchpy/orchpy/serialization.py

@@ -20,19 +20,20 @@ def from_primitive(primitive_obj):
     obj.deserialize(primitive_obj[1])
   return obj
 
-def serialize(obj):
+def serialize(worker_capsule, obj):
   primitive_obj = to_primitive(obj)
-  return orchpy.lib.serialize_object(primitive_obj)
+  obj_capsule, contained_objrefs = orchpy.lib.serialize_object(worker_capsule, primitive_obj) # contained_objrefs is a list of the objrefs contained in obj
+  return obj_capsule, contained_objrefs
 
-def deserialize(capsule):
-  primitive_obj = orchpy.lib.deserialize_object(capsule)
+def deserialize(worker_capsule, capsule):
+  primitive_obj = orchpy.lib.deserialize_object(worker_capsule, capsule)
   return from_primitive(primitive_obj)
 
-def serialize_call(func_name, args):
+def serialize_call(worker_capsule, 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)
+  return orchpy.lib.serialize_call(worker_capsule, func_name, primitive_args)
 
-def deserialize_call(call):
-  func_name, primitive_args, return_objrefs = orchpy.lib.deserialize_call(call)
+def deserialize_call(worker_capsule, call):
+  func_name, primitive_args, return_objrefs = orchpy.lib.deserialize_call(worker_capsule, call)
   args = [(arg if isinstance(arg, orchpy.lib.ObjRef) else from_primitive(arg)) for arg in primitive_args]
   return func_name, args, return_objrefs

lib/orchpy/orchpy/services.py

@@ -9,6 +9,7 @@ import orchpy.worker as worker
 _services_path = os.path.dirname(os.path.abspath(__file__))
 
 all_processes = []
+driver = None
 
 IP_ADDRESS = "127.0.0.1"
 TIMEOUT_SECONDS = 5
@@ -55,7 +56,14 @@ def cleanup():
     print "Termination attempt failed, giving up."
   all_processes = []
 
-atexit.register(cleanup)
+  global driver
+  if driver is not None:
+    orchpy.disconnect(driver)
+  else:
+    orchpy.disconnect()
+  driver = None
+
+# atexit.register(cleanup)
 
 def start_scheduler(scheduler_address):
   p = subprocess.Popen([os.path.join(_services_path, "scheduler"), scheduler_address])
@@ -74,6 +82,7 @@ def start_worker(test_path, scheduler_address, objstore_address, worker_address)
   all_processes.append((p, worker_address))
 
 def start_cluster(driver_worker=None, num_workers=0, worker_path=None):
+  global driver
   if num_workers > 0 and worker_path is None:
     raise Exception("Attempting to start a cluster with some workers, but `worker_path` is None.")
   scheduler_address = address(IP_ADDRESS, new_scheduler_port())
@@ -84,6 +93,7 @@ def start_cluster(driver_worker=None, num_workers=0, worker_path=None):
   time.sleep(0.2)
   if driver_worker is not None:
     orchpy.connect(scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()), driver_worker)
+    driver = driver_worker
   else:
     orchpy.connect(scheduler_address, objstore_address, address(IP_ADDRESS, new_worker_port()))
   for _ in range(num_workers):

lib/orchpy/orchpy/worker.py

@@ -10,7 +10,6 @@ class Worker(object):
 
   def __init__(self):
     self.functions = {}
-    self.connected = False
     self.handle = None
 
   def put_object(self, objref, value):
@@ -18,8 +17,8 @@ class Worker(object):
     if type(value) == np.ndarray:
       orchpy.lib.put_arrow(self.handle, objref, value)
     else:
-      object_capsule = serialization.serialize(value)
-      orchpy.lib.put_object(self.handle, objref, object_capsule)
+      object_capsule, contained_objrefs = serialization.serialize(self.handle, value) # contained_objrefs is a list of the objrefs contained in object_capsule
+      orchpy.lib.put_object(self.handle, objref, object_capsule, contained_objrefs)
 
   def get_object(self, objref):
     """
@@ -32,7 +31,7 @@ class Worker(object):
       return orchpy.lib.get_arrow(self.handle, objref)
     else:
       object_capsule = orchpy.lib.get_object(self.handle, objref)
-      return serialization.deserialize(object_capsule)
+      return serialization.deserialize(self.handle, object_capsule)
 
   def alias_objrefs(self, alias_objref, target_objref):
     """Make `alias_objref` refer to the same object that `target_objref` refers to."""
@@ -45,13 +44,16 @@ 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 = serialization.serialize_call(func_name, args)
+    call_capsule = serialization.serialize_call(self.handle, func_name, args)
     objrefs = orchpy.lib.remote_call(self.handle, call_capsule)
     return objrefs
 
 # We make `global_worker` a global variable so that there is one worker per worker process.
 global_worker = Worker()
 
+def scheduler_info(worker=global_worker):
+  return orchpy.lib.scheduler_info(worker.handle);
+
 def register_module(module, recursive=False, worker=global_worker):
   print "registering functions in module {}.".format(module.__name__)
   for name in dir(module):
@@ -63,10 +65,12 @@ def register_module(module, recursive=False, worker=global_worker):
     #   register_module(val, recursive, worker)
 
 def connect(scheduler_addr, objstore_addr, worker_addr, worker=global_worker):
-  if worker.connected:
-    del worker.handle # TODO(rkn): Make sure this actually deallocates (need a destructor for the capsule)
+  if hasattr(worker, "handle"):
+    del worker.handle
   worker.handle = orchpy.lib.create_worker(scheduler_addr, objstore_addr, worker_addr)
-  worker.connected = True
+
+def disconnect(worker=global_worker):
+  orchpy.lib.disconnect(worker.handle)
 
 def pull(objref, worker=global_worker):
   orchpy.lib.request_object(worker.handle, objref)
@@ -78,16 +82,18 @@ def push(value, worker=global_worker):
   return objref
 
 def main_loop(worker=global_worker):
-  if not worker.connected:
+  if not orchpy.lib.connected(worker.handle):
     raise Exception("Worker is attempting to enter main_loop but has not been connected yet.")
   orchpy.lib.start_worker_service(worker.handle)
-  while True:
-    call = orchpy.lib.wait_for_next_task(worker.handle)
-    func_name, args, return_objrefs = serialization.deserialize_call(call)
+  def process_call(call): # wrapping these calls in a function should cause the local variables to go out of scope more quickly, which is useful for inspecting reference counts
+    func_name, args, return_objrefs = serialization.deserialize_call(worker.handle, 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
     orchpy.lib.notify_task_completed(worker.handle) # notify the scheduler that the task has completed
+  while True:
+    call = orchpy.lib.wait_for_next_task(worker.handle)
+    process_call(call)
 
 def distributed(arg_types, return_types, worker=global_worker):
   def distributed_decorator(func):