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)