def execute_request(self, stream, ident, parent):
self.log.debug('execute request %s'%parent)
try:
code = parent['content']['code']
except:
self.log.error("Got bad msg: %s"%parent, exc_info=True)
return
self.session.send(self.iopub_stream, 'pyin', {'code':code},parent=parent,
ident=asbytes('%s.pyin'%self.prefix))
started = datetime.now()
try:
comp_code = self.compiler(code, '<zmq-kernel>')
# allow for not overriding displayhook
if hasattr(sys.displayhook, 'set_parent'):
sys.displayhook.set_parent(parent)
sys.stdout.set_parent(parent)
sys.stderr.set_parent(parent)
exec(comp_code, self.user_ns, self.user_ns)
except:
exc_content = self._wrap_exception('execute')
# exc_msg = self.session.msg(u'pyerr', exc_content, parent)
self.session.send(self.iopub_stream, 'pyerr', exc_content, parent=parent,
ident=asbytes('%s.pyerr'%self.prefix))
reply_content = exc_content
else:
reply_content = {'status' : 'ok'}
reply_msg = self.session.send(stream, 'execute_reply', reply_content, parent=parent,
ident=ident, subheader = dict(started=started))
self.log.debug(str(reply_msg))
if reply_msg['content']['status'] == 'error':
self.abort_queues()
def execute_request(self, stream, ident, parent):
self.log.debug('execute request %s'%parent)
try:
code = parent[u'content'][u'code']
except:
self.log.error("Got bad msg: %s"%parent, exc_info=True)
return
self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent,
ident=asbytes('%s.pyin'%self.prefix))
started = datetime.now()
try:
comp_code = self.compiler(code, '<zmq-kernel>')
# allow for not overriding displayhook
if hasattr(sys.displayhook, 'set_parent'):
sys.displayhook.set_parent(parent)
sys.stdout.set_parent(parent)
sys.stderr.set_parent(parent)
exec comp_code in self.user_ns, self.user_ns
except:
exc_content = self._wrap_exception('execute')
# exc_msg = self.session.msg(u'pyerr', exc_content, parent)
self.session.send(self.iopub_stream, u'pyerr', exc_content, parent=parent,
ident=asbytes('%s.pyerr'%self.prefix))
reply_content = exc_content
else:
reply_content = {'status' : 'ok'}
reply_msg = self.session.send(stream, u'execute_reply', reply_content, parent=parent,
ident=ident, subheader = dict(started=started))
self.log.debug(str(reply_msg))
if reply_msg['content']['status'] == u'error':
self.abort_queues()
def handle_pong(self, msg):
"a heart just beat"
current = asbytes(str(self.lifetime))
last = asbytes(str(self.last_ping))
if msg[1] == current:
delta = time.time()-self.tic
# self.log.debug("heartbeat::heart %r took %.2f ms to respond"%(msg[0], 1000*delta))
self.responses.add(msg[0])
elif msg[1] == last:
delta = time.time()-self.tic + (self.lifetime-self.last_ping)
self.log.warn("heartbeat::heart %r missed a beat, and took %.2f ms to respond", msg[0], 1000*delta)
self.responses.add(msg[0])
else:
self.log.warn("heartbeat::got bad heartbeat (possibly old?): %s (current=%.3f)", msg[1], self.lifetime)
def handle_pong(self, msg):
"a heart just beat"
current = asbytes(str(self.lifetime))
last = asbytes(str(self.last_ping))
if msg[1] == current:
delta = time.time() - self.tic
# self.log.debug("heartbeat::heart %r took %.2f ms to respond"%(msg[0], 1000*delta))
self.responses.add(msg[0])
elif msg[1] == last:
delta = time.time() - self.tic + (self.lifetime - self.last_ping)
self.log.warn("heartbeat::heart %r missed a beat, and took %.2f ms to respond", msg[0], 1000 * delta)
self.responses.add(msg[0])
else:
self.log.warn("heartbeat::got bad heartbeat (possibly old?): %s (current=%.3f)", msg[1], self.lifetime)
def load_config_from_json(self):
"""load config from existing json connector files."""
c = self.config
self.log.debug("loading config from JSON")
# load from engine config
fname = os.path.join(self.profile_dir.security_dir, self.engine_json_file)
self.log.info("loading connection info from %s", fname)
with open(fname) as f:
cfg = json.loads(f.read())
key = c.Session.key = asbytes(cfg['exec_key'])
xport,addr = cfg['url'].split('://')
c.HubFactory.engine_transport = xport
ip,ports = addr.split(':')
c.HubFactory.engine_ip = ip
c.HubFactory.regport = int(ports)
self.location = cfg['location']
if not self.engine_ssh_server:
self.engine_ssh_server = cfg['ssh']
# load client config
fname = os.path.join(self.profile_dir.security_dir, self.client_json_file)
self.log.info("loading connection info from %s", fname)
with open(fname) as f:
cfg = json.loads(f.read())
assert key == cfg['exec_key'], "exec_key mismatch between engine and client keys"
xport,addr = cfg['url'].split('://')
c.HubFactory.client_transport = xport
ip,ports = addr.split(':')
c.HubFactory.client_ip = ip
if not self.ssh_server:
self.ssh_server = cfg['ssh']
assert int(ports) == c.HubFactory.regport, "regport mismatch"
def dispatch_notification(self, msg):
"""dispatch register/unregister events."""
try:
idents, msg = self.session.feed_identities(msg)
except ValueError:
self.log.warn("task::Invalid Message: %r", msg)
return
try:
msg = self.session.unserialize(msg)
except ValueError:
self.log.warn("task::Unauthorized message from: %r" % idents)
return
msg_type = msg['header']['msg_type']
handler = self._notification_handlers.get(msg_type, None)
if handler is None:
self.log.error("Unhandled message type: %r" % msg_type)
else:
try:
handler(asbytes(msg['content']['queue']))
except Exception:
self.log.error("task::Invalid notification msg: %r",
msg,
exc_info=True)
def load_config_from_json(self):
"""load config from existing json connector files."""
c = self.config
self.log.debug("loading config from JSON")
# load from engine config
with open(
os.path.join(self.profile_dir.security_dir,
self.engine_json_file)) as f:
cfg = json.loads(f.read())
key = c.Session.key = asbytes(cfg['exec_key'])
xport, addr = cfg['url'].split('://')
c.HubFactory.engine_transport = xport
ip, ports = addr.split(':')
c.HubFactory.engine_ip = ip
c.HubFactory.regport = int(ports)
self.location = cfg['location']
if not self.engine_ssh_server:
self.engine_ssh_server = cfg['ssh']
# load client config
with open(
os.path.join(self.profile_dir.security_dir,
self.client_json_file)) as f:
cfg = json.loads(f.read())
assert key == cfg[
'exec_key'], "exec_key mismatch between engine and client keys"
xport, addr = cfg['url'].split('://')
c.HubFactory.client_transport = xport
ip, ports = addr.split(':')
c.HubFactory.client_ip = ip
if not self.ssh_server:
self.ssh_server = cfg['ssh']
assert int(ports) == c.HubFactory.regport, "regport mismatch"
def load_connector_file(self):
"""load config from a JSON connector file,
at a *lower* priority than command-line/config files.
"""
self.log.info("Loading url_file %r", self.url_file)
config = self.config
with open(self.url_file) as f:
d = json.loads(f.read())
if 'exec_key' in d:
config.Session.key = asbytes(d['exec_key'])
try:
config.EngineFactory.location
except AttributeError:
config.EngineFactory.location = d['location']
d['url'] = disambiguate_url(d['url'], config.EngineFactory.location)
try:
config.EngineFactory.url
except AttributeError:
config.EngineFactory.url = d['url']
try:
config.EngineFactory.sshserver
except AttributeError:
config.EngineFactory.sshserver = d['ssh']
def load_connector_file(self):
"""load config from a JSON connector file,
at a *lower* priority than command-line/config files.
"""
self.log.info("Loading url_file %r"%self.url_file)
config = self.config
with open(self.url_file) as f:
d = json.loads(f.read())
if 'exec_key' in d:
config.Session.key = asbytes(d['exec_key'])
try:
config.EngineFactory.location
except AttributeError:
config.EngineFactory.location = d['location']
d['url'] = disambiguate_url(d['url'], config.EngineFactory.location)
try:
config.EngineFactory.url
except AttributeError:
config.EngineFactory.url = d['url']
try:
config.EngineFactory.sshserver
except AttributeError:
config.EngineFactory.sshserver = d['ssh']
def init_hub(self):
c = self.config
self.do_import_statements()
reusing = self.reuse_files
if reusing:
try:
self.load_config_from_json()
except (AssertionError, IOError):
reusing = False
# check again, because reusing may have failed:
if reusing:
pass
elif self.secure:
key = str(uuid.uuid4())
# keyfile = os.path.join(self.profile_dir.security_dir, self.exec_key)
# with open(keyfile, 'w') as f:
# f.write(key)
# os.chmod(keyfile, stat.S_IRUSR|stat.S_IWUSR)
c.Session.key = asbytes(key)
else:
key = c.Session.key = b''
try:
self.factory = HubFactory(config=c, log=self.log)
# self.start_logging()
self.factory.init_hub()
except:
self.log.error("Couldn't construct the Controller", exc_info=True)
self.exit(1)
if not reusing:
# save to new json config files
f = self.factory
cdict = {
'exec_key': key,
'ssh': self.ssh_server,
'url':
"%s://%s:%s" % (f.client_transport, f.client_ip, f.regport),
'location': self.location
}
self.save_connection_dict('ipcontroller-client.json', cdict)
edict = cdict
edict['url'] = "%s://%s:%s" % (
(f.client_transport, f.client_ip, f.regport))
self.save_connection_dict('ipcontroller-engine.json', edict)
def init_hub(self):
c = self.config
self.do_import_statements()
reusing = self.reuse_files
if reusing:
try:
self.load_config_from_json()
except (AssertionError,IOError):
reusing=False
# check again, because reusing may have failed:
if reusing:
pass
elif self.secure:
key = str(uuid.uuid4())
# keyfile = os.path.join(self.profile_dir.security_dir, self.exec_key)
# with open(keyfile, 'w') as f:
# f.write(key)
# os.chmod(keyfile, stat.S_IRUSR|stat.S_IWUSR)
c.Session.key = asbytes(key)
else:
key = c.Session.key = b''
try:
self.factory = HubFactory(config=c, log=self.log)
# self.start_logging()
self.factory.init_hub()
except:
self.log.error("Couldn't construct the Controller", exc_info=True)
self.exit(1)
if not reusing:
# save to new json config files
f = self.factory
cdict = {'exec_key' : key,
'ssh' : self.ssh_server,
'url' : "%s://%s:%s"%(f.client_transport, f.client_ip, f.regport),
'location' : self.location
}
self.save_connection_dict('ipcontroller-client.json', cdict)
edict = cdict
edict['url']="%s://%s:%s"%((f.client_transport, f.client_ip, f.regport))
edict['ssh'] = self.engine_ssh_server
self.save_connection_dict('ipcontroller-engine.json', edict)
def beat(self):
self.pongstream.flush()
self.last_ping = self.lifetime
toc = time.time()
self.lifetime += toc - self.tic
self.tic = toc
# self.log.debug("heartbeat::%s"%self.lifetime)
goodhearts = self.hearts.intersection(self.responses)
missed_beats = self.hearts.difference(goodhearts)
heartfailures = self.on_probation.intersection(missed_beats)
newhearts = self.responses.difference(goodhearts)
map(self.handle_new_heart, newhearts)
map(self.handle_heart_failure, heartfailures)
self.on_probation = missed_beats.intersection(self.hearts)
self.responses = set()
# print self.on_probation, self.hearts
# self.log.debug("heartbeat::beat %.3f, %i beating hearts"%(self.lifetime, len(self.hearts)))
self.pingstream.send(asbytes(str(self.lifetime)))
def beat(self):
self.pongstream.flush()
self.last_ping = self.lifetime
toc = time.time()
self.lifetime += toc-self.tic
self.tic = toc
# self.log.debug("heartbeat::%s"%self.lifetime)
goodhearts = self.hearts.intersection(self.responses)
missed_beats = self.hearts.difference(goodhearts)
heartfailures = self.on_probation.intersection(missed_beats)
newhearts = self.responses.difference(goodhearts)
map(self.handle_new_heart, newhearts)
map(self.handle_heart_failure, heartfailures)
self.on_probation = missed_beats.intersection(self.hearts)
self.responses = set()
# print self.on_probation, self.hearts
# self.log.debug("heartbeat::beat %.3f, %i beating hearts"%(self.lifetime, len(self.hearts)))
self.pingstream.send(asbytes(str(self.lifetime)))
def load_config_from_json(self):
"""load config from existing json connector files."""
c = self.config
# load from engine config
with open(os.path.join(self.profile_dir.security_dir, 'ipcontroller-engine.json')) as f:
cfg = json.loads(f.read())
key = c.Session.key = asbytes(cfg['exec_key'])
xport,addr = cfg['url'].split('://')
c.HubFactory.engine_transport = xport
ip,ports = addr.split(':')
c.HubFactory.engine_ip = ip
c.HubFactory.regport = int(ports)
self.location = cfg['location']
# load client config
with open(os.path.join(self.profile_dir.security_dir, 'ipcontroller-client.json')) as f:
cfg = json.loads(f.read())
assert key == cfg['exec_key'], "exec_key mismatch between engine and client keys"
xport,addr = cfg['url'].split('://')
c.HubFactory.client_transport = xport
ip,ports = addr.split(':')
c.HubFactory.client_ip = ip
self.ssh_server = cfg['ssh']
assert int(ports) == c.HubFactory.regport, "regport mismatch"
def dispatch_notification(self, msg):
"""dispatch register/unregister events."""
try:
idents,msg = self.session.feed_identities(msg)
except ValueError:
self.log.warn("task::Invalid Message: %r",msg)
return
try:
msg = self.session.unserialize(msg)
except ValueError:
self.log.warn("task::Unauthorized message from: %r"%idents)
return
msg_type = msg['header']['msg_type']
handler = self._notification_handlers.get(msg_type, None)
if handler is None:
self.log.error("Unhandled message type: %r"%msg_type)
else:
try:
handler(asbytes(msg['content']['queue']))
except Exception:
self.log.error("task::Invalid notification msg: %r", msg, exc_info=True)
def save_task_destination(self, idents, msg):
try:
msg = self.session.unserialize(msg, content=True)
except Exception:
self.log.error("task::invalid task tracking message", exc_info=True)
return
content = msg['content']
# print (content)
msg_id = content['msg_id']
engine_uuid = content['engine_id']
eid = self.by_ident[util.asbytes(engine_uuid)]
self.log.info("task::task %r arrived on %r", msg_id, eid)
if msg_id in self.unassigned:
self.unassigned.remove(msg_id)
# else:
# self.log.debug("task::task %r not listed as MIA?!"%(msg_id))
self.tasks[eid].append(msg_id)
# self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
try:
self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
except Exception:
self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)
def save_task_destination(self, idents, msg):
try:
msg = self.session.unserialize(msg, content=True)
except Exception:
self.log.error("task::invalid task tracking message", exc_info=True)
return
content = msg['content']
# print (content)
msg_id = content['msg_id']
engine_uuid = content['engine_id']
eid = self.by_ident[util.asbytes(engine_uuid)]
self.log.info("task::task %r arrived on %r", msg_id, eid)
if msg_id in self.unassigned:
self.unassigned.remove(msg_id)
# else:
# self.log.debug("task::task %r not listed as MIA?!"%(msg_id))
self.tasks[eid].append(msg_id)
# self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
try:
self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
except Exception:
self.log.error("DB Error saving task destination %r", msg_id, exc_info=True)
def _ident_default(self):
return asbytes(self.session.session)
def _ident_default(self):
return asbytes(self.session.session)
def init_engine(self):
# This is the working dir by now.
sys.path.insert(0, '')
config = self.config
# print config
self.find_url_file()
# was the url manually specified?
keys = set(self.config.EngineFactory.keys())
keys = keys.union(set(self.config.RegistrationFactory.keys()))
if keys.intersection(set(['ip', 'url', 'port'])):
# Connection info was specified, don't wait for the file
url_specified = True
self.wait_for_url_file = 0
else:
url_specified = False
if self.wait_for_url_file and not os.path.exists(self.url_file):
self.log.warn("url_file %r not found" % self.url_file)
self.log.warn("Waiting up to %.1f seconds for it to arrive." %
self.wait_for_url_file)
tic = time.time()
while not os.path.exists(self.url_file) and (
time.time() - tic < self.wait_for_url_file):
# wait for url_file to exist, for up to 10 seconds
time.sleep(0.1)
if os.path.exists(self.url_file):
self.log.info("Loading url_file %r" % self.url_file)
with open(self.url_file) as f:
d = json.loads(f.read())
if d['exec_key']:
config.Session.key = asbytes(d['exec_key'])
d['url'] = disambiguate_url(d['url'], d['location'])
config.EngineFactory.url = d['url']
config.EngineFactory.location = d['location']
elif not url_specified:
self.log.critical("Fatal: url file never arrived: %s" %
self.url_file)
self.exit(1)
try:
exec_lines = config.Kernel.exec_lines
except AttributeError:
config.Kernel.exec_lines = []
exec_lines = config.Kernel.exec_lines
if self.startup_script:
enc = sys.getfilesystemencoding() or 'utf8'
cmd = "execfile(%r)" % self.startup_script.encode(enc)
exec_lines.append(cmd)
if self.startup_command:
exec_lines.append(self.startup_command)
# Create the underlying shell class and Engine
# shell_class = import_item(self.master_config.Global.shell_class)
# print self.config
try:
self.engine = EngineFactory(config=config, log=self.log)
except:
self.log.error("Couldn't start the Engine", exc_info=True)
self.exit(1)
def register_engine(self, reg, msg):
"""Register a new engine."""
content = msg['content']
try:
queue = util.asbytes(content['queue'])
except KeyError:
self.log.error("registration::queue not specified", exc_info=True)
return
heart = content.get('heartbeat', None)
if heart:
heart = util.asbytes(heart)
"""register a new engine, and create the socket(s) necessary"""
eid = self._next_id
# print (eid, queue, reg, heart)
self.log.debug("registration::register_engine(%i, %r, %r, %r)", eid, queue, reg, heart)
content = dict(id=eid,status='ok')
content.update(self.engine_info)
# check if requesting available IDs:
if queue in self.by_ident:
try:
raise KeyError("queue_id %r in use" % queue)
except:
content = error.wrap_exception()
self.log.error("queue_id %r in use", queue, exc_info=True)
elif heart in self.hearts: # need to check unique hearts?
try:
raise KeyError("heart_id %r in use" % heart)
except:
self.log.error("heart_id %r in use", heart, exc_info=True)
content = error.wrap_exception()
else:
for h, pack in self.incoming_registrations.iteritems():
if heart == h:
try:
raise KeyError("heart_id %r in use" % heart)
except:
self.log.error("heart_id %r in use", heart, exc_info=True)
content = error.wrap_exception()
break
elif queue == pack[1]:
try:
raise KeyError("queue_id %r in use" % queue)
except:
self.log.error("queue_id %r in use", queue, exc_info=True)
content = error.wrap_exception()
break
msg = self.session.send(self.query, "registration_reply",
content=content,
ident=reg)
if content['status'] == 'ok':
if heart in self.heartmonitor.hearts:
# already beating
self.incoming_registrations[heart] = (eid,queue,reg[0],None)
self.finish_registration(heart)
else:
purge = lambda : self._purge_stalled_registration(heart)
dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
dc.start()
self.incoming_registrations[heart] = (eid,queue,reg[0],dc)
else:
self.log.error("registration::registration %i failed: %r", eid, content['evalue'])
return eid
def apply_request(self, stream, ident, parent):
# flush previous reply, so this request won't block it
stream.flush(zmq.POLLOUT)
try:
content = parent[u'content']
bufs = parent[u'buffers']
msg_id = parent['header']['msg_id']
# bound = parent['header'].get('bound', False)
except:
self.log.error("Got bad msg: %s" % parent, exc_info=True)
return
# pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
# self.iopub_stream.send(pyin_msg)
# self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent)
sub = {
'dependencies_met': True,
'engine': self.ident,
'started': datetime.now()
}
try:
# allow for not overriding displayhook
if hasattr(sys.displayhook, 'set_parent'):
sys.displayhook.set_parent(parent)
sys.stdout.set_parent(parent)
sys.stderr.set_parent(parent)
# exec "f(*args,**kwargs)" in self.user_ns, self.user_ns
working = self.user_ns
# suffix =
prefix = "_" + str(msg_id).replace("-", "") + "_"
f, args, kwargs = unpack_apply_message(bufs, working, copy=False)
# if bound:
# bound_ns = Namespace(working)
# args = [bound_ns]+list(args)
fname = getattr(f, '__name__', 'f')
fname = prefix + "f"
argname = prefix + "args"
kwargname = prefix + "kwargs"
resultname = prefix + "result"
ns = {fname: f, argname: args, kwargname: kwargs, resultname: None}
# print ns
working.update(ns)
code = "%s=%s(*%s,**%s)" % (resultname, fname, argname, kwargname)
try:
exec code in working, working
result = working.get(resultname)
finally:
for key in ns.iterkeys():
working.pop(key)
# if bound:
# working.update(bound_ns)
packed_result, buf = serialize_object(result)
result_buf = [packed_result] + buf
except:
exc_content = self._wrap_exception('apply')
# exc_msg = self.session.msg(u'pyerr', exc_content, parent)
self.session.send(self.iopub_stream,
u'pyerr',
exc_content,
parent=parent,
ident=asbytes('%s.pyerr' % self.prefix))
reply_content = exc_content
result_buf = []
if exc_content['ename'] == 'UnmetDependency':
sub['dependencies_met'] = False
else:
reply_content = {'status': 'ok'}
# put 'ok'/'error' status in header, for scheduler introspection:
sub['status'] = reply_content['status']
reply_msg = self.session.send(stream,
u'apply_reply',
reply_content,
parent=parent,
ident=ident,
buffers=result_buf,
subheader=sub)
# flush i/o
# should this be before reply_msg is sent, like in the single-kernel code,
# or should nothing get in the way of real results?
sys.stdout.flush()
sys.stderr.flush()
def init_hub(self):
"""construct"""
client_iface = "%s://%s:"%(self.client_transport, self.client_ip) + "%i"
engine_iface = "%s://%s:"%(self.engine_transport, self.engine_ip) + "%i"
ctx = self.context
loop = self.loop
# Registrar socket
q = ZMQStream(ctx.socket(zmq.XREP), loop)
q.bind(client_iface % self.regport)
self.log.info("Hub listening on %s for registration."%(client_iface%self.regport))
if self.client_ip != self.engine_ip:
q.bind(engine_iface % self.regport)
self.log.info("Hub listening on %s for registration."%(engine_iface%self.regport))
### Engine connections ###
# heartbeat
hpub = ctx.socket(zmq.PUB)
hpub.bind(engine_iface % self.hb[0])
hrep = ctx.socket(zmq.XREP)
hrep.bind(engine_iface % self.hb[1])
self.heartmonitor = HeartMonitor(loop=loop, config=self.config, log=self.log,
pingstream=ZMQStream(hpub,loop),
pongstream=ZMQStream(hrep,loop)
)
### Client connections ###
# Notifier socket
n = ZMQStream(ctx.socket(zmq.PUB), loop)
n.bind(client_iface%self.notifier_port)
### build and launch the queues ###
# monitor socket
sub = ctx.socket(zmq.SUB)
sub.setsockopt(zmq.SUBSCRIBE, b"")
sub.bind(self.monitor_url)
sub.bind('inproc://monitor')
sub = ZMQStream(sub, loop)
# connect the db
self.log.info('Hub using DB backend: %r'%(self.db_class.split()[-1]))
# cdir = self.config.Global.cluster_dir
self.db = import_item(str(self.db_class))(session=self.session.session,
config=self.config, log=self.log)
time.sleep(.25)
try:
scheme = self.config.TaskScheduler.scheme_name
except AttributeError:
from .scheduler import TaskScheduler
scheme = TaskScheduler.scheme_name.get_default_value()
# build connection dicts
self.engine_info = {
'control' : engine_iface%self.control[1],
'mux': engine_iface%self.mux[1],
'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]),
'task' : engine_iface%self.task[1],
'iopub' : engine_iface%self.iopub[1],
# 'monitor' : engine_iface%self.mon_port,
}
self.client_info = {
'control' : client_iface%self.control[0],
'mux': client_iface%self.mux[0],
'task' : (scheme, client_iface%self.task[0]),
'iopub' : client_iface%self.iopub[0],
'notification': client_iface%self.notifier_port
}
self.log.debug("Hub engine addrs: %s"%self.engine_info)
self.log.debug("Hub client addrs: %s"%self.client_info)
# resubmit stream
r = ZMQStream(ctx.socket(zmq.XREQ), loop)
url = util.disambiguate_url(self.client_info['task'][-1])
r.setsockopt(zmq.IDENTITY, util.asbytes(self.session.session))
r.connect(url)
self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
query=q, notifier=n, resubmit=r, db=self.db,
engine_info=self.engine_info, client_info=self.client_info,
log=self.log)
def init_engine(self):
# This is the working dir by now.
sys.path.insert(0, '')
config = self.config
# print config
self.find_url_file()
# was the url manually specified?
keys = set(self.config.EngineFactory.keys())
keys = keys.union(set(self.config.RegistrationFactory.keys()))
if keys.intersection(set(['ip', 'url', 'port'])):
# Connection info was specified, don't wait for the file
url_specified = True
self.wait_for_url_file = 0
else:
url_specified = False
if self.wait_for_url_file and not os.path.exists(self.url_file):
self.log.warn("url_file %r not found"%self.url_file)
self.log.warn("Waiting up to %.1f seconds for it to arrive."%self.wait_for_url_file)
tic = time.time()
while not os.path.exists(self.url_file) and (time.time()-tic < self.wait_for_url_file):
# wait for url_file to exist, for up to 10 seconds
time.sleep(0.1)
if os.path.exists(self.url_file):
self.log.info("Loading url_file %r"%self.url_file)
with open(self.url_file) as f:
d = json.loads(f.read())
if d['exec_key']:
config.Session.key = asbytes(d['exec_key'])
d['url'] = disambiguate_url(d['url'], d['location'])
config.EngineFactory.url = d['url']
config.EngineFactory.location = d['location']
elif not url_specified:
self.log.critical("Fatal: url file never arrived: %s"%self.url_file)
self.exit(1)
try:
exec_lines = config.Kernel.exec_lines
except AttributeError:
config.Kernel.exec_lines = []
exec_lines = config.Kernel.exec_lines
if self.startup_script:
enc = sys.getfilesystemencoding() or 'utf8'
cmd="execfile(%r)"%self.startup_script.encode(enc)
exec_lines.append(cmd)
if self.startup_command:
exec_lines.append(self.startup_command)
# Create the underlying shell class and Engine
# shell_class = import_item(self.master_config.Global.shell_class)
# print self.config
try:
self.engine = EngineFactory(config=config, log=self.log)
except:
self.log.error("Couldn't start the Engine", exc_info=True)
self.exit(1)
def _ident_changed(self, name, old, new):
self.bident = asbytes(new)
def apply_request(self, stream, ident, parent):
# flush previous reply, so this request won't block it
stream.flush(zmq.POLLOUT)
try:
content = parent['content']
bufs = parent['buffers']
msg_id = parent['header']['msg_id']
# bound = parent['header'].get('bound', False)
except:
self.log.error("Got bad msg: %s"%parent, exc_info=True)
return
# pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
# self.iopub_stream.send(pyin_msg)
# self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent)
sub = {'dependencies_met' : True, 'engine' : self.ident,
'started': datetime.now()}
try:
# allow for not overriding displayhook
if hasattr(sys.displayhook, 'set_parent'):
sys.displayhook.set_parent(parent)
sys.stdout.set_parent(parent)
sys.stderr.set_parent(parent)
# exec "f(*args,**kwargs)" in self.user_ns, self.user_ns
working = self.user_ns
# suffix =
prefix = "_"+str(msg_id).replace("-","")+"_"
f,args,kwargs = unpack_apply_message(bufs, working, copy=False)
# if bound:
# bound_ns = Namespace(working)
# args = [bound_ns]+list(args)
fname = getattr(f, '__name__', 'f')
fname = prefix+"f"
argname = prefix+"args"
kwargname = prefix+"kwargs"
resultname = prefix+"result"
ns = { fname : f, argname : args, kwargname : kwargs , resultname : None }
# print ns
working.update(ns)
code = "%s=%s(*%s,**%s)"%(resultname, fname, argname, kwargname)
try:
exec(code, working,working)
result = working.get(resultname)
finally:
for key in ns.keys():
working.pop(key)
# if bound:
# working.update(bound_ns)
packed_result,buf = serialize_object(result)
result_buf = [packed_result]+buf
except:
exc_content = self._wrap_exception('apply')
# exc_msg = self.session.msg(u'pyerr', exc_content, parent)
self.session.send(self.iopub_stream, 'pyerr', exc_content, parent=parent,
ident=asbytes('%s.pyerr'%self.prefix))
reply_content = exc_content
result_buf = []
if exc_content['ename'] == 'UnmetDependency':
sub['dependencies_met'] = False
else:
reply_content = {'status' : 'ok'}
# put 'ok'/'error' status in header, for scheduler introspection:
sub['status'] = reply_content['status']
reply_msg = self.session.send(stream, 'apply_reply', reply_content,
parent=parent, ident=ident,buffers=result_buf, subheader=sub)
# flush i/o
# should this be before reply_msg is sent, like in the single-kernel code,
# or should nothing get in the way of real results?
sys.stdout.flush()
sys.stderr.flush()
def _ident_changed(self, name, old, new):
self.bident = asbytes(new)
def register_engine(self, reg, msg):
"""Register a new engine."""
content = msg['content']
try:
queue = util.asbytes(content['queue'])
except KeyError:
self.log.error("registration::queue not specified", exc_info=True)
return
heart = content.get('heartbeat', None)
if heart:
heart = util.asbytes(heart)
"""register a new engine, and create the socket(s) necessary"""
eid = self._next_id
# print (eid, queue, reg, heart)
self.log.debug("registration::register_engine(%i, %r, %r, %r)", eid, queue, reg, heart)
content = dict(id=eid,status='ok')
content.update(self.engine_info)
# check if requesting available IDs:
if queue in self.by_ident:
try:
raise KeyError("queue_id %r in use" % queue)
except:
content = error.wrap_exception()
self.log.error("queue_id %r in use", queue, exc_info=True)
elif heart in self.hearts: # need to check unique hearts?
try:
raise KeyError("heart_id %r in use" % heart)
except:
self.log.error("heart_id %r in use", heart, exc_info=True)
content = error.wrap_exception()
else:
for h, pack in self.incoming_registrations.iteritems():
if heart == h:
try:
raise KeyError("heart_id %r in use" % heart)
except:
self.log.error("heart_id %r in use", heart, exc_info=True)
content = error.wrap_exception()
break
elif queue == pack[1]:
try:
raise KeyError("queue_id %r in use" % queue)
except:
self.log.error("queue_id %r in use", queue, exc_info=True)
content = error.wrap_exception()
break
msg = self.session.send(self.query, "registration_reply",
content=content,
ident=reg)
if content['status'] == 'ok':
if heart in self.heartmonitor.hearts:
# already beating
self.incoming_registrations[heart] = (eid,queue,reg[0],None)
self.finish_registration(heart)
else:
purge = lambda : self._purge_stalled_registration(heart)
dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
dc.start()
self.incoming_registrations[heart] = (eid,queue,reg[0],dc)
else:
self.log.error("registration::registration %i failed: %r", eid, content['evalue'])
return eid
