/
remotefunction.py
276 lines (223 loc) · 8.45 KB
/
remotefunction.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
"""Remote Functions and decorators for Views."""
# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.
from __future__ import division
import sys
import warnings
from IPython.external.decorator import decorator
from IPython.testing.skipdoctest import skip_doctest
from . import map as Map
from .asyncresult import AsyncMapResult
#-----------------------------------------------------------------------------
# Functions and Decorators
#-----------------------------------------------------------------------------
@skip_doctest
def remote(view, block=None, **flags):
"""Turn a function into a remote function.
This method can be used for map:
In [1]: @remote(view,block=True)
...: def func(a):
...: pass
"""
def remote_function(f):
return RemoteFunction(view, f, block=block, **flags)
return remote_function
@skip_doctest
def parallel(view, dist='b', block=None, ordered=True, **flags):
"""Turn a function into a parallel remote function.
This method can be used for map:
In [1]: @parallel(view, block=True)
...: def func(a):
...: pass
"""
def parallel_function(f):
return ParallelFunction(view, f, dist=dist, block=block, ordered=ordered, **flags)
return parallel_function
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)
@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
#--------------------------------------------------------------------------
# Classes
#--------------------------------------------------------------------------
class RemoteFunction(object):
"""Turn an existing function into a remote function.
Parameters
----------
view : View instance
The view to be used for execution
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
to use the current `block` attribute of `view`
**flags : remaining kwargs are passed to View.temp_flags
"""
view = None # the remote connection
func = None # the wrapped function
block = None # whether to block
flags = None # dict of extra kwargs for temp_flags
def __init__(self, view, f, block=None, **flags):
self.view = view
self.func = f
self.block=block
self.flags=flags
def __call__(self, *args, **kwargs):
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)
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
----------
view : View instance
The view to be used for execution
f : callable
The function to be wrapped into a remote function
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
block : bool [default: None]
Whether to wait for results or not. The default behavior is
to use the current `block` attribute of `view`
chunksize : int or None
The size of chunk to use when breaking up sequences in a load-balanced manner
ordered : bool [default: True]
Whether the result should be kept in order. If False,
results become available as they arrive, regardless of submission order.
**flags
remaining kwargs are passed to View.temp_flags
"""
chunksize = None
ordered = None
mapObject = None
_mapping = False
def __init__(self, view, f, dist='b', block=None, chunksize=None, ordered=True, **flags):
super(ParallelFunction, self).__init__(view, f, block=block, **flags)
self.chunksize = chunksize
self.ordered = ordered
mapClass = Map.dists[dist]
self.mapObject = mapClass()
@sync_view_results
def __call__(self, *sequences):
client = self.view.client
lens = []
maxlen = minlen = -1
for i, seq in enumerate(sequences):
try:
n = len(seq)
except Exception:
seq = list(seq)
if isinstance(sequences, tuple):
# can't alter a tuple
sequences = list(sequences)
sequences[i] = seq
n = len(seq)
if n > maxlen:
maxlen = n
if minlen == -1 or n < minlen:
minlen = n
lens.append(n)
if maxlen == 0:
# nothing to iterate over
return []
# check that the length of sequences match
if not self._mapping and minlen != maxlen:
msg = 'all sequences must have equal length, but have %s' % lens
raise ValueError(msg)
balanced = 'Balanced' in self.view.__class__.__name__
if balanced:
if self.chunksize:
nparts = maxlen // self.chunksize + int(maxlen % self.chunksize > 0)
else:
nparts = maxlen
targets = [None]*nparts
else:
if self.chunksize:
warnings.warn("`chunksize` is ignored unless load balancing", UserWarning)
# multiplexed:
targets = self.view.targets
# 'all' is lazily evaluated at execution time, which is now:
if targets == 'all':
targets = client._build_targets(targets)[1]
elif isinstance(targets, int):
# single-engine view, targets must be iterable
targets = [targets]
nparts = len(targets)
msg_ids = []
for index, t in enumerate(targets):
args = []
for seq in sequences:
part = self.mapObject.getPartition(seq, index, nparts, maxlen)
args.append(part)
if sum([len(arg) for arg in args]) == 0:
continue
if self._mapping:
if sys.version_info[0] >= 3:
f = lambda f, *sequences: list(map(f, *sequences))
else:
f = map
args = [self.func] + args
else:
f=self.func
view = self.view if balanced else client[t]
with view.temp_flags(block=False, **self.flags):
ar = view.apply(f, *args)
msg_ids.extend(ar.msg_ids)
r = AsyncMapResult(self.view.client, msg_ids, self.mapObject,
fname=getname(self.func),
ordered=self.ordered
)
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 one or more sequence(s) remotely.
This should behave very much like the builtin map, but return an AsyncMapResult
if self.block is False.
That means it can take generators (will be cast to lists locally),
and mismatched sequence lengths will be padded with None.
"""
# set _mapping as a flag for use inside self.__call__
self._mapping = True
try:
ret = self(*sequences)
finally:
self._mapping = False
return ret
__all__ = ['remote', 'parallel', 'RemoteFunction', 'ParallelFunction']