Changed how ray treats deserialization of custom classes (#333)

lib/python/ray/array/distributed/core.py

@@ -9,7 +9,7 @@ __all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy",
 BLOCK_SIZE = 10
 
 class DistArray(object):
-  def construct(self, shape, objectids=None):
+  def __init__(self, shape, objectids=None):
     self.shape = shape
     self.ndim = len(shape)
     self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape]
@@ -17,17 +17,14 @@ class DistArray(object):
     if self.num_blocks != list(self.objectids.shape):
       raise Exception("The fields `num_blocks` and `objectids` are inconsistent, `num_blocks` is {} and `objectids` has shape {}".format(self.num_blocks, list(self.objectids.shape)))
 
-  def deserialize(self, primitives):
+  @staticmethod
+  def deserialize(primitives):
     (shape, objectids) = primitives
-    self.construct(shape, objectids)
+    return DistArray(shape, objectids)
 
   def serialize(self):
     return (self.shape, self.objectids)
 
-  def __init__(self, shape=None):
-    if shape is not None:
-      self.construct(shape)
-
   @staticmethod
   def compute_block_lower(index, shape):
     if len(index) != len(shape):

lib/python/ray/array/distributed/linalg.py

@@ -46,8 +46,6 @@ def tsqr(a):
     current_rs = new_rs
   assert len(current_rs) == 1, "len(current_rs) = " + str(len(current_rs))
 
-  q_result = DistArray()
-
   # handle the special case in which the whole DistArray "a" fits in one block
   # and has fewer rows than columns, this is a bit ugly so think about how to
   # remove it
@@ -56,9 +54,8 @@ def tsqr(a):
   else:
     q_shape = [a.shape[0], a.shape[0]]
   q_num_blocks = DistArray.compute_num_blocks(q_shape)
-  q_result = DistArray()
   q_objectids = np.empty(q_num_blocks, dtype=object)
-  q_result.construct(q_shape, q_objectids)
+  q_result = DistArray(q_shape, q_objectids)
 
   # reconstruct output
   for i in range(num_blocks):
@@ -145,8 +142,7 @@ def qr(a):
   k = min(m, n)
 
   # we will store our scratch work in a_work
-  a_work = DistArray()
-  a_work.construct(a.shape, np.copy(a.objectids))
+  a_work = DistArray(a.shape, np.copy(a.objectids))
 
   result_dtype = np.linalg.qr(ray.get(a.objectids[0, 0]))[0].dtype.name
   r_res = ray.get(zeros.remote([k, n], result_dtype)) # TODO(rkn): It would be preferable not to get this right after creating it.

lib/python/ray/array/distributed/random.py

@@ -12,6 +12,5 @@ def normal(shape):
   objectids = np.empty(num_blocks, dtype=object)
   for index in np.ndindex(*num_blocks):
     objectids[index] = ra.random.normal.remote(DistArray.compute_block_shape(index, shape))
-  result = DistArray()
-  result.construct(shape, objectids)
+  result = DistArray(shape, objectids)
   return result