def test_io_api():
"""Test that wrapped stdout has the same API as a normal TextIO object"""
session = Session()
ctx = zmq.Context()
pub = ctx.socket(zmq.PUB)
thread = IOPubThread(pub)
thread.start()
stream = OutStream(session, thread, 'stdout')
# cleanup unused zmq objects before we start testing
thread.stop()
thread.close()
ctx.term()
assert stream.errors is None
assert not stream.isatty()
with pytest.raises(io.UnsupportedOperation):
stream.detach()
with pytest.raises(io.UnsupportedOperation):
next(stream)
with pytest.raises(io.UnsupportedOperation):
stream.read()
with pytest.raises(io.UnsupportedOperation):
stream.readline()
with pytest.raises(io.UnsupportedOperation):
stream.seek(0)
with pytest.raises(io.UnsupportedOperation):
stream.tell()
with pytest.raises(TypeError):
stream.write(b'')
def test_io_isatty():
session = Session()
ctx = zmq.Context()
pub = ctx.socket(zmq.PUB)
thread = IOPubThread(pub)
thread.start()
stream = OutStream(session, thread, 'stdout', isatty=True)
assert stream.isatty()
def init_iopub(self, context):
self.iopub_socket = context.socket(zmq.PUB)
self.iopub_socket.linger = 1000
self.iopub_port = self._bind_socket(self.iopub_socket, self.iopub_port)
self.log.debug("iopub PUB Channel on port: %i" % self.iopub_port)
self.configure_tornado_logger()
self.iopub_thread = IOPubThread(self.iopub_socket, pipe=True)
self.iopub_thread.start()
# backward-compat: wrap iopub socket API in background thread
self.iopub_socket = self.iopub_thread.background_socket
def complete_registration(self, msg, connect, maybe_tunnel):
# print msg
self.loop.remove_timeout(self._abort_timeout)
ctx = self.context
loop = self.loop
identity = self.bident
idents, msg = self.session.feed_identities(msg)
msg = self.session.deserialize(msg)
content = msg['content']
info = self.connection_info
def url(key):
"""get zmq url for given channel"""
return str(info["interface"] + ":%i" % info[key])
if content['status'] == 'ok':
if self.id is not None and content['id'] != self.id:
self.log.warning("Did not get the requested id: %i != %i",
content['id'], self.id)
self.id = content['id']
# launch heartbeat
# possibly forward hb ports with tunnels
hb_ping = maybe_tunnel(url('hb_ping'))
hb_pong = maybe_tunnel(url('hb_pong'))
hb_monitor = None
if self.max_heartbeat_misses > 0:
# Add a monitor socket which will record the last time a ping was seen
mon = self.context.socket(zmq.SUB)
mport = mon.bind_to_random_port('tcp://%s' % localhost())
mon.setsockopt(zmq.SUBSCRIBE, b"")
self._hb_listener = zmqstream.ZMQStream(mon, self.loop)
self._hb_listener.on_recv(self._report_ping)
hb_monitor = "tcp://%s:%i" % (localhost(), mport)
heart = Heart(hb_ping, hb_pong, hb_monitor, heart_id=identity)
heart.start()
# create Shell Connections (MUX, Task, etc.):
shell_addrs = url('mux'), url('task')
# Use only one shell stream for mux and tasks
stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop)
stream.setsockopt(zmq.IDENTITY, identity)
shell_streams = [stream]
for addr in shell_addrs:
connect(stream, addr)
# control stream:
control_addr = url('control')
control_stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop)
control_stream.setsockopt(zmq.IDENTITY, identity)
connect(control_stream, control_addr)
# create iopub stream:
iopub_addr = url('iopub')
iopub_socket = ctx.socket(zmq.PUB)
iopub_socket.setsockopt(zmq.IDENTITY, identity)
connect(iopub_socket, iopub_addr)
try:
from ipykernel.iostream import IOPubThread
except ImportError:
pass
else:
iopub_socket = IOPubThread(iopub_socket)
iopub_socket.start()
# disable history:
self.config.HistoryManager.hist_file = ':memory:'
# Redirect input streams and set a display hook.
if self.out_stream_factory:
sys.stdout = self.out_stream_factory(self.session,
iopub_socket, u'stdout')
sys.stdout.topic = cast_bytes('engine.%i.stdout' % self.id)
sys.stderr = self.out_stream_factory(self.session,
iopub_socket, u'stderr')
sys.stderr.topic = cast_bytes('engine.%i.stderr' % self.id)
if self.display_hook_factory:
sys.displayhook = self.display_hook_factory(
self.session, iopub_socket)
sys.displayhook.topic = cast_bytes('engine.%i.execute_result' %
self.id)
self.kernel = Kernel(parent=self,
engine_id=self.id,
ident=self.ident,
session=self.session,
control_stream=control_stream,
shell_streams=shell_streams,
iopub_socket=iopub_socket,
loop=loop,
user_ns=self.user_ns,
log=self.log)
self.kernel.shell.display_pub.topic = cast_bytes(
'engine.%i.displaypub' % self.id)
# periodically check the heartbeat pings of the controller
# Should be started here and not in "start()" so that the right period can be taken
# from the hubs HeartBeatMonitor.period
if self.max_heartbeat_misses > 0:
# Use a slightly bigger check period than the hub signal period to not warn unnecessary
self.hb_check_period = int(content['hb_period']) + 10
self.log.info(
"Starting to monitor the heartbeat signal from the hub every %i ms.",
self.hb_check_period)
self._hb_reporter = ioloop.PeriodicCallback(
self._hb_monitor, self.hb_check_period)
self._hb_reporter.start()
else:
self.log.info(
"Monitoring of the heartbeat signal from the hub is not enabled."
)
# FIXME: This is a hack until IPKernelApp and IPEngineApp can be fully merged
app = IPKernelApp(parent=self,
shell=self.kernel.shell,
kernel=self.kernel,
log=self.log)
if self.use_mpi and self.init_mpi:
app.exec_lines.insert(0, self.init_mpi)
app.init_profile_dir()
app.init_code()
self.kernel.start()
else:
self.log.fatal("Registration Failed: %s" % msg)
raise Exception("Registration Failed: %s" % msg)
self.log.info("Completed registration with id %i" % self.id)
def _complete_registration(self, msg, connect, maybe_tunnel):
ctx = self.context
loop = self.loop
identity = self.bident
idents, msg = self.session.feed_identities(msg)
msg = self.session.deserialize(msg)
content = msg['content']
info = self.connection_info
def url(key):
"""get zmq url for given channel"""
return str(info["interface"] + ":%i" % info[key])
def urls(key):
return [f'{info["interface"]}:{port}' for port in info[key]]
if content['status'] == 'ok':
requested_id = self.id
self.id = content['id']
if requested_id is not None and self.id != requested_id:
self.log.warning(
f"Did not get the requested id: {self.id} != {requested_id}"
)
self.log.name = self.log.name.rsplit(".", 1)[0] + f".{self.id}"
elif self.id is None:
self.log.name += f".{self.id}"
# create Shell Connections (MUX, Task, etc.):
# select which broadcast endpoint to connect to
# use rank % len(broadcast_leaves)
broadcast_urls = urls('broadcast')
broadcast_leaves = len(broadcast_urls)
broadcast_index = self.id % len(broadcast_urls)
broadcast_url = broadcast_urls[broadcast_index]
shell_addrs = [url('mux'), url('task'), broadcast_url]
self.log.info(f'Shell_addrs: {shell_addrs}')
# Use only one shell stream for mux and tasks
stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop)
stream.setsockopt(zmq.IDENTITY, identity)
# TODO: enable PROBE_ROUTER when schedulers can handle the empty message
# stream.setsockopt(zmq.PROBE_ROUTER, 1)
self.log.debug("Setting shell identity %r", identity)
shell_streams = [stream]
for addr in shell_addrs:
self.log.info("Connecting shell to %s", addr)
connect(stream, addr)
# control stream:
control_url = url('control')
curve_serverkey = self.curve_serverkey
if self.enable_nanny:
nanny_url, self.nanny_pipe = self.start_nanny(
control_url=control_url, )
control_url = nanny_url
# nanny uses our curve_publickey, not the controller's publickey
curve_serverkey = self.curve_publickey
control_stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop)
control_stream.setsockopt(zmq.IDENTITY, identity)
connect(control_stream,
control_url,
curve_serverkey=curve_serverkey)
# create iopub stream:
iopub_addr = url('iopub')
iopub_socket = ctx.socket(zmq.PUB)
iopub_socket.SNDHWM = 0
iopub_socket.setsockopt(zmq.IDENTITY, identity)
connect(iopub_socket, iopub_addr)
try:
from ipykernel.iostream import IOPubThread
except ImportError:
pass
else:
iopub_socket = IOPubThread(iopub_socket)
iopub_socket.start()
# disable history:
self.config.HistoryManager.hist_file = ':memory:'
# Redirect input streams and set a display hook.
if self.out_stream_factory:
sys.stdout = self.out_stream_factory(self.session,
iopub_socket, 'stdout')
sys.stdout.topic = f"engine.{self.id}.stdout".encode("ascii")
sys.stderr = self.out_stream_factory(self.session,
iopub_socket, 'stderr')
sys.stderr.topic = f"engine.{self.id}.stderr".encode("ascii")
# copied from ipykernel 6, which captures sys.__stderr__ at the FD-level
if getattr(sys.stderr, "_original_stdstream_copy",
None) is not None:
for handler in self.log.handlers:
if isinstance(handler, StreamHandler) and (
handler.stream.buffer.fileno() == 2):
self.log.debug(
"Seeing logger to stderr, rerouting to raw filedescriptor."
)
handler.stream = TextIOWrapper(
FileIO(sys.stderr._original_stdstream_copy,
"w"))
if self.display_hook_factory:
sys.displayhook = self.display_hook_factory(
self.session, iopub_socket)
sys.displayhook.topic = f"engine.{self.id}.execute_result".encode(
"ascii")
# patch Session to always send engine uuid metadata
original_send = self.session.send
def send_with_metadata(
stream,
msg_or_type,
content=None,
parent=None,
ident=None,
buffers=None,
track=False,
header=None,
metadata=None,
**kwargs,
):
"""Ensure all messages set engine uuid metadata"""
metadata = metadata or {}
metadata.setdefault("engine", self.ident)
return original_send(
stream,
msg_or_type,
content=content,
parent=parent,
ident=ident,
buffers=buffers,
track=track,
header=header,
metadata=metadata,
**kwargs,
)
self.session.send = send_with_metadata
self.kernel = Kernel.instance(
parent=self,
engine_id=self.id,
ident=self.ident,
session=self.session,
control_stream=control_stream,
shell_streams=shell_streams,
iopub_socket=iopub_socket,
loop=loop,
user_ns=self.user_ns,
log=self.log,
)
self.kernel.shell.display_pub.topic = f"engine.{self.id}.displaypub".encode(
"ascii")
# FIXME: This is a hack until IPKernelApp and IPEngineApp can be fully merged
self.init_signal()
app = IPKernelApp(parent=self,
shell=self.kernel.shell,
kernel=self.kernel,
log=self.log)
if self.use_mpi and self.init_mpi:
app.exec_lines.insert(0, self.init_mpi)
app.init_profile_dir()
app.init_code()
self.kernel.start()
else:
self.log.fatal("Registration Failed: %s" % msg)
raise Exception("Registration Failed: %s" % msg)
self.start_heartbeat(
maybe_tunnel(url('hb_ping')),
maybe_tunnel(url('hb_pong')),
content['hb_period'],
identity,
)
self.log.info("Completed registration with id %i" % self.id)
self.loop.remove_timeout(self._abort_timeout)
