diff --git a/distarray/cleanup.py b/distarray/cleanup.py index 333ad2bc..cd18e1e0 100644 --- a/distarray/cleanup.py +++ b/distarray/cleanup.py @@ -8,30 +8,13 @@ from distarray.ipython_utils import IPythonClient +def cleanup(view, module_name, context_name): + """ Delete Context object with the given name from the given module""" + def _cleanup(module_name, context_name): + ns = __import__(module_name) + delattr(ns, context_name) -def engine_cleanup(module_name, prefix): - """ Remove variables with ``prefix`` prefix from the namespace of the - module with ``module_name``. - """ - mod = __import__(module_name) - ns = mod.__dict__ - keys = tuple(ns.keys()) - deleted = 0 - for k in keys: - if k.startswith(prefix): - del ns[k] - deleted += 1 - count = 0 - for k in ns: - if k.startswith(prefix): - count += 1 - return (deleted, count) - - -def cleanup(view, module_name, prefix): - """ Delete keys with prefix from client's engines. """ - remaining = view.apply_async(engine_cleanup, module_name, prefix).get_dict() - return remaining + view.apply_async(_cleanup, module_name, context_name) def cleanup_all(module_name, prefix): diff --git a/distarray/client_map.py b/distarray/client_map.py index a0a35c90..372b26f1 100644 --- a/distarray/client_map.py +++ b/distarray/client_map.py @@ -33,6 +33,7 @@ from distarray.metadata_utils import (normalize_dist, normalize_grid_shape, make_grid_shape, + positivify, validate_grid_shape, _start_stop_block) @@ -469,6 +470,7 @@ def owning_ranks(self, idxs): If the `idxs` tuple is out of bounds, raises `IndexError`. """ + idxs = map(positivify, idxs, self.shape) # positivify and check dim_coord_hits = [m.owners(idx) for (m, idx) in zip(self.maps, idxs)] all_coords = product(*dim_coord_hits) ranks = [self.rank_from_coords[c] for c in all_coords] diff --git a/distarray/context.py b/distarray/context.py index 489fb2e8..fe8f70e7 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -82,8 +82,10 @@ def _setup_context_key(self): Create a dict on the engines which will hold everything from this context. """ - context_key = self.uid() - cmd = '%s = {}' % (context_key) + context_key = DISTARRAY_BASE_NAME + self.uid() + cmd = ("import types, sys;" + "%s = types.ModuleType('%s');") + cmd %= (context_key, context_key) self._execute(cmd, targets=range(len(self.view))) return context_key @@ -151,7 +153,7 @@ def uid(): def _generate_key(self): """ Generate a unique key name for this context. """ - key = "%s['%s']" % (self.context_key, self.uid()) + key = "%s.%s" % (self.context_key, 'key_' + self.uid()) return key def _key_and_push(self, *values): @@ -166,7 +168,7 @@ def delete_key(self, key): def cleanup(self): """ Delete keys that this context created from all the engines. """ - cleanup.cleanup(view=self.view, module_name='__main__', prefix=self._key_prefix()) + cleanup.cleanup(view=self.view, module_name='__main__', context_name=self.context_key) def close(self): self.cleanup() diff --git a/distarray/local/maps.py b/distarray/local/maps.py index cdcb7cd9..302a0be9 100644 --- a/distarray/local/maps.py +++ b/distarray/local/maps.py @@ -29,7 +29,7 @@ from distarray.local import construct from distarray.metadata_utils import (validate_grid_shape, make_grid_shape, normalize_grid_shape, normalize_dist, - distribute_indices) + distribute_indices, positivify) class Distribution(object): @@ -136,6 +136,7 @@ def rank_from_coords(self, coords): def local_from_global(self, *global_ind): """ Given `global_ind` indices, translate into local indices.""" + global_ind = tuple(map(positivify, global_ind, self.global_shape)) return tuple(self._maps[dim].local_from_global(global_ind[dim]) for dim in range(self.ndim)) diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index 7668cafa..1bb60fb0 100644 --- a/distarray/metadata_utils.py +++ b/distarray/metadata_utils.py @@ -192,3 +192,12 @@ def distribute_indices(dd): except KeyError: msg = "dist_type %r not supported." raise TypeError(msg % dist_type) + + +def positivify(index, size): + if 0 <= index < size: + return index + elif -size <= index < 0: + return size + index + else: + raise IndexError("Index %s out of bounds" % index) diff --git a/distarray/tests/test_client.py b/distarray/tests/test_client.py index c9c626a9..3b54925c 100644 --- a/distarray/tests/test_client.py +++ b/distarray/tests/test_client.py @@ -41,6 +41,10 @@ def test_set_and_getitem_block_dist(self): for val in range(size): self.assertEqual(dap[val], val) + for i in range(1, size + 1): + dap[-i] = i + self.assertEqual(dap[-i], i) + def test_set_and_getitem_nd_block_dist(self): size = 5 dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'}) @@ -49,31 +53,39 @@ def test_set_and_getitem_nd_block_dist(self): for col in range(size): val = size*row + col dap[row, col] = val - - for row in range(size): - for col in range(size): - val = size*row + col self.assertEqual(dap[row, col], val) + for row in range(1 ,size + 1): + for col in range(1, size + 1): + dap[-row, -col] = row + col + self.assertEqual(dap[-row, -col], row + col) + + def test_set_and_getitem_cyclic_dist(self): size = 10 dap = self.dac.empty((size,), dist={0: 'c'}) for val in range(size): dap[val] = val - - for val in range(size): self.assertEqual(dap[val], val) + for i in range(1, size + 1): + dap[-i] = i + self.assertEqual(dap[-i], i) + def test_get_index_error(self): dap = self.dac.empty((10,), dist={0: 'c'}) with self.assertRaises(IndexError): dap[11] + with self.assertRaises(IndexError): + dap[-11] def test_set_index_error(self): dap = self.dac.empty((10,), dist={0: 'c'}) with self.assertRaises(IndexError): dap[11] = 55 + with self.assertRaises(IndexError): + dap[-11] = 55 def test_iteration(self): size = 10