remotefunction.py
241 lines
| 7.4 KiB
| text/x-python
|
PythonLexer
|
r4018 | """Remote Functions and decorators for Views. | ||
Authors: | ||||
* Brian Granger | ||||
* Min RK | ||||
""" | ||||
|
r3588 | #----------------------------------------------------------------------------- | ||
|
r4018 | # Copyright (C) 2010-2011 The IPython Development Team | ||
|
r3588 | # | ||
# Distributed under the terms of the BSD License. The full license is in | ||||
# the file COPYING, distributed as part of this software. | ||||
#----------------------------------------------------------------------------- | ||||
#----------------------------------------------------------------------------- | ||||
# Imports | ||||
#----------------------------------------------------------------------------- | ||||
|
r4155 | from __future__ import division | ||
|
r4159 | import sys | ||
|
r3636 | import warnings | ||
|
r3886 | from IPython.testing.skipdoctest import skip_doctest | ||
|
r3641 | |||
|
r3642 | from . import map as Map | ||
from .asyncresult import AsyncMapResult | ||||
|
r3588 | |||
#----------------------------------------------------------------------------- | ||||
|
r6360 | # 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 | ||
|
r6360 | 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) | ||||
|
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: | ||||
* '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] | ||
Whether | ||||
|
r3664 | **flags : remaining kwargs are passed to View.temp_flags | ||
|
r3644 | """ | ||
|
r4872 | |||
|
r3664 | chunksize=None | ||
|
r5171 | ordered=None | ||
|
r3664 | mapObject=None | ||
|
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() | ||||
|
r4872 | |||
|
r3588 | def __call__(self, *sequences): | ||
|
r5290 | client = self.view.client | ||
|
r3664 | # check that the length of sequences match | ||
|
r3588 | len_0 = len(sequences[0]) | ||
for s in sequences: | ||||
if len(s)!=len_0: | ||||
|
r3635 | msg = 'all sequences must have equal length, but %i!=%i'%(len_0,len(s)) | ||
raise ValueError(msg) | ||||
|
r3664 | balanced = 'Balanced' in self.view.__class__.__name__ | ||
if balanced: | ||||
if self.chunksize: | ||||
|
r4155 | nparts = len_0//self.chunksize + int(len_0%self.chunksize > 0) | ||
|
r3636 | else: | ||
nparts = len_0 | ||||
|
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] | ||||
|
r3636 | nparts = len(targets) | ||
|
r4872 | |||
|
r3588 | msg_ids = [] | ||
|
r3635 | for index, t in enumerate(targets): | ||
|
r3588 | args = [] | ||
for seq in sequences: | ||||
part = self.mapObject.getPartition(seq, index, nparts) | ||||
|
r3639 | if len(part) == 0: | ||
|
r3588 | continue | ||
else: | ||||
args.append(part) | ||||
if not args: | ||||
continue | ||||
|
r4872 | |||
|
r3588 | # print (args) | ||
if hasattr(self, '_map'): | ||||
|
r4155 | if sys.version_info[0] >= 3: | ||
f = lambda f, *sequences: list(map(f, *sequences)) | ||||
else: | ||||
f = map | ||||
|
r3588 | args = [self.func]+args | ||
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 | |||
|
r3635 | msg_ids.append(ar.msg_ids[0]) | ||
|
r4872 | |||
|
r5171 | r = AsyncMapResult(self.view.client, msg_ids, self.mapObject, | ||
|
r6360 | 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): | ||
|
r4872 | """call a function on each element of a sequence remotely. | ||
|
r3644 | This should behave very much like the builtin map, but return an AsyncMapResult | ||
if self.block is False. | ||||
""" | ||||
# set _map as a flag for use inside self.__call__ | ||||
|
r3588 | self._map = True | ||
|
r3635 | try: | ||
ret = self.__call__(*sequences) | ||||
finally: | ||||
del self._map | ||||
|
r3588 | return ret | ||
|
r3886 | __all__ = ['remote', 'parallel', 'RemoteFunction', 'ParallelFunction'] | ||