# coding: utf-8 """Wrappers for forwarding stdout/stderr over zmq""" # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. from __future__ import print_function import atexit from binascii import b2a_hex from collections import deque try: from importlib import lock_held as import_lock_held except ImportError: from imp import lock_held as import_lock_held import os import sys import threading import warnings from io import StringIO, TextIOBase import zmq from zmq.eventloop.ioloop import IOLoop from zmq.eventloop.zmqstream import ZMQStream from jupyter_client.session import extract_header from ipython_genutils import py3compat from ipython_genutils.py3compat import unicode_type #----------------------------------------------------------------------------- # Globals #----------------------------------------------------------------------------- MASTER = 0 CHILD = 1 #----------------------------------------------------------------------------- # IO classes #----------------------------------------------------------------------------- class IOPubThread(object): """An object for sending IOPub messages in a background thread Prevents a blocking main thread from delaying output from threads. IOPubThread(pub_socket).background_socket is a Socket-API-providing object whose IO is always run in a thread. """ def __init__(self, socket, pipe=False): """Create IOPub thread Parameters ---------- socket: zmq.PUB Socket the socket on which messages will be sent. pipe: bool Whether this process should listen for IOPub messages piped from subprocesses. """ self.socket = socket self.background_socket = BackgroundSocket(self) self._master_pid = os.getpid() self._pipe_flag = pipe self.io_loop = IOLoop(make_current=False) if pipe: self._setup_pipe_in() self._local = threading.local() self._events = deque() self._setup_event_pipe() self.thread = threading.Thread(target=self._thread_main) self.thread.daemon = True def _thread_main(self): """The inner loop that's actually run in a thread""" self.io_loop.make_current() self.io_loop.start() self.io_loop.close(all_fds=True) def _setup_event_pipe(self): """Create the PULL socket listening for events that should fire in this thread.""" ctx = self.socket.context pipe_in = ctx.socket(zmq.PULL) pipe_in.linger = 0 _uuid = b2a_hex(os.urandom(16)).decode('ascii') iface = self._event_interface = 'inproc://%s' % _uuid pipe_in.bind(iface) self._event_puller = ZMQStream(pipe_in, self.io_loop) self._event_puller.on_recv(self._handle_event) @property def _event_pipe(self): """thread-local event pipe for signaling events that should be processed in the thread""" try: event_pipe = self._local.event_pipe except AttributeError: # new thread, new event pipe ctx = self.socket.context event_pipe = ctx.socket(zmq.PUSH) event_pipe.linger = 0 event_pipe.connect(self._event_interface) self._local.event_pipe = event_pipe return event_pipe def _handle_event(self, msg): """Handle an event on the event pipe Content of the message is ignored. Whenever *an* event arrives on the event stream, *all* waiting events are processed in order. """ # freeze event count so new writes don't extend the queue # while we are processing n_events = len(self._events) for i in range(n_events): event_f = self._events.popleft() event_f() def _setup_pipe_in(self): """setup listening pipe for IOPub from forked subprocesses""" ctx = self.socket.context # use UUID to authenticate pipe messages self._pipe_uuid = os.urandom(16) pipe_in = ctx.socket(zmq.PULL) pipe_in.linger = 0 try: self._pipe_port = pipe_in.bind_to_random_port("tcp://127.0.0.1") except zmq.ZMQError as e: warnings.warn("Couldn't bind IOPub Pipe to 127.0.0.1: %s" % e + "\nsubprocess output will be unavailable." ) self._pipe_flag = False pipe_in.close() return self._pipe_in = ZMQStream(pipe_in, self.io_loop) self._pipe_in.on_recv(self._handle_pipe_msg) def _handle_pipe_msg(self, msg): """handle a pipe message from a subprocess""" if not self._pipe_flag or not self._is_master_process(): return if msg[0] != self._pipe_uuid: print("Bad pipe message: %s", msg, file=sys.__stderr__) return self.send_multipart(msg[1:]) def _setup_pipe_out(self): # must be new context after fork ctx = zmq.Context() pipe_out = ctx.socket(zmq.PUSH) pipe_out.linger = 3000 # 3s timeout for pipe_out sends before discarding the message pipe_out.connect("tcp://127.0.0.1:%i" % self._pipe_port) return ctx, pipe_out def _is_master_process(self): return os.getpid() == self._master_pid def _check_mp_mode(self): """check for forks, and switch to zmq pipeline if necessary""" if not self._pipe_flag or self._is_master_process(): return MASTER else: return CHILD def start(self): """Start the IOPub thread""" self.thread.start() # make sure we don't prevent process exit # I'm not sure why setting daemon=True above isn't enough, but it doesn't appear to be. atexit.register(self.stop) def stop(self): """Stop the IOPub thread""" if not self.thread.is_alive(): return self.io_loop.add_callback(self.io_loop.stop) self.thread.join() if hasattr(self._local, 'event_pipe'): self._local.event_pipe.close() def close(self): if self.closed: return self.socket.close() self.socket = None @property def closed(self): return self.socket is None def schedule(self, f): """Schedule a function to be called in our IO thread. If the thread is not running, call immediately. """ if self.thread.is_alive(): self._events.append(f) # wake event thread (message content is ignored) self._event_pipe.send(b'') else: f() def send_multipart(self, *args, **kwargs): """send_multipart schedules actual zmq send in my thread. If my thread isn't running (e.g. forked process), send immediately. """ self.schedule(lambda : self._really_send(*args, **kwargs)) def _really_send(self, msg, *args, **kwargs): """The callback that actually sends messages""" mp_mode = self._check_mp_mode() if mp_mode != CHILD: # we are master, do a regular send self.socket.send_multipart(msg, *args, **kwargs) else: # we are a child, pipe to master # new context/socket for every pipe-out # since forks don't teardown politely, use ctx.term to ensure send has completed ctx, pipe_out = self._setup_pipe_out() pipe_out.send_multipart([self._pipe_uuid] + msg, *args, **kwargs) pipe_out.close() ctx.term() class BackgroundSocket(object): """Wrapper around IOPub thread that provides zmq send[_multipart]""" io_thread = None def __init__(self, io_thread): self.io_thread = io_thread def __getattr__(self, attr): """Wrap socket attr access for backward-compatibility""" if attr.startswith('__') and attr.endswith('__'): # don't wrap magic methods super(BackgroundSocket, self).__getattr__(attr) if hasattr(self.io_thread.socket, attr): warnings.warn("Accessing zmq Socket attribute %s on BackgroundSocket" % attr, DeprecationWarning, stacklevel=2) return getattr(self.io_thread.socket, attr) super(BackgroundSocket, self).__getattr__(attr) def __setattr__(self, attr, value): if attr == 'io_thread' or (attr.startswith('__' and attr.endswith('__'))): super(BackgroundSocket, self).__setattr__(attr, value) else: warnings.warn("Setting zmq Socket attribute %s on BackgroundSocket" % attr, DeprecationWarning, stacklevel=2) setattr(self.io_thread.socket, attr, value) def send(self, msg, *args, **kwargs): return self.send_multipart([msg], *args, **kwargs) def send_multipart(self, *args, **kwargs): """Schedule send in IO thread""" return self.io_thread.send_multipart(*args, **kwargs) class OutStream(TextIOBase): """A file like object that publishes the stream to a 0MQ PUB socket. Output is handed off to an IO Thread """ # timeout for flush to avoid infinite hang # in case of misbehavior flush_timeout = 10 # The time interval between automatic flushes, in seconds. flush_interval = 0.2 topic = None encoding = 'UTF-8' def __init__(self, session, pub_thread, name, pipe=None, echo=None): if pipe is not None: warnings.warn("pipe argument to OutStream is deprecated and ignored", DeprecationWarning) # This is necessary for compatibility with Python built-in streams self.session = session if not isinstance(pub_thread, IOPubThread): # Backward-compat: given socket, not thread. Wrap in a thread. warnings.warn("OutStream should be created with IOPubThread, not %r" % pub_thread, DeprecationWarning, stacklevel=2) pub_thread = IOPubThread(pub_thread) pub_thread.start() self.pub_thread = pub_thread self.name = name self.topic = b'stream.' + py3compat.cast_bytes(name) self.parent_header = {} self._master_pid = os.getpid() self._flush_pending = False self._subprocess_flush_pending = False self._io_loop = pub_thread.io_loop self._new_buffer() self.echo = None if echo: if hasattr(echo, 'read') and hasattr(echo, 'write'): self.echo = echo else: raise ValueError("echo argument must be a file like object") def _is_master_process(self): return os.getpid() == self._master_pid def set_parent(self, parent): self.parent_header = extract_header(parent) def close(self): self.pub_thread = None @property def closed(self): return self.pub_thread is None def _schedule_flush(self): """schedule a flush in the IO thread call this on write, to indicate that flush should be called soon. """ if self._flush_pending: return self._flush_pending = True # add_timeout has to be handed to the io thread via event pipe def _schedule_in_thread(): self._io_loop.call_later(self.flush_interval, self._flush) self.pub_thread.schedule(_schedule_in_thread) def flush(self): """trigger actual zmq send send will happen in the background thread """ if self.pub_thread and self.pub_thread.thread is not None and self.pub_thread.thread.is_alive(): # request flush on the background thread self.pub_thread.schedule(self._flush) # wait for flush to actually get through, if we can. # waiting across threads during import can cause deadlocks # so only wait if import lock is not held if not import_lock_held(): evt = threading.Event() self.pub_thread.schedule(evt.set) # and give a timeout to avoid if not evt.wait(self.flush_timeout): # write directly to __stderr__ instead of warning because # if this is happening sys.stderr may be the problem. print("IOStream.flush timed out", file=sys.__stderr__) else: self._flush() def _flush(self): """This is where the actual send happens. _flush should generally be called in the IO thread, unless the thread has been destroyed (e.g. forked subprocess). """ self._flush_pending = False self._subprocess_flush_pending = False if self.echo is not None: try: self.echo.flush() except OSError as e: if self.echo is not sys.__stderr__: print("Flush failed: {}".format(e), file=sys.__stderr__) data = self._flush_buffer() if data: # FIXME: this disables Session's fork-safe check, # since pub_thread is itself fork-safe. # There should be a better way to do this. self.session.pid = os.getpid() content = {u'name':self.name, u'text':data} self.session.send(self.pub_thread, u'stream', content=content, parent=self.parent_header, ident=self.topic) def write(self, string): if self.echo is not None: try: self.echo.write(string) except OSError as e: if self.echo is not sys.__stderr__: print("Write failed: {}".format(e), file=sys.__stderr__) if self.pub_thread is None: raise ValueError('I/O operation on closed file') else: # Make sure that we're handling unicode if not isinstance(string, unicode_type): string = string.decode(self.encoding, 'replace') is_child = (not self._is_master_process()) # only touch the buffer in the IO thread to avoid races self.pub_thread.schedule(lambda : self._buffer.write(string)) if is_child: # mp.Pool cannot be trusted to flush promptly (or ever), # and this helps. if self._subprocess_flush_pending: return self._subprocess_flush_pending = True # We can not rely on self._io_loop.call_later from a subprocess self.pub_thread.schedule(self._flush) else: self._schedule_flush() def writelines(self, sequence): if self.pub_thread is None: raise ValueError('I/O operation on closed file') else: for string in sequence: self.write(string) def writable(self): return True def _flush_buffer(self): """clear the current buffer and return the current buffer data. This should only be called in the IO thread. """ data = u'' if self._buffer is not None: buf = self._buffer self._new_buffer() data = buf.getvalue() buf.close() return data def _new_buffer(self): self._buffer = StringIO()