diff --git a/lib/orchpy/arrays/dist/__init__.py b/lib/orchpy/arrays/dist/__init__.py index e1d2e30fd..3200379ef 100644 --- a/lib/orchpy/arrays/dist/__init__.py +++ b/lib/orchpy/arrays/dist/__init__.py @@ -1 +1,2 @@ -from core import DistArray, BLOCK_SIZE +import random +from core import * diff --git a/lib/orchpy/arrays/dist/core.py b/lib/orchpy/arrays/dist/core.py index 024faa2cd..08d379b31 100644 --- a/lib/orchpy/arrays/dist/core.py +++ b/lib/orchpy/arrays/dist/core.py @@ -3,12 +3,14 @@ import numpy as np import arrays.single as single import orchpy as op +__all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy", + "eye", "triu", "tril", "blockwise_dot", "dot", "block_column", "block_row"] + BLOCK_SIZE = 10 class DistArray(object): - def construct(self, shape, dtype, objrefs): + def construct(self, shape, 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] @@ -16,41 +18,49 @@ class DistArray(object): 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) + (shape, objrefs) = primitives + self.construct(shape, objrefs) def serialize(self): - return (self.shape, self.dtype.__name__, self.objrefs) + return (self.shape, 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)) + @staticmethod + def compute_block_lower(index, shape): + # TODO(rkn): Check that the entries of index are in the correct range. + # TODO(rkn): Check that len(index) == len(shape). 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)) + @staticmethod + def compute_block_upper(index, shape): + # TODO(rkn): Check that the entries of index are in the correct range. + # TODO(rkn): Check that len(index) == len(shape). upper = [] - for i in range(self.ndim): - upper.append(min((index[i] + 1) * BLOCK_SIZE, self.shape[i])) + for i in range(len(shape)): + upper.append(min((index[i] + 1) * BLOCK_SIZE, shape[i])) return upper - def compute_block_shape(self, index): - lower = self.compute_block_lower(index) - upper = self.compute_block_upper(index) + @staticmethod + def compute_block_shape(index, shape): + lower = DistArray.compute_block_lower(index, shape) + upper = DistArray.compute_block_upper(index, shape) return [u - l for (l, u) in zip(lower, upper)] + @staticmethod + def compute_num_blocks(shape): + return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape] + def assemble(self): """Assemble an array on this node from a distributed array object reference.""" - result = np.zeros(self.shape) + first_block = op.pull(self.objrefs[(0,) * self.ndim]) + dtype = first_block.dtype + result = np.zeros(self.shape, dtype=dtype) for index in np.ndindex(*self.num_blocks): - lower = self.compute_block_lower(index) - upper = self.compute_block_upper(index) + lower = DistArray.compute_block_lower(index, self.shape) + upper = DistArray.compute_block_upper(index, self.shape) result[[slice(l, u) for (l, u) in zip(lower, upper)]] = op.pull(self.objrefs[index]) return result @@ -58,3 +68,137 @@ class DistArray(object): # TODO(rkn): fix this, this is just a placeholder that should work but is inefficient a = self.assemble() return a[sliced] + +@op.distributed([DistArray], [np.ndarray]) +def assemble(a): + return a.assemble() + +@op.distributed([List[int], str], [DistArray]) +def zeros(shape, dtype_name): + num_blocks = DistArray.compute_num_blocks(shape) + objrefs = np.empty(num_blocks, dtype=object) + for index in np.ndindex(*num_blocks): + objrefs[index] = single.zeros(DistArray.compute_block_shape(index, shape), dtype_name) + result = DistArray() + result.construct(shape, objrefs) + return result + +@op.distributed([List[int], str], [DistArray]) +def ones(shape, dtype_name): + num_blocks = DistArray.compute_num_blocks(shape) + objrefs = np.empty(num_blocks, dtype=object) + for index in np.ndindex(*num_blocks): + objrefs[index] = single.ones(DistArray.compute_block_shape(index, shape), dtype_name) + result = DistArray() + result.construct(shape, objrefs) + return result + +@op.distributed([DistArray], [DistArray]) +def copy(a): + num_blocks = DistArray.compute_num_blocks(a.shape) + objrefs = np.empty(num_blocks, dtype=object) + for index in np.ndindex(*num_blocks): + objrefs[index] = single.copy(a.objrefs[index]) + result = DistArray() + result.construct(a.shape, objrefs) + return result + +@op.distributed([int, str], [DistArray]) +def eye(dim, dtype_name): + shape = [dim, dim] + num_blocks = DistArray.compute_num_blocks(shape) + objrefs = np.empty(num_blocks, dtype=object) + for (i, j) in np.ndindex(*num_blocks): + if i == j: + objrefs[i, j] = single.eye(DistArray.compute_block_shape([i, j], shape)[0], dtype_name) + else: + objrefs[i, j] = single.zeros(DistArray.compute_block_shape([i, j], shape), dtype_name) + result = DistArray() + result.construct(shape, objrefs) + return result + +@op.distributed([DistArray], [DistArray]) +def triu(a): + if a.ndim != 2: + raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim)) + objrefs = np.empty(a.num_blocks, dtype=object) + for i in range(a.num_blocks[0]): + for j in range(a.num_blocks[1]): + if i < j: + objrefs[i, j] = single.copy(a.objrefs[i, j]) + elif i == j: + objrefs[i, j] = single.triu(a.objrefs[i, j]) + else: + objrefs[i, j] = single.zeros_like(a.objrefs[i, j]) + result = DistArray() + result.construct(a.shape, objrefs) + return result + +@op.distributed([DistArray], [DistArray]) +def tril(a): + if a.ndim != 2: + raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim)) + objrefs = np.empty(a.num_blocks, dtype=object) + for i in range(a.num_blocks[0]): + for j in range(a.num_blocks[1]): + if i > j: + objrefs[i, j] = single.copy(a.objrefs[i, j]) + elif i == j: + objrefs[i, j] = single.tril(a.objrefs[i, j]) + else: + objrefs[i, j] = single.zeros_like(a.objrefs[i, j]) + result = DistArray() + result.construct(a.shape, objrefs) + return result + +@op.distributed([np.ndarray, None], [np.ndarray]) +def blockwise_dot(*matrices): + n = len(matrices) + if n % 2 != 0: + raise Exception("blockwise_dot expects an even number of arguments, but len(matrices) is {}.".format(n)) + shape = (matrices[0].shape[0], matrices[n / 2].shape[1]) + result = np.zeros(shape) + for i in range(n / 2): + result += np.dot(matrices[i], matrices[n / 2 + i]) + return result + +@op.distributed([DistArray, DistArray], [DistArray]) +def dot(a, b): + if a.ndim != 2: + raise Exception("dot expects its arguments to be 2-dimensional, but a.ndim = {}.".format(a.ndim)) + if b.ndim != 2: + raise Exception("dot expects its arguments to be 2-dimensional, but b.ndim = {}.".format(b.ndim)) + if a.shape[1] != b.shape[0]: + raise Exception("dot expects a.shape[1] to equal b.shape[0], but a.shape = {} and b.shape = {}.".format(a.shape, b.shape)) + shape = [a.shape[0], b.shape[1]] + num_blocks = DistArray.compute_num_blocks(shape) + objrefs = np.empty(num_blocks, dtype=object) + for i in range(num_blocks[0]): + for j in range(num_blocks[1]): + args = list(a.objrefs[i, :]) + list(b.objrefs[:, j]) + objrefs[i, j] = blockwise_dot(*args) + result = DistArray() + result.construct(shape, objrefs) + return result + +# This is not in numpy, should we expose this? +@op.distributed([DistArray], [DistArray]) +def block_column(a, col): + if a.ndim != 2: + raise Exception("block_column expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape)) + top_block_shape = DistArray.compute_block_shape([0, col]) + shape = [a.shape[0], top_block_shape[1]] + result = DistArray() + result.construct(shape, a.objrefs[:, col]) + return result + +# This is not in numpy, should we expose this? +@op.distributed([DistArray], [DistArray]) +def block_row(a, row): + if a.ndim != 2: + raise Exception("block_row expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape)) + left_block_shape = DistArray.compute_block_shape([row, 0]) + shape = [left_block_shape[0], a.shape[1]] + result = DistArray() + result.construct(shape, a.objrefs[row, :]) + return result diff --git a/lib/orchpy/arrays/dist/random.py b/lib/orchpy/arrays/dist/random.py new file mode 100644 index 000000000..9e0464884 --- /dev/null +++ b/lib/orchpy/arrays/dist/random.py @@ -0,0 +1,17 @@ +from typing import List + +import numpy as np +import arrays.single as single +import orchpy as op + +from core import * + +@op.distributed([List[int]], [DistArray]) +def normal(shape): + num_blocks = DistArray.compute_num_blocks(shape) + objrefs = np.empty(num_blocks, dtype=object) + for index in np.ndindex(*num_blocks): + objrefs[index] = single.random.normal(DistArray.compute_block_shape(index, shape)) + result = DistArray() + result.construct(shape, objrefs) + return result diff --git a/lib/orchpy/arrays/single/__init__.py b/lib/orchpy/arrays/single/__init__.py index 2890639cb..967afbb62 100644 --- a/lib/orchpy/arrays/single/__init__.py +++ b/lib/orchpy/arrays/single/__init__.py @@ -1,2 +1,2 @@ import random, linalg -from core import zeros, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu +from core import zeros, zeros_like, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu diff --git a/lib/orchpy/arrays/single/core.py b/lib/orchpy/arrays/single/core.py index 914c51ec0..e634dd8fb 100644 --- a/lib/orchpy/arrays/single/core.py +++ b/lib/orchpy/arrays/single/core.py @@ -2,17 +2,21 @@ from typing import List import numpy as np import orchpy as op -@op.distributed([List[int]], [np.ndarray]) -def zeros(shape): - return np.zeros(shape) +@op.distributed([List[int], str], [np.ndarray]) +def zeros(shape, dtype_name): + return np.zeros(shape, dtype=np.dtype(dtype_name)) -@op.distributed([List[int]], [np.ndarray]) -def ones(shape): - return np.ones(shape) +@op.distributed([np.ndarray], [np.ndarray]) +def zeros_like(x): + return np.zeros_like(x) -@op.distributed([int], [np.ndarray]) -def eye(dim): - return np.eye(dim) +@op.distributed([List[int], str], [np.ndarray]) +def ones(shape, dtype_name): + return np.ones(shape, dtype=np.dtype(dtype_name)) + +@op.distributed([int, str], [np.ndarray]) +def eye(dim, dtype_name): + return np.eye(dim, dtype=np.dtype(dtype_name)) @op.distributed([np.ndarray, np.ndarray], [np.ndarray]) def dot(a, b): diff --git a/lib/orchpy/arrays/single/linalg.py b/lib/orchpy/arrays/single/linalg.py index 19d3de6f4..1c969247d 100644 --- a/lib/orchpy/arrays/single/linalg.py +++ b/lib/orchpy/arrays/single/linalg.py @@ -2,21 +2,95 @@ from typing import List import numpy as np import orchpy as op -# TODO(rkn): this should take the same optional "mode" argument as np.linalg.qr, except that the different options sometimes have different numbers of return values, which could be a problem +__all__ = ["matrix_power", "solve", "tensorsolve", "tensorinv", "inv", + "cholesky", "eigvals", "eigvalsh", "pinv", "slogdet", "det", + "svd", "eig", "eigh", "lstsq", "norm", "qr", "cond", "matrix_rank", + "LinAlgError", "multi_dot"] + +@op.distributed([np.ndarray, int], [np.ndarray]) +def matrix_power(M, n): + return np.linalg.matrix_power(M, n) + +@op.distributed([np.ndarray, np.ndarray], [np.ndarray]) +def solve(a, b): + return np.linalg.solve(a, b) + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray]) +def tensorsolve(a): + raise NotImplementedError + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray]) +def tensorinv(a): + raise NotImplementedError + +@op.distributed([np.ndarray], [np.ndarray]) +def inv(a): + return np.linalg.inv(a) + +@op.distributed([np.ndarray], [np.ndarray]) +def cholesky(a): + return np.linalg.cholesky(a) + +@op.distributed([np.ndarray], [np.ndarray]) +def eigvals(a): + return np.linalg.eigvals(a) + +@op.distributed([np.ndarray], [np.ndarray]) +def eigvalsh(a): + raise NotImplementedError + +@op.distributed([np.ndarray], [np.ndarray]) +def pinv(a): + return np.linalg.pinv(a) + +@op.distributed([np.ndarray], [int]) +def slogdet(a): + raise NotImplementedError + +@op.distributed([np.ndarray], [float]) +def det(a): + return np.linalg.det(a) + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray]) +def svd(a): + return np.linalg.svd(a) + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray]) +def eig(a): + return np.linalg.eig(a) + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray]) +def eigh(a): + return np.linalg.eigh(a) + +@op.distributed([np.ndarray], [np.ndarray, np.ndarray, int, np.ndarray]) +def lstsq(a, b): + return np.linalg.lstsq(a) + +@op.distributed([np.ndarray], [float]) +def norm(x): + return np.linalg.norm(x) + @op.distributed([np.ndarray], [np.ndarray, np.ndarray]) def qr(a): - """ - Suppose (n, m) = a.shape - If n >= m: - q.shape == (n, m) - r.shape == (m, m) - If n < m: - q.shape == (n, n) - r.shape == (n, m) - """ return np.linalg.qr(a) -#@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray]) +@op.distributed([np.ndarray], [float]) +def cond(x): + return np.linalg.cond(x) + +@op.distributed([np.ndarray], [int]) +def matrix_rank(M): + return np.linalg.matrix_rank(M) + +@op.distributed([np.ndarray, None], [np.ndarray]) +def multi_dot(a): + raise NotImplementedError + + + +# This isn't in numpy, should we expose it? +@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray]) def modified_lu(q): """ Algorithm 5 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf diff --git a/lib/orchpy/orchpy/worker.py b/lib/orchpy/orchpy/worker.py index 6cf787221..c7dc77680 100644 --- a/lib/orchpy/orchpy/worker.py +++ b/lib/orchpy/orchpy/worker.py @@ -78,8 +78,8 @@ def distributed(arg_types, return_types, worker=global_worker): """This is what gets executed remotely on a worker after a distributed function is scheduled by the scheduler.""" print "Calling function {} with arguments {}".format(func.__name__, arguments) result = func(*arguments) - if len(return_types) != 1 and len(result) != len(return_types): - raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(func.__name__, len(return_types), return_types, func.__name__, len(result))) + check_return_values(func_call, result) # throws an exception if result is invalid + print "Finished executing function {} with arguments {}".format(func.__name__, arguments) return result def func_call(*args): """This is what gets run immediately when a worker calls a distributed function.""" @@ -94,6 +94,18 @@ def distributed(arg_types, return_types, worker=global_worker): return func_call return distributed_decorator +# helper method, this should not be called by the user +def check_return_values(function, result): + if len(function.return_types) == 1: + if not isinstance(result, function.return_types[0]): + raise Exception("The @distributed decorator for function {} expects one return value with type {}, but {} returned a {}.".format(function.__name__, function.return_types[0], function.__name__, type(result))) + else: + if len(result) != len(function.return_types): + raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(function.__name__, len(function.return_types), function.return_types, function.__name__, len(result))) + for i in range(len(result)): + if not isinstance(result[i], function.return_types[i]): + raise Exception("The {}th return value for function {} has type {}, but the @distributed decorator expected a return value of type {}.".format(i, function.__name__, type(result[i]), function.return_types[i])) + # helper method, this should not be called by the user def check_arguments(function, args): # check the number of args @@ -107,7 +119,7 @@ def check_arguments(function, args): expected_type = function.arg_types[i] elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None: expected_type = function.arg_types[-1] - elif function.arg_types[-1] is None and len(function.arg_types > 1): + elif function.arg_types[-1] is None and len(function.arg_types) > 1: expected_type = function.arg_types[-2] else: assert False, "This code should be unreachable." @@ -136,7 +148,7 @@ def get_arguments_for_execution(function, args, worker=global_worker): expected_type = function.arg_types[i] elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None: expected_type = function.arg_types[-1] - elif function.arg_types[-1] is None and len(function.arg_types > 1): + elif function.arg_types[-1] is None and len(function.arg_types) > 1: expected_type = function.arg_types[-2] else: assert False, "This code should be unreachable." @@ -145,6 +157,7 @@ def get_arguments_for_execution(function, args, worker=global_worker): # get the object from the local object store print "Getting argument {} for function {}.".format(i, function.__name__) argument = worker.get_object(arg) + print "Successfully retrieved argument {} for function {}.".format(i, function.__name__) else: # pass the argument by value argument = arg diff --git a/src/scheduler.cc b/src/scheduler.cc index e8c1e8503..9b0d1b96d 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -7,6 +7,12 @@ Status SchedulerService::RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) { std::unique_ptr task(new Call(request->call())); // need to copy, because request is const fntable_lock_.lock(); + + if (fntable_.find(task->name()) == fntable_.end()) { + // TODO(rkn): In the future, this should probably not be fatal. + ORCH_LOG(ORCH_FATAL, "The function " << task->name() << " has not been registered by any worker."); + } + size_t num_return_vals = fntable_[task->name()].num_return_vals(); fntable_lock_.unlock(); diff --git a/test/arrays_test.py b/test/arrays_test.py index ccc8e90a6..f2b6e972d 100644 --- a/test/arrays_test.py +++ b/test/arrays_test.py @@ -73,12 +73,12 @@ class ArraysSingleTest(unittest.TestCase): time.sleep(0.2) # test eye - ref = single.eye(3) + ref = single.eye(3, "float") val = orchpy.pull(ref) self.assertTrue(np.alltrue(val == np.eye(3))) # test zeros - ref = single.zeros([3, 4, 5]) + ref = single.zeros([3, 4, 5], "float") val = orchpy.pull(ref) self.assertTrue(np.alltrue(val == np.zeros([3, 4, 5]))) @@ -101,13 +101,12 @@ class ArraysSingleTest(unittest.TestCase): class ArraysDistTest(unittest.TestCase): - def testMethods(self): + def testSerialization(self): x = dist.DistArray() - x.construct([2, 3, 4], float, np.array([[[orchpy.lib.ObjRef(0)]]])) + x.construct([2, 3, 4], np.array([[[orchpy.lib.ObjRef(0)]]])) capsule = serialization.serialize(x) y = serialization.deserialize(capsule) self.assertEqual(x.shape, y.shape) - self.assertEqual(x.dtype, y.dtype) self.assertEqual(x.objrefs[0, 0, 0].val, y.objrefs[0, 0, 0].val) def testAssemble(self): @@ -132,13 +131,67 @@ class ArraysDistTest(unittest.TestCase): time.sleep(0.2) - a = single.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE]) - b = single.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE]) + a = single.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE], "float") + b = single.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE], "float") x = dist.DistArray() - x.construct([2 * dist.BLOCK_SIZE, dist.BLOCK_SIZE], float, np.array([[a], [b]])) + x.construct([2 * dist.BLOCK_SIZE, dist.BLOCK_SIZE], np.array([[a], [b]])) self.assertTrue(np.alltrue(x.assemble() == np.vstack([np.ones([dist.BLOCK_SIZE, dist.BLOCK_SIZE]), np.zeros([dist.BLOCK_SIZE, dist.BLOCK_SIZE])]))) services.cleanup() + def testMethods(self): + scheduler_port = new_scheduler_port() + objstore_port = new_objstore_port() + worker1_port = new_worker_port() + worker2_port = new_worker_port() + worker3_port = new_worker_port() + worker4_port = new_worker_port() + + services.start_scheduler(address(IP_ADDRESS, scheduler_port)) + + time.sleep(0.1) + + services.start_objstore(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port)) + + time.sleep(0.2) + + orchpy.connect(address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port), address(IP_ADDRESS, worker1_port)) + + test_dir = os.path.dirname(os.path.abspath(__file__)) + test_path = os.path.join(test_dir, "testrecv.py") + services.start_worker(test_path, address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port), address(IP_ADDRESS, worker2_port)) + services.start_worker(test_path, address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port), address(IP_ADDRESS, worker3_port)) + services.start_worker(test_path, address(IP_ADDRESS, scheduler_port), address(IP_ADDRESS, objstore_port), address(IP_ADDRESS, worker4_port)) + + time.sleep(0.2) + + x = dist.zeros([9, 25, 51], "float") + self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == 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]))) + + 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)))) + + x = dist.eye(25, "float") + self.assertTrue(np.alltrue(orchpy.pull(dist.assemble(x)) == 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))))) + + 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))))) + + 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))))) + + services.cleanup() + if __name__ == '__main__': unittest.main() diff --git a/test/shell.py b/test/shell.py index 13c58ee18..1a0e67806 100644 --- a/test/shell.py +++ b/test/shell.py @@ -5,7 +5,7 @@ import orchpy.services as services import orchpy.worker as worker import arrays.single as single -# import arrays.dist as dist +import arrays.dist as dist from grpc.beta import implementations import orchestra_pb2 diff --git a/test/testrecv.py b/test/testrecv.py index 938505297..7768000cf 100644 --- a/test/testrecv.py +++ b/test/testrecv.py @@ -2,7 +2,7 @@ import sys import argparse import arrays.single as single -# import arrays.dist as dist +import arrays.dist as dist import orchpy import orchpy.services as services @@ -31,7 +31,9 @@ if __name__ == '__main__': orchpy.register_module(single) orchpy.register_module(single.random) orchpy.register_module(single.linalg) - # orchpy.register_module(dist) + orchpy.register_module(dist) + orchpy.register_module(dist.random) + # orchpy.register_module(dist.linalg) orchpy.register_module(sys.modules[__name__]) worker.main_loop()