remotefunction.py
276 lines
| 8.4 KiB
| text/x-python
|
PythonLexer
|
r16634 | """Remote Functions and decorators for Views.""" | ||
|
r4018 | |||
|
r16634 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
|
r3588 | |||
|
r4155 | from __future__ import division | ||
|
r4159 | import sys | ||
|
r3636 | import warnings | ||
|
r10073 | from IPython.external.decorator import decorator | ||
|
r3886 | from IPython.testing.skipdoctest import skip_doctest | ||
|
r3641 | |||
|
r3642 | from . import map as Map | ||
from .asyncresult import AsyncMapResult | ||||
|
r3588 | |||
#----------------------------------------------------------------------------- | ||||
|
r5821 | # Functions and Decorators | ||
|
r3588 | #----------------------------------------------------------------------------- | ||
|
r3886 | @skip_doctest | ||
|
r3664 | def remote(view, block=None, **flags): | ||
|
r3588 | """Turn a function into a remote function. | ||
|
r4872 | |||
|
r3588 | This method can be used for map: | ||
|
r4872 | |||
|
r3664 | In [1]: @remote(view,block=True) | ||
|
r3641 | ...: def func(a): | ||
...: pass | ||||
|
r3588 | """ | ||
|
r4872 | |||
|
r3588 | def remote_function(f): | ||
|
r3664 | return RemoteFunction(view, f, block=block, **flags) | ||
|
r3588 | return remote_function | ||
|
r3886 | @skip_doctest | ||
|
r5171 | def parallel(view, dist='b', block=None, ordered=True, **flags): | ||
|
r3588 | """Turn a function into a parallel remote function. | ||
|
r4872 | |||
|
r3588 | This method can be used for map: | ||
|
r4872 | |||
|
r3664 | In [1]: @parallel(view, block=True) | ||
|
r3641 | ...: def func(a): | ||
...: pass | ||||
|
r3588 | """ | ||
|
r4872 | |||
|
r3588 | def parallel_function(f): | ||
|
r5171 | return ParallelFunction(view, f, dist=dist, block=block, ordered=ordered, **flags) | ||
|
r3588 | return parallel_function | ||
|
r5821 | def getname(f): | ||
"""Get the name of an object. | ||||
For use in case of callables that are not functions, and | ||||
thus may not have __name__ defined. | ||||
Order: f.__name__ > f.name > str(f) | ||||
""" | ||||
try: | ||||
return f.__name__ | ||||
except: | ||||
pass | ||||
try: | ||||
return f.name | ||||
except: | ||||
pass | ||||
return str(f) | ||||
|
r10073 | @decorator | ||
def sync_view_results(f, self, *args, **kwargs): | ||||
"""sync relevant results from self.client to our results attribute. | ||||
This is a clone of view.sync_results, but for remote functions | ||||
""" | ||||
view = self.view | ||||
if view._in_sync_results: | ||||
return f(self, *args, **kwargs) | ||||
view._in_sync_results = True | ||||
try: | ||||
ret = f(self, *args, **kwargs) | ||||
finally: | ||||
view._in_sync_results = False | ||||
view._sync_results() | ||||
return ret | ||||
|
r3588 | #-------------------------------------------------------------------------- | ||
# Classes | ||||
#-------------------------------------------------------------------------- | ||||
class RemoteFunction(object): | ||||
"""Turn an existing function into a remote function. | ||||
|
r4872 | |||
|
r3588 | Parameters | ||
---------- | ||||
|
r4872 | |||
|
r3664 | view : View instance | ||
The view to be used for execution | ||||
|
r3588 | f : callable | ||
The function to be wrapped into a remote function | ||||
block : bool [default: None] | ||||
Whether to wait for results or not. The default behavior is | ||||
|
r3664 | to use the current `block` attribute of `view` | ||
|
r4872 | |||
|
r3664 | **flags : remaining kwargs are passed to View.temp_flags | ||
|
r3588 | """ | ||
|
r4872 | |||
|
r3664 | view = None # the remote connection | ||
|
r3588 | func = None # the wrapped function | ||
block = None # whether to block | ||||
|
r3664 | flags = None # dict of extra kwargs for temp_flags | ||
|
r4872 | |||
|
r3664 | def __init__(self, view, f, block=None, **flags): | ||
self.view = view | ||||
|
r3588 | self.func = f | ||
self.block=block | ||||
|
r3664 | self.flags=flags | ||
|
r4872 | |||
|
r3588 | def __call__(self, *args, **kwargs): | ||
|
r3664 | block = self.view.block if self.block is None else self.block | ||
with self.view.temp_flags(block=block, **self.flags): | ||||
return self.view.apply(self.func, *args, **kwargs) | ||||
|
r4872 | |||
|
r3588 | |||
class ParallelFunction(RemoteFunction): | ||||
|
r3644 | """Class for mapping a function to sequences. | ||
|
r4872 | |||
|
r3644 | This will distribute the sequences according the a mapper, and call | ||
the function on each sub-sequence. If called via map, then the function | ||||
will be called once on each element, rather that each sub-sequence. | ||||
|
r4872 | |||
|
r3644 | Parameters | ||
---------- | ||||
|
r4872 | |||
|
r3664 | view : View instance | ||
The view to be used for execution | ||||
|
r3644 | f : callable | ||
The function to be wrapped into a remote function | ||||
|
r3664 | dist : str [default: 'b'] | ||
The key for which mapObject to use to distribute sequences | ||||
options are: | ||||
|
r13595 | |||
* 'b' : use contiguous chunks in order | ||||
* 'r' : use round-robin striping | ||||
|
r3644 | block : bool [default: None] | ||
Whether to wait for results or not. The default behavior is | ||||
|
r3664 | to use the current `block` attribute of `view` | ||
chunksize : int or None | ||||
|
r3644 | The size of chunk to use when breaking up sequences in a load-balanced manner | ||
|
r5171 | ordered : bool [default: True] | ||
|
r10568 | Whether the result should be kept in order. If False, | ||
results become available as they arrive, regardless of submission order. | ||||
|
r13595 | **flags | ||
remaining kwargs are passed to View.temp_flags | ||||
|
r3644 | """ | ||
|
r4872 | |||
|
r10568 | chunksize = None | ||
ordered = None | ||||
mapObject = None | ||||
|
r10566 | _mapping = False | ||
|
r4872 | |||
|
r5171 | def __init__(self, view, f, dist='b', block=None, chunksize=None, ordered=True, **flags): | ||
|
r3664 | super(ParallelFunction, self).__init__(view, f, block=block, **flags) | ||
self.chunksize = chunksize | ||||
|
r5171 | self.ordered = ordered | ||
|
r4872 | |||
|
r3588 | mapClass = Map.dists[dist] | ||
self.mapObject = mapClass() | ||||
|
r10566 | |||
|
r10073 | @sync_view_results | ||
|
r3588 | def __call__(self, *sequences): | ||
|
r5290 | client = self.view.client | ||
|
r10566 | lens = [] | ||
|
r10568 | maxlen = minlen = -1 | ||
|
r10566 | for i, seq in enumerate(sequences): | ||
try: | ||||
n = len(seq) | ||||
except Exception: | ||||
seq = list(seq) | ||||
|
r10570 | if isinstance(sequences, tuple): | ||
# can't alter a tuple | ||||
sequences = list(sequences) | ||||
|
r10566 | sequences[i] = seq | ||
n = len(seq) | ||||
|
r10568 | if n > maxlen: | ||
maxlen = n | ||||
if minlen == -1 or n < minlen: | ||||
minlen = n | ||||
|
r10566 | lens.append(n) | ||
|
r16634 | if maxlen == 0: | ||
# nothing to iterate over | ||||
return [] | ||||
|
r3664 | # check that the length of sequences match | ||
|
r10568 | if not self._mapping and minlen != maxlen: | ||
|
r10566 | msg = 'all sequences must have equal length, but have %s' % lens | ||
raise ValueError(msg) | ||||
|
r3664 | balanced = 'Balanced' in self.view.__class__.__name__ | ||
if balanced: | ||||
if self.chunksize: | ||||
|
r10568 | nparts = maxlen // self.chunksize + int(maxlen % self.chunksize > 0) | ||
|
r3636 | else: | ||
|
r10568 | nparts = maxlen | ||
|
r3664 | targets = [None]*nparts | ||
|
r3588 | else: | ||
|
r3664 | if self.chunksize: | ||
warnings.warn("`chunksize` is ignored unless load balancing", UserWarning) | ||||
|
r3588 | # multiplexed: | ||
|
r3664 | targets = self.view.targets | ||
|
r5290 | # 'all' is lazily evaluated at execution time, which is now: | ||
if targets == 'all': | ||||
targets = client._build_targets(targets)[1] | ||||
|
r6513 | elif isinstance(targets, int): | ||
# single-engine view, targets must be iterable | ||||
targets = [targets] | ||||
|
r3636 | nparts = len(targets) | ||
|
r4872 | |||
|
r3588 | msg_ids = [] | ||
|
r3635 | for index, t in enumerate(targets): | ||
|
r3588 | args = [] | ||
for seq in sequences: | ||||
|
r10568 | part = self.mapObject.getPartition(seq, index, nparts, maxlen) | ||
args.append(part) | ||||
|
r12174 | |||
if sum([len(arg) for arg in args]) == 0: | ||||
|
r3588 | continue | ||
|
r4872 | |||
|
r10566 | if self._mapping: | ||
|
r4155 | if sys.version_info[0] >= 3: | ||
f = lambda f, *sequences: list(map(f, *sequences)) | ||||
else: | ||||
f = map | ||||
|
r10568 | args = [self.func] + args | ||
|
r3588 | else: | ||
f=self.func | ||||
|
r4872 | |||
|
r3664 | view = self.view if balanced else client[t] | ||
with view.temp_flags(block=False, **self.flags): | ||||
ar = view.apply(f, *args) | ||||
|
r4872 | |||
|
r10568 | msg_ids.extend(ar.msg_ids) | ||
|
r4872 | |||
|
r10568 | r = AsyncMapResult(self.view.client, msg_ids, self.mapObject, | ||
|
r5821 | fname=getname(self.func), | ||
|
r5171 | ordered=self.ordered | ||
) | ||||
|
r4872 | |||
|
r3588 | if self.block: | ||
|
r3635 | try: | ||
return r.get() | ||||
except KeyboardInterrupt: | ||||
return r | ||||
|
r3588 | else: | ||
return r | ||||
|
r4872 | |||
|
r3588 | def map(self, *sequences): | ||
|
r10604 | """call a function on each element of one or more sequence(s) remotely. | ||
|
r3644 | This should behave very much like the builtin map, but return an AsyncMapResult | ||
if self.block is False. | ||||
|
r10604 | |||
That means it can take generators (will be cast to lists locally), | ||||
and mismatched sequence lengths will be padded with None. | ||||
|
r3644 | """ | ||
|
r10566 | # set _mapping as a flag for use inside self.__call__ | ||
self._mapping = True | ||||
|
r3635 | try: | ||
|
r10566 | ret = self(*sequences) | ||
|
r3635 | finally: | ||
|
r10566 | self._mapping = False | ||
|
r3588 | return ret | ||
|
r3886 | __all__ = ['remote', 'parallel', 'RemoteFunction', 'ParallelFunction'] | ||