##// END OF EJS Templates
add 'received' timestamp to DB...
MinRK -
Show More
@@ -1,1452 +1,1454 b''
1 """A semi-synchronous Client for the ZMQ cluster
1 """A semi-synchronous Client for the ZMQ cluster
2
2
3 Authors:
3 Authors:
4
4
5 * MinRK
5 * MinRK
6 """
6 """
7 #-----------------------------------------------------------------------------
7 #-----------------------------------------------------------------------------
8 # Copyright (C) 2010-2011 The IPython Development Team
8 # Copyright (C) 2010-2011 The IPython Development Team
9 #
9 #
10 # Distributed under the terms of the BSD License. The full license is in
10 # Distributed under the terms of the BSD License. The full license is in
11 # the file COPYING, distributed as part of this software.
11 # the file COPYING, distributed as part of this software.
12 #-----------------------------------------------------------------------------
12 #-----------------------------------------------------------------------------
13
13
14 #-----------------------------------------------------------------------------
14 #-----------------------------------------------------------------------------
15 # Imports
15 # Imports
16 #-----------------------------------------------------------------------------
16 #-----------------------------------------------------------------------------
17
17
18 import os
18 import os
19 import json
19 import json
20 import sys
20 import sys
21 import time
21 import time
22 import warnings
22 import warnings
23 from datetime import datetime
23 from datetime import datetime
24 from getpass import getpass
24 from getpass import getpass
25 from pprint import pprint
25 from pprint import pprint
26
26
27 pjoin = os.path.join
27 pjoin = os.path.join
28
28
29 import zmq
29 import zmq
30 # from zmq.eventloop import ioloop, zmqstream
30 # from zmq.eventloop import ioloop, zmqstream
31
31
32 from IPython.config.configurable import MultipleInstanceError
32 from IPython.config.configurable import MultipleInstanceError
33 from IPython.core.application import BaseIPythonApplication
33 from IPython.core.application import BaseIPythonApplication
34
34
35 from IPython.utils.jsonutil import rekey
35 from IPython.utils.jsonutil import rekey
36 from IPython.utils.localinterfaces import LOCAL_IPS
36 from IPython.utils.localinterfaces import LOCAL_IPS
37 from IPython.utils.path import get_ipython_dir
37 from IPython.utils.path import get_ipython_dir
38 from IPython.utils.traitlets import (HasTraits, Integer, Instance, Unicode,
38 from IPython.utils.traitlets import (HasTraits, Integer, Instance, Unicode,
39 Dict, List, Bool, Set)
39 Dict, List, Bool, Set)
40 from IPython.external.decorator import decorator
40 from IPython.external.decorator import decorator
41 from IPython.external.ssh import tunnel
41 from IPython.external.ssh import tunnel
42
42
43 from IPython.parallel import Reference
43 from IPython.parallel import Reference
44 from IPython.parallel import error
44 from IPython.parallel import error
45 from IPython.parallel import util
45 from IPython.parallel import util
46
46
47 from IPython.zmq.session import Session, Message
47 from IPython.zmq.session import Session, Message
48
48
49 from .asyncresult import AsyncResult, AsyncHubResult
49 from .asyncresult import AsyncResult, AsyncHubResult
50 from IPython.core.profiledir import ProfileDir, ProfileDirError
50 from IPython.core.profiledir import ProfileDir, ProfileDirError
51 from .view import DirectView, LoadBalancedView
51 from .view import DirectView, LoadBalancedView
52
52
53 if sys.version_info[0] >= 3:
53 if sys.version_info[0] >= 3:
54 # xrange is used in a couple 'isinstance' tests in py2
54 # xrange is used in a couple 'isinstance' tests in py2
55 # should be just 'range' in 3k
55 # should be just 'range' in 3k
56 xrange = range
56 xrange = range
57
57
58 #--------------------------------------------------------------------------
58 #--------------------------------------------------------------------------
59 # Decorators for Client methods
59 # Decorators for Client methods
60 #--------------------------------------------------------------------------
60 #--------------------------------------------------------------------------
61
61
62 @decorator
62 @decorator
63 def spin_first(f, self, *args, **kwargs):
63 def spin_first(f, self, *args, **kwargs):
64 """Call spin() to sync state prior to calling the method."""
64 """Call spin() to sync state prior to calling the method."""
65 self.spin()
65 self.spin()
66 return f(self, *args, **kwargs)
66 return f(self, *args, **kwargs)
67
67
68
68
69 #--------------------------------------------------------------------------
69 #--------------------------------------------------------------------------
70 # Classes
70 # Classes
71 #--------------------------------------------------------------------------
71 #--------------------------------------------------------------------------
72
72
73 class Metadata(dict):
73 class Metadata(dict):
74 """Subclass of dict for initializing metadata values.
74 """Subclass of dict for initializing metadata values.
75
75
76 Attribute access works on keys.
76 Attribute access works on keys.
77
77
78 These objects have a strict set of keys - errors will raise if you try
78 These objects have a strict set of keys - errors will raise if you try
79 to add new keys.
79 to add new keys.
80 """
80 """
81 def __init__(self, *args, **kwargs):
81 def __init__(self, *args, **kwargs):
82 dict.__init__(self)
82 dict.__init__(self)
83 md = {'msg_id' : None,
83 md = {'msg_id' : None,
84 'submitted' : None,
84 'submitted' : None,
85 'started' : None,
85 'started' : None,
86 'completed' : None,
86 'completed' : None,
87 'received' : None,
87 'received' : None,
88 'engine_uuid' : None,
88 'engine_uuid' : None,
89 'engine_id' : None,
89 'engine_id' : None,
90 'follow' : None,
90 'follow' : None,
91 'after' : None,
91 'after' : None,
92 'status' : None,
92 'status' : None,
93
93
94 'pyin' : None,
94 'pyin' : None,
95 'pyout' : None,
95 'pyout' : None,
96 'pyerr' : None,
96 'pyerr' : None,
97 'stdout' : '',
97 'stdout' : '',
98 'stderr' : '',
98 'stderr' : '',
99 }
99 }
100 self.update(md)
100 self.update(md)
101 self.update(dict(*args, **kwargs))
101 self.update(dict(*args, **kwargs))
102
102
103 def __getattr__(self, key):
103 def __getattr__(self, key):
104 """getattr aliased to getitem"""
104 """getattr aliased to getitem"""
105 if key in self.iterkeys():
105 if key in self.iterkeys():
106 return self[key]
106 return self[key]
107 else:
107 else:
108 raise AttributeError(key)
108 raise AttributeError(key)
109
109
110 def __setattr__(self, key, value):
110 def __setattr__(self, key, value):
111 """setattr aliased to setitem, with strict"""
111 """setattr aliased to setitem, with strict"""
112 if key in self.iterkeys():
112 if key in self.iterkeys():
113 self[key] = value
113 self[key] = value
114 else:
114 else:
115 raise AttributeError(key)
115 raise AttributeError(key)
116
116
117 def __setitem__(self, key, value):
117 def __setitem__(self, key, value):
118 """strict static key enforcement"""
118 """strict static key enforcement"""
119 if key in self.iterkeys():
119 if key in self.iterkeys():
120 dict.__setitem__(self, key, value)
120 dict.__setitem__(self, key, value)
121 else:
121 else:
122 raise KeyError(key)
122 raise KeyError(key)
123
123
124
124
125 class Client(HasTraits):
125 class Client(HasTraits):
126 """A semi-synchronous client to the IPython ZMQ cluster
126 """A semi-synchronous client to the IPython ZMQ cluster
127
127
128 Parameters
128 Parameters
129 ----------
129 ----------
130
130
131 url_or_file : bytes or unicode; zmq url or path to ipcontroller-client.json
131 url_or_file : bytes or unicode; zmq url or path to ipcontroller-client.json
132 Connection information for the Hub's registration. If a json connector
132 Connection information for the Hub's registration. If a json connector
133 file is given, then likely no further configuration is necessary.
133 file is given, then likely no further configuration is necessary.
134 [Default: use profile]
134 [Default: use profile]
135 profile : bytes
135 profile : bytes
136 The name of the Cluster profile to be used to find connector information.
136 The name of the Cluster profile to be used to find connector information.
137 If run from an IPython application, the default profile will be the same
137 If run from an IPython application, the default profile will be the same
138 as the running application, otherwise it will be 'default'.
138 as the running application, otherwise it will be 'default'.
139 context : zmq.Context
139 context : zmq.Context
140 Pass an existing zmq.Context instance, otherwise the client will create its own.
140 Pass an existing zmq.Context instance, otherwise the client will create its own.
141 debug : bool
141 debug : bool
142 flag for lots of message printing for debug purposes
142 flag for lots of message printing for debug purposes
143 timeout : int/float
143 timeout : int/float
144 time (in seconds) to wait for connection replies from the Hub
144 time (in seconds) to wait for connection replies from the Hub
145 [Default: 10]
145 [Default: 10]
146
146
147 #-------------- session related args ----------------
147 #-------------- session related args ----------------
148
148
149 config : Config object
149 config : Config object
150 If specified, this will be relayed to the Session for configuration
150 If specified, this will be relayed to the Session for configuration
151 username : str
151 username : str
152 set username for the session object
152 set username for the session object
153 packer : str (import_string) or callable
153 packer : str (import_string) or callable
154 Can be either the simple keyword 'json' or 'pickle', or an import_string to a
154 Can be either the simple keyword 'json' or 'pickle', or an import_string to a
155 function to serialize messages. Must support same input as
155 function to serialize messages. Must support same input as
156 JSON, and output must be bytes.
156 JSON, and output must be bytes.
157 You can pass a callable directly as `pack`
157 You can pass a callable directly as `pack`
158 unpacker : str (import_string) or callable
158 unpacker : str (import_string) or callable
159 The inverse of packer. Only necessary if packer is specified as *not* one
159 The inverse of packer. Only necessary if packer is specified as *not* one
160 of 'json' or 'pickle'.
160 of 'json' or 'pickle'.
161
161
162 #-------------- ssh related args ----------------
162 #-------------- ssh related args ----------------
163 # These are args for configuring the ssh tunnel to be used
163 # These are args for configuring the ssh tunnel to be used
164 # credentials are used to forward connections over ssh to the Controller
164 # credentials are used to forward connections over ssh to the Controller
165 # Note that the ip given in `addr` needs to be relative to sshserver
165 # Note that the ip given in `addr` needs to be relative to sshserver
166 # The most basic case is to leave addr as pointing to localhost (127.0.0.1),
166 # The most basic case is to leave addr as pointing to localhost (127.0.0.1),
167 # and set sshserver as the same machine the Controller is on. However,
167 # and set sshserver as the same machine the Controller is on. However,
168 # the only requirement is that sshserver is able to see the Controller
168 # the only requirement is that sshserver is able to see the Controller
169 # (i.e. is within the same trusted network).
169 # (i.e. is within the same trusted network).
170
170
171 sshserver : str
171 sshserver : str
172 A string of the form passed to ssh, i.e. 'server.tld' or 'user@server.tld:port'
172 A string of the form passed to ssh, i.e. 'server.tld' or 'user@server.tld:port'
173 If keyfile or password is specified, and this is not, it will default to
173 If keyfile or password is specified, and this is not, it will default to
174 the ip given in addr.
174 the ip given in addr.
175 sshkey : str; path to ssh private key file
175 sshkey : str; path to ssh private key file
176 This specifies a key to be used in ssh login, default None.
176 This specifies a key to be used in ssh login, default None.
177 Regular default ssh keys will be used without specifying this argument.
177 Regular default ssh keys will be used without specifying this argument.
178 password : str
178 password : str
179 Your ssh password to sshserver. Note that if this is left None,
179 Your ssh password to sshserver. Note that if this is left None,
180 you will be prompted for it if passwordless key based login is unavailable.
180 you will be prompted for it if passwordless key based login is unavailable.
181 paramiko : bool
181 paramiko : bool
182 flag for whether to use paramiko instead of shell ssh for tunneling.
182 flag for whether to use paramiko instead of shell ssh for tunneling.
183 [default: True on win32, False else]
183 [default: True on win32, False else]
184
184
185 ------- exec authentication args -------
185 ------- exec authentication args -------
186 If even localhost is untrusted, you can have some protection against
186 If even localhost is untrusted, you can have some protection against
187 unauthorized execution by signing messages with HMAC digests.
187 unauthorized execution by signing messages with HMAC digests.
188 Messages are still sent as cleartext, so if someone can snoop your
188 Messages are still sent as cleartext, so if someone can snoop your
189 loopback traffic this will not protect your privacy, but will prevent
189 loopback traffic this will not protect your privacy, but will prevent
190 unauthorized execution.
190 unauthorized execution.
191
191
192 exec_key : str
192 exec_key : str
193 an authentication key or file containing a key
193 an authentication key or file containing a key
194 default: None
194 default: None
195
195
196
196
197 Attributes
197 Attributes
198 ----------
198 ----------
199
199
200 ids : list of int engine IDs
200 ids : list of int engine IDs
201 requesting the ids attribute always synchronizes
201 requesting the ids attribute always synchronizes
202 the registration state. To request ids without synchronization,
202 the registration state. To request ids without synchronization,
203 use semi-private _ids attributes.
203 use semi-private _ids attributes.
204
204
205 history : list of msg_ids
205 history : list of msg_ids
206 a list of msg_ids, keeping track of all the execution
206 a list of msg_ids, keeping track of all the execution
207 messages you have submitted in order.
207 messages you have submitted in order.
208
208
209 outstanding : set of msg_ids
209 outstanding : set of msg_ids
210 a set of msg_ids that have been submitted, but whose
210 a set of msg_ids that have been submitted, but whose
211 results have not yet been received.
211 results have not yet been received.
212
212
213 results : dict
213 results : dict
214 a dict of all our results, keyed by msg_id
214 a dict of all our results, keyed by msg_id
215
215
216 block : bool
216 block : bool
217 determines default behavior when block not specified
217 determines default behavior when block not specified
218 in execution methods
218 in execution methods
219
219
220 Methods
220 Methods
221 -------
221 -------
222
222
223 spin
223 spin
224 flushes incoming results and registration state changes
224 flushes incoming results and registration state changes
225 control methods spin, and requesting `ids` also ensures up to date
225 control methods spin, and requesting `ids` also ensures up to date
226
226
227 wait
227 wait
228 wait on one or more msg_ids
228 wait on one or more msg_ids
229
229
230 execution methods
230 execution methods
231 apply
231 apply
232 legacy: execute, run
232 legacy: execute, run
233
233
234 data movement
234 data movement
235 push, pull, scatter, gather
235 push, pull, scatter, gather
236
236
237 query methods
237 query methods
238 queue_status, get_result, purge, result_status
238 queue_status, get_result, purge, result_status
239
239
240 control methods
240 control methods
241 abort, shutdown
241 abort, shutdown
242
242
243 """
243 """
244
244
245
245
246 block = Bool(False)
246 block = Bool(False)
247 outstanding = Set()
247 outstanding = Set()
248 results = Instance('collections.defaultdict', (dict,))
248 results = Instance('collections.defaultdict', (dict,))
249 metadata = Instance('collections.defaultdict', (Metadata,))
249 metadata = Instance('collections.defaultdict', (Metadata,))
250 history = List()
250 history = List()
251 debug = Bool(False)
251 debug = Bool(False)
252
252
253 profile=Unicode()
253 profile=Unicode()
254 def _profile_default(self):
254 def _profile_default(self):
255 if BaseIPythonApplication.initialized():
255 if BaseIPythonApplication.initialized():
256 # an IPython app *might* be running, try to get its profile
256 # an IPython app *might* be running, try to get its profile
257 try:
257 try:
258 return BaseIPythonApplication.instance().profile
258 return BaseIPythonApplication.instance().profile
259 except (AttributeError, MultipleInstanceError):
259 except (AttributeError, MultipleInstanceError):
260 # could be a *different* subclass of config.Application,
260 # could be a *different* subclass of config.Application,
261 # which would raise one of these two errors.
261 # which would raise one of these two errors.
262 return u'default'
262 return u'default'
263 else:
263 else:
264 return u'default'
264 return u'default'
265
265
266
266
267 _outstanding_dict = Instance('collections.defaultdict', (set,))
267 _outstanding_dict = Instance('collections.defaultdict', (set,))
268 _ids = List()
268 _ids = List()
269 _connected=Bool(False)
269 _connected=Bool(False)
270 _ssh=Bool(False)
270 _ssh=Bool(False)
271 _context = Instance('zmq.Context')
271 _context = Instance('zmq.Context')
272 _config = Dict()
272 _config = Dict()
273 _engines=Instance(util.ReverseDict, (), {})
273 _engines=Instance(util.ReverseDict, (), {})
274 # _hub_socket=Instance('zmq.Socket')
274 # _hub_socket=Instance('zmq.Socket')
275 _query_socket=Instance('zmq.Socket')
275 _query_socket=Instance('zmq.Socket')
276 _control_socket=Instance('zmq.Socket')
276 _control_socket=Instance('zmq.Socket')
277 _iopub_socket=Instance('zmq.Socket')
277 _iopub_socket=Instance('zmq.Socket')
278 _notification_socket=Instance('zmq.Socket')
278 _notification_socket=Instance('zmq.Socket')
279 _mux_socket=Instance('zmq.Socket')
279 _mux_socket=Instance('zmq.Socket')
280 _task_socket=Instance('zmq.Socket')
280 _task_socket=Instance('zmq.Socket')
281 _task_scheme=Unicode()
281 _task_scheme=Unicode()
282 _closed = False
282 _closed = False
283 _ignored_control_replies=Integer(0)
283 _ignored_control_replies=Integer(0)
284 _ignored_hub_replies=Integer(0)
284 _ignored_hub_replies=Integer(0)
285
285
286 def __new__(self, *args, **kw):
286 def __new__(self, *args, **kw):
287 # don't raise on positional args
287 # don't raise on positional args
288 return HasTraits.__new__(self, **kw)
288 return HasTraits.__new__(self, **kw)
289
289
290 def __init__(self, url_or_file=None, profile=None, profile_dir=None, ipython_dir=None,
290 def __init__(self, url_or_file=None, profile=None, profile_dir=None, ipython_dir=None,
291 context=None, debug=False, exec_key=None,
291 context=None, debug=False, exec_key=None,
292 sshserver=None, sshkey=None, password=None, paramiko=None,
292 sshserver=None, sshkey=None, password=None, paramiko=None,
293 timeout=10, **extra_args
293 timeout=10, **extra_args
294 ):
294 ):
295 if profile:
295 if profile:
296 super(Client, self).__init__(debug=debug, profile=profile)
296 super(Client, self).__init__(debug=debug, profile=profile)
297 else:
297 else:
298 super(Client, self).__init__(debug=debug)
298 super(Client, self).__init__(debug=debug)
299 if context is None:
299 if context is None:
300 context = zmq.Context.instance()
300 context = zmq.Context.instance()
301 self._context = context
301 self._context = context
302
302
303 self._setup_profile_dir(self.profile, profile_dir, ipython_dir)
303 self._setup_profile_dir(self.profile, profile_dir, ipython_dir)
304 if self._cd is not None:
304 if self._cd is not None:
305 if url_or_file is None:
305 if url_or_file is None:
306 url_or_file = pjoin(self._cd.security_dir, 'ipcontroller-client.json')
306 url_or_file = pjoin(self._cd.security_dir, 'ipcontroller-client.json')
307 assert url_or_file is not None, "I can't find enough information to connect to a hub!"\
307 assert url_or_file is not None, "I can't find enough information to connect to a hub!"\
308 " Please specify at least one of url_or_file or profile."
308 " Please specify at least one of url_or_file or profile."
309
309
310 if not util.is_url(url_or_file):
310 if not util.is_url(url_or_file):
311 # it's not a url, try for a file
311 # it's not a url, try for a file
312 if not os.path.exists(url_or_file):
312 if not os.path.exists(url_or_file):
313 if self._cd:
313 if self._cd:
314 url_or_file = os.path.join(self._cd.security_dir, url_or_file)
314 url_or_file = os.path.join(self._cd.security_dir, url_or_file)
315 assert os.path.exists(url_or_file), "Not a valid connection file or url: %r"%url_or_file
315 assert os.path.exists(url_or_file), "Not a valid connection file or url: %r"%url_or_file
316 with open(url_or_file) as f:
316 with open(url_or_file) as f:
317 cfg = json.loads(f.read())
317 cfg = json.loads(f.read())
318 else:
318 else:
319 cfg = {'url':url_or_file}
319 cfg = {'url':url_or_file}
320
320
321 # sync defaults from args, json:
321 # sync defaults from args, json:
322 if sshserver:
322 if sshserver:
323 cfg['ssh'] = sshserver
323 cfg['ssh'] = sshserver
324 if exec_key:
324 if exec_key:
325 cfg['exec_key'] = exec_key
325 cfg['exec_key'] = exec_key
326 exec_key = cfg['exec_key']
326 exec_key = cfg['exec_key']
327 location = cfg.setdefault('location', None)
327 location = cfg.setdefault('location', None)
328 cfg['url'] = util.disambiguate_url(cfg['url'], location)
328 cfg['url'] = util.disambiguate_url(cfg['url'], location)
329 url = cfg['url']
329 url = cfg['url']
330 proto,addr,port = util.split_url(url)
330 proto,addr,port = util.split_url(url)
331 if location is not None and addr == '127.0.0.1':
331 if location is not None and addr == '127.0.0.1':
332 # location specified, and connection is expected to be local
332 # location specified, and connection is expected to be local
333 if location not in LOCAL_IPS and not sshserver:
333 if location not in LOCAL_IPS and not sshserver:
334 # load ssh from JSON *only* if the controller is not on
334 # load ssh from JSON *only* if the controller is not on
335 # this machine
335 # this machine
336 sshserver=cfg['ssh']
336 sshserver=cfg['ssh']
337 if location not in LOCAL_IPS and not sshserver:
337 if location not in LOCAL_IPS and not sshserver:
338 # warn if no ssh specified, but SSH is probably needed
338 # warn if no ssh specified, but SSH is probably needed
339 # This is only a warning, because the most likely cause
339 # This is only a warning, because the most likely cause
340 # is a local Controller on a laptop whose IP is dynamic
340 # is a local Controller on a laptop whose IP is dynamic
341 warnings.warn("""
341 warnings.warn("""
342 Controller appears to be listening on localhost, but not on this machine.
342 Controller appears to be listening on localhost, but not on this machine.
343 If this is true, you should specify Client(...,sshserver='you@%s')
343 If this is true, you should specify Client(...,sshserver='you@%s')
344 or instruct your controller to listen on an external IP."""%location,
344 or instruct your controller to listen on an external IP."""%location,
345 RuntimeWarning)
345 RuntimeWarning)
346 elif not sshserver:
346 elif not sshserver:
347 # otherwise sync with cfg
347 # otherwise sync with cfg
348 sshserver = cfg['ssh']
348 sshserver = cfg['ssh']
349
349
350 self._config = cfg
350 self._config = cfg
351
351
352 self._ssh = bool(sshserver or sshkey or password)
352 self._ssh = bool(sshserver or sshkey or password)
353 if self._ssh and sshserver is None:
353 if self._ssh and sshserver is None:
354 # default to ssh via localhost
354 # default to ssh via localhost
355 sshserver = url.split('://')[1].split(':')[0]
355 sshserver = url.split('://')[1].split(':')[0]
356 if self._ssh and password is None:
356 if self._ssh and password is None:
357 if tunnel.try_passwordless_ssh(sshserver, sshkey, paramiko):
357 if tunnel.try_passwordless_ssh(sshserver, sshkey, paramiko):
358 password=False
358 password=False
359 else:
359 else:
360 password = getpass("SSH Password for %s: "%sshserver)
360 password = getpass("SSH Password for %s: "%sshserver)
361 ssh_kwargs = dict(keyfile=sshkey, password=password, paramiko=paramiko)
361 ssh_kwargs = dict(keyfile=sshkey, password=password, paramiko=paramiko)
362
362
363 # configure and construct the session
363 # configure and construct the session
364 if exec_key is not None:
364 if exec_key is not None:
365 if os.path.isfile(exec_key):
365 if os.path.isfile(exec_key):
366 extra_args['keyfile'] = exec_key
366 extra_args['keyfile'] = exec_key
367 else:
367 else:
368 exec_key = util.asbytes(exec_key)
368 exec_key = util.asbytes(exec_key)
369 extra_args['key'] = exec_key
369 extra_args['key'] = exec_key
370 self.session = Session(**extra_args)
370 self.session = Session(**extra_args)
371
371
372 self._query_socket = self._context.socket(zmq.DEALER)
372 self._query_socket = self._context.socket(zmq.DEALER)
373 self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
373 self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
374 if self._ssh:
374 if self._ssh:
375 tunnel.tunnel_connection(self._query_socket, url, sshserver, **ssh_kwargs)
375 tunnel.tunnel_connection(self._query_socket, url, sshserver, **ssh_kwargs)
376 else:
376 else:
377 self._query_socket.connect(url)
377 self._query_socket.connect(url)
378
378
379 self.session.debug = self.debug
379 self.session.debug = self.debug
380
380
381 self._notification_handlers = {'registration_notification' : self._register_engine,
381 self._notification_handlers = {'registration_notification' : self._register_engine,
382 'unregistration_notification' : self._unregister_engine,
382 'unregistration_notification' : self._unregister_engine,
383 'shutdown_notification' : lambda msg: self.close(),
383 'shutdown_notification' : lambda msg: self.close(),
384 }
384 }
385 self._queue_handlers = {'execute_reply' : self._handle_execute_reply,
385 self._queue_handlers = {'execute_reply' : self._handle_execute_reply,
386 'apply_reply' : self._handle_apply_reply}
386 'apply_reply' : self._handle_apply_reply}
387 self._connect(sshserver, ssh_kwargs, timeout)
387 self._connect(sshserver, ssh_kwargs, timeout)
388
388
389 def __del__(self):
389 def __del__(self):
390 """cleanup sockets, but _not_ context."""
390 """cleanup sockets, but _not_ context."""
391 self.close()
391 self.close()
392
392
393 def _setup_profile_dir(self, profile, profile_dir, ipython_dir):
393 def _setup_profile_dir(self, profile, profile_dir, ipython_dir):
394 if ipython_dir is None:
394 if ipython_dir is None:
395 ipython_dir = get_ipython_dir()
395 ipython_dir = get_ipython_dir()
396 if profile_dir is not None:
396 if profile_dir is not None:
397 try:
397 try:
398 self._cd = ProfileDir.find_profile_dir(profile_dir)
398 self._cd = ProfileDir.find_profile_dir(profile_dir)
399 return
399 return
400 except ProfileDirError:
400 except ProfileDirError:
401 pass
401 pass
402 elif profile is not None:
402 elif profile is not None:
403 try:
403 try:
404 self._cd = ProfileDir.find_profile_dir_by_name(
404 self._cd = ProfileDir.find_profile_dir_by_name(
405 ipython_dir, profile)
405 ipython_dir, profile)
406 return
406 return
407 except ProfileDirError:
407 except ProfileDirError:
408 pass
408 pass
409 self._cd = None
409 self._cd = None
410
410
411 def _update_engines(self, engines):
411 def _update_engines(self, engines):
412 """Update our engines dict and _ids from a dict of the form: {id:uuid}."""
412 """Update our engines dict and _ids from a dict of the form: {id:uuid}."""
413 for k,v in engines.iteritems():
413 for k,v in engines.iteritems():
414 eid = int(k)
414 eid = int(k)
415 self._engines[eid] = v
415 self._engines[eid] = v
416 self._ids.append(eid)
416 self._ids.append(eid)
417 self._ids = sorted(self._ids)
417 self._ids = sorted(self._ids)
418 if sorted(self._engines.keys()) != range(len(self._engines)) and \
418 if sorted(self._engines.keys()) != range(len(self._engines)) and \
419 self._task_scheme == 'pure' and self._task_socket:
419 self._task_scheme == 'pure' and self._task_socket:
420 self._stop_scheduling_tasks()
420 self._stop_scheduling_tasks()
421
421
422 def _stop_scheduling_tasks(self):
422 def _stop_scheduling_tasks(self):
423 """Stop scheduling tasks because an engine has been unregistered
423 """Stop scheduling tasks because an engine has been unregistered
424 from a pure ZMQ scheduler.
424 from a pure ZMQ scheduler.
425 """
425 """
426 self._task_socket.close()
426 self._task_socket.close()
427 self._task_socket = None
427 self._task_socket = None
428 msg = "An engine has been unregistered, and we are using pure " +\
428 msg = "An engine has been unregistered, and we are using pure " +\
429 "ZMQ task scheduling. Task farming will be disabled."
429 "ZMQ task scheduling. Task farming will be disabled."
430 if self.outstanding:
430 if self.outstanding:
431 msg += " If you were running tasks when this happened, " +\
431 msg += " If you were running tasks when this happened, " +\
432 "some `outstanding` msg_ids may never resolve."
432 "some `outstanding` msg_ids may never resolve."
433 warnings.warn(msg, RuntimeWarning)
433 warnings.warn(msg, RuntimeWarning)
434
434
435 def _build_targets(self, targets):
435 def _build_targets(self, targets):
436 """Turn valid target IDs or 'all' into two lists:
436 """Turn valid target IDs or 'all' into two lists:
437 (int_ids, uuids).
437 (int_ids, uuids).
438 """
438 """
439 if not self._ids:
439 if not self._ids:
440 # flush notification socket if no engines yet, just in case
440 # flush notification socket if no engines yet, just in case
441 if not self.ids:
441 if not self.ids:
442 raise error.NoEnginesRegistered("Can't build targets without any engines")
442 raise error.NoEnginesRegistered("Can't build targets without any engines")
443
443
444 if targets is None:
444 if targets is None:
445 targets = self._ids
445 targets = self._ids
446 elif isinstance(targets, basestring):
446 elif isinstance(targets, basestring):
447 if targets.lower() == 'all':
447 if targets.lower() == 'all':
448 targets = self._ids
448 targets = self._ids
449 else:
449 else:
450 raise TypeError("%r not valid str target, must be 'all'"%(targets))
450 raise TypeError("%r not valid str target, must be 'all'"%(targets))
451 elif isinstance(targets, int):
451 elif isinstance(targets, int):
452 if targets < 0:
452 if targets < 0:
453 targets = self.ids[targets]
453 targets = self.ids[targets]
454 if targets not in self._ids:
454 if targets not in self._ids:
455 raise IndexError("No such engine: %i"%targets)
455 raise IndexError("No such engine: %i"%targets)
456 targets = [targets]
456 targets = [targets]
457
457
458 if isinstance(targets, slice):
458 if isinstance(targets, slice):
459 indices = range(len(self._ids))[targets]
459 indices = range(len(self._ids))[targets]
460 ids = self.ids
460 ids = self.ids
461 targets = [ ids[i] for i in indices ]
461 targets = [ ids[i] for i in indices ]
462
462
463 if not isinstance(targets, (tuple, list, xrange)):
463 if not isinstance(targets, (tuple, list, xrange)):
464 raise TypeError("targets by int/slice/collection of ints only, not %s"%(type(targets)))
464 raise TypeError("targets by int/slice/collection of ints only, not %s"%(type(targets)))
465
465
466 return [util.asbytes(self._engines[t]) for t in targets], list(targets)
466 return [util.asbytes(self._engines[t]) for t in targets], list(targets)
467
467
468 def _connect(self, sshserver, ssh_kwargs, timeout):
468 def _connect(self, sshserver, ssh_kwargs, timeout):
469 """setup all our socket connections to the cluster. This is called from
469 """setup all our socket connections to the cluster. This is called from
470 __init__."""
470 __init__."""
471
471
472 # Maybe allow reconnecting?
472 # Maybe allow reconnecting?
473 if self._connected:
473 if self._connected:
474 return
474 return
475 self._connected=True
475 self._connected=True
476
476
477 def connect_socket(s, url):
477 def connect_socket(s, url):
478 url = util.disambiguate_url(url, self._config['location'])
478 url = util.disambiguate_url(url, self._config['location'])
479 if self._ssh:
479 if self._ssh:
480 return tunnel.tunnel_connection(s, url, sshserver, **ssh_kwargs)
480 return tunnel.tunnel_connection(s, url, sshserver, **ssh_kwargs)
481 else:
481 else:
482 return s.connect(url)
482 return s.connect(url)
483
483
484 self.session.send(self._query_socket, 'connection_request')
484 self.session.send(self._query_socket, 'connection_request')
485 # use Poller because zmq.select has wrong units in pyzmq 2.1.7
485 # use Poller because zmq.select has wrong units in pyzmq 2.1.7
486 poller = zmq.Poller()
486 poller = zmq.Poller()
487 poller.register(self._query_socket, zmq.POLLIN)
487 poller.register(self._query_socket, zmq.POLLIN)
488 # poll expects milliseconds, timeout is seconds
488 # poll expects milliseconds, timeout is seconds
489 evts = poller.poll(timeout*1000)
489 evts = poller.poll(timeout*1000)
490 if not evts:
490 if not evts:
491 raise error.TimeoutError("Hub connection request timed out")
491 raise error.TimeoutError("Hub connection request timed out")
492 idents,msg = self.session.recv(self._query_socket,mode=0)
492 idents,msg = self.session.recv(self._query_socket,mode=0)
493 if self.debug:
493 if self.debug:
494 pprint(msg)
494 pprint(msg)
495 msg = Message(msg)
495 msg = Message(msg)
496 content = msg.content
496 content = msg.content
497 self._config['registration'] = dict(content)
497 self._config['registration'] = dict(content)
498 if content.status == 'ok':
498 if content.status == 'ok':
499 ident = self.session.bsession
499 ident = self.session.bsession
500 if content.mux:
500 if content.mux:
501 self._mux_socket = self._context.socket(zmq.DEALER)
501 self._mux_socket = self._context.socket(zmq.DEALER)
502 self._mux_socket.setsockopt(zmq.IDENTITY, ident)
502 self._mux_socket.setsockopt(zmq.IDENTITY, ident)
503 connect_socket(self._mux_socket, content.mux)
503 connect_socket(self._mux_socket, content.mux)
504 if content.task:
504 if content.task:
505 self._task_scheme, task_addr = content.task
505 self._task_scheme, task_addr = content.task
506 self._task_socket = self._context.socket(zmq.DEALER)
506 self._task_socket = self._context.socket(zmq.DEALER)
507 self._task_socket.setsockopt(zmq.IDENTITY, ident)
507 self._task_socket.setsockopt(zmq.IDENTITY, ident)
508 connect_socket(self._task_socket, task_addr)
508 connect_socket(self._task_socket, task_addr)
509 if content.notification:
509 if content.notification:
510 self._notification_socket = self._context.socket(zmq.SUB)
510 self._notification_socket = self._context.socket(zmq.SUB)
511 connect_socket(self._notification_socket, content.notification)
511 connect_socket(self._notification_socket, content.notification)
512 self._notification_socket.setsockopt(zmq.SUBSCRIBE, b'')
512 self._notification_socket.setsockopt(zmq.SUBSCRIBE, b'')
513 # if content.query:
513 # if content.query:
514 # self._query_socket = self._context.socket(zmq.DEALER)
514 # self._query_socket = self._context.socket(zmq.DEALER)
515 # self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
515 # self._query_socket.setsockopt(zmq.IDENTITY, self.session.bsession)
516 # connect_socket(self._query_socket, content.query)
516 # connect_socket(self._query_socket, content.query)
517 if content.control:
517 if content.control:
518 self._control_socket = self._context.socket(zmq.DEALER)
518 self._control_socket = self._context.socket(zmq.DEALER)
519 self._control_socket.setsockopt(zmq.IDENTITY, ident)
519 self._control_socket.setsockopt(zmq.IDENTITY, ident)
520 connect_socket(self._control_socket, content.control)
520 connect_socket(self._control_socket, content.control)
521 if content.iopub:
521 if content.iopub:
522 self._iopub_socket = self._context.socket(zmq.SUB)
522 self._iopub_socket = self._context.socket(zmq.SUB)
523 self._iopub_socket.setsockopt(zmq.SUBSCRIBE, b'')
523 self._iopub_socket.setsockopt(zmq.SUBSCRIBE, b'')
524 self._iopub_socket.setsockopt(zmq.IDENTITY, ident)
524 self._iopub_socket.setsockopt(zmq.IDENTITY, ident)
525 connect_socket(self._iopub_socket, content.iopub)
525 connect_socket(self._iopub_socket, content.iopub)
526 self._update_engines(dict(content.engines))
526 self._update_engines(dict(content.engines))
527 else:
527 else:
528 self._connected = False
528 self._connected = False
529 raise Exception("Failed to connect!")
529 raise Exception("Failed to connect!")
530
530
531 #--------------------------------------------------------------------------
531 #--------------------------------------------------------------------------
532 # handlers and callbacks for incoming messages
532 # handlers and callbacks for incoming messages
533 #--------------------------------------------------------------------------
533 #--------------------------------------------------------------------------
534
534
535 def _unwrap_exception(self, content):
535 def _unwrap_exception(self, content):
536 """unwrap exception, and remap engine_id to int."""
536 """unwrap exception, and remap engine_id to int."""
537 e = error.unwrap_exception(content)
537 e = error.unwrap_exception(content)
538 # print e.traceback
538 # print e.traceback
539 if e.engine_info:
539 if e.engine_info:
540 e_uuid = e.engine_info['engine_uuid']
540 e_uuid = e.engine_info['engine_uuid']
541 eid = self._engines[e_uuid]
541 eid = self._engines[e_uuid]
542 e.engine_info['engine_id'] = eid
542 e.engine_info['engine_id'] = eid
543 return e
543 return e
544
544
545 def _extract_metadata(self, header, parent, content):
545 def _extract_metadata(self, header, parent, content):
546 md = {'msg_id' : parent['msg_id'],
546 md = {'msg_id' : parent['msg_id'],
547 'received' : datetime.now(),
547 'received' : datetime.now(),
548 'engine_uuid' : header.get('engine', None),
548 'engine_uuid' : header.get('engine', None),
549 'follow' : parent.get('follow', []),
549 'follow' : parent.get('follow', []),
550 'after' : parent.get('after', []),
550 'after' : parent.get('after', []),
551 'status' : content['status'],
551 'status' : content['status'],
552 }
552 }
553
553
554 if md['engine_uuid'] is not None:
554 if md['engine_uuid'] is not None:
555 md['engine_id'] = self._engines.get(md['engine_uuid'], None)
555 md['engine_id'] = self._engines.get(md['engine_uuid'], None)
556
556
557 if 'date' in parent:
557 if 'date' in parent:
558 md['submitted'] = parent['date']
558 md['submitted'] = parent['date']
559 if 'started' in header:
559 if 'started' in header:
560 md['started'] = header['started']
560 md['started'] = header['started']
561 if 'date' in header:
561 if 'date' in header:
562 md['completed'] = header['date']
562 md['completed'] = header['date']
563 return md
563 return md
564
564
565 def _register_engine(self, msg):
565 def _register_engine(self, msg):
566 """Register a new engine, and update our connection info."""
566 """Register a new engine, and update our connection info."""
567 content = msg['content']
567 content = msg['content']
568 eid = content['id']
568 eid = content['id']
569 d = {eid : content['queue']}
569 d = {eid : content['queue']}
570 self._update_engines(d)
570 self._update_engines(d)
571
571
572 def _unregister_engine(self, msg):
572 def _unregister_engine(self, msg):
573 """Unregister an engine that has died."""
573 """Unregister an engine that has died."""
574 content = msg['content']
574 content = msg['content']
575 eid = int(content['id'])
575 eid = int(content['id'])
576 if eid in self._ids:
576 if eid in self._ids:
577 self._ids.remove(eid)
577 self._ids.remove(eid)
578 uuid = self._engines.pop(eid)
578 uuid = self._engines.pop(eid)
579
579
580 self._handle_stranded_msgs(eid, uuid)
580 self._handle_stranded_msgs(eid, uuid)
581
581
582 if self._task_socket and self._task_scheme == 'pure':
582 if self._task_socket and self._task_scheme == 'pure':
583 self._stop_scheduling_tasks()
583 self._stop_scheduling_tasks()
584
584
585 def _handle_stranded_msgs(self, eid, uuid):
585 def _handle_stranded_msgs(self, eid, uuid):
586 """Handle messages known to be on an engine when the engine unregisters.
586 """Handle messages known to be on an engine when the engine unregisters.
587
587
588 It is possible that this will fire prematurely - that is, an engine will
588 It is possible that this will fire prematurely - that is, an engine will
589 go down after completing a result, and the client will be notified
589 go down after completing a result, and the client will be notified
590 of the unregistration and later receive the successful result.
590 of the unregistration and later receive the successful result.
591 """
591 """
592
592
593 outstanding = self._outstanding_dict[uuid]
593 outstanding = self._outstanding_dict[uuid]
594
594
595 for msg_id in list(outstanding):
595 for msg_id in list(outstanding):
596 if msg_id in self.results:
596 if msg_id in self.results:
597 # we already
597 # we already
598 continue
598 continue
599 try:
599 try:
600 raise error.EngineError("Engine %r died while running task %r"%(eid, msg_id))
600 raise error.EngineError("Engine %r died while running task %r"%(eid, msg_id))
601 except:
601 except:
602 content = error.wrap_exception()
602 content = error.wrap_exception()
603 # build a fake message:
603 # build a fake message:
604 parent = {}
604 parent = {}
605 header = {}
605 header = {}
606 parent['msg_id'] = msg_id
606 parent['msg_id'] = msg_id
607 header['engine'] = uuid
607 header['engine'] = uuid
608 header['date'] = datetime.now()
608 header['date'] = datetime.now()
609 msg = dict(parent_header=parent, header=header, content=content)
609 msg = dict(parent_header=parent, header=header, content=content)
610 self._handle_apply_reply(msg)
610 self._handle_apply_reply(msg)
611
611
612 def _handle_execute_reply(self, msg):
612 def _handle_execute_reply(self, msg):
613 """Save the reply to an execute_request into our results.
613 """Save the reply to an execute_request into our results.
614
614
615 execute messages are never actually used. apply is used instead.
615 execute messages are never actually used. apply is used instead.
616 """
616 """
617
617
618 parent = msg['parent_header']
618 parent = msg['parent_header']
619 msg_id = parent['msg_id']
619 msg_id = parent['msg_id']
620 if msg_id not in self.outstanding:
620 if msg_id not in self.outstanding:
621 if msg_id in self.history:
621 if msg_id in self.history:
622 print ("got stale result: %s"%msg_id)
622 print ("got stale result: %s"%msg_id)
623 else:
623 else:
624 print ("got unknown result: %s"%msg_id)
624 print ("got unknown result: %s"%msg_id)
625 else:
625 else:
626 self.outstanding.remove(msg_id)
626 self.outstanding.remove(msg_id)
627 self.results[msg_id] = self._unwrap_exception(msg['content'])
627 self.results[msg_id] = self._unwrap_exception(msg['content'])
628
628
629 def _handle_apply_reply(self, msg):
629 def _handle_apply_reply(self, msg):
630 """Save the reply to an apply_request into our results."""
630 """Save the reply to an apply_request into our results."""
631 parent = msg['parent_header']
631 parent = msg['parent_header']
632 msg_id = parent['msg_id']
632 msg_id = parent['msg_id']
633 if msg_id not in self.outstanding:
633 if msg_id not in self.outstanding:
634 if msg_id in self.history:
634 if msg_id in self.history:
635 print ("got stale result: %s"%msg_id)
635 print ("got stale result: %s"%msg_id)
636 print self.results[msg_id]
636 print self.results[msg_id]
637 print msg
637 print msg
638 else:
638 else:
639 print ("got unknown result: %s"%msg_id)
639 print ("got unknown result: %s"%msg_id)
640 else:
640 else:
641 self.outstanding.remove(msg_id)
641 self.outstanding.remove(msg_id)
642 content = msg['content']
642 content = msg['content']
643 header = msg['header']
643 header = msg['header']
644
644
645 # construct metadata:
645 # construct metadata:
646 md = self.metadata[msg_id]
646 md = self.metadata[msg_id]
647 md.update(self._extract_metadata(header, parent, content))
647 md.update(self._extract_metadata(header, parent, content))
648 # is this redundant?
648 # is this redundant?
649 self.metadata[msg_id] = md
649 self.metadata[msg_id] = md
650
650
651 e_outstanding = self._outstanding_dict[md['engine_uuid']]
651 e_outstanding = self._outstanding_dict[md['engine_uuid']]
652 if msg_id in e_outstanding:
652 if msg_id in e_outstanding:
653 e_outstanding.remove(msg_id)
653 e_outstanding.remove(msg_id)
654
654
655 # construct result:
655 # construct result:
656 if content['status'] == 'ok':
656 if content['status'] == 'ok':
657 self.results[msg_id] = util.unserialize_object(msg['buffers'])[0]
657 self.results[msg_id] = util.unserialize_object(msg['buffers'])[0]
658 elif content['status'] == 'aborted':
658 elif content['status'] == 'aborted':
659 self.results[msg_id] = error.TaskAborted(msg_id)
659 self.results[msg_id] = error.TaskAborted(msg_id)
660 elif content['status'] == 'resubmitted':
660 elif content['status'] == 'resubmitted':
661 # TODO: handle resubmission
661 # TODO: handle resubmission
662 pass
662 pass
663 else:
663 else:
664 self.results[msg_id] = self._unwrap_exception(content)
664 self.results[msg_id] = self._unwrap_exception(content)
665
665
666 def _flush_notifications(self):
666 def _flush_notifications(self):
667 """Flush notifications of engine registrations waiting
667 """Flush notifications of engine registrations waiting
668 in ZMQ queue."""
668 in ZMQ queue."""
669 idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
669 idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
670 while msg is not None:
670 while msg is not None:
671 if self.debug:
671 if self.debug:
672 pprint(msg)
672 pprint(msg)
673 msg_type = msg['header']['msg_type']
673 msg_type = msg['header']['msg_type']
674 handler = self._notification_handlers.get(msg_type, None)
674 handler = self._notification_handlers.get(msg_type, None)
675 if handler is None:
675 if handler is None:
676 raise Exception("Unhandled message type: %s"%msg.msg_type)
676 raise Exception("Unhandled message type: %s"%msg.msg_type)
677 else:
677 else:
678 handler(msg)
678 handler(msg)
679 idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
679 idents,msg = self.session.recv(self._notification_socket, mode=zmq.NOBLOCK)
680
680
681 def _flush_results(self, sock):
681 def _flush_results(self, sock):
682 """Flush task or queue results waiting in ZMQ queue."""
682 """Flush task or queue results waiting in ZMQ queue."""
683 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
683 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
684 while msg is not None:
684 while msg is not None:
685 if self.debug:
685 if self.debug:
686 pprint(msg)
686 pprint(msg)
687 msg_type = msg['header']['msg_type']
687 msg_type = msg['header']['msg_type']
688 handler = self._queue_handlers.get(msg_type, None)
688 handler = self._queue_handlers.get(msg_type, None)
689 if handler is None:
689 if handler is None:
690 raise Exception("Unhandled message type: %s"%msg.msg_type)
690 raise Exception("Unhandled message type: %s"%msg.msg_type)
691 else:
691 else:
692 handler(msg)
692 handler(msg)
693 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
693 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
694
694
695 def _flush_control(self, sock):
695 def _flush_control(self, sock):
696 """Flush replies from the control channel waiting
696 """Flush replies from the control channel waiting
697 in the ZMQ queue.
697 in the ZMQ queue.
698
698
699 Currently: ignore them."""
699 Currently: ignore them."""
700 if self._ignored_control_replies <= 0:
700 if self._ignored_control_replies <= 0:
701 return
701 return
702 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
702 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
703 while msg is not None:
703 while msg is not None:
704 self._ignored_control_replies -= 1
704 self._ignored_control_replies -= 1
705 if self.debug:
705 if self.debug:
706 pprint(msg)
706 pprint(msg)
707 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
707 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
708
708
709 def _flush_ignored_control(self):
709 def _flush_ignored_control(self):
710 """flush ignored control replies"""
710 """flush ignored control replies"""
711 while self._ignored_control_replies > 0:
711 while self._ignored_control_replies > 0:
712 self.session.recv(self._control_socket)
712 self.session.recv(self._control_socket)
713 self._ignored_control_replies -= 1
713 self._ignored_control_replies -= 1
714
714
715 def _flush_ignored_hub_replies(self):
715 def _flush_ignored_hub_replies(self):
716 ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
716 ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
717 while msg is not None:
717 while msg is not None:
718 ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
718 ident,msg = self.session.recv(self._query_socket, mode=zmq.NOBLOCK)
719
719
720 def _flush_iopub(self, sock):
720 def _flush_iopub(self, sock):
721 """Flush replies from the iopub channel waiting
721 """Flush replies from the iopub channel waiting
722 in the ZMQ queue.
722 in the ZMQ queue.
723 """
723 """
724 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
724 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
725 while msg is not None:
725 while msg is not None:
726 if self.debug:
726 if self.debug:
727 pprint(msg)
727 pprint(msg)
728 parent = msg['parent_header']
728 parent = msg['parent_header']
729 # ignore IOPub messages with no parent.
729 # ignore IOPub messages with no parent.
730 # Caused by print statements or warnings from before the first execution.
730 # Caused by print statements or warnings from before the first execution.
731 if not parent:
731 if not parent:
732 continue
732 continue
733 msg_id = parent['msg_id']
733 msg_id = parent['msg_id']
734 content = msg['content']
734 content = msg['content']
735 header = msg['header']
735 header = msg['header']
736 msg_type = msg['header']['msg_type']
736 msg_type = msg['header']['msg_type']
737
737
738 # init metadata:
738 # init metadata:
739 md = self.metadata[msg_id]
739 md = self.metadata[msg_id]
740
740
741 if msg_type == 'stream':
741 if msg_type == 'stream':
742 name = content['name']
742 name = content['name']
743 s = md[name] or ''
743 s = md[name] or ''
744 md[name] = s + content['data']
744 md[name] = s + content['data']
745 elif msg_type == 'pyerr':
745 elif msg_type == 'pyerr':
746 md.update({'pyerr' : self._unwrap_exception(content)})
746 md.update({'pyerr' : self._unwrap_exception(content)})
747 elif msg_type == 'pyin':
747 elif msg_type == 'pyin':
748 md.update({'pyin' : content['code']})
748 md.update({'pyin' : content['code']})
749 else:
749 else:
750 md.update({msg_type : content.get('data', '')})
750 md.update({msg_type : content.get('data', '')})
751
751
752 # reduntant?
752 # reduntant?
753 self.metadata[msg_id] = md
753 self.metadata[msg_id] = md
754
754
755 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
755 idents,msg = self.session.recv(sock, mode=zmq.NOBLOCK)
756
756
757 #--------------------------------------------------------------------------
757 #--------------------------------------------------------------------------
758 # len, getitem
758 # len, getitem
759 #--------------------------------------------------------------------------
759 #--------------------------------------------------------------------------
760
760
761 def __len__(self):
761 def __len__(self):
762 """len(client) returns # of engines."""
762 """len(client) returns # of engines."""
763 return len(self.ids)
763 return len(self.ids)
764
764
765 def __getitem__(self, key):
765 def __getitem__(self, key):
766 """index access returns DirectView multiplexer objects
766 """index access returns DirectView multiplexer objects
767
767
768 Must be int, slice, or list/tuple/xrange of ints"""
768 Must be int, slice, or list/tuple/xrange of ints"""
769 if not isinstance(key, (int, slice, tuple, list, xrange)):
769 if not isinstance(key, (int, slice, tuple, list, xrange)):
770 raise TypeError("key by int/slice/iterable of ints only, not %s"%(type(key)))
770 raise TypeError("key by int/slice/iterable of ints only, not %s"%(type(key)))
771 else:
771 else:
772 return self.direct_view(key)
772 return self.direct_view(key)
773
773
774 #--------------------------------------------------------------------------
774 #--------------------------------------------------------------------------
775 # Begin public methods
775 # Begin public methods
776 #--------------------------------------------------------------------------
776 #--------------------------------------------------------------------------
777
777
778 @property
778 @property
779 def ids(self):
779 def ids(self):
780 """Always up-to-date ids property."""
780 """Always up-to-date ids property."""
781 self._flush_notifications()
781 self._flush_notifications()
782 # always copy:
782 # always copy:
783 return list(self._ids)
783 return list(self._ids)
784
784
785 def close(self):
785 def close(self):
786 if self._closed:
786 if self._closed:
787 return
787 return
788 snames = filter(lambda n: n.endswith('socket'), dir(self))
788 snames = filter(lambda n: n.endswith('socket'), dir(self))
789 for socket in map(lambda name: getattr(self, name), snames):
789 for socket in map(lambda name: getattr(self, name), snames):
790 if isinstance(socket, zmq.Socket) and not socket.closed:
790 if isinstance(socket, zmq.Socket) and not socket.closed:
791 socket.close()
791 socket.close()
792 self._closed = True
792 self._closed = True
793
793
794 def spin(self):
794 def spin(self):
795 """Flush any registration notifications and execution results
795 """Flush any registration notifications and execution results
796 waiting in the ZMQ queue.
796 waiting in the ZMQ queue.
797 """
797 """
798 if self._notification_socket:
798 if self._notification_socket:
799 self._flush_notifications()
799 self._flush_notifications()
800 if self._mux_socket:
800 if self._mux_socket:
801 self._flush_results(self._mux_socket)
801 self._flush_results(self._mux_socket)
802 if self._task_socket:
802 if self._task_socket:
803 self._flush_results(self._task_socket)
803 self._flush_results(self._task_socket)
804 if self._control_socket:
804 if self._control_socket:
805 self._flush_control(self._control_socket)
805 self._flush_control(self._control_socket)
806 if self._iopub_socket:
806 if self._iopub_socket:
807 self._flush_iopub(self._iopub_socket)
807 self._flush_iopub(self._iopub_socket)
808 if self._query_socket:
808 if self._query_socket:
809 self._flush_ignored_hub_replies()
809 self._flush_ignored_hub_replies()
810
810
811 def wait(self, jobs=None, timeout=-1):
811 def wait(self, jobs=None, timeout=-1):
812 """waits on one or more `jobs`, for up to `timeout` seconds.
812 """waits on one or more `jobs`, for up to `timeout` seconds.
813
813
814 Parameters
814 Parameters
815 ----------
815 ----------
816
816
817 jobs : int, str, or list of ints and/or strs, or one or more AsyncResult objects
817 jobs : int, str, or list of ints and/or strs, or one or more AsyncResult objects
818 ints are indices to self.history
818 ints are indices to self.history
819 strs are msg_ids
819 strs are msg_ids
820 default: wait on all outstanding messages
820 default: wait on all outstanding messages
821 timeout : float
821 timeout : float
822 a time in seconds, after which to give up.
822 a time in seconds, after which to give up.
823 default is -1, which means no timeout
823 default is -1, which means no timeout
824
824
825 Returns
825 Returns
826 -------
826 -------
827
827
828 True : when all msg_ids are done
828 True : when all msg_ids are done
829 False : timeout reached, some msg_ids still outstanding
829 False : timeout reached, some msg_ids still outstanding
830 """
830 """
831 tic = time.time()
831 tic = time.time()
832 if jobs is None:
832 if jobs is None:
833 theids = self.outstanding
833 theids = self.outstanding
834 else:
834 else:
835 if isinstance(jobs, (int, basestring, AsyncResult)):
835 if isinstance(jobs, (int, basestring, AsyncResult)):
836 jobs = [jobs]
836 jobs = [jobs]
837 theids = set()
837 theids = set()
838 for job in jobs:
838 for job in jobs:
839 if isinstance(job, int):
839 if isinstance(job, int):
840 # index access
840 # index access
841 job = self.history[job]
841 job = self.history[job]
842 elif isinstance(job, AsyncResult):
842 elif isinstance(job, AsyncResult):
843 map(theids.add, job.msg_ids)
843 map(theids.add, job.msg_ids)
844 continue
844 continue
845 theids.add(job)
845 theids.add(job)
846 if not theids.intersection(self.outstanding):
846 if not theids.intersection(self.outstanding):
847 return True
847 return True
848 self.spin()
848 self.spin()
849 while theids.intersection(self.outstanding):
849 while theids.intersection(self.outstanding):
850 if timeout >= 0 and ( time.time()-tic ) > timeout:
850 if timeout >= 0 and ( time.time()-tic ) > timeout:
851 break
851 break
852 time.sleep(1e-3)
852 time.sleep(1e-3)
853 self.spin()
853 self.spin()
854 return len(theids.intersection(self.outstanding)) == 0
854 return len(theids.intersection(self.outstanding)) == 0
855
855
856 #--------------------------------------------------------------------------
856 #--------------------------------------------------------------------------
857 # Control methods
857 # Control methods
858 #--------------------------------------------------------------------------
858 #--------------------------------------------------------------------------
859
859
860 @spin_first
860 @spin_first
861 def clear(self, targets=None, block=None):
861 def clear(self, targets=None, block=None):
862 """Clear the namespace in target(s)."""
862 """Clear the namespace in target(s)."""
863 block = self.block if block is None else block
863 block = self.block if block is None else block
864 targets = self._build_targets(targets)[0]
864 targets = self._build_targets(targets)[0]
865 for t in targets:
865 for t in targets:
866 self.session.send(self._control_socket, 'clear_request', content={}, ident=t)
866 self.session.send(self._control_socket, 'clear_request', content={}, ident=t)
867 error = False
867 error = False
868 if block:
868 if block:
869 self._flush_ignored_control()
869 self._flush_ignored_control()
870 for i in range(len(targets)):
870 for i in range(len(targets)):
871 idents,msg = self.session.recv(self._control_socket,0)
871 idents,msg = self.session.recv(self._control_socket,0)
872 if self.debug:
872 if self.debug:
873 pprint(msg)
873 pprint(msg)
874 if msg['content']['status'] != 'ok':
874 if msg['content']['status'] != 'ok':
875 error = self._unwrap_exception(msg['content'])
875 error = self._unwrap_exception(msg['content'])
876 else:
876 else:
877 self._ignored_control_replies += len(targets)
877 self._ignored_control_replies += len(targets)
878 if error:
878 if error:
879 raise error
879 raise error
880
880
881
881
882 @spin_first
882 @spin_first
883 def abort(self, jobs=None, targets=None, block=None):
883 def abort(self, jobs=None, targets=None, block=None):
884 """Abort specific jobs from the execution queues of target(s).
884 """Abort specific jobs from the execution queues of target(s).
885
885
886 This is a mechanism to prevent jobs that have already been submitted
886 This is a mechanism to prevent jobs that have already been submitted
887 from executing.
887 from executing.
888
888
889 Parameters
889 Parameters
890 ----------
890 ----------
891
891
892 jobs : msg_id, list of msg_ids, or AsyncResult
892 jobs : msg_id, list of msg_ids, or AsyncResult
893 The jobs to be aborted
893 The jobs to be aborted
894
894
895 If unspecified/None: abort all outstanding jobs.
895 If unspecified/None: abort all outstanding jobs.
896
896
897 """
897 """
898 block = self.block if block is None else block
898 block = self.block if block is None else block
899 jobs = jobs if jobs is not None else list(self.outstanding)
899 jobs = jobs if jobs is not None else list(self.outstanding)
900 targets = self._build_targets(targets)[0]
900 targets = self._build_targets(targets)[0]
901
901
902 msg_ids = []
902 msg_ids = []
903 if isinstance(jobs, (basestring,AsyncResult)):
903 if isinstance(jobs, (basestring,AsyncResult)):
904 jobs = [jobs]
904 jobs = [jobs]
905 bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
905 bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
906 if bad_ids:
906 if bad_ids:
907 raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
907 raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
908 for j in jobs:
908 for j in jobs:
909 if isinstance(j, AsyncResult):
909 if isinstance(j, AsyncResult):
910 msg_ids.extend(j.msg_ids)
910 msg_ids.extend(j.msg_ids)
911 else:
911 else:
912 msg_ids.append(j)
912 msg_ids.append(j)
913 content = dict(msg_ids=msg_ids)
913 content = dict(msg_ids=msg_ids)
914 for t in targets:
914 for t in targets:
915 self.session.send(self._control_socket, 'abort_request',
915 self.session.send(self._control_socket, 'abort_request',
916 content=content, ident=t)
916 content=content, ident=t)
917 error = False
917 error = False
918 if block:
918 if block:
919 self._flush_ignored_control()
919 self._flush_ignored_control()
920 for i in range(len(targets)):
920 for i in range(len(targets)):
921 idents,msg = self.session.recv(self._control_socket,0)
921 idents,msg = self.session.recv(self._control_socket,0)
922 if self.debug:
922 if self.debug:
923 pprint(msg)
923 pprint(msg)
924 if msg['content']['status'] != 'ok':
924 if msg['content']['status'] != 'ok':
925 error = self._unwrap_exception(msg['content'])
925 error = self._unwrap_exception(msg['content'])
926 else:
926 else:
927 self._ignored_control_replies += len(targets)
927 self._ignored_control_replies += len(targets)
928 if error:
928 if error:
929 raise error
929 raise error
930
930
931 @spin_first
931 @spin_first
932 def shutdown(self, targets=None, restart=False, hub=False, block=None):
932 def shutdown(self, targets=None, restart=False, hub=False, block=None):
933 """Terminates one or more engine processes, optionally including the hub."""
933 """Terminates one or more engine processes, optionally including the hub."""
934 block = self.block if block is None else block
934 block = self.block if block is None else block
935 if hub:
935 if hub:
936 targets = 'all'
936 targets = 'all'
937 targets = self._build_targets(targets)[0]
937 targets = self._build_targets(targets)[0]
938 for t in targets:
938 for t in targets:
939 self.session.send(self._control_socket, 'shutdown_request',
939 self.session.send(self._control_socket, 'shutdown_request',
940 content={'restart':restart},ident=t)
940 content={'restart':restart},ident=t)
941 error = False
941 error = False
942 if block or hub:
942 if block or hub:
943 self._flush_ignored_control()
943 self._flush_ignored_control()
944 for i in range(len(targets)):
944 for i in range(len(targets)):
945 idents,msg = self.session.recv(self._control_socket, 0)
945 idents,msg = self.session.recv(self._control_socket, 0)
946 if self.debug:
946 if self.debug:
947 pprint(msg)
947 pprint(msg)
948 if msg['content']['status'] != 'ok':
948 if msg['content']['status'] != 'ok':
949 error = self._unwrap_exception(msg['content'])
949 error = self._unwrap_exception(msg['content'])
950 else:
950 else:
951 self._ignored_control_replies += len(targets)
951 self._ignored_control_replies += len(targets)
952
952
953 if hub:
953 if hub:
954 time.sleep(0.25)
954 time.sleep(0.25)
955 self.session.send(self._query_socket, 'shutdown_request')
955 self.session.send(self._query_socket, 'shutdown_request')
956 idents,msg = self.session.recv(self._query_socket, 0)
956 idents,msg = self.session.recv(self._query_socket, 0)
957 if self.debug:
957 if self.debug:
958 pprint(msg)
958 pprint(msg)
959 if msg['content']['status'] != 'ok':
959 if msg['content']['status'] != 'ok':
960 error = self._unwrap_exception(msg['content'])
960 error = self._unwrap_exception(msg['content'])
961
961
962 if error:
962 if error:
963 raise error
963 raise error
964
964
965 #--------------------------------------------------------------------------
965 #--------------------------------------------------------------------------
966 # Execution related methods
966 # Execution related methods
967 #--------------------------------------------------------------------------
967 #--------------------------------------------------------------------------
968
968
969 def _maybe_raise(self, result):
969 def _maybe_raise(self, result):
970 """wrapper for maybe raising an exception if apply failed."""
970 """wrapper for maybe raising an exception if apply failed."""
971 if isinstance(result, error.RemoteError):
971 if isinstance(result, error.RemoteError):
972 raise result
972 raise result
973
973
974 return result
974 return result
975
975
976 def send_apply_message(self, socket, f, args=None, kwargs=None, subheader=None, track=False,
976 def send_apply_message(self, socket, f, args=None, kwargs=None, subheader=None, track=False,
977 ident=None):
977 ident=None):
978 """construct and send an apply message via a socket.
978 """construct and send an apply message via a socket.
979
979
980 This is the principal method with which all engine execution is performed by views.
980 This is the principal method with which all engine execution is performed by views.
981 """
981 """
982
982
983 assert not self._closed, "cannot use me anymore, I'm closed!"
983 assert not self._closed, "cannot use me anymore, I'm closed!"
984 # defaults:
984 # defaults:
985 args = args if args is not None else []
985 args = args if args is not None else []
986 kwargs = kwargs if kwargs is not None else {}
986 kwargs = kwargs if kwargs is not None else {}
987 subheader = subheader if subheader is not None else {}
987 subheader = subheader if subheader is not None else {}
988
988
989 # validate arguments
989 # validate arguments
990 if not callable(f) and not isinstance(f, Reference):
990 if not callable(f) and not isinstance(f, Reference):
991 raise TypeError("f must be callable, not %s"%type(f))
991 raise TypeError("f must be callable, not %s"%type(f))
992 if not isinstance(args, (tuple, list)):
992 if not isinstance(args, (tuple, list)):
993 raise TypeError("args must be tuple or list, not %s"%type(args))
993 raise TypeError("args must be tuple or list, not %s"%type(args))
994 if not isinstance(kwargs, dict):
994 if not isinstance(kwargs, dict):
995 raise TypeError("kwargs must be dict, not %s"%type(kwargs))
995 raise TypeError("kwargs must be dict, not %s"%type(kwargs))
996 if not isinstance(subheader, dict):
996 if not isinstance(subheader, dict):
997 raise TypeError("subheader must be dict, not %s"%type(subheader))
997 raise TypeError("subheader must be dict, not %s"%type(subheader))
998
998
999 bufs = util.pack_apply_message(f,args,kwargs)
999 bufs = util.pack_apply_message(f,args,kwargs)
1000
1000
1001 msg = self.session.send(socket, "apply_request", buffers=bufs, ident=ident,
1001 msg = self.session.send(socket, "apply_request", buffers=bufs, ident=ident,
1002 subheader=subheader, track=track)
1002 subheader=subheader, track=track)
1003
1003
1004 msg_id = msg['header']['msg_id']
1004 msg_id = msg['header']['msg_id']
1005 self.outstanding.add(msg_id)
1005 self.outstanding.add(msg_id)
1006 if ident:
1006 if ident:
1007 # possibly routed to a specific engine
1007 # possibly routed to a specific engine
1008 if isinstance(ident, list):
1008 if isinstance(ident, list):
1009 ident = ident[-1]
1009 ident = ident[-1]
1010 if ident in self._engines.values():
1010 if ident in self._engines.values():
1011 # save for later, in case of engine death
1011 # save for later, in case of engine death
1012 self._outstanding_dict[ident].add(msg_id)
1012 self._outstanding_dict[ident].add(msg_id)
1013 self.history.append(msg_id)
1013 self.history.append(msg_id)
1014 self.metadata[msg_id]['submitted'] = datetime.now()
1014 self.metadata[msg_id]['submitted'] = datetime.now()
1015
1015
1016 return msg
1016 return msg
1017
1017
1018 #--------------------------------------------------------------------------
1018 #--------------------------------------------------------------------------
1019 # construct a View object
1019 # construct a View object
1020 #--------------------------------------------------------------------------
1020 #--------------------------------------------------------------------------
1021
1021
1022 def load_balanced_view(self, targets=None):
1022 def load_balanced_view(self, targets=None):
1023 """construct a DirectView object.
1023 """construct a DirectView object.
1024
1024
1025 If no arguments are specified, create a LoadBalancedView
1025 If no arguments are specified, create a LoadBalancedView
1026 using all engines.
1026 using all engines.
1027
1027
1028 Parameters
1028 Parameters
1029 ----------
1029 ----------
1030
1030
1031 targets: list,slice,int,etc. [default: use all engines]
1031 targets: list,slice,int,etc. [default: use all engines]
1032 The subset of engines across which to load-balance
1032 The subset of engines across which to load-balance
1033 """
1033 """
1034 if targets == 'all':
1034 if targets == 'all':
1035 targets = None
1035 targets = None
1036 if targets is not None:
1036 if targets is not None:
1037 targets = self._build_targets(targets)[1]
1037 targets = self._build_targets(targets)[1]
1038 return LoadBalancedView(client=self, socket=self._task_socket, targets=targets)
1038 return LoadBalancedView(client=self, socket=self._task_socket, targets=targets)
1039
1039
1040 def direct_view(self, targets='all'):
1040 def direct_view(self, targets='all'):
1041 """construct a DirectView object.
1041 """construct a DirectView object.
1042
1042
1043 If no targets are specified, create a DirectView using all engines.
1043 If no targets are specified, create a DirectView using all engines.
1044
1044
1045 rc.direct_view('all') is distinguished from rc[:] in that 'all' will
1045 rc.direct_view('all') is distinguished from rc[:] in that 'all' will
1046 evaluate the target engines at each execution, whereas rc[:] will connect to
1046 evaluate the target engines at each execution, whereas rc[:] will connect to
1047 all *current* engines, and that list will not change.
1047 all *current* engines, and that list will not change.
1048
1048
1049 That is, 'all' will always use all engines, whereas rc[:] will not use
1049 That is, 'all' will always use all engines, whereas rc[:] will not use
1050 engines added after the DirectView is constructed.
1050 engines added after the DirectView is constructed.
1051
1051
1052 Parameters
1052 Parameters
1053 ----------
1053 ----------
1054
1054
1055 targets: list,slice,int,etc. [default: use all engines]
1055 targets: list,slice,int,etc. [default: use all engines]
1056 The engines to use for the View
1056 The engines to use for the View
1057 """
1057 """
1058 single = isinstance(targets, int)
1058 single = isinstance(targets, int)
1059 # allow 'all' to be lazily evaluated at each execution
1059 # allow 'all' to be lazily evaluated at each execution
1060 if targets != 'all':
1060 if targets != 'all':
1061 targets = self._build_targets(targets)[1]
1061 targets = self._build_targets(targets)[1]
1062 if single:
1062 if single:
1063 targets = targets[0]
1063 targets = targets[0]
1064 return DirectView(client=self, socket=self._mux_socket, targets=targets)
1064 return DirectView(client=self, socket=self._mux_socket, targets=targets)
1065
1065
1066 #--------------------------------------------------------------------------
1066 #--------------------------------------------------------------------------
1067 # Query methods
1067 # Query methods
1068 #--------------------------------------------------------------------------
1068 #--------------------------------------------------------------------------
1069
1069
1070 @spin_first
1070 @spin_first
1071 def get_result(self, indices_or_msg_ids=None, block=None):
1071 def get_result(self, indices_or_msg_ids=None, block=None):
1072 """Retrieve a result by msg_id or history index, wrapped in an AsyncResult object.
1072 """Retrieve a result by msg_id or history index, wrapped in an AsyncResult object.
1073
1073
1074 If the client already has the results, no request to the Hub will be made.
1074 If the client already has the results, no request to the Hub will be made.
1075
1075
1076 This is a convenient way to construct AsyncResult objects, which are wrappers
1076 This is a convenient way to construct AsyncResult objects, which are wrappers
1077 that include metadata about execution, and allow for awaiting results that
1077 that include metadata about execution, and allow for awaiting results that
1078 were not submitted by this Client.
1078 were not submitted by this Client.
1079
1079
1080 It can also be a convenient way to retrieve the metadata associated with
1080 It can also be a convenient way to retrieve the metadata associated with
1081 blocking execution, since it always retrieves
1081 blocking execution, since it always retrieves
1082
1082
1083 Examples
1083 Examples
1084 --------
1084 --------
1085 ::
1085 ::
1086
1086
1087 In [10]: r = client.apply()
1087 In [10]: r = client.apply()
1088
1088
1089 Parameters
1089 Parameters
1090 ----------
1090 ----------
1091
1091
1092 indices_or_msg_ids : integer history index, str msg_id, or list of either
1092 indices_or_msg_ids : integer history index, str msg_id, or list of either
1093 The indices or msg_ids of indices to be retrieved
1093 The indices or msg_ids of indices to be retrieved
1094
1094
1095 block : bool
1095 block : bool
1096 Whether to wait for the result to be done
1096 Whether to wait for the result to be done
1097
1097
1098 Returns
1098 Returns
1099 -------
1099 -------
1100
1100
1101 AsyncResult
1101 AsyncResult
1102 A single AsyncResult object will always be returned.
1102 A single AsyncResult object will always be returned.
1103
1103
1104 AsyncHubResult
1104 AsyncHubResult
1105 A subclass of AsyncResult that retrieves results from the Hub
1105 A subclass of AsyncResult that retrieves results from the Hub
1106
1106
1107 """
1107 """
1108 block = self.block if block is None else block
1108 block = self.block if block is None else block
1109 if indices_or_msg_ids is None:
1109 if indices_or_msg_ids is None:
1110 indices_or_msg_ids = -1
1110 indices_or_msg_ids = -1
1111
1111
1112 if not isinstance(indices_or_msg_ids, (list,tuple)):
1112 if not isinstance(indices_or_msg_ids, (list,tuple)):
1113 indices_or_msg_ids = [indices_or_msg_ids]
1113 indices_or_msg_ids = [indices_or_msg_ids]
1114
1114
1115 theids = []
1115 theids = []
1116 for id in indices_or_msg_ids:
1116 for id in indices_or_msg_ids:
1117 if isinstance(id, int):
1117 if isinstance(id, int):
1118 id = self.history[id]
1118 id = self.history[id]
1119 if not isinstance(id, basestring):
1119 if not isinstance(id, basestring):
1120 raise TypeError("indices must be str or int, not %r"%id)
1120 raise TypeError("indices must be str or int, not %r"%id)
1121 theids.append(id)
1121 theids.append(id)
1122
1122
1123 local_ids = filter(lambda msg_id: msg_id in self.history or msg_id in self.results, theids)
1123 local_ids = filter(lambda msg_id: msg_id in self.history or msg_id in self.results, theids)
1124 remote_ids = filter(lambda msg_id: msg_id not in local_ids, theids)
1124 remote_ids = filter(lambda msg_id: msg_id not in local_ids, theids)
1125
1125
1126 if remote_ids:
1126 if remote_ids:
1127 ar = AsyncHubResult(self, msg_ids=theids)
1127 ar = AsyncHubResult(self, msg_ids=theids)
1128 else:
1128 else:
1129 ar = AsyncResult(self, msg_ids=theids)
1129 ar = AsyncResult(self, msg_ids=theids)
1130
1130
1131 if block:
1131 if block:
1132 ar.wait()
1132 ar.wait()
1133
1133
1134 return ar
1134 return ar
1135
1135
1136 @spin_first
1136 @spin_first
1137 def resubmit(self, indices_or_msg_ids=None, subheader=None, block=None):
1137 def resubmit(self, indices_or_msg_ids=None, subheader=None, block=None):
1138 """Resubmit one or more tasks.
1138 """Resubmit one or more tasks.
1139
1139
1140 in-flight tasks may not be resubmitted.
1140 in-flight tasks may not be resubmitted.
1141
1141
1142 Parameters
1142 Parameters
1143 ----------
1143 ----------
1144
1144
1145 indices_or_msg_ids : integer history index, str msg_id, or list of either
1145 indices_or_msg_ids : integer history index, str msg_id, or list of either
1146 The indices or msg_ids of indices to be retrieved
1146 The indices or msg_ids of indices to be retrieved
1147
1147
1148 block : bool
1148 block : bool
1149 Whether to wait for the result to be done
1149 Whether to wait for the result to be done
1150
1150
1151 Returns
1151 Returns
1152 -------
1152 -------
1153
1153
1154 AsyncHubResult
1154 AsyncHubResult
1155 A subclass of AsyncResult that retrieves results from the Hub
1155 A subclass of AsyncResult that retrieves results from the Hub
1156
1156
1157 """
1157 """
1158 block = self.block if block is None else block
1158 block = self.block if block is None else block
1159 if indices_or_msg_ids is None:
1159 if indices_or_msg_ids is None:
1160 indices_or_msg_ids = -1
1160 indices_or_msg_ids = -1
1161
1161
1162 if not isinstance(indices_or_msg_ids, (list,tuple)):
1162 if not isinstance(indices_or_msg_ids, (list,tuple)):
1163 indices_or_msg_ids = [indices_or_msg_ids]
1163 indices_or_msg_ids = [indices_or_msg_ids]
1164
1164
1165 theids = []
1165 theids = []
1166 for id in indices_or_msg_ids:
1166 for id in indices_or_msg_ids:
1167 if isinstance(id, int):
1167 if isinstance(id, int):
1168 id = self.history[id]
1168 id = self.history[id]
1169 if not isinstance(id, basestring):
1169 if not isinstance(id, basestring):
1170 raise TypeError("indices must be str or int, not %r"%id)
1170 raise TypeError("indices must be str or int, not %r"%id)
1171 theids.append(id)
1171 theids.append(id)
1172
1172
1173 for msg_id in theids:
1173 for msg_id in theids:
1174 self.outstanding.discard(msg_id)
1174 self.outstanding.discard(msg_id)
1175 if msg_id in self.history:
1175 if msg_id in self.history:
1176 self.history.remove(msg_id)
1176 self.history.remove(msg_id)
1177 self.results.pop(msg_id, None)
1177 self.results.pop(msg_id, None)
1178 self.metadata.pop(msg_id, None)
1178 self.metadata.pop(msg_id, None)
1179 content = dict(msg_ids = theids)
1179 content = dict(msg_ids = theids)
1180
1180
1181 self.session.send(self._query_socket, 'resubmit_request', content)
1181 self.session.send(self._query_socket, 'resubmit_request', content)
1182
1182
1183 zmq.select([self._query_socket], [], [])
1183 zmq.select([self._query_socket], [], [])
1184 idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
1184 idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
1185 if self.debug:
1185 if self.debug:
1186 pprint(msg)
1186 pprint(msg)
1187 content = msg['content']
1187 content = msg['content']
1188 if content['status'] != 'ok':
1188 if content['status'] != 'ok':
1189 raise self._unwrap_exception(content)
1189 raise self._unwrap_exception(content)
1190
1190
1191 ar = AsyncHubResult(self, msg_ids=theids)
1191 ar = AsyncHubResult(self, msg_ids=theids)
1192
1192
1193 if block:
1193 if block:
1194 ar.wait()
1194 ar.wait()
1195
1195
1196 return ar
1196 return ar
1197
1197
1198 @spin_first
1198 @spin_first
1199 def result_status(self, msg_ids, status_only=True):
1199 def result_status(self, msg_ids, status_only=True):
1200 """Check on the status of the result(s) of the apply request with `msg_ids`.
1200 """Check on the status of the result(s) of the apply request with `msg_ids`.
1201
1201
1202 If status_only is False, then the actual results will be retrieved, else
1202 If status_only is False, then the actual results will be retrieved, else
1203 only the status of the results will be checked.
1203 only the status of the results will be checked.
1204
1204
1205 Parameters
1205 Parameters
1206 ----------
1206 ----------
1207
1207
1208 msg_ids : list of msg_ids
1208 msg_ids : list of msg_ids
1209 if int:
1209 if int:
1210 Passed as index to self.history for convenience.
1210 Passed as index to self.history for convenience.
1211 status_only : bool (default: True)
1211 status_only : bool (default: True)
1212 if False:
1212 if False:
1213 Retrieve the actual results of completed tasks.
1213 Retrieve the actual results of completed tasks.
1214
1214
1215 Returns
1215 Returns
1216 -------
1216 -------
1217
1217
1218 results : dict
1218 results : dict
1219 There will always be the keys 'pending' and 'completed', which will
1219 There will always be the keys 'pending' and 'completed', which will
1220 be lists of msg_ids that are incomplete or complete. If `status_only`
1220 be lists of msg_ids that are incomplete or complete. If `status_only`
1221 is False, then completed results will be keyed by their `msg_id`.
1221 is False, then completed results will be keyed by their `msg_id`.
1222 """
1222 """
1223 if not isinstance(msg_ids, (list,tuple)):
1223 if not isinstance(msg_ids, (list,tuple)):
1224 msg_ids = [msg_ids]
1224 msg_ids = [msg_ids]
1225
1225
1226 theids = []
1226 theids = []
1227 for msg_id in msg_ids:
1227 for msg_id in msg_ids:
1228 if isinstance(msg_id, int):
1228 if isinstance(msg_id, int):
1229 msg_id = self.history[msg_id]
1229 msg_id = self.history[msg_id]
1230 if not isinstance(msg_id, basestring):
1230 if not isinstance(msg_id, basestring):
1231 raise TypeError("msg_ids must be str, not %r"%msg_id)
1231 raise TypeError("msg_ids must be str, not %r"%msg_id)
1232 theids.append(msg_id)
1232 theids.append(msg_id)
1233
1233
1234 completed = []
1234 completed = []
1235 local_results = {}
1235 local_results = {}
1236
1236
1237 # comment this block out to temporarily disable local shortcut:
1237 # comment this block out to temporarily disable local shortcut:
1238 for msg_id in theids:
1238 for msg_id in theids:
1239 if msg_id in self.results:
1239 if msg_id in self.results:
1240 completed.append(msg_id)
1240 completed.append(msg_id)
1241 local_results[msg_id] = self.results[msg_id]
1241 local_results[msg_id] = self.results[msg_id]
1242 theids.remove(msg_id)
1242 theids.remove(msg_id)
1243
1243
1244 if theids: # some not locally cached
1244 if theids: # some not locally cached
1245 content = dict(msg_ids=theids, status_only=status_only)
1245 content = dict(msg_ids=theids, status_only=status_only)
1246 msg = self.session.send(self._query_socket, "result_request", content=content)
1246 msg = self.session.send(self._query_socket, "result_request", content=content)
1247 zmq.select([self._query_socket], [], [])
1247 zmq.select([self._query_socket], [], [])
1248 idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
1248 idents,msg = self.session.recv(self._query_socket, zmq.NOBLOCK)
1249 if self.debug:
1249 if self.debug:
1250 pprint(msg)
1250 pprint(msg)
1251 content = msg['content']
1251 content = msg['content']
1252 if content['status'] != 'ok':
1252 if content['status'] != 'ok':
1253 raise self._unwrap_exception(content)
1253 raise self._unwrap_exception(content)
1254 buffers = msg['buffers']
1254 buffers = msg['buffers']
1255 else:
1255 else:
1256 content = dict(completed=[],pending=[])
1256 content = dict(completed=[],pending=[])
1257
1257
1258 content['completed'].extend(completed)
1258 content['completed'].extend(completed)
1259
1259
1260 if status_only:
1260 if status_only:
1261 return content
1261 return content
1262
1262
1263 failures = []
1263 failures = []
1264 # load cached results into result:
1264 # load cached results into result:
1265 content.update(local_results)
1265 content.update(local_results)
1266
1266
1267 # update cache with results:
1267 # update cache with results:
1268 for msg_id in sorted(theids):
1268 for msg_id in sorted(theids):
1269 if msg_id in content['completed']:
1269 if msg_id in content['completed']:
1270 rec = content[msg_id]
1270 rec = content[msg_id]
1271 parent = rec['header']
1271 parent = rec['header']
1272 header = rec['result_header']
1272 header = rec['result_header']
1273 rcontent = rec['result_content']
1273 rcontent = rec['result_content']
1274 iodict = rec['io']
1274 iodict = rec['io']
1275 if isinstance(rcontent, str):
1275 if isinstance(rcontent, str):
1276 rcontent = self.session.unpack(rcontent)
1276 rcontent = self.session.unpack(rcontent)
1277
1277
1278 md = self.metadata[msg_id]
1278 md = self.metadata[msg_id]
1279 md.update(self._extract_metadata(header, parent, rcontent))
1279 md.update(self._extract_metadata(header, parent, rcontent))
1280 if rec.get('received'):
1281 md['received'] = rec['received']
1280 md.update(iodict)
1282 md.update(iodict)
1281
1283
1282 if rcontent['status'] == 'ok':
1284 if rcontent['status'] == 'ok':
1283 res,buffers = util.unserialize_object(buffers)
1285 res,buffers = util.unserialize_object(buffers)
1284 else:
1286 else:
1285 print rcontent
1287 print rcontent
1286 res = self._unwrap_exception(rcontent)
1288 res = self._unwrap_exception(rcontent)
1287 failures.append(res)
1289 failures.append(res)
1288
1290
1289 self.results[msg_id] = res
1291 self.results[msg_id] = res
1290 content[msg_id] = res
1292 content[msg_id] = res
1291
1293
1292 if len(theids) == 1 and failures:
1294 if len(theids) == 1 and failures:
1293 raise failures[0]
1295 raise failures[0]
1294
1296
1295 error.collect_exceptions(failures, "result_status")
1297 error.collect_exceptions(failures, "result_status")
1296 return content
1298 return content
1297
1299
1298 @spin_first
1300 @spin_first
1299 def queue_status(self, targets='all', verbose=False):
1301 def queue_status(self, targets='all', verbose=False):
1300 """Fetch the status of engine queues.
1302 """Fetch the status of engine queues.
1301
1303
1302 Parameters
1304 Parameters
1303 ----------
1305 ----------
1304
1306
1305 targets : int/str/list of ints/strs
1307 targets : int/str/list of ints/strs
1306 the engines whose states are to be queried.
1308 the engines whose states are to be queried.
1307 default : all
1309 default : all
1308 verbose : bool
1310 verbose : bool
1309 Whether to return lengths only, or lists of ids for each element
1311 Whether to return lengths only, or lists of ids for each element
1310 """
1312 """
1311 if targets == 'all':
1313 if targets == 'all':
1312 # allow 'all' to be evaluated on the engine
1314 # allow 'all' to be evaluated on the engine
1313 engine_ids = None
1315 engine_ids = None
1314 else:
1316 else:
1315 engine_ids = self._build_targets(targets)[1]
1317 engine_ids = self._build_targets(targets)[1]
1316 content = dict(targets=engine_ids, verbose=verbose)
1318 content = dict(targets=engine_ids, verbose=verbose)
1317 self.session.send(self._query_socket, "queue_request", content=content)
1319 self.session.send(self._query_socket, "queue_request", content=content)
1318 idents,msg = self.session.recv(self._query_socket, 0)
1320 idents,msg = self.session.recv(self._query_socket, 0)
1319 if self.debug:
1321 if self.debug:
1320 pprint(msg)
1322 pprint(msg)
1321 content = msg['content']
1323 content = msg['content']
1322 status = content.pop('status')
1324 status = content.pop('status')
1323 if status != 'ok':
1325 if status != 'ok':
1324 raise self._unwrap_exception(content)
1326 raise self._unwrap_exception(content)
1325 content = rekey(content)
1327 content = rekey(content)
1326 if isinstance(targets, int):
1328 if isinstance(targets, int):
1327 return content[targets]
1329 return content[targets]
1328 else:
1330 else:
1329 return content
1331 return content
1330
1332
1331 @spin_first
1333 @spin_first
1332 def purge_results(self, jobs=[], targets=[]):
1334 def purge_results(self, jobs=[], targets=[]):
1333 """Tell the Hub to forget results.
1335 """Tell the Hub to forget results.
1334
1336
1335 Individual results can be purged by msg_id, or the entire
1337 Individual results can be purged by msg_id, or the entire
1336 history of specific targets can be purged.
1338 history of specific targets can be purged.
1337
1339
1338 Use `purge_results('all')` to scrub everything from the Hub's db.
1340 Use `purge_results('all')` to scrub everything from the Hub's db.
1339
1341
1340 Parameters
1342 Parameters
1341 ----------
1343 ----------
1342
1344
1343 jobs : str or list of str or AsyncResult objects
1345 jobs : str or list of str or AsyncResult objects
1344 the msg_ids whose results should be forgotten.
1346 the msg_ids whose results should be forgotten.
1345 targets : int/str/list of ints/strs
1347 targets : int/str/list of ints/strs
1346 The targets, by int_id, whose entire history is to be purged.
1348 The targets, by int_id, whose entire history is to be purged.
1347
1349
1348 default : None
1350 default : None
1349 """
1351 """
1350 if not targets and not jobs:
1352 if not targets and not jobs:
1351 raise ValueError("Must specify at least one of `targets` and `jobs`")
1353 raise ValueError("Must specify at least one of `targets` and `jobs`")
1352 if targets:
1354 if targets:
1353 targets = self._build_targets(targets)[1]
1355 targets = self._build_targets(targets)[1]
1354
1356
1355 # construct msg_ids from jobs
1357 # construct msg_ids from jobs
1356 if jobs == 'all':
1358 if jobs == 'all':
1357 msg_ids = jobs
1359 msg_ids = jobs
1358 else:
1360 else:
1359 msg_ids = []
1361 msg_ids = []
1360 if isinstance(jobs, (basestring,AsyncResult)):
1362 if isinstance(jobs, (basestring,AsyncResult)):
1361 jobs = [jobs]
1363 jobs = [jobs]
1362 bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
1364 bad_ids = filter(lambda obj: not isinstance(obj, (basestring, AsyncResult)), jobs)
1363 if bad_ids:
1365 if bad_ids:
1364 raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
1366 raise TypeError("Invalid msg_id type %r, expected str or AsyncResult"%bad_ids[0])
1365 for j in jobs:
1367 for j in jobs:
1366 if isinstance(j, AsyncResult):
1368 if isinstance(j, AsyncResult):
1367 msg_ids.extend(j.msg_ids)
1369 msg_ids.extend(j.msg_ids)
1368 else:
1370 else:
1369 msg_ids.append(j)
1371 msg_ids.append(j)
1370
1372
1371 content = dict(engine_ids=targets, msg_ids=msg_ids)
1373 content = dict(engine_ids=targets, msg_ids=msg_ids)
1372 self.session.send(self._query_socket, "purge_request", content=content)
1374 self.session.send(self._query_socket, "purge_request", content=content)
1373 idents, msg = self.session.recv(self._query_socket, 0)
1375 idents, msg = self.session.recv(self._query_socket, 0)
1374 if self.debug:
1376 if self.debug:
1375 pprint(msg)
1377 pprint(msg)
1376 content = msg['content']
1378 content = msg['content']
1377 if content['status'] != 'ok':
1379 if content['status'] != 'ok':
1378 raise self._unwrap_exception(content)
1380 raise self._unwrap_exception(content)
1379
1381
1380 @spin_first
1382 @spin_first
1381 def hub_history(self):
1383 def hub_history(self):
1382 """Get the Hub's history
1384 """Get the Hub's history
1383
1385
1384 Just like the Client, the Hub has a history, which is a list of msg_ids.
1386 Just like the Client, the Hub has a history, which is a list of msg_ids.
1385 This will contain the history of all clients, and, depending on configuration,
1387 This will contain the history of all clients, and, depending on configuration,
1386 may contain history across multiple cluster sessions.
1388 may contain history across multiple cluster sessions.
1387
1389
1388 Any msg_id returned here is a valid argument to `get_result`.
1390 Any msg_id returned here is a valid argument to `get_result`.
1389
1391
1390 Returns
1392 Returns
1391 -------
1393 -------
1392
1394
1393 msg_ids : list of strs
1395 msg_ids : list of strs
1394 list of all msg_ids, ordered by task submission time.
1396 list of all msg_ids, ordered by task submission time.
1395 """
1397 """
1396
1398
1397 self.session.send(self._query_socket, "history_request", content={})
1399 self.session.send(self._query_socket, "history_request", content={})
1398 idents, msg = self.session.recv(self._query_socket, 0)
1400 idents, msg = self.session.recv(self._query_socket, 0)
1399
1401
1400 if self.debug:
1402 if self.debug:
1401 pprint(msg)
1403 pprint(msg)
1402 content = msg['content']
1404 content = msg['content']
1403 if content['status'] != 'ok':
1405 if content['status'] != 'ok':
1404 raise self._unwrap_exception(content)
1406 raise self._unwrap_exception(content)
1405 else:
1407 else:
1406 return content['history']
1408 return content['history']
1407
1409
1408 @spin_first
1410 @spin_first
1409 def db_query(self, query, keys=None):
1411 def db_query(self, query, keys=None):
1410 """Query the Hub's TaskRecord database
1412 """Query the Hub's TaskRecord database
1411
1413
1412 This will return a list of task record dicts that match `query`
1414 This will return a list of task record dicts that match `query`
1413
1415
1414 Parameters
1416 Parameters
1415 ----------
1417 ----------
1416
1418
1417 query : mongodb query dict
1419 query : mongodb query dict
1418 The search dict. See mongodb query docs for details.
1420 The search dict. See mongodb query docs for details.
1419 keys : list of strs [optional]
1421 keys : list of strs [optional]
1420 The subset of keys to be returned. The default is to fetch everything but buffers.
1422 The subset of keys to be returned. The default is to fetch everything but buffers.
1421 'msg_id' will *always* be included.
1423 'msg_id' will *always* be included.
1422 """
1424 """
1423 if isinstance(keys, basestring):
1425 if isinstance(keys, basestring):
1424 keys = [keys]
1426 keys = [keys]
1425 content = dict(query=query, keys=keys)
1427 content = dict(query=query, keys=keys)
1426 self.session.send(self._query_socket, "db_request", content=content)
1428 self.session.send(self._query_socket, "db_request", content=content)
1427 idents, msg = self.session.recv(self._query_socket, 0)
1429 idents, msg = self.session.recv(self._query_socket, 0)
1428 if self.debug:
1430 if self.debug:
1429 pprint(msg)
1431 pprint(msg)
1430 content = msg['content']
1432 content = msg['content']
1431 if content['status'] != 'ok':
1433 if content['status'] != 'ok':
1432 raise self._unwrap_exception(content)
1434 raise self._unwrap_exception(content)
1433
1435
1434 records = content['records']
1436 records = content['records']
1435
1437
1436 buffer_lens = content['buffer_lens']
1438 buffer_lens = content['buffer_lens']
1437 result_buffer_lens = content['result_buffer_lens']
1439 result_buffer_lens = content['result_buffer_lens']
1438 buffers = msg['buffers']
1440 buffers = msg['buffers']
1439 has_bufs = buffer_lens is not None
1441 has_bufs = buffer_lens is not None
1440 has_rbufs = result_buffer_lens is not None
1442 has_rbufs = result_buffer_lens is not None
1441 for i,rec in enumerate(records):
1443 for i,rec in enumerate(records):
1442 # relink buffers
1444 # relink buffers
1443 if has_bufs:
1445 if has_bufs:
1444 blen = buffer_lens[i]
1446 blen = buffer_lens[i]
1445 rec['buffers'], buffers = buffers[:blen],buffers[blen:]
1447 rec['buffers'], buffers = buffers[:blen],buffers[blen:]
1446 if has_rbufs:
1448 if has_rbufs:
1447 blen = result_buffer_lens[i]
1449 blen = result_buffer_lens[i]
1448 rec['result_buffers'], buffers = buffers[:blen],buffers[blen:]
1450 rec['result_buffers'], buffers = buffers[:blen],buffers[blen:]
1449
1451
1450 return records
1452 return records
1451
1453
1452 __all__ = [ 'Client' ]
1454 __all__ = [ 'Client' ]
@@ -1,1293 +1,1298 b''
1 """The IPython Controller Hub with 0MQ
1 """The IPython Controller Hub with 0MQ
2 This is the master object that handles connections from engines and clients,
2 This is the master object that handles connections from engines and clients,
3 and monitors traffic through the various queues.
3 and monitors traffic through the various queues.
4
4
5 Authors:
5 Authors:
6
6
7 * Min RK
7 * Min RK
8 """
8 """
9 #-----------------------------------------------------------------------------
9 #-----------------------------------------------------------------------------
10 # Copyright (C) 2010-2011 The IPython Development Team
10 # Copyright (C) 2010-2011 The IPython Development Team
11 #
11 #
12 # Distributed under the terms of the BSD License. The full license is in
12 # Distributed under the terms of the BSD License. The full license is in
13 # the file COPYING, distributed as part of this software.
13 # the file COPYING, distributed as part of this software.
14 #-----------------------------------------------------------------------------
14 #-----------------------------------------------------------------------------
15
15
16 #-----------------------------------------------------------------------------
16 #-----------------------------------------------------------------------------
17 # Imports
17 # Imports
18 #-----------------------------------------------------------------------------
18 #-----------------------------------------------------------------------------
19 from __future__ import print_function
19 from __future__ import print_function
20
20
21 import sys
21 import sys
22 import time
22 import time
23 from datetime import datetime
23 from datetime import datetime
24
24
25 import zmq
25 import zmq
26 from zmq.eventloop import ioloop
26 from zmq.eventloop import ioloop
27 from zmq.eventloop.zmqstream import ZMQStream
27 from zmq.eventloop.zmqstream import ZMQStream
28
28
29 # internal:
29 # internal:
30 from IPython.utils.importstring import import_item
30 from IPython.utils.importstring import import_item
31 from IPython.utils.traitlets import (
31 from IPython.utils.traitlets import (
32 HasTraits, Instance, Integer, Unicode, Dict, Set, Tuple, CBytes, DottedObjectName
32 HasTraits, Instance, Integer, Unicode, Dict, Set, Tuple, CBytes, DottedObjectName
33 )
33 )
34
34
35 from IPython.parallel import error, util
35 from IPython.parallel import error, util
36 from IPython.parallel.factory import RegistrationFactory
36 from IPython.parallel.factory import RegistrationFactory
37
37
38 from IPython.zmq.session import SessionFactory
38 from IPython.zmq.session import SessionFactory
39
39
40 from .heartmonitor import HeartMonitor
40 from .heartmonitor import HeartMonitor
41
41
42 #-----------------------------------------------------------------------------
42 #-----------------------------------------------------------------------------
43 # Code
43 # Code
44 #-----------------------------------------------------------------------------
44 #-----------------------------------------------------------------------------
45
45
46 def _passer(*args, **kwargs):
46 def _passer(*args, **kwargs):
47 return
47 return
48
48
49 def _printer(*args, **kwargs):
49 def _printer(*args, **kwargs):
50 print (args)
50 print (args)
51 print (kwargs)
51 print (kwargs)
52
52
53 def empty_record():
53 def empty_record():
54 """Return an empty dict with all record keys."""
54 """Return an empty dict with all record keys."""
55 return {
55 return {
56 'msg_id' : None,
56 'msg_id' : None,
57 'header' : None,
57 'header' : None,
58 'content': None,
58 'content': None,
59 'buffers': None,
59 'buffers': None,
60 'submitted': None,
60 'submitted': None,
61 'client_uuid' : None,
61 'client_uuid' : None,
62 'engine_uuid' : None,
62 'engine_uuid' : None,
63 'started': None,
63 'started': None,
64 'completed': None,
64 'completed': None,
65 'resubmitted': None,
65 'resubmitted': None,
66 'received': None,
66 'result_header' : None,
67 'result_header' : None,
67 'result_content' : None,
68 'result_content' : None,
68 'result_buffers' : None,
69 'result_buffers' : None,
69 'queue' : None,
70 'queue' : None,
70 'pyin' : None,
71 'pyin' : None,
71 'pyout': None,
72 'pyout': None,
72 'pyerr': None,
73 'pyerr': None,
73 'stdout': '',
74 'stdout': '',
74 'stderr': '',
75 'stderr': '',
75 }
76 }
76
77
77 def init_record(msg):
78 def init_record(msg):
78 """Initialize a TaskRecord based on a request."""
79 """Initialize a TaskRecord based on a request."""
79 header = msg['header']
80 header = msg['header']
80 return {
81 return {
81 'msg_id' : header['msg_id'],
82 'msg_id' : header['msg_id'],
82 'header' : header,
83 'header' : header,
83 'content': msg['content'],
84 'content': msg['content'],
84 'buffers': msg['buffers'],
85 'buffers': msg['buffers'],
85 'submitted': header['date'],
86 'submitted': header['date'],
86 'client_uuid' : None,
87 'client_uuid' : None,
87 'engine_uuid' : None,
88 'engine_uuid' : None,
88 'started': None,
89 'started': None,
89 'completed': None,
90 'completed': None,
90 'resubmitted': None,
91 'resubmitted': None,
92 'received': None,
91 'result_header' : None,
93 'result_header' : None,
92 'result_content' : None,
94 'result_content' : None,
93 'result_buffers' : None,
95 'result_buffers' : None,
94 'queue' : None,
96 'queue' : None,
95 'pyin' : None,
97 'pyin' : None,
96 'pyout': None,
98 'pyout': None,
97 'pyerr': None,
99 'pyerr': None,
98 'stdout': '',
100 'stdout': '',
99 'stderr': '',
101 'stderr': '',
100 }
102 }
101
103
102
104
103 class EngineConnector(HasTraits):
105 class EngineConnector(HasTraits):
104 """A simple object for accessing the various zmq connections of an object.
106 """A simple object for accessing the various zmq connections of an object.
105 Attributes are:
107 Attributes are:
106 id (int): engine ID
108 id (int): engine ID
107 uuid (str): uuid (unused?)
109 uuid (str): uuid (unused?)
108 queue (str): identity of queue's XREQ socket
110 queue (str): identity of queue's XREQ socket
109 registration (str): identity of registration XREQ socket
111 registration (str): identity of registration XREQ socket
110 heartbeat (str): identity of heartbeat XREQ socket
112 heartbeat (str): identity of heartbeat XREQ socket
111 """
113 """
112 id=Integer(0)
114 id=Integer(0)
113 queue=CBytes()
115 queue=CBytes()
114 control=CBytes()
116 control=CBytes()
115 registration=CBytes()
117 registration=CBytes()
116 heartbeat=CBytes()
118 heartbeat=CBytes()
117 pending=Set()
119 pending=Set()
118
120
119 class HubFactory(RegistrationFactory):
121 class HubFactory(RegistrationFactory):
120 """The Configurable for setting up a Hub."""
122 """The Configurable for setting up a Hub."""
121
123
122 # port-pairs for monitoredqueues:
124 # port-pairs for monitoredqueues:
123 hb = Tuple(Integer,Integer,config=True,
125 hb = Tuple(Integer,Integer,config=True,
124 help="""XREQ/SUB Port pair for Engine heartbeats""")
126 help="""XREQ/SUB Port pair for Engine heartbeats""")
125 def _hb_default(self):
127 def _hb_default(self):
126 return tuple(util.select_random_ports(2))
128 return tuple(util.select_random_ports(2))
127
129
128 mux = Tuple(Integer,Integer,config=True,
130 mux = Tuple(Integer,Integer,config=True,
129 help="""Engine/Client Port pair for MUX queue""")
131 help="""Engine/Client Port pair for MUX queue""")
130
132
131 def _mux_default(self):
133 def _mux_default(self):
132 return tuple(util.select_random_ports(2))
134 return tuple(util.select_random_ports(2))
133
135
134 task = Tuple(Integer,Integer,config=True,
136 task = Tuple(Integer,Integer,config=True,
135 help="""Engine/Client Port pair for Task queue""")
137 help="""Engine/Client Port pair for Task queue""")
136 def _task_default(self):
138 def _task_default(self):
137 return tuple(util.select_random_ports(2))
139 return tuple(util.select_random_ports(2))
138
140
139 control = Tuple(Integer,Integer,config=True,
141 control = Tuple(Integer,Integer,config=True,
140 help="""Engine/Client Port pair for Control queue""")
142 help="""Engine/Client Port pair for Control queue""")
141
143
142 def _control_default(self):
144 def _control_default(self):
143 return tuple(util.select_random_ports(2))
145 return tuple(util.select_random_ports(2))
144
146
145 iopub = Tuple(Integer,Integer,config=True,
147 iopub = Tuple(Integer,Integer,config=True,
146 help="""Engine/Client Port pair for IOPub relay""")
148 help="""Engine/Client Port pair for IOPub relay""")
147
149
148 def _iopub_default(self):
150 def _iopub_default(self):
149 return tuple(util.select_random_ports(2))
151 return tuple(util.select_random_ports(2))
150
152
151 # single ports:
153 # single ports:
152 mon_port = Integer(config=True,
154 mon_port = Integer(config=True,
153 help="""Monitor (SUB) port for queue traffic""")
155 help="""Monitor (SUB) port for queue traffic""")
154
156
155 def _mon_port_default(self):
157 def _mon_port_default(self):
156 return util.select_random_ports(1)[0]
158 return util.select_random_ports(1)[0]
157
159
158 notifier_port = Integer(config=True,
160 notifier_port = Integer(config=True,
159 help="""PUB port for sending engine status notifications""")
161 help="""PUB port for sending engine status notifications""")
160
162
161 def _notifier_port_default(self):
163 def _notifier_port_default(self):
162 return util.select_random_ports(1)[0]
164 return util.select_random_ports(1)[0]
163
165
164 engine_ip = Unicode('127.0.0.1', config=True,
166 engine_ip = Unicode('127.0.0.1', config=True,
165 help="IP on which to listen for engine connections. [default: loopback]")
167 help="IP on which to listen for engine connections. [default: loopback]")
166 engine_transport = Unicode('tcp', config=True,
168 engine_transport = Unicode('tcp', config=True,
167 help="0MQ transport for engine connections. [default: tcp]")
169 help="0MQ transport for engine connections. [default: tcp]")
168
170
169 client_ip = Unicode('127.0.0.1', config=True,
171 client_ip = Unicode('127.0.0.1', config=True,
170 help="IP on which to listen for client connections. [default: loopback]")
172 help="IP on which to listen for client connections. [default: loopback]")
171 client_transport = Unicode('tcp', config=True,
173 client_transport = Unicode('tcp', config=True,
172 help="0MQ transport for client connections. [default : tcp]")
174 help="0MQ transport for client connections. [default : tcp]")
173
175
174 monitor_ip = Unicode('127.0.0.1', config=True,
176 monitor_ip = Unicode('127.0.0.1', config=True,
175 help="IP on which to listen for monitor messages. [default: loopback]")
177 help="IP on which to listen for monitor messages. [default: loopback]")
176 monitor_transport = Unicode('tcp', config=True,
178 monitor_transport = Unicode('tcp', config=True,
177 help="0MQ transport for monitor messages. [default : tcp]")
179 help="0MQ transport for monitor messages. [default : tcp]")
178
180
179 monitor_url = Unicode('')
181 monitor_url = Unicode('')
180
182
181 db_class = DottedObjectName('IPython.parallel.controller.dictdb.DictDB',
183 db_class = DottedObjectName('IPython.parallel.controller.dictdb.DictDB',
182 config=True, help="""The class to use for the DB backend""")
184 config=True, help="""The class to use for the DB backend""")
183
185
184 # not configurable
186 # not configurable
185 db = Instance('IPython.parallel.controller.dictdb.BaseDB')
187 db = Instance('IPython.parallel.controller.dictdb.BaseDB')
186 heartmonitor = Instance('IPython.parallel.controller.heartmonitor.HeartMonitor')
188 heartmonitor = Instance('IPython.parallel.controller.heartmonitor.HeartMonitor')
187
189
188 def _ip_changed(self, name, old, new):
190 def _ip_changed(self, name, old, new):
189 self.engine_ip = new
191 self.engine_ip = new
190 self.client_ip = new
192 self.client_ip = new
191 self.monitor_ip = new
193 self.monitor_ip = new
192 self._update_monitor_url()
194 self._update_monitor_url()
193
195
194 def _update_monitor_url(self):
196 def _update_monitor_url(self):
195 self.monitor_url = "%s://%s:%i" % (self.monitor_transport, self.monitor_ip, self.mon_port)
197 self.monitor_url = "%s://%s:%i" % (self.monitor_transport, self.monitor_ip, self.mon_port)
196
198
197 def _transport_changed(self, name, old, new):
199 def _transport_changed(self, name, old, new):
198 self.engine_transport = new
200 self.engine_transport = new
199 self.client_transport = new
201 self.client_transport = new
200 self.monitor_transport = new
202 self.monitor_transport = new
201 self._update_monitor_url()
203 self._update_monitor_url()
202
204
203 def __init__(self, **kwargs):
205 def __init__(self, **kwargs):
204 super(HubFactory, self).__init__(**kwargs)
206 super(HubFactory, self).__init__(**kwargs)
205 self._update_monitor_url()
207 self._update_monitor_url()
206
208
207
209
208 def construct(self):
210 def construct(self):
209 self.init_hub()
211 self.init_hub()
210
212
211 def start(self):
213 def start(self):
212 self.heartmonitor.start()
214 self.heartmonitor.start()
213 self.log.info("Heartmonitor started")
215 self.log.info("Heartmonitor started")
214
216
215 def init_hub(self):
217 def init_hub(self):
216 """construct"""
218 """construct"""
217 client_iface = "%s://%s:" % (self.client_transport, self.client_ip) + "%i"
219 client_iface = "%s://%s:" % (self.client_transport, self.client_ip) + "%i"
218 engine_iface = "%s://%s:" % (self.engine_transport, self.engine_ip) + "%i"
220 engine_iface = "%s://%s:" % (self.engine_transport, self.engine_ip) + "%i"
219
221
220 ctx = self.context
222 ctx = self.context
221 loop = self.loop
223 loop = self.loop
222
224
223 # Registrar socket
225 # Registrar socket
224 q = ZMQStream(ctx.socket(zmq.ROUTER), loop)
226 q = ZMQStream(ctx.socket(zmq.ROUTER), loop)
225 q.bind(client_iface % self.regport)
227 q.bind(client_iface % self.regport)
226 self.log.info("Hub listening on %s for registration.", client_iface % self.regport)
228 self.log.info("Hub listening on %s for registration.", client_iface % self.regport)
227 if self.client_ip != self.engine_ip:
229 if self.client_ip != self.engine_ip:
228 q.bind(engine_iface % self.regport)
230 q.bind(engine_iface % self.regport)
229 self.log.info("Hub listening on %s for registration.", engine_iface % self.regport)
231 self.log.info("Hub listening on %s for registration.", engine_iface % self.regport)
230
232
231 ### Engine connections ###
233 ### Engine connections ###
232
234
233 # heartbeat
235 # heartbeat
234 hpub = ctx.socket(zmq.PUB)
236 hpub = ctx.socket(zmq.PUB)
235 hpub.bind(engine_iface % self.hb[0])
237 hpub.bind(engine_iface % self.hb[0])
236 hrep = ctx.socket(zmq.ROUTER)
238 hrep = ctx.socket(zmq.ROUTER)
237 hrep.bind(engine_iface % self.hb[1])
239 hrep.bind(engine_iface % self.hb[1])
238 self.heartmonitor = HeartMonitor(loop=loop, config=self.config, log=self.log,
240 self.heartmonitor = HeartMonitor(loop=loop, config=self.config, log=self.log,
239 pingstream=ZMQStream(hpub,loop),
241 pingstream=ZMQStream(hpub,loop),
240 pongstream=ZMQStream(hrep,loop)
242 pongstream=ZMQStream(hrep,loop)
241 )
243 )
242
244
243 ### Client connections ###
245 ### Client connections ###
244 # Notifier socket
246 # Notifier socket
245 n = ZMQStream(ctx.socket(zmq.PUB), loop)
247 n = ZMQStream(ctx.socket(zmq.PUB), loop)
246 n.bind(client_iface%self.notifier_port)
248 n.bind(client_iface%self.notifier_port)
247
249
248 ### build and launch the queues ###
250 ### build and launch the queues ###
249
251
250 # monitor socket
252 # monitor socket
251 sub = ctx.socket(zmq.SUB)
253 sub = ctx.socket(zmq.SUB)
252 sub.setsockopt(zmq.SUBSCRIBE, b"")
254 sub.setsockopt(zmq.SUBSCRIBE, b"")
253 sub.bind(self.monitor_url)
255 sub.bind(self.monitor_url)
254 sub.bind('inproc://monitor')
256 sub.bind('inproc://monitor')
255 sub = ZMQStream(sub, loop)
257 sub = ZMQStream(sub, loop)
256
258
257 # connect the db
259 # connect the db
258 self.log.info('Hub using DB backend: %r'%(self.db_class.split()[-1]))
260 self.log.info('Hub using DB backend: %r'%(self.db_class.split()[-1]))
259 # cdir = self.config.Global.cluster_dir
261 # cdir = self.config.Global.cluster_dir
260 self.db = import_item(str(self.db_class))(session=self.session.session,
262 self.db = import_item(str(self.db_class))(session=self.session.session,
261 config=self.config, log=self.log)
263 config=self.config, log=self.log)
262 time.sleep(.25)
264 time.sleep(.25)
263 try:
265 try:
264 scheme = self.config.TaskScheduler.scheme_name
266 scheme = self.config.TaskScheduler.scheme_name
265 except AttributeError:
267 except AttributeError:
266 from .scheduler import TaskScheduler
268 from .scheduler import TaskScheduler
267 scheme = TaskScheduler.scheme_name.get_default_value()
269 scheme = TaskScheduler.scheme_name.get_default_value()
268 # build connection dicts
270 # build connection dicts
269 self.engine_info = {
271 self.engine_info = {
270 'control' : engine_iface%self.control[1],
272 'control' : engine_iface%self.control[1],
271 'mux': engine_iface%self.mux[1],
273 'mux': engine_iface%self.mux[1],
272 'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]),
274 'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]),
273 'task' : engine_iface%self.task[1],
275 'task' : engine_iface%self.task[1],
274 'iopub' : engine_iface%self.iopub[1],
276 'iopub' : engine_iface%self.iopub[1],
275 # 'monitor' : engine_iface%self.mon_port,
277 # 'monitor' : engine_iface%self.mon_port,
276 }
278 }
277
279
278 self.client_info = {
280 self.client_info = {
279 'control' : client_iface%self.control[0],
281 'control' : client_iface%self.control[0],
280 'mux': client_iface%self.mux[0],
282 'mux': client_iface%self.mux[0],
281 'task' : (scheme, client_iface%self.task[0]),
283 'task' : (scheme, client_iface%self.task[0]),
282 'iopub' : client_iface%self.iopub[0],
284 'iopub' : client_iface%self.iopub[0],
283 'notification': client_iface%self.notifier_port
285 'notification': client_iface%self.notifier_port
284 }
286 }
285 self.log.debug("Hub engine addrs: %s", self.engine_info)
287 self.log.debug("Hub engine addrs: %s", self.engine_info)
286 self.log.debug("Hub client addrs: %s", self.client_info)
288 self.log.debug("Hub client addrs: %s", self.client_info)
287
289
288 # resubmit stream
290 # resubmit stream
289 r = ZMQStream(ctx.socket(zmq.DEALER), loop)
291 r = ZMQStream(ctx.socket(zmq.DEALER), loop)
290 url = util.disambiguate_url(self.client_info['task'][-1])
292 url = util.disambiguate_url(self.client_info['task'][-1])
291 r.setsockopt(zmq.IDENTITY, self.session.bsession)
293 r.setsockopt(zmq.IDENTITY, self.session.bsession)
292 r.connect(url)
294 r.connect(url)
293
295
294 self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
296 self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
295 query=q, notifier=n, resubmit=r, db=self.db,
297 query=q, notifier=n, resubmit=r, db=self.db,
296 engine_info=self.engine_info, client_info=self.client_info,
298 engine_info=self.engine_info, client_info=self.client_info,
297 log=self.log)
299 log=self.log)
298
300
299
301
300 class Hub(SessionFactory):
302 class Hub(SessionFactory):
301 """The IPython Controller Hub with 0MQ connections
303 """The IPython Controller Hub with 0MQ connections
302
304
303 Parameters
305 Parameters
304 ==========
306 ==========
305 loop: zmq IOLoop instance
307 loop: zmq IOLoop instance
306 session: Session object
308 session: Session object
307 <removed> context: zmq context for creating new connections (?)
309 <removed> context: zmq context for creating new connections (?)
308 queue: ZMQStream for monitoring the command queue (SUB)
310 queue: ZMQStream for monitoring the command queue (SUB)
309 query: ZMQStream for engine registration and client queries requests (XREP)
311 query: ZMQStream for engine registration and client queries requests (XREP)
310 heartbeat: HeartMonitor object checking the pulse of the engines
312 heartbeat: HeartMonitor object checking the pulse of the engines
311 notifier: ZMQStream for broadcasting engine registration changes (PUB)
313 notifier: ZMQStream for broadcasting engine registration changes (PUB)
312 db: connection to db for out of memory logging of commands
314 db: connection to db for out of memory logging of commands
313 NotImplemented
315 NotImplemented
314 engine_info: dict of zmq connection information for engines to connect
316 engine_info: dict of zmq connection information for engines to connect
315 to the queues.
317 to the queues.
316 client_info: dict of zmq connection information for engines to connect
318 client_info: dict of zmq connection information for engines to connect
317 to the queues.
319 to the queues.
318 """
320 """
319 # internal data structures:
321 # internal data structures:
320 ids=Set() # engine IDs
322 ids=Set() # engine IDs
321 keytable=Dict()
323 keytable=Dict()
322 by_ident=Dict()
324 by_ident=Dict()
323 engines=Dict()
325 engines=Dict()
324 clients=Dict()
326 clients=Dict()
325 hearts=Dict()
327 hearts=Dict()
326 pending=Set()
328 pending=Set()
327 queues=Dict() # pending msg_ids keyed by engine_id
329 queues=Dict() # pending msg_ids keyed by engine_id
328 tasks=Dict() # pending msg_ids submitted as tasks, keyed by client_id
330 tasks=Dict() # pending msg_ids submitted as tasks, keyed by client_id
329 completed=Dict() # completed msg_ids keyed by engine_id
331 completed=Dict() # completed msg_ids keyed by engine_id
330 all_completed=Set() # completed msg_ids keyed by engine_id
332 all_completed=Set() # completed msg_ids keyed by engine_id
331 dead_engines=Set() # completed msg_ids keyed by engine_id
333 dead_engines=Set() # completed msg_ids keyed by engine_id
332 unassigned=Set() # set of task msg_ds not yet assigned a destination
334 unassigned=Set() # set of task msg_ds not yet assigned a destination
333 incoming_registrations=Dict()
335 incoming_registrations=Dict()
334 registration_timeout=Integer()
336 registration_timeout=Integer()
335 _idcounter=Integer(0)
337 _idcounter=Integer(0)
336
338
337 # objects from constructor:
339 # objects from constructor:
338 query=Instance(ZMQStream)
340 query=Instance(ZMQStream)
339 monitor=Instance(ZMQStream)
341 monitor=Instance(ZMQStream)
340 notifier=Instance(ZMQStream)
342 notifier=Instance(ZMQStream)
341 resubmit=Instance(ZMQStream)
343 resubmit=Instance(ZMQStream)
342 heartmonitor=Instance(HeartMonitor)
344 heartmonitor=Instance(HeartMonitor)
343 db=Instance(object)
345 db=Instance(object)
344 client_info=Dict()
346 client_info=Dict()
345 engine_info=Dict()
347 engine_info=Dict()
346
348
347
349
348 def __init__(self, **kwargs):
350 def __init__(self, **kwargs):
349 """
351 """
350 # universal:
352 # universal:
351 loop: IOLoop for creating future connections
353 loop: IOLoop for creating future connections
352 session: streamsession for sending serialized data
354 session: streamsession for sending serialized data
353 # engine:
355 # engine:
354 queue: ZMQStream for monitoring queue messages
356 queue: ZMQStream for monitoring queue messages
355 query: ZMQStream for engine+client registration and client requests
357 query: ZMQStream for engine+client registration and client requests
356 heartbeat: HeartMonitor object for tracking engines
358 heartbeat: HeartMonitor object for tracking engines
357 # extra:
359 # extra:
358 db: ZMQStream for db connection (NotImplemented)
360 db: ZMQStream for db connection (NotImplemented)
359 engine_info: zmq address/protocol dict for engine connections
361 engine_info: zmq address/protocol dict for engine connections
360 client_info: zmq address/protocol dict for client connections
362 client_info: zmq address/protocol dict for client connections
361 """
363 """
362
364
363 super(Hub, self).__init__(**kwargs)
365 super(Hub, self).__init__(**kwargs)
364 self.registration_timeout = max(5000, 2*self.heartmonitor.period)
366 self.registration_timeout = max(5000, 2*self.heartmonitor.period)
365
367
366 # validate connection dicts:
368 # validate connection dicts:
367 for k,v in self.client_info.iteritems():
369 for k,v in self.client_info.iteritems():
368 if k == 'task':
370 if k == 'task':
369 util.validate_url_container(v[1])
371 util.validate_url_container(v[1])
370 else:
372 else:
371 util.validate_url_container(v)
373 util.validate_url_container(v)
372 # util.validate_url_container(self.client_info)
374 # util.validate_url_container(self.client_info)
373 util.validate_url_container(self.engine_info)
375 util.validate_url_container(self.engine_info)
374
376
375 # register our callbacks
377 # register our callbacks
376 self.query.on_recv(self.dispatch_query)
378 self.query.on_recv(self.dispatch_query)
377 self.monitor.on_recv(self.dispatch_monitor_traffic)
379 self.monitor.on_recv(self.dispatch_monitor_traffic)
378
380
379 self.heartmonitor.add_heart_failure_handler(self.handle_heart_failure)
381 self.heartmonitor.add_heart_failure_handler(self.handle_heart_failure)
380 self.heartmonitor.add_new_heart_handler(self.handle_new_heart)
382 self.heartmonitor.add_new_heart_handler(self.handle_new_heart)
381
383
382 self.monitor_handlers = {b'in' : self.save_queue_request,
384 self.monitor_handlers = {b'in' : self.save_queue_request,
383 b'out': self.save_queue_result,
385 b'out': self.save_queue_result,
384 b'intask': self.save_task_request,
386 b'intask': self.save_task_request,
385 b'outtask': self.save_task_result,
387 b'outtask': self.save_task_result,
386 b'tracktask': self.save_task_destination,
388 b'tracktask': self.save_task_destination,
387 b'incontrol': _passer,
389 b'incontrol': _passer,
388 b'outcontrol': _passer,
390 b'outcontrol': _passer,
389 b'iopub': self.save_iopub_message,
391 b'iopub': self.save_iopub_message,
390 }
392 }
391
393
392 self.query_handlers = {'queue_request': self.queue_status,
394 self.query_handlers = {'queue_request': self.queue_status,
393 'result_request': self.get_results,
395 'result_request': self.get_results,
394 'history_request': self.get_history,
396 'history_request': self.get_history,
395 'db_request': self.db_query,
397 'db_request': self.db_query,
396 'purge_request': self.purge_results,
398 'purge_request': self.purge_results,
397 'load_request': self.check_load,
399 'load_request': self.check_load,
398 'resubmit_request': self.resubmit_task,
400 'resubmit_request': self.resubmit_task,
399 'shutdown_request': self.shutdown_request,
401 'shutdown_request': self.shutdown_request,
400 'registration_request' : self.register_engine,
402 'registration_request' : self.register_engine,
401 'unregistration_request' : self.unregister_engine,
403 'unregistration_request' : self.unregister_engine,
402 'connection_request': self.connection_request,
404 'connection_request': self.connection_request,
403 }
405 }
404
406
405 # ignore resubmit replies
407 # ignore resubmit replies
406 self.resubmit.on_recv(lambda msg: None, copy=False)
408 self.resubmit.on_recv(lambda msg: None, copy=False)
407
409
408 self.log.info("hub::created hub")
410 self.log.info("hub::created hub")
409
411
410 @property
412 @property
411 def _next_id(self):
413 def _next_id(self):
412 """gemerate a new ID.
414 """gemerate a new ID.
413
415
414 No longer reuse old ids, just count from 0."""
416 No longer reuse old ids, just count from 0."""
415 newid = self._idcounter
417 newid = self._idcounter
416 self._idcounter += 1
418 self._idcounter += 1
417 return newid
419 return newid
418 # newid = 0
420 # newid = 0
419 # incoming = [id[0] for id in self.incoming_registrations.itervalues()]
421 # incoming = [id[0] for id in self.incoming_registrations.itervalues()]
420 # # print newid, self.ids, self.incoming_registrations
422 # # print newid, self.ids, self.incoming_registrations
421 # while newid in self.ids or newid in incoming:
423 # while newid in self.ids or newid in incoming:
422 # newid += 1
424 # newid += 1
423 # return newid
425 # return newid
424
426
425 #-----------------------------------------------------------------------------
427 #-----------------------------------------------------------------------------
426 # message validation
428 # message validation
427 #-----------------------------------------------------------------------------
429 #-----------------------------------------------------------------------------
428
430
429 def _validate_targets(self, targets):
431 def _validate_targets(self, targets):
430 """turn any valid targets argument into a list of integer ids"""
432 """turn any valid targets argument into a list of integer ids"""
431 if targets is None:
433 if targets is None:
432 # default to all
434 # default to all
433 return self.ids
435 return self.ids
434
436
435 if isinstance(targets, (int,str,unicode)):
437 if isinstance(targets, (int,str,unicode)):
436 # only one target specified
438 # only one target specified
437 targets = [targets]
439 targets = [targets]
438 _targets = []
440 _targets = []
439 for t in targets:
441 for t in targets:
440 # map raw identities to ids
442 # map raw identities to ids
441 if isinstance(t, (str,unicode)):
443 if isinstance(t, (str,unicode)):
442 t = self.by_ident.get(t, t)
444 t = self.by_ident.get(t, t)
443 _targets.append(t)
445 _targets.append(t)
444 targets = _targets
446 targets = _targets
445 bad_targets = [ t for t in targets if t not in self.ids ]
447 bad_targets = [ t for t in targets if t not in self.ids ]
446 if bad_targets:
448 if bad_targets:
447 raise IndexError("No Such Engine: %r" % bad_targets)
449 raise IndexError("No Such Engine: %r" % bad_targets)
448 if not targets:
450 if not targets:
449 raise IndexError("No Engines Registered")
451 raise IndexError("No Engines Registered")
450 return targets
452 return targets
451
453
452 #-----------------------------------------------------------------------------
454 #-----------------------------------------------------------------------------
453 # dispatch methods (1 per stream)
455 # dispatch methods (1 per stream)
454 #-----------------------------------------------------------------------------
456 #-----------------------------------------------------------------------------
455
457
456
458
457 def dispatch_monitor_traffic(self, msg):
459 def dispatch_monitor_traffic(self, msg):
458 """all ME and Task queue messages come through here, as well as
460 """all ME and Task queue messages come through here, as well as
459 IOPub traffic."""
461 IOPub traffic."""
460 self.log.debug("monitor traffic: %r", msg[0])
462 self.log.debug("monitor traffic: %r", msg[0])
461 switch = msg[0]
463 switch = msg[0]
462 try:
464 try:
463 idents, msg = self.session.feed_identities(msg[1:])
465 idents, msg = self.session.feed_identities(msg[1:])
464 except ValueError:
466 except ValueError:
465 idents=[]
467 idents=[]
466 if not idents:
468 if not idents:
467 self.log.error("Bad Monitor Message: %r", msg)
469 self.log.error("Bad Monitor Message: %r", msg)
468 return
470 return
469 handler = self.monitor_handlers.get(switch, None)
471 handler = self.monitor_handlers.get(switch, None)
470 if handler is not None:
472 if handler is not None:
471 handler(idents, msg)
473 handler(idents, msg)
472 else:
474 else:
473 self.log.error("Invalid monitor topic: %r", switch)
475 self.log.error("Invalid monitor topic: %r", switch)
474
476
475
477
476 def dispatch_query(self, msg):
478 def dispatch_query(self, msg):
477 """Route registration requests and queries from clients."""
479 """Route registration requests and queries from clients."""
478 try:
480 try:
479 idents, msg = self.session.feed_identities(msg)
481 idents, msg = self.session.feed_identities(msg)
480 except ValueError:
482 except ValueError:
481 idents = []
483 idents = []
482 if not idents:
484 if not idents:
483 self.log.error("Bad Query Message: %r", msg)
485 self.log.error("Bad Query Message: %r", msg)
484 return
486 return
485 client_id = idents[0]
487 client_id = idents[0]
486 try:
488 try:
487 msg = self.session.unserialize(msg, content=True)
489 msg = self.session.unserialize(msg, content=True)
488 except Exception:
490 except Exception:
489 content = error.wrap_exception()
491 content = error.wrap_exception()
490 self.log.error("Bad Query Message: %r", msg, exc_info=True)
492 self.log.error("Bad Query Message: %r", msg, exc_info=True)
491 self.session.send(self.query, "hub_error", ident=client_id,
493 self.session.send(self.query, "hub_error", ident=client_id,
492 content=content)
494 content=content)
493 return
495 return
494 # print client_id, header, parent, content
496 # print client_id, header, parent, content
495 #switch on message type:
497 #switch on message type:
496 msg_type = msg['header']['msg_type']
498 msg_type = msg['header']['msg_type']
497 self.log.info("client::client %r requested %r", client_id, msg_type)
499 self.log.info("client::client %r requested %r", client_id, msg_type)
498 handler = self.query_handlers.get(msg_type, None)
500 handler = self.query_handlers.get(msg_type, None)
499 try:
501 try:
500 assert handler is not None, "Bad Message Type: %r" % msg_type
502 assert handler is not None, "Bad Message Type: %r" % msg_type
501 except:
503 except:
502 content = error.wrap_exception()
504 content = error.wrap_exception()
503 self.log.error("Bad Message Type: %r", msg_type, exc_info=True)
505 self.log.error("Bad Message Type: %r", msg_type, exc_info=True)
504 self.session.send(self.query, "hub_error", ident=client_id,
506 self.session.send(self.query, "hub_error", ident=client_id,
505 content=content)
507 content=content)
506 return
508 return
507
509
508 else:
510 else:
509 handler(idents, msg)
511 handler(idents, msg)
510
512
511 def dispatch_db(self, msg):
513 def dispatch_db(self, msg):
512 """"""
514 """"""
513 raise NotImplementedError
515 raise NotImplementedError
514
516
515 #---------------------------------------------------------------------------
517 #---------------------------------------------------------------------------
516 # handler methods (1 per event)
518 # handler methods (1 per event)
517 #---------------------------------------------------------------------------
519 #---------------------------------------------------------------------------
518
520
519 #----------------------- Heartbeat --------------------------------------
521 #----------------------- Heartbeat --------------------------------------
520
522
521 def handle_new_heart(self, heart):
523 def handle_new_heart(self, heart):
522 """handler to attach to heartbeater.
524 """handler to attach to heartbeater.
523 Called when a new heart starts to beat.
525 Called when a new heart starts to beat.
524 Triggers completion of registration."""
526 Triggers completion of registration."""
525 self.log.debug("heartbeat::handle_new_heart(%r)", heart)
527 self.log.debug("heartbeat::handle_new_heart(%r)", heart)
526 if heart not in self.incoming_registrations:
528 if heart not in self.incoming_registrations:
527 self.log.info("heartbeat::ignoring new heart: %r", heart)
529 self.log.info("heartbeat::ignoring new heart: %r", heart)
528 else:
530 else:
529 self.finish_registration(heart)
531 self.finish_registration(heart)
530
532
531
533
532 def handle_heart_failure(self, heart):
534 def handle_heart_failure(self, heart):
533 """handler to attach to heartbeater.
535 """handler to attach to heartbeater.
534 called when a previously registered heart fails to respond to beat request.
536 called when a previously registered heart fails to respond to beat request.
535 triggers unregistration"""
537 triggers unregistration"""
536 self.log.debug("heartbeat::handle_heart_failure(%r)", heart)
538 self.log.debug("heartbeat::handle_heart_failure(%r)", heart)
537 eid = self.hearts.get(heart, None)
539 eid = self.hearts.get(heart, None)
538 queue = self.engines[eid].queue
540 queue = self.engines[eid].queue
539 if eid is None or self.keytable[eid] in self.dead_engines:
541 if eid is None or self.keytable[eid] in self.dead_engines:
540 self.log.info("heartbeat::ignoring heart failure %r (not an engine or already dead)", heart)
542 self.log.info("heartbeat::ignoring heart failure %r (not an engine or already dead)", heart)
541 else:
543 else:
542 self.unregister_engine(heart, dict(content=dict(id=eid, queue=queue)))
544 self.unregister_engine(heart, dict(content=dict(id=eid, queue=queue)))
543
545
544 #----------------------- MUX Queue Traffic ------------------------------
546 #----------------------- MUX Queue Traffic ------------------------------
545
547
546 def save_queue_request(self, idents, msg):
548 def save_queue_request(self, idents, msg):
547 if len(idents) < 2:
549 if len(idents) < 2:
548 self.log.error("invalid identity prefix: %r", idents)
550 self.log.error("invalid identity prefix: %r", idents)
549 return
551 return
550 queue_id, client_id = idents[:2]
552 queue_id, client_id = idents[:2]
551 try:
553 try:
552 msg = self.session.unserialize(msg)
554 msg = self.session.unserialize(msg)
553 except Exception:
555 except Exception:
554 self.log.error("queue::client %r sent invalid message to %r: %r", client_id, queue_id, msg, exc_info=True)
556 self.log.error("queue::client %r sent invalid message to %r: %r", client_id, queue_id, msg, exc_info=True)
555 return
557 return
556
558
557 eid = self.by_ident.get(queue_id, None)
559 eid = self.by_ident.get(queue_id, None)
558 if eid is None:
560 if eid is None:
559 self.log.error("queue::target %r not registered", queue_id)
561 self.log.error("queue::target %r not registered", queue_id)
560 self.log.debug("queue:: valid are: %r", self.by_ident.keys())
562 self.log.debug("queue:: valid are: %r", self.by_ident.keys())
561 return
563 return
562 record = init_record(msg)
564 record = init_record(msg)
563 msg_id = record['msg_id']
565 msg_id = record['msg_id']
564 self.log.info("queue::client %r submitted request %r to %s", client_id, msg_id, eid)
566 self.log.info("queue::client %r submitted request %r to %s", client_id, msg_id, eid)
565 # Unicode in records
567 # Unicode in records
566 record['engine_uuid'] = queue_id.decode('ascii')
568 record['engine_uuid'] = queue_id.decode('ascii')
567 record['client_uuid'] = client_id.decode('ascii')
569 record['client_uuid'] = client_id.decode('ascii')
568 record['queue'] = 'mux'
570 record['queue'] = 'mux'
569
571
570 try:
572 try:
571 # it's posible iopub arrived first:
573 # it's posible iopub arrived first:
572 existing = self.db.get_record(msg_id)
574 existing = self.db.get_record(msg_id)
573 for key,evalue in existing.iteritems():
575 for key,evalue in existing.iteritems():
574 rvalue = record.get(key, None)
576 rvalue = record.get(key, None)
575 if evalue and rvalue and evalue != rvalue:
577 if evalue and rvalue and evalue != rvalue:
576 self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
578 self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
577 elif evalue and not rvalue:
579 elif evalue and not rvalue:
578 record[key] = evalue
580 record[key] = evalue
579 try:
581 try:
580 self.db.update_record(msg_id, record)
582 self.db.update_record(msg_id, record)
581 except Exception:
583 except Exception:
582 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
584 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
583 except KeyError:
585 except KeyError:
584 try:
586 try:
585 self.db.add_record(msg_id, record)
587 self.db.add_record(msg_id, record)
586 except Exception:
588 except Exception:
587 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
589 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
588
590
589
591
590 self.pending.add(msg_id)
592 self.pending.add(msg_id)
591 self.queues[eid].append(msg_id)
593 self.queues[eid].append(msg_id)
592
594
593 def save_queue_result(self, idents, msg):
595 def save_queue_result(self, idents, msg):
594 if len(idents) < 2:
596 if len(idents) < 2:
595 self.log.error("invalid identity prefix: %r", idents)
597 self.log.error("invalid identity prefix: %r", idents)
596 return
598 return
597
599
598 client_id, queue_id = idents[:2]
600 client_id, queue_id = idents[:2]
599 try:
601 try:
600 msg = self.session.unserialize(msg)
602 msg = self.session.unserialize(msg)
601 except Exception:
603 except Exception:
602 self.log.error("queue::engine %r sent invalid message to %r: %r",
604 self.log.error("queue::engine %r sent invalid message to %r: %r",
603 queue_id, client_id, msg, exc_info=True)
605 queue_id, client_id, msg, exc_info=True)
604 return
606 return
605
607
606 eid = self.by_ident.get(queue_id, None)
608 eid = self.by_ident.get(queue_id, None)
607 if eid is None:
609 if eid is None:
608 self.log.error("queue::unknown engine %r is sending a reply: ", queue_id)
610 self.log.error("queue::unknown engine %r is sending a reply: ", queue_id)
609 return
611 return
610
612
611 parent = msg['parent_header']
613 parent = msg['parent_header']
612 if not parent:
614 if not parent:
613 return
615 return
614 msg_id = parent['msg_id']
616 msg_id = parent['msg_id']
615 if msg_id in self.pending:
617 if msg_id in self.pending:
616 self.pending.remove(msg_id)
618 self.pending.remove(msg_id)
617 self.all_completed.add(msg_id)
619 self.all_completed.add(msg_id)
618 self.queues[eid].remove(msg_id)
620 self.queues[eid].remove(msg_id)
619 self.completed[eid].append(msg_id)
621 self.completed[eid].append(msg_id)
620 self.log.info("queue::request %r completed on %s", msg_id, eid)
622 self.log.info("queue::request %r completed on %s", msg_id, eid)
621 elif msg_id not in self.all_completed:
623 elif msg_id not in self.all_completed:
622 # it could be a result from a dead engine that died before delivering the
624 # it could be a result from a dead engine that died before delivering the
623 # result
625 # result
624 self.log.warn("queue:: unknown msg finished %r", msg_id)
626 self.log.warn("queue:: unknown msg finished %r", msg_id)
625 return
627 return
626 # update record anyway, because the unregistration could have been premature
628 # update record anyway, because the unregistration could have been premature
627 rheader = msg['header']
629 rheader = msg['header']
628 completed = rheader['date']
630 completed = rheader['date']
629 started = rheader.get('started', None)
631 started = rheader.get('started', None)
630 result = {
632 result = {
631 'result_header' : rheader,
633 'result_header' : rheader,
632 'result_content': msg['content'],
634 'result_content': msg['content'],
635 'received': datetime.now(),
633 'started' : started,
636 'started' : started,
634 'completed' : completed
637 'completed' : completed
635 }
638 }
636
639
637 result['result_buffers'] = msg['buffers']
640 result['result_buffers'] = msg['buffers']
638 try:
641 try:
639 self.db.update_record(msg_id, result)
642 self.db.update_record(msg_id, result)
640 except Exception:
643 except Exception:
641 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
644 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
642
645
643
646
644 #--------------------- Task Queue Traffic ------------------------------
647 #--------------------- Task Queue Traffic ------------------------------
645
648
646 def save_task_request(self, idents, msg):
649 def save_task_request(self, idents, msg):
647 """Save the submission of a task."""
650 """Save the submission of a task."""
648 client_id = idents[0]
651 client_id = idents[0]
649
652
650 try:
653 try:
651 msg = self.session.unserialize(msg)
654 msg = self.session.unserialize(msg)
652 except Exception:
655 except Exception:
653 self.log.error("task::client %r sent invalid task message: %r",
656 self.log.error("task::client %r sent invalid task message: %r",
654 client_id, msg, exc_info=True)
657 client_id, msg, exc_info=True)
655 return
658 return
656 record = init_record(msg)
659 record = init_record(msg)
657
660
658 record['client_uuid'] = client_id.decode('ascii')
661 record['client_uuid'] = client_id.decode('ascii')
659 record['queue'] = 'task'
662 record['queue'] = 'task'
660 header = msg['header']
663 header = msg['header']
661 msg_id = header['msg_id']
664 msg_id = header['msg_id']
662 self.pending.add(msg_id)
665 self.pending.add(msg_id)
663 self.unassigned.add(msg_id)
666 self.unassigned.add(msg_id)
664 try:
667 try:
665 # it's posible iopub arrived first:
668 # it's posible iopub arrived first:
666 existing = self.db.get_record(msg_id)
669 existing = self.db.get_record(msg_id)
667 if existing['resubmitted']:
670 if existing['resubmitted']:
668 for key in ('submitted', 'client_uuid', 'buffers'):
671 for key in ('submitted', 'client_uuid', 'buffers'):
669 # don't clobber these keys on resubmit
672 # don't clobber these keys on resubmit
670 # submitted and client_uuid should be different
673 # submitted and client_uuid should be different
671 # and buffers might be big, and shouldn't have changed
674 # and buffers might be big, and shouldn't have changed
672 record.pop(key)
675 record.pop(key)
673 # still check content,header which should not change
676 # still check content,header which should not change
674 # but are not expensive to compare as buffers
677 # but are not expensive to compare as buffers
675
678
676 for key,evalue in existing.iteritems():
679 for key,evalue in existing.iteritems():
677 if key.endswith('buffers'):
680 if key.endswith('buffers'):
678 # don't compare buffers
681 # don't compare buffers
679 continue
682 continue
680 rvalue = record.get(key, None)
683 rvalue = record.get(key, None)
681 if evalue and rvalue and evalue != rvalue:
684 if evalue and rvalue and evalue != rvalue:
682 self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
685 self.log.warn("conflicting initial state for record: %r:%r <%r> %r", msg_id, rvalue, key, evalue)
683 elif evalue and not rvalue:
686 elif evalue and not rvalue:
684 record[key] = evalue
687 record[key] = evalue
685 try:
688 try:
686 self.db.update_record(msg_id, record)
689 self.db.update_record(msg_id, record)
687 except Exception:
690 except Exception:
688 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
691 self.log.error("DB Error updating record %r", msg_id, exc_info=True)
689 except KeyError:
692 except KeyError:
690 try:
693 try:
691 self.db.add_record(msg_id, record)
694 self.db.add_record(msg_id, record)
692 except Exception:
695 except Exception:
693 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
696 self.log.error("DB Error adding record %r", msg_id, exc_info=True)
694 except Exception:
697 except Exception:
695 self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
698 self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
696
699
697 def save_task_result(self, idents, msg):
700 def save_task_result(self, idents, msg):
698 """save the result of a completed task."""
701 """save the result of a completed task."""
699 client_id = idents[0]
702 client_id = idents[0]
700 try:
703 try:
701 msg = self.session.unserialize(msg)
704 msg = self.session.unserialize(msg)
702 except Exception:
705 except Exception:
703 self.log.error("task::invalid task result message send to %r: %r",
706 self.log.error("task::invalid task result message send to %r: %r",
704 client_id, msg, exc_info=True)
707 client_id, msg, exc_info=True)
705 return
708 return
706
709
707 parent = msg['parent_header']
710 parent = msg['parent_header']
708 if not parent:
711 if not parent:
709 # print msg
712 # print msg
710 self.log.warn("Task %r had no parent!", msg)
713 self.log.warn("Task %r had no parent!", msg)
711 return
714 return
712 msg_id = parent['msg_id']
715 msg_id = parent['msg_id']
713 if msg_id in self.unassigned:
716 if msg_id in self.unassigned:
714 self.unassigned.remove(msg_id)
717 self.unassigned.remove(msg_id)
715
718
716 header = msg['header']
719 header = msg['header']
717 engine_uuid = header.get('engine', None)
720 engine_uuid = header.get('engine', None)
718 eid = self.by_ident.get(engine_uuid, None)
721 eid = self.by_ident.get(engine_uuid, None)
719
722
720 if msg_id in self.pending:
723 if msg_id in self.pending:
721 self.log.info("task::task %r finished on %s", msg_id, eid)
724 self.log.info("task::task %r finished on %s", msg_id, eid)
722 self.pending.remove(msg_id)
725 self.pending.remove(msg_id)
723 self.all_completed.add(msg_id)
726 self.all_completed.add(msg_id)
724 if eid is not None:
727 if eid is not None:
725 self.completed[eid].append(msg_id)
728 self.completed[eid].append(msg_id)
726 if msg_id in self.tasks[eid]:
729 if msg_id in self.tasks[eid]:
727 self.tasks[eid].remove(msg_id)
730 self.tasks[eid].remove(msg_id)
728 completed = header['date']
731 completed = header['date']
729 started = header.get('started', None)
732 started = header.get('started', None)
730 result = {
733 result = {
731 'result_header' : header,
734 'result_header' : header,
732 'result_content': msg['content'],
735 'result_content': msg['content'],
733 'started' : started,
736 'started' : started,
734 'completed' : completed,
737 'completed' : completed,
735 'engine_uuid': engine_uuid
738 'received' : datetime.now(),
739 'engine_uuid': engine_uuid,
736 }
740 }
737
741
738 result['result_buffers'] = msg['buffers']
742 result['result_buffers'] = msg['buffers']
739 try:
743 try:
740 self.db.update_record(msg_id, result)
744 self.db.update_record(msg_id, result)
741 except Exception:
745 except Exception:
742 self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
746 self.log.error("DB Error saving task request %r", msg_id, exc_info=True)
743
747
744 else:
748 else:
745 self.log.debug("task::unknown task %r finished", msg_id)
749 self.log.debug("task::unknown task %r finished", msg_id)
746
750
747 def save_task_destination(self, idents, msg):
751 def save_task_destination(self, idents, msg):
748 try:
752 try:
749 msg = self.session.unserialize(msg, content=True)
753 msg = self.session.unserialize(msg, content=True)
750 except Exception:
754 except Exception:
751 self.log.error("task::invalid task tracking message", exc_info=True)
755 self.log.error("task::invalid task tracking message", exc_info=True)
752 return
756 return
753 content = msg['content']
757 content = msg['content']
754 # print (content)
758 # print (content)
755 msg_id = content['msg_id']
759 msg_id = content['msg_id']
756 engine_uuid = content['engine_id']
760 engine_uuid = content['engine_id']
757 eid = self.by_ident[util.asbytes(engine_uuid)]
761 eid = self.by_ident[util.asbytes(engine_uuid)]
758
762
759 self.log.info("task::task %r arrived on %r", msg_id, eid)
763 self.log.info("task::task %r arrived on %r", msg_id, eid)
760 if msg_id in self.unassigned:
764 if msg_id in self.unassigned:
761 self.unassigned.remove(msg_id)
765 self.unassigned.remove(msg_id)
762 # else:
766 # else:
763 # self.log.debug("task::task %r not listed as MIA?!"%(msg_id))
767 # self.log.debug("task::task %r not listed as MIA?!"%(msg_id))
764
768
765 self.tasks[eid].append(msg_id)
769 self.tasks[eid].append(msg_id)
766 # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
770 # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
767 try:
771 try:
768 self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
772 self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
769 except Exception:
773 except Exception:
770 self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)
774 self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)
771
775
772
776
773 def mia_task_request(self, idents, msg):
777 def mia_task_request(self, idents, msg):
774 raise NotImplementedError
778 raise NotImplementedError
775 client_id = idents[0]
779 client_id = idents[0]
776 # content = dict(mia=self.mia,status='ok')
780 # content = dict(mia=self.mia,status='ok')
777 # self.session.send('mia_reply', content=content, idents=client_id)
781 # self.session.send('mia_reply', content=content, idents=client_id)
778
782
779
783
780 #--------------------- IOPub Traffic ------------------------------
784 #--------------------- IOPub Traffic ------------------------------
781
785
782 def save_iopub_message(self, topics, msg):
786 def save_iopub_message(self, topics, msg):
783 """save an iopub message into the db"""
787 """save an iopub message into the db"""
784 # print (topics)
788 # print (topics)
785 try:
789 try:
786 msg = self.session.unserialize(msg, content=True)
790 msg = self.session.unserialize(msg, content=True)
787 except Exception:
791 except Exception:
788 self.log.error("iopub::invalid IOPub message", exc_info=True)
792 self.log.error("iopub::invalid IOPub message", exc_info=True)
789 return
793 return
790
794
791 parent = msg['parent_header']
795 parent = msg['parent_header']
792 if not parent:
796 if not parent:
793 self.log.error("iopub::invalid IOPub message: %r", msg)
797 self.log.error("iopub::invalid IOPub message: %r", msg)
794 return
798 return
795 msg_id = parent['msg_id']
799 msg_id = parent['msg_id']
796 msg_type = msg['header']['msg_type']
800 msg_type = msg['header']['msg_type']
797 content = msg['content']
801 content = msg['content']
798
802
799 # ensure msg_id is in db
803 # ensure msg_id is in db
800 try:
804 try:
801 rec = self.db.get_record(msg_id)
805 rec = self.db.get_record(msg_id)
802 except KeyError:
806 except KeyError:
803 rec = empty_record()
807 rec = empty_record()
804 rec['msg_id'] = msg_id
808 rec['msg_id'] = msg_id
805 self.db.add_record(msg_id, rec)
809 self.db.add_record(msg_id, rec)
806 # stream
810 # stream
807 d = {}
811 d = {}
808 if msg_type == 'stream':
812 if msg_type == 'stream':
809 name = content['name']
813 name = content['name']
810 s = rec[name] or ''
814 s = rec[name] or ''
811 d[name] = s + content['data']
815 d[name] = s + content['data']
812
816
813 elif msg_type == 'pyerr':
817 elif msg_type == 'pyerr':
814 d['pyerr'] = content
818 d['pyerr'] = content
815 elif msg_type == 'pyin':
819 elif msg_type == 'pyin':
816 d['pyin'] = content['code']
820 d['pyin'] = content['code']
817 else:
821 else:
818 d[msg_type] = content.get('data', '')
822 d[msg_type] = content.get('data', '')
819
823
820 try:
824 try:
821 self.db.update_record(msg_id, d)
825 self.db.update_record(msg_id, d)
822 except Exception:
826 except Exception:
823 self.log.error("DB Error saving iopub message %r", msg_id, exc_info=True)
827 self.log.error("DB Error saving iopub message %r", msg_id, exc_info=True)
824
828
825
829
826
830
827 #-------------------------------------------------------------------------
831 #-------------------------------------------------------------------------
828 # Registration requests
832 # Registration requests
829 #-------------------------------------------------------------------------
833 #-------------------------------------------------------------------------
830
834
831 def connection_request(self, client_id, msg):
835 def connection_request(self, client_id, msg):
832 """Reply with connection addresses for clients."""
836 """Reply with connection addresses for clients."""
833 self.log.info("client::client %r connected", client_id)
837 self.log.info("client::client %r connected", client_id)
834 content = dict(status='ok')
838 content = dict(status='ok')
835 content.update(self.client_info)
839 content.update(self.client_info)
836 jsonable = {}
840 jsonable = {}
837 for k,v in self.keytable.iteritems():
841 for k,v in self.keytable.iteritems():
838 if v not in self.dead_engines:
842 if v not in self.dead_engines:
839 jsonable[str(k)] = v.decode('ascii')
843 jsonable[str(k)] = v.decode('ascii')
840 content['engines'] = jsonable
844 content['engines'] = jsonable
841 self.session.send(self.query, 'connection_reply', content, parent=msg, ident=client_id)
845 self.session.send(self.query, 'connection_reply', content, parent=msg, ident=client_id)
842
846
843 def register_engine(self, reg, msg):
847 def register_engine(self, reg, msg):
844 """Register a new engine."""
848 """Register a new engine."""
845 content = msg['content']
849 content = msg['content']
846 try:
850 try:
847 queue = util.asbytes(content['queue'])
851 queue = util.asbytes(content['queue'])
848 except KeyError:
852 except KeyError:
849 self.log.error("registration::queue not specified", exc_info=True)
853 self.log.error("registration::queue not specified", exc_info=True)
850 return
854 return
851 heart = content.get('heartbeat', None)
855 heart = content.get('heartbeat', None)
852 if heart:
856 if heart:
853 heart = util.asbytes(heart)
857 heart = util.asbytes(heart)
854 """register a new engine, and create the socket(s) necessary"""
858 """register a new engine, and create the socket(s) necessary"""
855 eid = self._next_id
859 eid = self._next_id
856 # print (eid, queue, reg, heart)
860 # print (eid, queue, reg, heart)
857
861
858 self.log.debug("registration::register_engine(%i, %r, %r, %r)", eid, queue, reg, heart)
862 self.log.debug("registration::register_engine(%i, %r, %r, %r)", eid, queue, reg, heart)
859
863
860 content = dict(id=eid,status='ok')
864 content = dict(id=eid,status='ok')
861 content.update(self.engine_info)
865 content.update(self.engine_info)
862 # check if requesting available IDs:
866 # check if requesting available IDs:
863 if queue in self.by_ident:
867 if queue in self.by_ident:
864 try:
868 try:
865 raise KeyError("queue_id %r in use" % queue)
869 raise KeyError("queue_id %r in use" % queue)
866 except:
870 except:
867 content = error.wrap_exception()
871 content = error.wrap_exception()
868 self.log.error("queue_id %r in use", queue, exc_info=True)
872 self.log.error("queue_id %r in use", queue, exc_info=True)
869 elif heart in self.hearts: # need to check unique hearts?
873 elif heart in self.hearts: # need to check unique hearts?
870 try:
874 try:
871 raise KeyError("heart_id %r in use" % heart)
875 raise KeyError("heart_id %r in use" % heart)
872 except:
876 except:
873 self.log.error("heart_id %r in use", heart, exc_info=True)
877 self.log.error("heart_id %r in use", heart, exc_info=True)
874 content = error.wrap_exception()
878 content = error.wrap_exception()
875 else:
879 else:
876 for h, pack in self.incoming_registrations.iteritems():
880 for h, pack in self.incoming_registrations.iteritems():
877 if heart == h:
881 if heart == h:
878 try:
882 try:
879 raise KeyError("heart_id %r in use" % heart)
883 raise KeyError("heart_id %r in use" % heart)
880 except:
884 except:
881 self.log.error("heart_id %r in use", heart, exc_info=True)
885 self.log.error("heart_id %r in use", heart, exc_info=True)
882 content = error.wrap_exception()
886 content = error.wrap_exception()
883 break
887 break
884 elif queue == pack[1]:
888 elif queue == pack[1]:
885 try:
889 try:
886 raise KeyError("queue_id %r in use" % queue)
890 raise KeyError("queue_id %r in use" % queue)
887 except:
891 except:
888 self.log.error("queue_id %r in use", queue, exc_info=True)
892 self.log.error("queue_id %r in use", queue, exc_info=True)
889 content = error.wrap_exception()
893 content = error.wrap_exception()
890 break
894 break
891
895
892 msg = self.session.send(self.query, "registration_reply",
896 msg = self.session.send(self.query, "registration_reply",
893 content=content,
897 content=content,
894 ident=reg)
898 ident=reg)
895
899
896 if content['status'] == 'ok':
900 if content['status'] == 'ok':
897 if heart in self.heartmonitor.hearts:
901 if heart in self.heartmonitor.hearts:
898 # already beating
902 # already beating
899 self.incoming_registrations[heart] = (eid,queue,reg[0],None)
903 self.incoming_registrations[heart] = (eid,queue,reg[0],None)
900 self.finish_registration(heart)
904 self.finish_registration(heart)
901 else:
905 else:
902 purge = lambda : self._purge_stalled_registration(heart)
906 purge = lambda : self._purge_stalled_registration(heart)
903 dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
907 dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
904 dc.start()
908 dc.start()
905 self.incoming_registrations[heart] = (eid,queue,reg[0],dc)
909 self.incoming_registrations[heart] = (eid,queue,reg[0],dc)
906 else:
910 else:
907 self.log.error("registration::registration %i failed: %r", eid, content['evalue'])
911 self.log.error("registration::registration %i failed: %r", eid, content['evalue'])
908 return eid
912 return eid
909
913
910 def unregister_engine(self, ident, msg):
914 def unregister_engine(self, ident, msg):
911 """Unregister an engine that explicitly requested to leave."""
915 """Unregister an engine that explicitly requested to leave."""
912 try:
916 try:
913 eid = msg['content']['id']
917 eid = msg['content']['id']
914 except:
918 except:
915 self.log.error("registration::bad engine id for unregistration: %r", ident, exc_info=True)
919 self.log.error("registration::bad engine id for unregistration: %r", ident, exc_info=True)
916 return
920 return
917 self.log.info("registration::unregister_engine(%r)", eid)
921 self.log.info("registration::unregister_engine(%r)", eid)
918 # print (eid)
922 # print (eid)
919 uuid = self.keytable[eid]
923 uuid = self.keytable[eid]
920 content=dict(id=eid, queue=uuid.decode('ascii'))
924 content=dict(id=eid, queue=uuid.decode('ascii'))
921 self.dead_engines.add(uuid)
925 self.dead_engines.add(uuid)
922 # self.ids.remove(eid)
926 # self.ids.remove(eid)
923 # uuid = self.keytable.pop(eid)
927 # uuid = self.keytable.pop(eid)
924 #
928 #
925 # ec = self.engines.pop(eid)
929 # ec = self.engines.pop(eid)
926 # self.hearts.pop(ec.heartbeat)
930 # self.hearts.pop(ec.heartbeat)
927 # self.by_ident.pop(ec.queue)
931 # self.by_ident.pop(ec.queue)
928 # self.completed.pop(eid)
932 # self.completed.pop(eid)
929 handleit = lambda : self._handle_stranded_msgs(eid, uuid)
933 handleit = lambda : self._handle_stranded_msgs(eid, uuid)
930 dc = ioloop.DelayedCallback(handleit, self.registration_timeout, self.loop)
934 dc = ioloop.DelayedCallback(handleit, self.registration_timeout, self.loop)
931 dc.start()
935 dc.start()
932 ############## TODO: HANDLE IT ################
936 ############## TODO: HANDLE IT ################
933
937
934 if self.notifier:
938 if self.notifier:
935 self.session.send(self.notifier, "unregistration_notification", content=content)
939 self.session.send(self.notifier, "unregistration_notification", content=content)
936
940
937 def _handle_stranded_msgs(self, eid, uuid):
941 def _handle_stranded_msgs(self, eid, uuid):
938 """Handle messages known to be on an engine when the engine unregisters.
942 """Handle messages known to be on an engine when the engine unregisters.
939
943
940 It is possible that this will fire prematurely - that is, an engine will
944 It is possible that this will fire prematurely - that is, an engine will
941 go down after completing a result, and the client will be notified
945 go down after completing a result, and the client will be notified
942 that the result failed and later receive the actual result.
946 that the result failed and later receive the actual result.
943 """
947 """
944
948
945 outstanding = self.queues[eid]
949 outstanding = self.queues[eid]
946
950
947 for msg_id in outstanding:
951 for msg_id in outstanding:
948 self.pending.remove(msg_id)
952 self.pending.remove(msg_id)
949 self.all_completed.add(msg_id)
953 self.all_completed.add(msg_id)
950 try:
954 try:
951 raise error.EngineError("Engine %r died while running task %r" % (eid, msg_id))
955 raise error.EngineError("Engine %r died while running task %r" % (eid, msg_id))
952 except:
956 except:
953 content = error.wrap_exception()
957 content = error.wrap_exception()
954 # build a fake header:
958 # build a fake header:
955 header = {}
959 header = {}
956 header['engine'] = uuid
960 header['engine'] = uuid
957 header['date'] = datetime.now()
961 header['date'] = datetime.now()
958 rec = dict(result_content=content, result_header=header, result_buffers=[])
962 rec = dict(result_content=content, result_header=header, result_buffers=[])
959 rec['completed'] = header['date']
963 rec['completed'] = header['date']
960 rec['engine_uuid'] = uuid
964 rec['engine_uuid'] = uuid
961 try:
965 try:
962 self.db.update_record(msg_id, rec)
966 self.db.update_record(msg_id, rec)
963 except Exception:
967 except Exception:
964 self.log.error("DB Error handling stranded msg %r", msg_id, exc_info=True)
968 self.log.error("DB Error handling stranded msg %r", msg_id, exc_info=True)
965
969
966
970
967 def finish_registration(self, heart):
971 def finish_registration(self, heart):
968 """Second half of engine registration, called after our HeartMonitor
972 """Second half of engine registration, called after our HeartMonitor
969 has received a beat from the Engine's Heart."""
973 has received a beat from the Engine's Heart."""
970 try:
974 try:
971 (eid,queue,reg,purge) = self.incoming_registrations.pop(heart)
975 (eid,queue,reg,purge) = self.incoming_registrations.pop(heart)
972 except KeyError:
976 except KeyError:
973 self.log.error("registration::tried to finish nonexistant registration", exc_info=True)
977 self.log.error("registration::tried to finish nonexistant registration", exc_info=True)
974 return
978 return
975 self.log.info("registration::finished registering engine %i:%r", eid, queue)
979 self.log.info("registration::finished registering engine %i:%r", eid, queue)
976 if purge is not None:
980 if purge is not None:
977 purge.stop()
981 purge.stop()
978 control = queue
982 control = queue
979 self.ids.add(eid)
983 self.ids.add(eid)
980 self.keytable[eid] = queue
984 self.keytable[eid] = queue
981 self.engines[eid] = EngineConnector(id=eid, queue=queue, registration=reg,
985 self.engines[eid] = EngineConnector(id=eid, queue=queue, registration=reg,
982 control=control, heartbeat=heart)
986 control=control, heartbeat=heart)
983 self.by_ident[queue] = eid
987 self.by_ident[queue] = eid
984 self.queues[eid] = list()
988 self.queues[eid] = list()
985 self.tasks[eid] = list()
989 self.tasks[eid] = list()
986 self.completed[eid] = list()
990 self.completed[eid] = list()
987 self.hearts[heart] = eid
991 self.hearts[heart] = eid
988 content = dict(id=eid, queue=self.engines[eid].queue.decode('ascii'))
992 content = dict(id=eid, queue=self.engines[eid].queue.decode('ascii'))
989 if self.notifier:
993 if self.notifier:
990 self.session.send(self.notifier, "registration_notification", content=content)
994 self.session.send(self.notifier, "registration_notification", content=content)
991 self.log.info("engine::Engine Connected: %i", eid)
995 self.log.info("engine::Engine Connected: %i", eid)
992
996
993 def _purge_stalled_registration(self, heart):
997 def _purge_stalled_registration(self, heart):
994 if heart in self.incoming_registrations:
998 if heart in self.incoming_registrations:
995 eid = self.incoming_registrations.pop(heart)[0]
999 eid = self.incoming_registrations.pop(heart)[0]
996 self.log.info("registration::purging stalled registration: %i", eid)
1000 self.log.info("registration::purging stalled registration: %i", eid)
997 else:
1001 else:
998 pass
1002 pass
999
1003
1000 #-------------------------------------------------------------------------
1004 #-------------------------------------------------------------------------
1001 # Client Requests
1005 # Client Requests
1002 #-------------------------------------------------------------------------
1006 #-------------------------------------------------------------------------
1003
1007
1004 def shutdown_request(self, client_id, msg):
1008 def shutdown_request(self, client_id, msg):
1005 """handle shutdown request."""
1009 """handle shutdown request."""
1006 self.session.send(self.query, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
1010 self.session.send(self.query, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
1007 # also notify other clients of shutdown
1011 # also notify other clients of shutdown
1008 self.session.send(self.notifier, 'shutdown_notice', content={'status': 'ok'})
1012 self.session.send(self.notifier, 'shutdown_notice', content={'status': 'ok'})
1009 dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
1013 dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
1010 dc.start()
1014 dc.start()
1011
1015
1012 def _shutdown(self):
1016 def _shutdown(self):
1013 self.log.info("hub::hub shutting down.")
1017 self.log.info("hub::hub shutting down.")
1014 time.sleep(0.1)
1018 time.sleep(0.1)
1015 sys.exit(0)
1019 sys.exit(0)
1016
1020
1017
1021
1018 def check_load(self, client_id, msg):
1022 def check_load(self, client_id, msg):
1019 content = msg['content']
1023 content = msg['content']
1020 try:
1024 try:
1021 targets = content['targets']
1025 targets = content['targets']
1022 targets = self._validate_targets(targets)
1026 targets = self._validate_targets(targets)
1023 except:
1027 except:
1024 content = error.wrap_exception()
1028 content = error.wrap_exception()
1025 self.session.send(self.query, "hub_error",
1029 self.session.send(self.query, "hub_error",
1026 content=content, ident=client_id)
1030 content=content, ident=client_id)
1027 return
1031 return
1028
1032
1029 content = dict(status='ok')
1033 content = dict(status='ok')
1030 # loads = {}
1034 # loads = {}
1031 for t in targets:
1035 for t in targets:
1032 content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
1036 content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
1033 self.session.send(self.query, "load_reply", content=content, ident=client_id)
1037 self.session.send(self.query, "load_reply", content=content, ident=client_id)
1034
1038
1035
1039
1036 def queue_status(self, client_id, msg):
1040 def queue_status(self, client_id, msg):
1037 """Return the Queue status of one or more targets.
1041 """Return the Queue status of one or more targets.
1038 if verbose: return the msg_ids
1042 if verbose: return the msg_ids
1039 else: return len of each type.
1043 else: return len of each type.
1040 keys: queue (pending MUX jobs)
1044 keys: queue (pending MUX jobs)
1041 tasks (pending Task jobs)
1045 tasks (pending Task jobs)
1042 completed (finished jobs from both queues)"""
1046 completed (finished jobs from both queues)"""
1043 content = msg['content']
1047 content = msg['content']
1044 targets = content['targets']
1048 targets = content['targets']
1045 try:
1049 try:
1046 targets = self._validate_targets(targets)
1050 targets = self._validate_targets(targets)
1047 except:
1051 except:
1048 content = error.wrap_exception()
1052 content = error.wrap_exception()
1049 self.session.send(self.query, "hub_error",
1053 self.session.send(self.query, "hub_error",
1050 content=content, ident=client_id)
1054 content=content, ident=client_id)
1051 return
1055 return
1052 verbose = content.get('verbose', False)
1056 verbose = content.get('verbose', False)
1053 content = dict(status='ok')
1057 content = dict(status='ok')
1054 for t in targets:
1058 for t in targets:
1055 queue = self.queues[t]
1059 queue = self.queues[t]
1056 completed = self.completed[t]
1060 completed = self.completed[t]
1057 tasks = self.tasks[t]
1061 tasks = self.tasks[t]
1058 if not verbose:
1062 if not verbose:
1059 queue = len(queue)
1063 queue = len(queue)
1060 completed = len(completed)
1064 completed = len(completed)
1061 tasks = len(tasks)
1065 tasks = len(tasks)
1062 content[str(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
1066 content[str(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
1063 content['unassigned'] = list(self.unassigned) if verbose else len(self.unassigned)
1067 content['unassigned'] = list(self.unassigned) if verbose else len(self.unassigned)
1064 # print (content)
1068 # print (content)
1065 self.session.send(self.query, "queue_reply", content=content, ident=client_id)
1069 self.session.send(self.query, "queue_reply", content=content, ident=client_id)
1066
1070
1067 def purge_results(self, client_id, msg):
1071 def purge_results(self, client_id, msg):
1068 """Purge results from memory. This method is more valuable before we move
1072 """Purge results from memory. This method is more valuable before we move
1069 to a DB based message storage mechanism."""
1073 to a DB based message storage mechanism."""
1070 content = msg['content']
1074 content = msg['content']
1071 self.log.info("Dropping records with %s", content)
1075 self.log.info("Dropping records with %s", content)
1072 msg_ids = content.get('msg_ids', [])
1076 msg_ids = content.get('msg_ids', [])
1073 reply = dict(status='ok')
1077 reply = dict(status='ok')
1074 if msg_ids == 'all':
1078 if msg_ids == 'all':
1075 try:
1079 try:
1076 self.db.drop_matching_records(dict(completed={'$ne':None}))
1080 self.db.drop_matching_records(dict(completed={'$ne':None}))
1077 except Exception:
1081 except Exception:
1078 reply = error.wrap_exception()
1082 reply = error.wrap_exception()
1079 else:
1083 else:
1080 pending = filter(lambda m: m in self.pending, msg_ids)
1084 pending = filter(lambda m: m in self.pending, msg_ids)
1081 if pending:
1085 if pending:
1082 try:
1086 try:
1083 raise IndexError("msg pending: %r" % pending[0])
1087 raise IndexError("msg pending: %r" % pending[0])
1084 except:
1088 except:
1085 reply = error.wrap_exception()
1089 reply = error.wrap_exception()
1086 else:
1090 else:
1087 try:
1091 try:
1088 self.db.drop_matching_records(dict(msg_id={'$in':msg_ids}))
1092 self.db.drop_matching_records(dict(msg_id={'$in':msg_ids}))
1089 except Exception:
1093 except Exception:
1090 reply = error.wrap_exception()
1094 reply = error.wrap_exception()
1091
1095
1092 if reply['status'] == 'ok':
1096 if reply['status'] == 'ok':
1093 eids = content.get('engine_ids', [])
1097 eids = content.get('engine_ids', [])
1094 for eid in eids:
1098 for eid in eids:
1095 if eid not in self.engines:
1099 if eid not in self.engines:
1096 try:
1100 try:
1097 raise IndexError("No such engine: %i" % eid)
1101 raise IndexError("No such engine: %i" % eid)
1098 except:
1102 except:
1099 reply = error.wrap_exception()
1103 reply = error.wrap_exception()
1100 break
1104 break
1101 uid = self.engines[eid].queue
1105 uid = self.engines[eid].queue
1102 try:
1106 try:
1103 self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
1107 self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
1104 except Exception:
1108 except Exception:
1105 reply = error.wrap_exception()
1109 reply = error.wrap_exception()
1106 break
1110 break
1107
1111
1108 self.session.send(self.query, 'purge_reply', content=reply, ident=client_id)
1112 self.session.send(self.query, 'purge_reply', content=reply, ident=client_id)
1109
1113
1110 def resubmit_task(self, client_id, msg):
1114 def resubmit_task(self, client_id, msg):
1111 """Resubmit one or more tasks."""
1115 """Resubmit one or more tasks."""
1112 def finish(reply):
1116 def finish(reply):
1113 self.session.send(self.query, 'resubmit_reply', content=reply, ident=client_id)
1117 self.session.send(self.query, 'resubmit_reply', content=reply, ident=client_id)
1114
1118
1115 content = msg['content']
1119 content = msg['content']
1116 msg_ids = content['msg_ids']
1120 msg_ids = content['msg_ids']
1117 reply = dict(status='ok')
1121 reply = dict(status='ok')
1118 try:
1122 try:
1119 records = self.db.find_records({'msg_id' : {'$in' : msg_ids}}, keys=[
1123 records = self.db.find_records({'msg_id' : {'$in' : msg_ids}}, keys=[
1120 'header', 'content', 'buffers'])
1124 'header', 'content', 'buffers'])
1121 except Exception:
1125 except Exception:
1122 self.log.error('db::db error finding tasks to resubmit', exc_info=True)
1126 self.log.error('db::db error finding tasks to resubmit', exc_info=True)
1123 return finish(error.wrap_exception())
1127 return finish(error.wrap_exception())
1124
1128
1125 # validate msg_ids
1129 # validate msg_ids
1126 found_ids = [ rec['msg_id'] for rec in records ]
1130 found_ids = [ rec['msg_id'] for rec in records ]
1127 invalid_ids = filter(lambda m: m in self.pending, found_ids)
1131 invalid_ids = filter(lambda m: m in self.pending, found_ids)
1128 if len(records) > len(msg_ids):
1132 if len(records) > len(msg_ids):
1129 try:
1133 try:
1130 raise RuntimeError("DB appears to be in an inconsistent state."
1134 raise RuntimeError("DB appears to be in an inconsistent state."
1131 "More matching records were found than should exist")
1135 "More matching records were found than should exist")
1132 except Exception:
1136 except Exception:
1133 return finish(error.wrap_exception())
1137 return finish(error.wrap_exception())
1134 elif len(records) < len(msg_ids):
1138 elif len(records) < len(msg_ids):
1135 missing = [ m for m in msg_ids if m not in found_ids ]
1139 missing = [ m for m in msg_ids if m not in found_ids ]
1136 try:
1140 try:
1137 raise KeyError("No such msg(s): %r" % missing)
1141 raise KeyError("No such msg(s): %r" % missing)
1138 except KeyError:
1142 except KeyError:
1139 return finish(error.wrap_exception())
1143 return finish(error.wrap_exception())
1140 elif invalid_ids:
1144 elif invalid_ids:
1141 msg_id = invalid_ids[0]
1145 msg_id = invalid_ids[0]
1142 try:
1146 try:
1143 raise ValueError("Task %r appears to be inflight" % msg_id)
1147 raise ValueError("Task %r appears to be inflight" % msg_id)
1144 except Exception:
1148 except Exception:
1145 return finish(error.wrap_exception())
1149 return finish(error.wrap_exception())
1146
1150
1147 # clear the existing records
1151 # clear the existing records
1148 now = datetime.now()
1152 now = datetime.now()
1149 rec = empty_record()
1153 rec = empty_record()
1150 map(rec.pop, ['msg_id', 'header', 'content', 'buffers', 'submitted'])
1154 map(rec.pop, ['msg_id', 'header', 'content', 'buffers', 'submitted'])
1151 rec['resubmitted'] = now
1155 rec['resubmitted'] = now
1152 rec['queue'] = 'task'
1156 rec['queue'] = 'task'
1153 rec['client_uuid'] = client_id[0]
1157 rec['client_uuid'] = client_id[0]
1154 try:
1158 try:
1155 for msg_id in msg_ids:
1159 for msg_id in msg_ids:
1156 self.all_completed.discard(msg_id)
1160 self.all_completed.discard(msg_id)
1157 self.db.update_record(msg_id, rec)
1161 self.db.update_record(msg_id, rec)
1158 except Exception:
1162 except Exception:
1159 self.log.error('db::db error upating record', exc_info=True)
1163 self.log.error('db::db error upating record', exc_info=True)
1160 reply = error.wrap_exception()
1164 reply = error.wrap_exception()
1161 else:
1165 else:
1162 # send the messages
1166 # send the messages
1163 for rec in records:
1167 for rec in records:
1164 header = rec['header']
1168 header = rec['header']
1165 # include resubmitted in header to prevent digest collision
1169 # include resubmitted in header to prevent digest collision
1166 header['resubmitted'] = now
1170 header['resubmitted'] = now
1167 msg = self.session.msg(header['msg_type'])
1171 msg = self.session.msg(header['msg_type'])
1168 msg['content'] = rec['content']
1172 msg['content'] = rec['content']
1169 msg['header'] = header
1173 msg['header'] = header
1170 msg['header']['msg_id'] = rec['msg_id']
1174 msg['header']['msg_id'] = rec['msg_id']
1171 self.session.send(self.resubmit, msg, buffers=rec['buffers'])
1175 self.session.send(self.resubmit, msg, buffers=rec['buffers'])
1172
1176
1173 finish(dict(status='ok'))
1177 finish(dict(status='ok'))
1174
1178
1175
1179
1176 def _extract_record(self, rec):
1180 def _extract_record(self, rec):
1177 """decompose a TaskRecord dict into subsection of reply for get_result"""
1181 """decompose a TaskRecord dict into subsection of reply for get_result"""
1178 io_dict = {}
1182 io_dict = {}
1179 for key in 'pyin pyout pyerr stdout stderr'.split():
1183 for key in ('pyin', 'pyout', 'pyerr', 'stdout', 'stderr'):
1180 io_dict[key] = rec[key]
1184 io_dict[key] = rec[key]
1181 content = { 'result_content': rec['result_content'],
1185 content = { 'result_content': rec['result_content'],
1182 'header': rec['header'],
1186 'header': rec['header'],
1183 'result_header' : rec['result_header'],
1187 'result_header' : rec['result_header'],
1188 'received' : rec['received'],
1184 'io' : io_dict,
1189 'io' : io_dict,
1185 }
1190 }
1186 if rec['result_buffers']:
1191 if rec['result_buffers']:
1187 buffers = map(bytes, rec['result_buffers'])
1192 buffers = map(bytes, rec['result_buffers'])
1188 else:
1193 else:
1189 buffers = []
1194 buffers = []
1190
1195
1191 return content, buffers
1196 return content, buffers
1192
1197
1193 def get_results(self, client_id, msg):
1198 def get_results(self, client_id, msg):
1194 """Get the result of 1 or more messages."""
1199 """Get the result of 1 or more messages."""
1195 content = msg['content']
1200 content = msg['content']
1196 msg_ids = sorted(set(content['msg_ids']))
1201 msg_ids = sorted(set(content['msg_ids']))
1197 statusonly = content.get('status_only', False)
1202 statusonly = content.get('status_only', False)
1198 pending = []
1203 pending = []
1199 completed = []
1204 completed = []
1200 content = dict(status='ok')
1205 content = dict(status='ok')
1201 content['pending'] = pending
1206 content['pending'] = pending
1202 content['completed'] = completed
1207 content['completed'] = completed
1203 buffers = []
1208 buffers = []
1204 if not statusonly:
1209 if not statusonly:
1205 try:
1210 try:
1206 matches = self.db.find_records(dict(msg_id={'$in':msg_ids}))
1211 matches = self.db.find_records(dict(msg_id={'$in':msg_ids}))
1207 # turn match list into dict, for faster lookup
1212 # turn match list into dict, for faster lookup
1208 records = {}
1213 records = {}
1209 for rec in matches:
1214 for rec in matches:
1210 records[rec['msg_id']] = rec
1215 records[rec['msg_id']] = rec
1211 except Exception:
1216 except Exception:
1212 content = error.wrap_exception()
1217 content = error.wrap_exception()
1213 self.session.send(self.query, "result_reply", content=content,
1218 self.session.send(self.query, "result_reply", content=content,
1214 parent=msg, ident=client_id)
1219 parent=msg, ident=client_id)
1215 return
1220 return
1216 else:
1221 else:
1217 records = {}
1222 records = {}
1218 for msg_id in msg_ids:
1223 for msg_id in msg_ids:
1219 if msg_id in self.pending:
1224 if msg_id in self.pending:
1220 pending.append(msg_id)
1225 pending.append(msg_id)
1221 elif msg_id in self.all_completed:
1226 elif msg_id in self.all_completed:
1222 completed.append(msg_id)
1227 completed.append(msg_id)
1223 if not statusonly:
1228 if not statusonly:
1224 c,bufs = self._extract_record(records[msg_id])
1229 c,bufs = self._extract_record(records[msg_id])
1225 content[msg_id] = c
1230 content[msg_id] = c
1226 buffers.extend(bufs)
1231 buffers.extend(bufs)
1227 elif msg_id in records:
1232 elif msg_id in records:
1228 if rec['completed']:
1233 if rec['completed']:
1229 completed.append(msg_id)
1234 completed.append(msg_id)
1230 c,bufs = self._extract_record(records[msg_id])
1235 c,bufs = self._extract_record(records[msg_id])
1231 content[msg_id] = c
1236 content[msg_id] = c
1232 buffers.extend(bufs)
1237 buffers.extend(bufs)
1233 else:
1238 else:
1234 pending.append(msg_id)
1239 pending.append(msg_id)
1235 else:
1240 else:
1236 try:
1241 try:
1237 raise KeyError('No such message: '+msg_id)
1242 raise KeyError('No such message: '+msg_id)
1238 except:
1243 except:
1239 content = error.wrap_exception()
1244 content = error.wrap_exception()
1240 break
1245 break
1241 self.session.send(self.query, "result_reply", content=content,
1246 self.session.send(self.query, "result_reply", content=content,
1242 parent=msg, ident=client_id,
1247 parent=msg, ident=client_id,
1243 buffers=buffers)
1248 buffers=buffers)
1244
1249
1245 def get_history(self, client_id, msg):
1250 def get_history(self, client_id, msg):
1246 """Get a list of all msg_ids in our DB records"""
1251 """Get a list of all msg_ids in our DB records"""
1247 try:
1252 try:
1248 msg_ids = self.db.get_history()
1253 msg_ids = self.db.get_history()
1249 except Exception as e:
1254 except Exception as e:
1250 content = error.wrap_exception()
1255 content = error.wrap_exception()
1251 else:
1256 else:
1252 content = dict(status='ok', history=msg_ids)
1257 content = dict(status='ok', history=msg_ids)
1253
1258
1254 self.session.send(self.query, "history_reply", content=content,
1259 self.session.send(self.query, "history_reply", content=content,
1255 parent=msg, ident=client_id)
1260 parent=msg, ident=client_id)
1256
1261
1257 def db_query(self, client_id, msg):
1262 def db_query(self, client_id, msg):
1258 """Perform a raw query on the task record database."""
1263 """Perform a raw query on the task record database."""
1259 content = msg['content']
1264 content = msg['content']
1260 query = content.get('query', {})
1265 query = content.get('query', {})
1261 keys = content.get('keys', None)
1266 keys = content.get('keys', None)
1262 buffers = []
1267 buffers = []
1263 empty = list()
1268 empty = list()
1264 try:
1269 try:
1265 records = self.db.find_records(query, keys)
1270 records = self.db.find_records(query, keys)
1266 except Exception as e:
1271 except Exception as e:
1267 content = error.wrap_exception()
1272 content = error.wrap_exception()
1268 else:
1273 else:
1269 # extract buffers from reply content:
1274 # extract buffers from reply content:
1270 if keys is not None:
1275 if keys is not None:
1271 buffer_lens = [] if 'buffers' in keys else None
1276 buffer_lens = [] if 'buffers' in keys else None
1272 result_buffer_lens = [] if 'result_buffers' in keys else None
1277 result_buffer_lens = [] if 'result_buffers' in keys else None
1273 else:
1278 else:
1274 buffer_lens = None
1279 buffer_lens = None
1275 result_buffer_lens = None
1280 result_buffer_lens = None
1276
1281
1277 for rec in records:
1282 for rec in records:
1278 # buffers may be None, so double check
1283 # buffers may be None, so double check
1279 b = rec.pop('buffers', empty) or empty
1284 b = rec.pop('buffers', empty) or empty
1280 if buffer_lens is not None:
1285 if buffer_lens is not None:
1281 buffer_lens.append(len(b))
1286 buffer_lens.append(len(b))
1282 buffers.extend(b)
1287 buffers.extend(b)
1283 rb = rec.pop('result_buffers', empty) or empty
1288 rb = rec.pop('result_buffers', empty) or empty
1284 if result_buffer_lens is not None:
1289 if result_buffer_lens is not None:
1285 result_buffer_lens.append(len(rb))
1290 result_buffer_lens.append(len(rb))
1286 buffers.extend(rb)
1291 buffers.extend(rb)
1287 content = dict(status='ok', records=records, buffer_lens=buffer_lens,
1292 content = dict(status='ok', records=records, buffer_lens=buffer_lens,
1288 result_buffer_lens=result_buffer_lens)
1293 result_buffer_lens=result_buffer_lens)
1289 # self.log.debug (content)
1294 # self.log.debug (content)
1290 self.session.send(self.query, "db_reply", content=content,
1295 self.session.send(self.query, "db_reply", content=content,
1291 parent=msg, ident=client_id,
1296 parent=msg, ident=client_id,
1292 buffers=buffers)
1297 buffers=buffers)
1293
1298
@@ -1,408 +1,411 b''
1 """A TaskRecord backend using sqlite3
1 """A TaskRecord backend using sqlite3
2
2
3 Authors:
3 Authors:
4
4
5 * Min RK
5 * Min RK
6 """
6 """
7 #-----------------------------------------------------------------------------
7 #-----------------------------------------------------------------------------
8 # Copyright (C) 2011 The IPython Development Team
8 # Copyright (C) 2011 The IPython Development Team
9 #
9 #
10 # Distributed under the terms of the BSD License. The full license is in
10 # Distributed under the terms of the BSD License. The full license is in
11 # the file COPYING, distributed as part of this software.
11 # the file COPYING, distributed as part of this software.
12 #-----------------------------------------------------------------------------
12 #-----------------------------------------------------------------------------
13
13
14 import json
14 import json
15 import os
15 import os
16 import cPickle as pickle
16 import cPickle as pickle
17 from datetime import datetime
17 from datetime import datetime
18
18
19 try:
19 try:
20 import sqlite3
20 import sqlite3
21 except ImportError:
21 except ImportError:
22 sqlite3 = None
22 sqlite3 = None
23
23
24 from zmq.eventloop import ioloop
24 from zmq.eventloop import ioloop
25
25
26 from IPython.utils.traitlets import Unicode, Instance, List, Dict
26 from IPython.utils.traitlets import Unicode, Instance, List, Dict
27 from .dictdb import BaseDB
27 from .dictdb import BaseDB
28 from IPython.utils.jsonutil import date_default, extract_dates, squash_dates
28 from IPython.utils.jsonutil import date_default, extract_dates, squash_dates
29
29
30 #-----------------------------------------------------------------------------
30 #-----------------------------------------------------------------------------
31 # SQLite operators, adapters, and converters
31 # SQLite operators, adapters, and converters
32 #-----------------------------------------------------------------------------
32 #-----------------------------------------------------------------------------
33
33
34 try:
34 try:
35 buffer
35 buffer
36 except NameError:
36 except NameError:
37 # py3k
37 # py3k
38 buffer = memoryview
38 buffer = memoryview
39
39
40 operators = {
40 operators = {
41 '$lt' : "<",
41 '$lt' : "<",
42 '$gt' : ">",
42 '$gt' : ">",
43 # null is handled weird with ==,!=
43 # null is handled weird with ==,!=
44 '$eq' : "=",
44 '$eq' : "=",
45 '$ne' : "!=",
45 '$ne' : "!=",
46 '$lte': "<=",
46 '$lte': "<=",
47 '$gte': ">=",
47 '$gte': ">=",
48 '$in' : ('=', ' OR '),
48 '$in' : ('=', ' OR '),
49 '$nin': ('!=', ' AND '),
49 '$nin': ('!=', ' AND '),
50 # '$all': None,
50 # '$all': None,
51 # '$mod': None,
51 # '$mod': None,
52 # '$exists' : None
52 # '$exists' : None
53 }
53 }
54 null_operators = {
54 null_operators = {
55 '=' : "IS NULL",
55 '=' : "IS NULL",
56 '!=' : "IS NOT NULL",
56 '!=' : "IS NOT NULL",
57 }
57 }
58
58
59 def _adapt_dict(d):
59 def _adapt_dict(d):
60 return json.dumps(d, default=date_default)
60 return json.dumps(d, default=date_default)
61
61
62 def _convert_dict(ds):
62 def _convert_dict(ds):
63 if ds is None:
63 if ds is None:
64 return ds
64 return ds
65 else:
65 else:
66 if isinstance(ds, bytes):
66 if isinstance(ds, bytes):
67 # If I understand the sqlite doc correctly, this will always be utf8
67 # If I understand the sqlite doc correctly, this will always be utf8
68 ds = ds.decode('utf8')
68 ds = ds.decode('utf8')
69 return extract_dates(json.loads(ds))
69 return extract_dates(json.loads(ds))
70
70
71 def _adapt_bufs(bufs):
71 def _adapt_bufs(bufs):
72 # this is *horrible*
72 # this is *horrible*
73 # copy buffers into single list and pickle it:
73 # copy buffers into single list and pickle it:
74 if bufs and isinstance(bufs[0], (bytes, buffer)):
74 if bufs and isinstance(bufs[0], (bytes, buffer)):
75 return sqlite3.Binary(pickle.dumps(map(bytes, bufs),-1))
75 return sqlite3.Binary(pickle.dumps(map(bytes, bufs),-1))
76 elif bufs:
76 elif bufs:
77 return bufs
77 return bufs
78 else:
78 else:
79 return None
79 return None
80
80
81 def _convert_bufs(bs):
81 def _convert_bufs(bs):
82 if bs is None:
82 if bs is None:
83 return []
83 return []
84 else:
84 else:
85 return pickle.loads(bytes(bs))
85 return pickle.loads(bytes(bs))
86
86
87 #-----------------------------------------------------------------------------
87 #-----------------------------------------------------------------------------
88 # SQLiteDB class
88 # SQLiteDB class
89 #-----------------------------------------------------------------------------
89 #-----------------------------------------------------------------------------
90
90
91 class SQLiteDB(BaseDB):
91 class SQLiteDB(BaseDB):
92 """SQLite3 TaskRecord backend."""
92 """SQLite3 TaskRecord backend."""
93
93
94 filename = Unicode('tasks.db', config=True,
94 filename = Unicode('tasks.db', config=True,
95 help="""The filename of the sqlite task database. [default: 'tasks.db']""")
95 help="""The filename of the sqlite task database. [default: 'tasks.db']""")
96 location = Unicode('', config=True,
96 location = Unicode('', config=True,
97 help="""The directory containing the sqlite task database. The default
97 help="""The directory containing the sqlite task database. The default
98 is to use the cluster_dir location.""")
98 is to use the cluster_dir location.""")
99 table = Unicode("", config=True,
99 table = Unicode("", config=True,
100 help="""The SQLite Table to use for storing tasks for this session. If unspecified,
100 help="""The SQLite Table to use for storing tasks for this session. If unspecified,
101 a new table will be created with the Hub's IDENT. Specifying the table will result
101 a new table will be created with the Hub's IDENT. Specifying the table will result
102 in tasks from previous sessions being available via Clients' db_query and
102 in tasks from previous sessions being available via Clients' db_query and
103 get_result methods.""")
103 get_result methods.""")
104
104
105 if sqlite3 is not None:
105 if sqlite3 is not None:
106 _db = Instance('sqlite3.Connection')
106 _db = Instance('sqlite3.Connection')
107 else:
107 else:
108 _db = None
108 _db = None
109 # the ordered list of column names
109 # the ordered list of column names
110 _keys = List(['msg_id' ,
110 _keys = List(['msg_id' ,
111 'header' ,
111 'header' ,
112 'content',
112 'content',
113 'buffers',
113 'buffers',
114 'submitted',
114 'submitted',
115 'client_uuid' ,
115 'client_uuid' ,
116 'engine_uuid' ,
116 'engine_uuid' ,
117 'started',
117 'started',
118 'completed',
118 'completed',
119 'resubmitted',
119 'resubmitted',
120 'received',
120 'result_header' ,
121 'result_header' ,
121 'result_content' ,
122 'result_content' ,
122 'result_buffers' ,
123 'result_buffers' ,
123 'queue' ,
124 'queue' ,
124 'pyin' ,
125 'pyin' ,
125 'pyout',
126 'pyout',
126 'pyerr',
127 'pyerr',
127 'stdout',
128 'stdout',
128 'stderr',
129 'stderr',
129 ])
130 ])
130 # sqlite datatypes for checking that db is current format
131 # sqlite datatypes for checking that db is current format
131 _types = Dict({'msg_id' : 'text' ,
132 _types = Dict({'msg_id' : 'text' ,
132 'header' : 'dict text',
133 'header' : 'dict text',
133 'content' : 'dict text',
134 'content' : 'dict text',
134 'buffers' : 'bufs blob',
135 'buffers' : 'bufs blob',
135 'submitted' : 'timestamp',
136 'submitted' : 'timestamp',
136 'client_uuid' : 'text',
137 'client_uuid' : 'text',
137 'engine_uuid' : 'text',
138 'engine_uuid' : 'text',
138 'started' : 'timestamp',
139 'started' : 'timestamp',
139 'completed' : 'timestamp',
140 'completed' : 'timestamp',
140 'resubmitted' : 'timestamp',
141 'resubmitted' : 'timestamp',
142 'received' : 'timestamp',
141 'result_header' : 'dict text',
143 'result_header' : 'dict text',
142 'result_content' : 'dict text',
144 'result_content' : 'dict text',
143 'result_buffers' : 'bufs blob',
145 'result_buffers' : 'bufs blob',
144 'queue' : 'text',
146 'queue' : 'text',
145 'pyin' : 'text',
147 'pyin' : 'text',
146 'pyout' : 'text',
148 'pyout' : 'text',
147 'pyerr' : 'text',
149 'pyerr' : 'text',
148 'stdout' : 'text',
150 'stdout' : 'text',
149 'stderr' : 'text',
151 'stderr' : 'text',
150 })
152 })
151
153
152 def __init__(self, **kwargs):
154 def __init__(self, **kwargs):
153 super(SQLiteDB, self).__init__(**kwargs)
155 super(SQLiteDB, self).__init__(**kwargs)
154 if sqlite3 is None:
156 if sqlite3 is None:
155 raise ImportError("SQLiteDB requires sqlite3")
157 raise ImportError("SQLiteDB requires sqlite3")
156 if not self.table:
158 if not self.table:
157 # use session, and prefix _, since starting with # is illegal
159 # use session, and prefix _, since starting with # is illegal
158 self.table = '_'+self.session.replace('-','_')
160 self.table = '_'+self.session.replace('-','_')
159 if not self.location:
161 if not self.location:
160 # get current profile
162 # get current profile
161 from IPython.core.application import BaseIPythonApplication
163 from IPython.core.application import BaseIPythonApplication
162 if BaseIPythonApplication.initialized():
164 if BaseIPythonApplication.initialized():
163 app = BaseIPythonApplication.instance()
165 app = BaseIPythonApplication.instance()
164 if app.profile_dir is not None:
166 if app.profile_dir is not None:
165 self.location = app.profile_dir.location
167 self.location = app.profile_dir.location
166 else:
168 else:
167 self.location = u'.'
169 self.location = u'.'
168 else:
170 else:
169 self.location = u'.'
171 self.location = u'.'
170 self._init_db()
172 self._init_db()
171
173
172 # register db commit as 2s periodic callback
174 # register db commit as 2s periodic callback
173 # to prevent clogging pipes
175 # to prevent clogging pipes
174 # assumes we are being run in a zmq ioloop app
176 # assumes we are being run in a zmq ioloop app
175 loop = ioloop.IOLoop.instance()
177 loop = ioloop.IOLoop.instance()
176 pc = ioloop.PeriodicCallback(self._db.commit, 2000, loop)
178 pc = ioloop.PeriodicCallback(self._db.commit, 2000, loop)
177 pc.start()
179 pc.start()
178
180
179 def _defaults(self, keys=None):
181 def _defaults(self, keys=None):
180 """create an empty record"""
182 """create an empty record"""
181 d = {}
183 d = {}
182 keys = self._keys if keys is None else keys
184 keys = self._keys if keys is None else keys
183 for key in keys:
185 for key in keys:
184 d[key] = None
186 d[key] = None
185 return d
187 return d
186
188
187 def _check_table(self):
189 def _check_table(self):
188 """Ensure that an incorrect table doesn't exist
190 """Ensure that an incorrect table doesn't exist
189
191
190 If a bad (old) table does exist, return False
192 If a bad (old) table does exist, return False
191 """
193 """
192 cursor = self._db.execute("PRAGMA table_info(%s)"%self.table)
194 cursor = self._db.execute("PRAGMA table_info(%s)"%self.table)
193 lines = cursor.fetchall()
195 lines = cursor.fetchall()
194 if not lines:
196 if not lines:
195 # table does not exist
197 # table does not exist
196 return True
198 return True
197 types = {}
199 types = {}
198 keys = []
200 keys = []
199 for line in lines:
201 for line in lines:
200 keys.append(line[1])
202 keys.append(line[1])
201 types[line[1]] = line[2]
203 types[line[1]] = line[2]
202 if self._keys != keys:
204 if self._keys != keys:
203 # key mismatch
205 # key mismatch
204 self.log.warn('keys mismatch')
206 self.log.warn('keys mismatch')
205 return False
207 return False
206 for key in self._keys:
208 for key in self._keys:
207 if types[key] != self._types[key]:
209 if types[key] != self._types[key]:
208 self.log.warn(
210 self.log.warn(
209 'type mismatch: %s: %s != %s'%(key,types[key],self._types[key])
211 'type mismatch: %s: %s != %s'%(key,types[key],self._types[key])
210 )
212 )
211 return False
213 return False
212 return True
214 return True
213
215
214 def _init_db(self):
216 def _init_db(self):
215 """Connect to the database and get new session number."""
217 """Connect to the database and get new session number."""
216 # register adapters
218 # register adapters
217 sqlite3.register_adapter(dict, _adapt_dict)
219 sqlite3.register_adapter(dict, _adapt_dict)
218 sqlite3.register_converter('dict', _convert_dict)
220 sqlite3.register_converter('dict', _convert_dict)
219 sqlite3.register_adapter(list, _adapt_bufs)
221 sqlite3.register_adapter(list, _adapt_bufs)
220 sqlite3.register_converter('bufs', _convert_bufs)
222 sqlite3.register_converter('bufs', _convert_bufs)
221 # connect to the db
223 # connect to the db
222 dbfile = os.path.join(self.location, self.filename)
224 dbfile = os.path.join(self.location, self.filename)
223 self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES,
225 self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES,
224 # isolation_level = None)#,
226 # isolation_level = None)#,
225 cached_statements=64)
227 cached_statements=64)
226 # print dir(self._db)
228 # print dir(self._db)
227 first_table = self.table
229 first_table = self.table
228 i=0
230 i=0
229 while not self._check_table():
231 while not self._check_table():
230 i+=1
232 i+=1
231 self.table = first_table+'_%i'%i
233 self.table = first_table+'_%i'%i
232 self.log.warn(
234 self.log.warn(
233 "Table %s exists and doesn't match db format, trying %s"%
235 "Table %s exists and doesn't match db format, trying %s"%
234 (first_table,self.table)
236 (first_table,self.table)
235 )
237 )
236
238
237 self._db.execute("""CREATE TABLE IF NOT EXISTS %s
239 self._db.execute("""CREATE TABLE IF NOT EXISTS %s
238 (msg_id text PRIMARY KEY,
240 (msg_id text PRIMARY KEY,
239 header dict text,
241 header dict text,
240 content dict text,
242 content dict text,
241 buffers bufs blob,
243 buffers bufs blob,
242 submitted timestamp,
244 submitted timestamp,
243 client_uuid text,
245 client_uuid text,
244 engine_uuid text,
246 engine_uuid text,
245 started timestamp,
247 started timestamp,
246 completed timestamp,
248 completed timestamp,
247 resubmitted timestamp,
249 resubmitted timestamp,
250 received timestamp,
248 result_header dict text,
251 result_header dict text,
249 result_content dict text,
252 result_content dict text,
250 result_buffers bufs blob,
253 result_buffers bufs blob,
251 queue text,
254 queue text,
252 pyin text,
255 pyin text,
253 pyout text,
256 pyout text,
254 pyerr text,
257 pyerr text,
255 stdout text,
258 stdout text,
256 stderr text)
259 stderr text)
257 """%self.table)
260 """%self.table)
258 self._db.commit()
261 self._db.commit()
259
262
260 def _dict_to_list(self, d):
263 def _dict_to_list(self, d):
261 """turn a mongodb-style record dict into a list."""
264 """turn a mongodb-style record dict into a list."""
262
265
263 return [ d[key] for key in self._keys ]
266 return [ d[key] for key in self._keys ]
264
267
265 def _list_to_dict(self, line, keys=None):
268 def _list_to_dict(self, line, keys=None):
266 """Inverse of dict_to_list"""
269 """Inverse of dict_to_list"""
267 keys = self._keys if keys is None else keys
270 keys = self._keys if keys is None else keys
268 d = self._defaults(keys)
271 d = self._defaults(keys)
269 for key,value in zip(keys, line):
272 for key,value in zip(keys, line):
270 d[key] = value
273 d[key] = value
271
274
272 return d
275 return d
273
276
274 def _render_expression(self, check):
277 def _render_expression(self, check):
275 """Turn a mongodb-style search dict into an SQL query."""
278 """Turn a mongodb-style search dict into an SQL query."""
276 expressions = []
279 expressions = []
277 args = []
280 args = []
278
281
279 skeys = set(check.keys())
282 skeys = set(check.keys())
280 skeys.difference_update(set(self._keys))
283 skeys.difference_update(set(self._keys))
281 skeys.difference_update(set(['buffers', 'result_buffers']))
284 skeys.difference_update(set(['buffers', 'result_buffers']))
282 if skeys:
285 if skeys:
283 raise KeyError("Illegal testing key(s): %s"%skeys)
286 raise KeyError("Illegal testing key(s): %s"%skeys)
284
287
285 for name,sub_check in check.iteritems():
288 for name,sub_check in check.iteritems():
286 if isinstance(sub_check, dict):
289 if isinstance(sub_check, dict):
287 for test,value in sub_check.iteritems():
290 for test,value in sub_check.iteritems():
288 try:
291 try:
289 op = operators[test]
292 op = operators[test]
290 except KeyError:
293 except KeyError:
291 raise KeyError("Unsupported operator: %r"%test)
294 raise KeyError("Unsupported operator: %r"%test)
292 if isinstance(op, tuple):
295 if isinstance(op, tuple):
293 op, join = op
296 op, join = op
294
297
295 if value is None and op in null_operators:
298 if value is None and op in null_operators:
296 expr = "%s %s" % (name, null_operators[op])
299 expr = "%s %s" % (name, null_operators[op])
297 else:
300 else:
298 expr = "%s %s ?"%(name, op)
301 expr = "%s %s ?"%(name, op)
299 if isinstance(value, (tuple,list)):
302 if isinstance(value, (tuple,list)):
300 if op in null_operators and any([v is None for v in value]):
303 if op in null_operators and any([v is None for v in value]):
301 # equality tests don't work with NULL
304 # equality tests don't work with NULL
302 raise ValueError("Cannot use %r test with NULL values on SQLite backend"%test)
305 raise ValueError("Cannot use %r test with NULL values on SQLite backend"%test)
303 expr = '( %s )'%( join.join([expr]*len(value)) )
306 expr = '( %s )'%( join.join([expr]*len(value)) )
304 args.extend(value)
307 args.extend(value)
305 else:
308 else:
306 args.append(value)
309 args.append(value)
307 expressions.append(expr)
310 expressions.append(expr)
308 else:
311 else:
309 # it's an equality check
312 # it's an equality check
310 if sub_check is None:
313 if sub_check is None:
311 expressions.append("%s IS NULL" % name)
314 expressions.append("%s IS NULL" % name)
312 else:
315 else:
313 expressions.append("%s = ?"%name)
316 expressions.append("%s = ?"%name)
314 args.append(sub_check)
317 args.append(sub_check)
315
318
316 expr = " AND ".join(expressions)
319 expr = " AND ".join(expressions)
317 return expr, args
320 return expr, args
318
321
319 def add_record(self, msg_id, rec):
322 def add_record(self, msg_id, rec):
320 """Add a new Task Record, by msg_id."""
323 """Add a new Task Record, by msg_id."""
321 d = self._defaults()
324 d = self._defaults()
322 d.update(rec)
325 d.update(rec)
323 d['msg_id'] = msg_id
326 d['msg_id'] = msg_id
324 line = self._dict_to_list(d)
327 line = self._dict_to_list(d)
325 tups = '(%s)'%(','.join(['?']*len(line)))
328 tups = '(%s)'%(','.join(['?']*len(line)))
326 self._db.execute("INSERT INTO %s VALUES %s"%(self.table, tups), line)
329 self._db.execute("INSERT INTO %s VALUES %s"%(self.table, tups), line)
327 # self._db.commit()
330 # self._db.commit()
328
331
329 def get_record(self, msg_id):
332 def get_record(self, msg_id):
330 """Get a specific Task Record, by msg_id."""
333 """Get a specific Task Record, by msg_id."""
331 cursor = self._db.execute("""SELECT * FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
334 cursor = self._db.execute("""SELECT * FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
332 line = cursor.fetchone()
335 line = cursor.fetchone()
333 if line is None:
336 if line is None:
334 raise KeyError("No such msg: %r"%msg_id)
337 raise KeyError("No such msg: %r"%msg_id)
335 return self._list_to_dict(line)
338 return self._list_to_dict(line)
336
339
337 def update_record(self, msg_id, rec):
340 def update_record(self, msg_id, rec):
338 """Update the data in an existing record."""
341 """Update the data in an existing record."""
339 query = "UPDATE %s SET "%self.table
342 query = "UPDATE %s SET "%self.table
340 sets = []
343 sets = []
341 keys = sorted(rec.keys())
344 keys = sorted(rec.keys())
342 values = []
345 values = []
343 for key in keys:
346 for key in keys:
344 sets.append('%s = ?'%key)
347 sets.append('%s = ?'%key)
345 values.append(rec[key])
348 values.append(rec[key])
346 query += ', '.join(sets)
349 query += ', '.join(sets)
347 query += ' WHERE msg_id == ?'
350 query += ' WHERE msg_id == ?'
348 values.append(msg_id)
351 values.append(msg_id)
349 self._db.execute(query, values)
352 self._db.execute(query, values)
350 # self._db.commit()
353 # self._db.commit()
351
354
352 def drop_record(self, msg_id):
355 def drop_record(self, msg_id):
353 """Remove a record from the DB."""
356 """Remove a record from the DB."""
354 self._db.execute("""DELETE FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
357 self._db.execute("""DELETE FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
355 # self._db.commit()
358 # self._db.commit()
356
359
357 def drop_matching_records(self, check):
360 def drop_matching_records(self, check):
358 """Remove a record from the DB."""
361 """Remove a record from the DB."""
359 expr,args = self._render_expression(check)
362 expr,args = self._render_expression(check)
360 query = "DELETE FROM %s WHERE %s"%(self.table, expr)
363 query = "DELETE FROM %s WHERE %s"%(self.table, expr)
361 self._db.execute(query,args)
364 self._db.execute(query,args)
362 # self._db.commit()
365 # self._db.commit()
363
366
364 def find_records(self, check, keys=None):
367 def find_records(self, check, keys=None):
365 """Find records matching a query dict, optionally extracting subset of keys.
368 """Find records matching a query dict, optionally extracting subset of keys.
366
369
367 Returns list of matching records.
370 Returns list of matching records.
368
371
369 Parameters
372 Parameters
370 ----------
373 ----------
371
374
372 check: dict
375 check: dict
373 mongodb-style query argument
376 mongodb-style query argument
374 keys: list of strs [optional]
377 keys: list of strs [optional]
375 if specified, the subset of keys to extract. msg_id will *always* be
378 if specified, the subset of keys to extract. msg_id will *always* be
376 included.
379 included.
377 """
380 """
378 if keys:
381 if keys:
379 bad_keys = [ key for key in keys if key not in self._keys ]
382 bad_keys = [ key for key in keys if key not in self._keys ]
380 if bad_keys:
383 if bad_keys:
381 raise KeyError("Bad record key(s): %s"%bad_keys)
384 raise KeyError("Bad record key(s): %s"%bad_keys)
382
385
383 if keys:
386 if keys:
384 # ensure msg_id is present and first:
387 # ensure msg_id is present and first:
385 if 'msg_id' in keys:
388 if 'msg_id' in keys:
386 keys.remove('msg_id')
389 keys.remove('msg_id')
387 keys.insert(0, 'msg_id')
390 keys.insert(0, 'msg_id')
388 req = ', '.join(keys)
391 req = ', '.join(keys)
389 else:
392 else:
390 req = '*'
393 req = '*'
391 expr,args = self._render_expression(check)
394 expr,args = self._render_expression(check)
392 query = """SELECT %s FROM %s WHERE %s"""%(req, self.table, expr)
395 query = """SELECT %s FROM %s WHERE %s"""%(req, self.table, expr)
393 cursor = self._db.execute(query, args)
396 cursor = self._db.execute(query, args)
394 matches = cursor.fetchall()
397 matches = cursor.fetchall()
395 records = []
398 records = []
396 for line in matches:
399 for line in matches:
397 rec = self._list_to_dict(line, keys)
400 rec = self._list_to_dict(line, keys)
398 records.append(rec)
401 records.append(rec)
399 return records
402 return records
400
403
401 def get_history(self):
404 def get_history(self):
402 """get all msg_ids, ordered by time submitted."""
405 """get all msg_ids, ordered by time submitted."""
403 query = """SELECT msg_id FROM %s ORDER by submitted ASC"""%self.table
406 query = """SELECT msg_id FROM %s ORDER by submitted ASC"""%self.table
404 cursor = self._db.execute(query)
407 cursor = self._db.execute(query)
405 # will be a list of length 1 tuples
408 # will be a list of length 1 tuples
406 return [ tup[0] for tup in cursor.fetchall()]
409 return [ tup[0] for tup in cursor.fetchall()]
407
410
408 __all__ = ['SQLiteDB'] No newline at end of file
411 __all__ = ['SQLiteDB']
General Comments 0
You need to be logged in to leave comments. Login now