"""Remote Functions and decorators for the client.""" #----------------------------------------------------------------------------- # Copyright (C) 2010 The IPython Development Team # # Distributed under the terms of the BSD License. The full license is in # the file COPYING, distributed as part of this software. #----------------------------------------------------------------------------- #----------------------------------------------------------------------------- # Imports #----------------------------------------------------------------------------- import warnings from IPython.testing import decorators as testdec from . import map as Map from .asyncresult import AsyncMapResult #----------------------------------------------------------------------------- # Decorators #----------------------------------------------------------------------------- @testdec.skip_doctest def remote(client, bound=False, block=None, targets=None, balanced=None): """Turn a function into a remote function. This method can be used for map: In [1]: @remote(client,block=True) ...: def func(a): ...: pass """ def remote_function(f): return RemoteFunction(client, f, bound, block, targets, balanced) return remote_function @testdec.skip_doctest def parallel(client, dist='b', bound=False, block=None, targets='all', balanced=None): """Turn a function into a parallel remote function. This method can be used for map: In [1]: @parallel(client,block=True) ...: def func(a): ...: pass """ def parallel_function(f): return ParallelFunction(client, f, dist, bound, block, targets, balanced) return parallel_function #-------------------------------------------------------------------------- # Classes #-------------------------------------------------------------------------- class RemoteFunction(object): """Turn an existing function into a remote function. Parameters ---------- client : Client instance The client to be used to connect to engines f : callable The function to be wrapped into a remote function bound : bool [default: False] Whether the affect the remote namespace when called block : bool [default: None] Whether to wait for results or not. The default behavior is to use the current `block` attribute of `client` targets : valid target list [default: all] The targets on which to execute. balanced : bool Whether to load-balance with the Task scheduler or not """ client = None # the remote connection func = None # the wrapped function block = None # whether to block bound = None # whether to affect the namespace targets = None # where to execute balanced = None # whether to load-balance def __init__(self, client, f, bound=False, block=None, targets=None, balanced=None): self.client = client self.func = f self.block=block self.bound=bound self.targets=targets if balanced is None: if targets is None: balanced = True else: balanced = False self.balanced = balanced def __call__(self, *args, **kwargs): return self.client.apply(self.func, args=args, kwargs=kwargs, block=self.block, targets=self.targets, bound=self.bound, balanced=self.balanced) class ParallelFunction(RemoteFunction): """Class for mapping a function to sequences. 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. Parameters ---------- client : Client instance The client to be used to connect to engines f : callable The function to be wrapped into a remote function bound : bool [default: False] Whether the affect the remote namespace when called block : bool [default: None] Whether to wait for results or not. The default behavior is to use the current `block` attribute of `client` targets : valid target list [default: all] The targets on which to execute. balanced : bool Whether to load-balance with the Task scheduler or not chunk_size : int or None The size of chunk to use when breaking up sequences in a load-balanced manner """ def __init__(self, client, f, dist='b', bound=False, block=None, targets='all', balanced=None, chunk_size=None): super(ParallelFunction, self).__init__(client,f,bound,block,targets,balanced) self.chunk_size = chunk_size mapClass = Map.dists[dist] self.mapObject = mapClass() def __call__(self, *sequences): len_0 = len(sequences[0]) for s in sequences: if len(s)!=len_0: msg = 'all sequences must have equal length, but %i!=%i'%(len_0,len(s)) raise ValueError(msg) if self.balanced: if self.chunk_size: nparts = len_0/self.chunk_size + int(len_0%self.chunk_size > 0) else: nparts = len_0 targets = [self.targets]*nparts else: if self.chunk_size: warnings.warn("`chunk_size` is ignored when `balanced=False", UserWarning) # multiplexed: targets = self.client._build_targets(self.targets)[-1] nparts = len(targets) msg_ids = [] # my_f = lambda *a: map(self.func, *a) for index, t in enumerate(targets): args = [] for seq in sequences: part = self.mapObject.getPartition(seq, index, nparts) if len(part) == 0: continue else: args.append(part) if not args: continue # print (args) if hasattr(self, '_map'): f = map args = [self.func]+args else: f=self.func ar = self.client.apply(f, args=args, block=False, bound=self.bound, targets=t, balanced=self.balanced) msg_ids.append(ar.msg_ids[0]) r = AsyncMapResult(self.client, msg_ids, self.mapObject, fname=self.func.__name__) if self.block: try: return r.get() except KeyboardInterrupt: return r else: return r def map(self, *sequences): """call a function on each element of a sequence remotely. 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__ self._map = True try: ret = self.__call__(*sequences) finally: del self._map return ret __all__ = ['remote', 'parallel', 'RemoteFunction', 'ParallelFunction']