improve serialization and add DistArray class

lib/orchpy/arrays/dist/__init__.py

@@ -0,0 +1 @@
+from core import DistArray, BLOCK_SIZE

lib/orchpy/arrays/dist/core.py

@@ -0,0 +1,60 @@
+from typing import List
+import numpy as np
+import arrays.single as single
+import orchpy as op
+
+BLOCK_SIZE = 10
+
+class DistArray(object):
+  def construct(self, shape, dtype, objrefs):
+    self.shape = shape
+    self.dtype = dtype
+    self.objrefs = objrefs
+    self.ndim = len(shape)
+    self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape]
+    if self.num_blocks != list(self.objrefs.shape):
+      raise Exception("The fields `num_blocks` and `objrefs` are inconsistent, `num_blocks` is {} and `objrefs` has shape {}".format(self.num_blocks, list(self.objrefs.shape)))
+
+  def deserialize(self, primitives):
+    (shape, dtype_name, objrefs) = primitives
+    self.construct(shape, np.dtype(dtype_name), objrefs)
+
+  def serialize(self):
+    return (self.shape, self.dtype.__name__, self.objrefs)
+
+  def __init__(self):
+    self.shape = None
+    self.dtype = None
+    self.objrefs = None
+
+  def compute_block_lower(self, index):
+    if len(index) != self.ndim:
+      raise Exception("The value `index` equals {}, but `ndim` is {}.".format(index, self.ndim))
+    return [elem * BLOCK_SIZE for elem in index]
+
+  def compute_block_upper(self, index):
+    if len(index) != self.ndim:
+      raise Exception("The value `index` equals {}, but `ndim` is {}.".format(index, self.ndim))
+    upper = []
+    for i in range(self.ndim):
+      upper.append(min((index[i] + 1) * BLOCK_SIZE, self.shape[i]))
+    return upper
+
+  def compute_block_shape(self, index):
+    lower = self.compute_block_lower(index)
+    upper = self.compute_block_upper(index)
+    return [u - l for (l, u) in zip(lower, upper)]
+
+  def assemble(self):
+    """Assemble an array on this node from a distributed array object reference."""
+    result = np.zeros(self.shape)
+    for index in np.ndindex(*self.num_blocks):
+      lower = self.compute_block_lower(index)
+      upper = self.compute_block_upper(index)
+      result[[slice(l, u) for (l, u) in zip(lower, upper)]] = op.pull(self.objrefs[index])
+    return result
+
+  def __getitem__(self, sliced):
+    # TODO(rkn): fix this, this is just a placeholder that should work but is inefficient
+    a = self.assemble()
+    return a[sliced]

lib/orchpy/arrays/single/__init__.py

@@ -1,2 +1,2 @@
 import random, linalg
-from core import zeros, eye, dot, vstack, hstack, subarray, copy, tril, triu
+from core import zeros, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu

lib/orchpy/arrays/single/core.py

@@ -6,6 +6,10 @@ import orchpy as op
 def zeros(shape):
   return np.zeros(shape)
 
+@op.distributed([List[int]], [np.ndarray])
+def ones(shape):
+  return np.ones(shape)
+
 @op.distributed([int], [np.ndarray])
 def eye(dim):
   return np.eye(dim)

lib/orchpy/orchpy/__init__.py

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

lib/orchpy/orchpy/serialization.py

@@ -0,0 +1,26 @@
+import importlib
+
+import orchpy
+
+def serialize(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 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

lib/orchpy/orchpy/worker.py

@@ -2,6 +2,7 @@ from types import ModuleType
 import typing
 
 import orchpy
+import serialization
 
 class Worker(object):
   """The methods in this class are considered unexposed to the user. The functions outside of this class are considered exposed."""
@@ -13,13 +14,13 @@ class Worker(object):
 
   def put_object(self, objref, value):
     """Put `value` in the local object store with objref `objref`. This assumes that the value for `objref` has not yet been placed in the local object store."""
-    object_capsule = orchpy.lib.serialize_object(value)
+    object_capsule = serialization.serialize(value)
     orchpy.lib.put_object(self.handle, objref, object_capsule)
 
   def get_object(self, objref):
     """Return the value from the local object store for objref `objref`. This will block until the value for `objref` has been written to the local object store."""
     object_capsule = orchpy.lib.get_object(self.handle, objref)
-    return orchpy.lib.deserialize_object(object_capsule)
+    return serialization.deserialize(object_capsule)
 
   def register_function(self, function):
     """Notify the scheduler that this worker can execute the function with name `func_name`. Store the function `function` locally."""
@@ -47,16 +48,16 @@ def register_module(module, recursive=False, worker=global_worker):
 
 def connect(scheduler_addr, objstore_addr, worker_addr, worker=global_worker):
   if worker.connected:
-    raise Exception("Worker called connect, but worker is already connected")
+    del worker.handle # TODO(rkn): Make sure this actually deallocates (need a destructor for the capsule)
   worker.handle = orchpy.lib.create_worker(scheduler_addr, objstore_addr, worker_addr)
   worker.connected = True
 
 def pull(objref, worker=global_worker):
   object_capsule = orchpy.lib.pull_object(worker.handle, objref)
-  return orchpy.lib.deserialize_object(object_capsule)
+  return serialization.deserialize(object_capsule)
 
 def push(value, worker=global_worker):
-  object_capsule = orchpy.lib.serialize_object(value)
+  object_capsule = serialization.serialize(value)
   return orchpy.lib.push_object(worker.handle, object_capsule)
 
 def main_loop(worker=global_worker):