implement reference counting and much more (#43)

test/arrays_test.py

@@ -2,6 +2,7 @@ import unittest
 import orchpy
 import orchpy.serialization as serialization
 import orchpy.services as services
+import orchpy.worker as worker
 import numpy as np
 import time
 import subprocess32 as subprocess
@@ -53,13 +54,18 @@ class ArraysSingleTest(unittest.TestCase):
 class ArraysDistTest(unittest.TestCase):
 
   def testSerialization(self):
+    w = worker.Worker()
+    services.start_cluster(driver_worker=w)
+
     x = dist.DistArray()
-    x.construct([2, 3, 4], np.array([[[orchpy.lib.ObjRef(0)]]]))
-    capsule = serialization.serialize(x)
-    y = serialization.deserialize(capsule)
+    x.construct([2, 3, 4], np.array([[[orchpy.push(0, w)]]]))
+    capsule, _ = serialization.serialize(w.handle, x) # TODO(rkn): THIS REQUIRES A WORKER_HANDLE
+    y = serialization.deserialize(w.handle, capsule) # TODO(rkn): THIS REQUIRES A WORKER_HANDLE
     self.assertEqual(x.shape, y.shape)
     self.assertEqual(x.objrefs[0, 0, 0].val, y.objrefs[0, 0, 0].val)
 
+    services.cleanup()
+
   def testAssemble(self):
     test_dir = os.path.dirname(os.path.abspath(__file__))
     test_path = os.path.join(test_dir, "testrecv.py")
@@ -76,33 +82,46 @@ class ArraysDistTest(unittest.TestCase):
   def testMethods(self):
     test_dir = os.path.dirname(os.path.abspath(__file__))
     test_path = os.path.join(test_dir, "testrecv.py")
-    services.start_cluster(num_workers=3, worker_path=test_path)
+    services.start_cluster(num_workers=4, worker_path=test_path)
 
     x = dist.zeros([9, 25, 51], "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.zeros([9, 25, 51])))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.zeros([9, 25, 51])))
 
     x = dist.ones([11, 25, 49], "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.ones([11, 25, 49])))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.ones([11, 25, 49])))
 
     x = dist.random.normal([11, 25, 49])
     y = dist.copy(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == orchpy.pull(dist.assemble(y))))
+    z = dist.assemble(x)
+    w = dist.assemble(y)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == orchpy.pull(w)))
 
     x = dist.eye(25, "float")
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == np.eye(25)))
+    y = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(y) == np.eye(25)))
 
     x = dist.random.normal([25, 49])
     y = dist.triu(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(y)) == np.triu(orchpy.pull(dist.assemble(x)))))
+    z = dist.assemble(y)
+    w = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == np.triu(orchpy.pull(w))))
 
     x = dist.random.normal([25, 49])
     y = dist.tril(x)
-    self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(y)) == np.tril(orchpy.pull(dist.assemble(x)))))
+    z = dist.assemble(y)
+    w = dist.assemble(x)
+    self.assertTrue(np.alltrue(orchpy.pull(z) == np.tril(orchpy.pull(w))))
 
     x = dist.random.normal([25, 49])
     y = dist.random.normal([49, 18])
     z = dist.dot(x, y)
-    self.assertTrue(np.allclose(orchpy.pull(dist.assemble(z)), np.dot(orchpy.pull(dist.assemble(x)), orchpy.pull(dist.assemble(y)))))
+    w = dist.assemble(z)
+    u = dist.assemble(x)
+    v = dist.assemble(y)
+    np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v)))
+    self.assertTrue(np.allclose(orchpy.pull(w), np.dot(orchpy.pull(u), orchpy.pull(v))))
 
     services.cleanup()
 

test/runtest.py

@@ -13,58 +13,58 @@ from google.protobuf.text_format import *
 import orchestra_pb2
 import types_pb2
 
+import test_functions
+import arrays.single as single
+import arrays.dist as dist
+
 class SerializationTest(unittest.TestCase):
 
-  def roundTripTest(self, data):
-    serialized = serialization.serialize(data)
-    result = serialization.deserialize(serialized)
+  def roundTripTest(self, worker, data):
+    serialized, _ = serialization.serialize(worker.handle, data)
+    result = serialization.deserialize(worker.handle, serialized)
     self.assertEqual(data, result)
 
-  def numpyTypeTest(self, typ):
+  def numpyTypeTest(self, worker, typ):
     a = np.random.randint(0, 10, size=(100, 100)).astype(typ)
-    b = serialization.serialize(a)
-    c = serialization.deserialize(b)
+    b, _ = serialization.serialize(worker.handle, a)
+    c = serialization.deserialize(worker.handle, b)
     self.assertTrue((a == c).all())
 
   def testSerialize(self):
-    self.roundTripTest([1, "hello", 3.0])
-    self.roundTripTest(42)
-    self.roundTripTest("hello world")
-    self.roundTripTest(42.0)
-    self.roundTripTest((1.0, "hi"))
+    w = worker.Worker()
+    services.start_cluster(driver_worker=w)
 
-    self.roundTripTest({"hello" : "world", 1: 42, 1.0: 45})
-    self.roundTripTest({})
+    self.roundTripTest(w, [1, "hello", 3.0])
+    self.roundTripTest(w, 42)
+    self.roundTripTest(w, "hello world")
+    self.roundTripTest(w, 42.0)
+    self.roundTripTest(w, (1.0, "hi"))
+
+    self.roundTripTest(w, {"hello" : "world", 1: 42, 1.0: 45})
+    self.roundTripTest(w, {})
 
     a = np.zeros((100, 100))
-    res = serialization.serialize(a)
-    b = serialization.deserialize(res)
+    res, _ = serialization.serialize(w.handle, a)
+    b = serialization.deserialize(w.handle, res)
     self.assertTrue((a == b).all())
 
-    self.numpyTypeTest('int8')
-    self.numpyTypeTest('uint8')
+    self.numpyTypeTest(w, 'int8')
+    self.numpyTypeTest(w, 'uint8')
     # self.numpyTypeTest('int16') # TODO(pcm): implement this
     # self.numpyTypeTest('int32') # TODO(pcm): implement this
-    self.numpyTypeTest('float32')
-    self.numpyTypeTest('float64')
+    self.numpyTypeTest(w, 'float32')
+    self.numpyTypeTest(w, 'float64')
 
-    a = np.array([[orchpy.lib.ObjRef(0), orchpy.lib.ObjRef(1)], [orchpy.lib.ObjRef(41), orchpy.lib.ObjRef(42)]])
-    capsule = serialization.serialize(a)
-    result = serialization.deserialize(capsule)
+    ref0 = orchpy.push(0, w)
+    ref1 = orchpy.push(0, w)
+    ref2 = orchpy.push(0, w)
+    ref3 = orchpy.push(0, w)
+    a = np.array([[ref0, ref1], [ref2, ref3]])
+    capsule, _ = serialization.serialize(w.handle, a)
+    result = serialization.deserialize(w.handle, capsule)
     self.assertTrue((a == result).all())
 
-class OrchPyLibTest(unittest.TestCase):
-
-    def testOrchPyLib(self):
-      w = worker.Worker()
-      services.start_cluster(driver_worker=w)
-
-      w.put_object(orchpy.lib.ObjRef(0), 'hello world')
-      result = w.get_object(orchpy.lib.ObjRef(0))
-
-      self.assertEqual(result, 'hello world')
-
-      services.cleanup()
+    services.cleanup()
 
 class ObjStoreTest(unittest.TestCase):
 
@@ -97,7 +97,7 @@ class SchedulerTest(unittest.TestCase):
     services.start_cluster(driver_worker=w, num_workers=1, worker_path=test_path)
 
     value_before = "test_string"
-    objref = w.remote_call("__main__.print_string", [value_before])
+    objref = w.remote_call("test_functions.print_string", [value_before])
 
     time.sleep(0.2)
 
@@ -148,12 +148,67 @@ class APITest(unittest.TestCase):
     test_path = os.path.join(test_dir, "testrecv.py")
     services.start_cluster(num_workers=3, worker_path=test_path, driver_worker=w)
 
-    objref = w.remote_call("__main__.test_alias_f", [])
+    objref = w.remote_call("test_functions.test_alias_f", [])
     self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))
-    objref = w.remote_call("__main__.test_alias_g", [])
+    objref = w.remote_call("test_functions.test_alias_g", [])
     self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))
-    objref = w.remote_call("__main__.test_alias_h", [])
+    objref = w.remote_call("test_functions.test_alias_h", [])
     self.assertTrue(np.alltrue(orchpy.pull(objref[0], w) == np.ones([3, 4, 5])))
 
+    services.cleanup()
+
+class ReferenceCountingTest(unittest.TestCase):
+
+  def testDeallocation(self):
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    test_path = os.path.join(test_dir, "testrecv.py")
+    services.start_cluster(num_workers=3, worker_path=test_path)
+
+    x = test_functions.test_alias_f()
+    orchpy.pull(x)
+    time.sleep(0.1)
+    objref_val = x.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val] == 1)
+
+    del x
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val] == -1) # -1 indicates deallocated
+
+    y = test_functions.test_alias_h()
+    orchpy.pull(y)
+    time.sleep(0.1)
+    objref_val = y.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 0, 0])
+
+    del y
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    z = dist.zeros([dist.BLOCK_SIZE, 2 * dist.BLOCK_SIZE], "float")
+    time.sleep(0.1)
+    objref_val = z.val
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 1, 1])
+
+    del z
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    x = single.zeros([10, 10], "float")
+    y = single.zeros([10, 10], "float")
+    z = single.dot(x, y)
+    objref_val = x.val
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [1, 1, 1])
+
+    del x
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, 1, 1])
+    del y
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, 1])
+    del z
+    time.sleep(0.1)
+    self.assertTrue(orchpy.scheduler_info()["reference_counts"][objref_val:(objref_val + 3)] == [-1, -1, -1])
+
+    services.cleanup()
+
 if __name__ == '__main__':
     unittest.main()

test/shell.py

@@ -5,6 +5,7 @@ import orchpy
 import orchpy.services as services
 import orchpy.worker as worker
 
+import test_functions
 import arrays.single as single
 import arrays.dist as dist
 
@@ -19,50 +20,9 @@ parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str,
 parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
 parser.add_argument("--worker-address", default="127.0.0.1:30001", type=str, help="the worker's address")
 
-@orchpy.distributed([], [np.ndarray])
-def test_alias_f():
-  return np.ones([3, 4, 5])
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_g():
-  return test_alias_f()
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_h():
-  return test_alias_g()
-
-@orchpy.distributed([str], [str])
-def print_string(string):
-  print "called print_string with", string
-  f = open("asdfasdf.txt", "w")
-  f.write("successfully called print_string with argument {}.".format(string))
-  return string
-
-@orchpy.distributed([int, int], [int, int])
-def handle_int(a, b):
-  return a + 1, b + 1
-
-def connect_to_scheduler(host, port):
-  channel = implementations.insecure_channel(host, port)
-  return orchestra_pb2.beta_create_Scheduler_stub(channel)
-
-def connect_to_objstore(host, port):
-  channel = implementations.insecure_channel(host, port)
-  return orchestra_pb2.beta_create_ObjStore_stub(channel)
-
 if __name__ == '__main__':
   args = parser.parse_args()
-  scheduler_ip_address, scheduler_port = args.scheduler_address.split(":")
-  scheduler_stub = connect_to_scheduler(scheduler_ip_address, int(scheduler_port))
-  objstore_ip_address, objstore_port = args.objstore_address.split(":")
-  objstore_stub = connect_to_objstore(objstore_ip_address, int(objstore_port))
   worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)
 
-  def scheduler_debug_info():
-    return scheduler_stub.SchedulerDebugInfo(orchestra_pb2.SchedulerDebugInfoRequest(), TIMEOUT_SECONDS)
-
-  def objstore_debug_info():
-    return objstore_stub.ObjStoreDebugInfo(orchestra_pb2.ObjStoreDebugInfoRequest(), TIMEOUT_SECONDS)
-
   import IPython
   IPython.embed()

test/test_functions.py

@@ -0,0 +1,32 @@
+import orchpy
+
+import numpy as np
+
+# Test simple functionality
+
+@orchpy.distributed([str], [str])
+def print_string(string):
+  print "called print_string with", string
+  f = open("asdfasdf.txt", "w")
+  f.write("successfully called print_string with argument {}.".format(string))
+  return string
+
+@orchpy.distributed([int, int], [int, int])
+def handle_int(a, b):
+  return a + 1, b + 1
+
+# Test aliasing
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_f():
+  return np.ones([3, 4, 5])
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_g():
+  return test_alias_f()
+
+@orchpy.distributed([], [np.ndarray])
+def test_alias_h():
+  return test_alias_g()
+
+# Test reference counting

test/testrecv.py

@@ -2,6 +2,7 @@ import sys
 import argparse
 import numpy as np
 
+import test_functions
 import arrays.single as single
 import arrays.dist as dist
 
@@ -14,33 +15,11 @@ parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str,
 parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
 parser.add_argument("--worker-address", default="127.0.0.1:40001", type=str, help="the worker's address")
 
-@orchpy.distributed([], [np.ndarray])
-def test_alias_f():
-  return np.ones([3, 4, 5])
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_g():
-  return test_alias_f()
-
-@orchpy.distributed([], [np.ndarray])
-def test_alias_h():
-  return test_alias_g()
-
-@orchpy.distributed([str], [str])
-def print_string(string):
-  print "called print_string with", string
-  f = open("asdfasdf.txt", "w")
-  f.write("successfully called print_string with argument {}.".format(string))
-  return string
-
-@orchpy.distributed([int, int], [int, int])
-def handle_int(a, b):
-  return a + 1, b + 1
-
 if __name__ == '__main__':
   args = parser.parse_args()
   worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)
 
+  orchpy.register_module(test_functions)
   orchpy.register_module(single)
   orchpy.register_module(single.random)
   orchpy.register_module(single.linalg)