mirror of
https://github.com/wassname/ray.git
synced 2026-06-29 17:04:56 +08:00
Change Python's ObjectID to ObjectRef (#9353)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user