Dual buffer with automatic switch. No performance hit for single process.

IPython/kernel/zmq/iostream.py

@@ -1,95 +1,157 b''
 import sys
 import time
+import os
 from io import StringIO
 
 from session import extract_header, Message
 
 from IPython.utils import io, text
 from IPython.utils import py3compat
 
+import multiprocessing as mp
+import multiprocessing.sharedctypes as mpshc
+from ctypes import c_bool
 #-----------------------------------------------------------------------------
 # Globals
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Stream classes
 #-----------------------------------------------------------------------------
 
 class OutStream(object):
     """A file like object that publishes the stream to a 0MQ PUB socket."""
 
     # The time interval between automatic flushes, in seconds.
     flush_interval = 0.05
     topic=None
 
     def __init__(self, session, pub_socket, name):
         self.encoding = 'UTF-8'
         self.session = session
         self.pub_socket = pub_socket
         self.name = name
         self.parent_header = {}
         self._new_buffer()
+        self._manager = mp.Manager()
+        #use sharectype here so it don't have to hit the manager
+        #no synchronize needed either(right?). Just a flag telling the master
+        #to switch the buffer to que
+        self._found_newprocess = mpshc.RawValue(c_bool, False)
+        self._que_buffer = self._manager.Queue()
+        self._que_lock = self._manager.Lock()
+        self._masterpid = os.getpid()
+        self._master_has_switched = False
+
+    def _switch_to_que(self):
+        #should only be called on master process
+        #don't clear the que before putting data in since
+        #child process might have put something in the que before the
+        #master know it.
+        self._que_buffer.put(self._buffer.getvalue())
+        self._new_buffer()
+        self._start = -1
+
+    def _is_master_process(self):
+        return os.getpid()==self._masterpid
+
+    def _debug_print(self,s):
+        sys.__stdout__.write(s+'\n')
+        sys.__stdout__.flush()
+
+    def _check_mp_mode(self):
+        """check multiprocess and switch to que if necessary"""
+        if not self._found_newprocess.value:
+            if not self._is_master_process():
+                self._found_newprocess.value = True
+        elif self._found_newprocess.value and not self._master_has_switched:
+
+            #switch to que if it has not been switch
+            if self._is_master_process():
+                self._switch_to_que()
+                self._master_has_switched = True
+
+        return self._found_newprocess.value
+
 
     def set_parent(self, parent):
         self.parent_header = extract_header(parent)
 
     def close(self):
         self.pub_socket = None
 
     def flush(self):
         #io.rprint('>>>flushing output buffer: %s<<<' % self.name)  # dbg
+
         if self.pub_socket is None:
             raise ValueError(u'I/O operation on closed file')
         else:
-            data = self._buffer.getvalue()
-            if data:
-                content = {u'name':self.name, u'data':data}
-                msg = self.session.send(self.pub_socket, u'stream', content=content,
-                                       parent=self.parent_header, ident=self.topic)
-
-                if hasattr(self.pub_socket, 'flush'):
-                    # socket itself has flush (presumably ZMQStream)
-                    self.pub_socket.flush()
-                self._buffer.close()
-                self._new_buffer()
+            if self._is_master_process():
+                data = u''
+                #obtain data
+                if self._check_mp_mode():#multiprocess
+                    with self._que_lock:
+                        while not self._que_buffer.empty():
+                            data += self._que_buffer.get()
+                else:#single process mode
+                    data = self._buffer.getvalue()
+
+                if data:
+                    content = {u'name':self.name, u'data':data}
+                    msg = self.session.send(self.pub_socket, u'stream', content=content,
+                                           parent=self.parent_header, ident=self.topic)
+
+                    if hasattr(self.pub_socket, 'flush'):
+                        # socket itself has flush (presumably ZMQStream)
+                        self.pub_socket.flush()
+                    self._buffer.close()
+                    self._new_buffer()
+            else:
+                pass
+
 
     def isatty(self):
         return False
 
     def __next__(self):
         raise IOError('Read not supported on a write only stream.')
 
     if not py3compat.PY3:
         next = __next__
 
     def read(self, size=-1):
         raise IOError('Read not supported on a write only stream.')
 
     def readline(self, size=-1):
         raise IOError('Read not supported on a write only stream.')
 
     def write(self, string):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             # Make sure that we're handling unicode
             if not isinstance(string, unicode):
                 string = string.decode(self.encoding, 'replace')
 
-            self._buffer.write(string)
+            if self._check_mp_mode(): #multi process mode
+                with self._que_lock:
+                    self._que_buffer.put(string)
+            else: #sigle process mode
+                self._buffer.write(string)
+
             current_time = time.time()
             if self._start <= 0:
                 self._start = current_time
             elif current_time - self._start > self.flush_interval:
                 self.flush()
 
     def writelines(self, sequence):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             for string in sequence:
                 self.write(string)
 
     def _new_buffer(self):
         self._buffer = StringIO()
         self._start = -1