Change Python's ObjectID to ObjectRef (#9353)

This commit is contained in:
Hao Chen
2020-07-10 17:49:04 +08:00
committed by GitHub
parent 6311e5a947
commit d49dadf891
91 changed files with 959 additions and 907 deletions
@@ -6,21 +6,22 @@ BLOCK_SIZE = 10
class DistArray:
def __init__(self, shape, objectids=None):
def __init__(self, shape, object_refs=None):
self.shape = shape
self.ndim = len(shape)
self.num_blocks = [
int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape
]
if objectids is not None:
self.objectids = objectids
if object_refs is not None:
self.object_refs = object_refs
else:
self.objectids = np.empty(self.num_blocks, dtype=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)))
self.object_refs = np.empty(self.num_blocks, dtype=object)
if self.num_blocks != list(self.object_refs.shape):
raise Exception(
"The fields `num_blocks` and `object_refs` are "
"inconsistent, `num_blocks` is {} and `object_refs` "
"has shape {}".format(self.num_blocks,
list(self.object_refs.shape)))
@staticmethod
def compute_block_lower(index, shape):
@@ -52,14 +53,14 @@ class DistArray:
return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape]
def assemble(self):
"""Assemble an array from a distributed array of object IDs."""
first_block = ray.get(self.objectids[(0, ) * self.ndim])
"""Assemble an array from a distributed array of object refs."""
first_block = ray.get(self.object_refs[(0, ) * self.ndim])
dtype = first_block.dtype
result = np.zeros(self.shape, dtype=dtype)
for index in np.ndindex(*self.num_blocks):
lower = DistArray.compute_block_lower(index, self.shape)
upper = DistArray.compute_block_upper(index, self.shape)
value = ray.get(self.objectids[index])
value = ray.get(self.object_refs[index])
result[tuple(slice(l, u) for (l, u) in zip(lower, upper))] = value
return result
@@ -83,7 +84,7 @@ def numpy_to_dist(a):
lower = DistArray.compute_block_lower(index, a.shape)
upper = DistArray.compute_block_upper(index, a.shape)
idx = tuple(slice(l, u) for (l, u) in zip(lower, upper))
result.objectids[index] = ray.put(a[idx])
result.object_refs[index] = ray.put(a[idx])
return result
@@ -91,7 +92,7 @@ def numpy_to_dist(a):
def zeros(shape, dtype_name="float"):
result = DistArray(shape)
for index in np.ndindex(*result.num_blocks):
result.objectids[index] = ra.zeros.remote(
result.object_refs[index] = ra.zeros.remote(
DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
return result
@@ -100,7 +101,7 @@ def zeros(shape, dtype_name="float"):
def ones(shape, dtype_name="float"):
result = DistArray(shape)
for index in np.ndindex(*result.num_blocks):
result.objectids[index] = ra.ones.remote(
result.object_refs[index] = ra.ones.remote(
DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
return result
@@ -111,7 +112,7 @@ def copy(a):
for index in np.ndindex(*result.num_blocks):
# We don't need to actually copy the objects because remote objects are
# immutable.
result.objectids[index] = a.objectids[index]
result.object_refs[index] = a.object_refs[index]
return result
@@ -123,10 +124,10 @@ def eye(dim1, dim2=-1, dtype_name="float"):
for (i, j) in np.ndindex(*result.num_blocks):
block_shape = DistArray.compute_block_shape([i, j], shape)
if i == j:
result.objectids[i, j] = ra.eye.remote(
result.object_refs[i, j] = ra.eye.remote(
block_shape[0], block_shape[1], dtype_name=dtype_name)
else:
result.objectids[i, j] = ra.zeros.remote(
result.object_refs[i, j] = ra.zeros.remote(
block_shape, dtype_name=dtype_name)
return result
@@ -139,11 +140,12 @@ def triu(a):
result = DistArray(a.shape)
for (i, j) in np.ndindex(*result.num_blocks):
if i < j:
result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.copy.remote(a.object_refs[i, j])
elif i == j:
result.objectids[i, j] = ra.triu.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.triu.remote(a.object_refs[i, j])
else:
result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.zeros_like.remote(
a.object_refs[i, j])
return result
@@ -155,11 +157,12 @@ def tril(a):
result = DistArray(a.shape)
for (i, j) in np.ndindex(*result.num_blocks):
if i > j:
result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.copy.remote(a.object_refs[i, j])
elif i == j:
result.objectids[i, j] = ra.tril.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.tril.remote(a.object_refs[i, j])
else:
result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
result.object_refs[i, j] = ra.zeros_like.remote(
a.object_refs[i, j])
return result
@@ -191,8 +194,8 @@ def dot(a, b):
shape = [a.shape[0], b.shape[1]]
result = DistArray(shape)
for (i, j) in np.ndindex(*result.num_blocks):
args = list(a.objectids[i, :]) + list(b.objectids[:, j])
result.objectids[i, j] = blockwise_dot.remote(*args)
args = list(a.object_refs[i, :]) + list(b.object_refs[:, j])
result.object_refs[i, j] = blockwise_dot.remote(*args)
return result
@@ -203,9 +206,9 @@ def subblocks(a, *ranges):
the `a`. The result and `a` will have the same number of dimensions. For
example,
subblocks(a, [0, 1], [2, 4])
will produce a DistArray whose objectids are
[[a.objectids[0, 2], a.objectids[0, 4]],
[a.objectids[1, 2], a.objectids[1, 4]]]
will produce a DistArray whose object_refs are
[[a.object_refs[0, 2], a.object_refs[0, 4]],
[a.object_refs[1, 2], a.object_refs[1, 4]]]
We allow the user to pass in an empty list [] to indicate the full range.
"""
ranges = list(ranges)
@@ -236,7 +239,7 @@ def subblocks(a, *ranges):
for i in range(a.ndim)]
result = DistArray(shape)
for index in np.ndindex(*result.num_blocks):
result.objectids[index] = a.objectids[tuple(
result.object_refs[index] = a.object_refs[tuple(
ranges[i][index[i]] for i in range(a.ndim))]
return result
@@ -250,7 +253,7 @@ def transpose(a):
result = DistArray([a.shape[1], a.shape[0]])
for i in range(result.num_blocks[0]):
for j in range(result.num_blocks[1]):
result.objectids[i, j] = ra.transpose.remote(a.objectids[j, i])
result.object_refs[i, j] = ra.transpose.remote(a.object_refs[j, i])
return result
@@ -263,8 +266,8 @@ def add(x1, x2):
x1.shape, x2.shape))
result = DistArray(x1.shape)
for index in np.ndindex(*result.num_blocks):
result.objectids[index] = ra.add.remote(x1.objectids[index],
x2.objectids[index])
result.object_refs[index] = ra.add.remote(x1.object_refs[index],
x2.object_refs[index])
return result
@@ -277,6 +280,6 @@ def subtract(x1, x2):
.format(x1.shape, x2.shape))
result = DistArray(x1.shape)
for index in np.ndindex(*result.num_blocks):
result.objectids[index] = ra.subtract.remote(x1.objectids[index],
x2.objectids[index])
result.object_refs[index] = ra.subtract.remote(x1.object_refs[index],
x2.object_refs[index])
return result
@@ -35,7 +35,7 @@ def tsqr(a):
q_tree = np.empty((num_blocks, K), dtype=object)
current_rs = []
for i in range(num_blocks):
block = a.objectids[i, 0]
block = a.object_refs[i, 0]
q, r = ra.linalg.qr.remote(block)
q_tree[i, 0] = q
current_rs.append(r)
@@ -57,8 +57,8 @@ def tsqr(a):
else:
q_shape = [a.shape[0], a.shape[0]]
q_num_blocks = core.DistArray.compute_num_blocks(q_shape)
q_objectids = np.empty(q_num_blocks, dtype=object)
q_result = core.DistArray(q_shape, q_objectids)
q_object_refs = np.empty(q_num_blocks, dtype=object)
q_result = core.DistArray(q_shape, q_object_refs)
# reconstruct output
for i in range(num_blocks):
@@ -75,7 +75,7 @@ def tsqr(a):
q_block_current = ra.dot.remote(
q_block_current,
ra.subarray.remote(q_tree[ith_index, j], lower, upper))
q_result.objectids[i] = q_block_current
q_result.object_refs[i] = q_block_current
r = current_rs[0]
return q_result, ray.get(r)
@@ -142,7 +142,7 @@ def tsqr_hr(a):
q, r_temp = tsqr.remote(a)
y, u, s = modified_lu.remote(q)
y_blocked = ray.get(y)
t, y_top = tsqr_hr_helper1.remote(u, s, y_blocked.objectids[0, 0],
t, y_top = tsqr_hr_helper1.remote(u, s, y_blocked.object_refs[0, 0],
a.shape[1])
r = tsqr_hr_helper2.remote(s, r_temp)
return ray.get(y), ray.get(t), ray.get(y_top), ray.get(r)
@@ -167,9 +167,9 @@ def qr(a):
k = min(m, n)
# we will store our scratch work in a_work
a_work = core.DistArray(a.shape, np.copy(a.objectids))
a_work = core.DistArray(a.shape, np.copy(a.object_refs))
result_dtype = np.linalg.qr(ray.get(a.objectids[0, 0]))[0].dtype.name
result_dtype = np.linalg.qr(ray.get(a.object_refs[0, 0]))[0].dtype.name
# TODO(rkn): It would be preferable not to get this right after creating
# it.
r_res = ray.get(core.zeros.remote([k, n], result_dtype))
@@ -188,28 +188,29 @@ def qr(a):
y_val = ray.get(y)
for j in range(i, a.num_blocks[0]):
y_res.objectids[j, i] = y_val.objectids[j - i, 0]
y_res.object_refs[j, i] = y_val.object_refs[j - i, 0]
if a.shape[0] > a.shape[1]:
# in this case, R needs to be square
R_shape = ray.get(ra.shape.remote(R))
eye_temp = ra.eye.remote(
R_shape[1], R_shape[0], dtype_name=result_dtype)
r_res.objectids[i, i] = ra.dot.remote(eye_temp, R)
r_res.object_refs[i, i] = ra.dot.remote(eye_temp, R)
else:
r_res.objectids[i, i] = R
r_res.object_refs[i, i] = R
Ts.append(core.numpy_to_dist.remote(t))
for c in range(i + 1, a.num_blocks[1]):
W_rcs = []
for r in range(i, a.num_blocks[0]):
y_ri = y_val.objectids[r - i, 0]
W_rcs.append(qr_helper2.remote(y_ri, a_work.objectids[r, c]))
y_ri = y_val.object_refs[r - i, 0]
W_rcs.append(qr_helper2.remote(y_ri, a_work.object_refs[r, c]))
W_c = ra.sum_list.remote(*W_rcs)
for r in range(i, a.num_blocks[0]):
y_ri = y_val.objectids[r - i, 0]
A_rc = qr_helper1.remote(a_work.objectids[r, c], y_ri, t, W_c)
a_work.objectids[r, c] = A_rc
r_res.objectids[i, c] = a_work.objectids[i, c]
y_ri = y_val.object_refs[r - i, 0]
A_rc = qr_helper1.remote(a_work.object_refs[r, c], y_ri, t,
W_c)
a_work.object_refs[r, c] = A_rc
r_res.object_refs[i, c] = a_work.object_refs[i, c]
# construct q_res from Ys and Ts
q = core.eye.remote(m, k, dtype_name=result_dtype)
@@ -8,9 +8,9 @@ from .core import DistArray
@ray.remote
def normal(shape):
num_blocks = DistArray.compute_num_blocks(shape)
objectids = np.empty(num_blocks, dtype=object)
object_refs = np.empty(num_blocks, dtype=object)
for index in np.ndindex(*num_blocks):
objectids[index] = ra.random.normal.remote(
object_refs[index] = ra.random.normal.remote(
DistArray.compute_block_shape(index, shape))
result = DistArray(shape, objectids)
result = DistArray(shape, object_refs)
return result