def region_join_mapper(ex, arrays, axes, local_user_fn, local_user_fn_kw,
target, region):
# FIXME: This should not be a special case. We need to think how to
# implement this for more applications not just for Cholesky.
# This is extremely ugly now. Please FIXME after paper submission.
old_ex = ex
ex = extent.change_partition_axis(ex, axes[0])
util.log_debug("old = %s, new = %s" % (str(old_ex), str(ex)))
tile = arrays[0].fetch(ex)
futures = rpc.FutureGroup()
for area in region:
intersection = extent.intersection(area, ex)
if intersection:
tile = tile.copy()
tiles = [tile]
join_extents = [ex]
subslice = extent.offset_slice(ex, intersection)
for i in range(1, len(arrays)):
ul = [0 for j in range(len(arrays[i].shape))]
lr = list(arrays[i].shape)
if axes[i] is not None:
ul[axes[i]] = ex.ul[axes[0][i - 1]]
lr[axes[i]] = ex.lr[axes[0][i - 1]]
join_extents.append(extent.create(ul, lr, arrays[i].shape))
tiles.append(arrays[i].fetch(join_extents[i]))
if local_user_fn_kw is None:
local_user_fn_kw = {}
_, tile[subslice] = local_user_fn(join_extents, tiles, **local_user_fn_kw)
break
futures = rpc.FutureGroup()
futures.append(target.update(ex, tile, wait=False))
return LocalKernelResult(result=[], futures=futures)
def target_mapper(ex, map_fn=None, source=None, target=None, fn_kw=None):
'''
Kernel function invoked during shuffle.
Runs ``map_fn`` over a single tile of the source array.
Args:
ex (Extent): Extent being processed.
map_fn (function): Function passed into `shuffle`.
source (DistArray): DistArray being mapped over.
target (DistArray): Array being written to.
fn_kw (dict): Keyword arguments for ``map_fn``.
Returns:
LocalKernelResult: No result data (all output is written to ``target``).
'''
result = list(map_fn(source, ex, **fn_kw))
futures = rpc.FutureGroup()
if result is not None:
for ex, v in result:
#update_time, _ = util.timeit(lambda: target.update(ex, v))
futures.append(target.update(ex, v, wait=False))
# util.log_warn('%s futures', len(futures))
return LocalKernelResult(result=[], futures=futures)
def _write_mapper(ex, source=None, sregion=None, dst_slice=None):
intersection = extent.intersection(ex, sregion)
futures = rpc.FutureGroup()
if intersection is not None:
dst_lr = np.asarray(intersection.lr) - np.asarray(sregion.ul)
dst_ul = np.asarray(intersection.ul) - np.asarray(sregion.ul)
dst_ex = extent.create(tuple(dst_ul), tuple(dst_lr), dst_slice.shape)
v = dst_slice.fetch(dst_ex)
futures.append(source.update(intersection, v, wait=False))
return LocalKernelResult(result=None, futures=futures)
def shutdown(self):
'''Shutdown all workers and halt.'''
if self._ctx.active is False:
return
self._ctx.active = False
futures = rpc.FutureGroup()
for id, w in self._workers.iteritems():
util.log_info('Shutting down worker %d', id)
futures.append(w.shutdown())
# Wait a second to let our shutdown request go out.
time.sleep(1)
self._server.shutdown()
def outer_mapper(ex, arrays, axes, local_user_fn, local_user_fn_kw, target):
# Fetch the tile for the first array
first_extent = extent.change_partition_axis(ex, axes[0])
first_tile = arrays[0].fetch(first_extent)
futures = rpc.FutureGroup()
if local_user_fn_kw is None:
local_user_fn_kw = {}
if axes[1] is None:
outer_extent = extent.from_shape(arrays[1].shape)
outer_tile = arrays[1].fetch(outer_extent)
result = local_user_fn(first_extent, first_tile, outer_extent,
outer_tile, **local_user_fn_kw)
if result is not None:
for ex, v in result:
futures.append(target.update(ex, v, wait=False))
else:
done_extent = {}
for key in arrays[1].tiles.iterkeys():
if hasattr(arrays[1], 'view_extent'):
key = arrays[1].view_extent(key)
outer_extent = extent.change_partition_axis(key, axes[1])
if done_extent.get(outer_extent, None) is not None:
# Usually, change_partition_axis won't return the same extents
# (if there is no bugs :) ). However, if the underline array
# is a vector and the new partition axis is 1, the API will
# return an extent that cover the whole vector.
# We need to avoid redo the extent.
continue
if outer_extent is None:
# It is possible that the return value of change_partition_axis
# is None if the dimension of new partition axis is smaller than
# the dimension of the original axis.
continue
outer_tile = arrays[1].fetch(outer_extent)
result = local_user_fn(first_extent, first_tile, outer_extent,
outer_tile, **local_user_fn_kw)
if result is not None:
for ex, v in result:
futures.append(target.update(ex, v, wait=False))
done_extent[outer_extent] = True
return LocalKernelResult(result=[], futures=futures)
def _initialize(self):
'''Sends an initialization request to all workers and waits
for their response.
'''
util.log_info('Initializing...')
req = core.InitializeReq(
peers=dict([(id, w.addr())
for id, w in self._workers.iteritems()]))
futures = rpc.FutureGroup()
for id, w in self._workers.iteritems():
req.id = id
futures.append(w.initialize(req))
futures.wait()
self._ctx = blob_ctx.BlobCtx(blob_ctx.MASTER_ID, self._workers, self)
self._initialized = True
util.log_info('done...')
def join_mapper(ex, arrays, axes, local_user_fn, local_user_fn_kw, target):
if len(axes) == 0:
tiles = []
# Fetch the extents
for i in range(0, len(arrays)):
tiles.append(arrays[i].fetch(ex))
join_extents = ex
else:
# First find out extents for all arrays
first_extent = extent.change_partition_axis(ex, axes[0])
# FIXME: I'm not sure if the following comment is really true for map2.
# assert first_extent is not None
if first_extent is None:
# It is possible that the return value of change_partition_axis
# is None if the dimension of new partition axis is smaller than
# the dimension of the original axis.
return LocalKernelResult(result=[])
keys = (first_extent.ul[axes[0]], first_extent.lr[axes[0]])
join_extents = [first_extent]
for i in range(1, len(arrays)):
ul = [0 for j in range(len(arrays[i].shape))]
lr = list(arrays[i].shape)
ul[axes[i]] = keys[0]
lr[axes[i]] = keys[1]
join_extents.append(extent.create(ul, lr, arrays[i].shape))
tiles = []
# Fetch the extents
for i in range(0, len(arrays)):
tiles.append(arrays[i].fetch(join_extents[i]))
if local_user_fn_kw is None:
local_user_fn_kw = {}
result = local_user_fn(join_extents, tiles, **local_user_fn_kw)
futures = rpc.FutureGroup()
if result is not None:
for ex, v in result:
futures.append(target.update(ex, v, wait=False))
return LocalKernelResult(result=[], futures=futures)
