Percolate a control connection throughout the code.

Author: jasongrout

jupyter_client/blocking/client.py

@@ -36,15 +36,15 @@
     TimeoutError = RuntimeError
 
 
-def reqrep(meth):
+def reqrep(meth, channel='shell'):
     def wrapped(self, *args, **kwargs):
         reply = kwargs.pop('reply', False)
         timeout = kwargs.pop('timeout', None)
         msg_id = meth(self, *args, **kwargs)
         if not reply:
             return msg_id
 
-        return self._recv_reply(msg_id, timeout=timeout)
+        return self._recv_reply(msg_id, timeout=timeout, channel=channel)
 
     if not meth.__doc__:
         # python -OO removes docstrings,
@@ -135,17 +135,21 @@ def wait_for_ready(self, timeout=None):
     iopub_channel_class = Type(ZMQSocketChannel)
     stdin_channel_class = Type(ZMQSocketChannel)
     hb_channel_class = Type(HBChannel)
+    control_channel_class = Type(ZMQSocketChannel)
 
 
-    def _recv_reply(self, msg_id, timeout=None):
+    def _recv_reply(self, msg_id, timeout=None, channel='shell'):
         """Receive and return the reply for a given request"""
         if timeout is not None:
             deadline = monotonic() + timeout
         while True:
             if timeout is not None:
                 timeout = max(0, deadline - monotonic())
             try:
-                reply = self.get_shell_msg(timeout=timeout)
+                if channel == 'control':
+                    reply = self.get_control_msg(timeout=timeout)
+                else:
+                    reply = self.get_shell_msg(timeout=timeout)
             except Empty:
                 raise TimeoutError("Timeout waiting for reply")
             if reply['parent_header'].get('msg_id') != msg_id:
@@ -154,13 +158,16 @@ def _recv_reply(self, msg_id, timeout=None):
             return reply
 
 
+    # replies come on the shell channel
     execute = reqrep(KernelClient.execute)
     history = reqrep(KernelClient.history)
     complete = reqrep(KernelClient.complete)
     inspect = reqrep(KernelClient.inspect)
     kernel_info = reqrep(KernelClient.kernel_info)
     comm_info = reqrep(KernelClient.comm_info)
-    shutdown = reqrep(KernelClient.shutdown)
+
+    # replies come on the control channel
+    shutdown = reqrep(KernelClient.shutdown, channel='control')
 
 
     def _stdin_hook_default(self, msg):

jupyter_client/clientabc.py

@@ -47,12 +47,16 @@ def hb_channel_class(self):
     def stdin_channel_class(self):
         pass
 
+    @abc.abstractproperty
+    def control_channel_class(self):
+        pass
+
     #--------------------------------------------------------------------------
     # Channel management methods
     #--------------------------------------------------------------------------
 
     @abc.abstractmethod
-    def start_channels(self, shell=True, iopub=True, stdin=True, hb=True):
+    def start_channels(self, shell=True, iopub=True, stdin=True, hb=True, control=True):
         pass
 
     @abc.abstractmethod
@@ -78,3 +82,7 @@ def stdin_channel(self):
     @abc.abstractproperty
     def hb_channel(self):
         pass
+
+    @abc.abstractproperty
+    def control_channel(self):
+        pass

jupyter_client/connect.py

@@ -226,7 +226,7 @@ def find_connection_file(filename='kernel-*.json', path=None, profile=None):
 def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
     """tunnel connections to a kernel via ssh
 
-    This will open four SSH tunnels from localhost on this machine to the
+    This will open five SSH tunnels from localhost on this machine to the
     ports associated with the kernel.  They can be either direct
     localhost-localhost tunnels, or if an intermediate server is necessary,
     the kernel must be listening on a public IP.
@@ -246,8 +246,8 @@ def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
     Returns
     -------
 
-    (shell, iopub, stdin, hb) : ints
-        The four ports on localhost that have been forwarded to the kernel.
+    (shell, iopub, stdin, hb, control) : ints
+        The five ports on localhost that have been forwarded to the kernel.
     """
     from .ssh import tunnel
     if isinstance(connection_info, string_types):
@@ -257,8 +257,8 @@ def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
 
     cf = connection_info
 
-    lports = tunnel.select_random_ports(4)
-    rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port']
+    lports = tunnel.select_random_ports(5)
+    rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port'], cf['control_port']
 
     remote_ip = cf['ip']
 

jupyter_client/consoleapp.py

@@ -72,6 +72,7 @@
     shell = 'JupyterConsoleApp.shell_port',
     iopub = 'JupyterConsoleApp.iopub_port',
     stdin = 'JupyterConsoleApp.stdin_port',
+    control = 'JupyterConsoleApp.control_port',
     existing = 'JupyterConsoleApp.existing',
     f = 'JupyterConsoleApp.connection_file',
 
@@ -222,7 +223,8 @@ def init_ssh(self):
                     shell_port=self.shell_port,
                     iopub_port=self.iopub_port,
                     stdin_port=self.stdin_port,
-                    hb_port=self.hb_port
+                    hb_port=self.hb_port,
+                    control_port=self.control_port
         )
 
         self.log.info("Forwarding connections to %s via %s"%(ip, self.sshserver))
@@ -236,7 +238,7 @@ def init_ssh(self):
             self.log.error("Could not setup tunnels", exc_info=True)
             self.exit(1)
 
-        self.shell_port, self.iopub_port, self.stdin_port, self.hb_port = newports
+        self.shell_port, self.iopub_port, self.stdin_port, self.hb_port, self.control_port = newports
 
         cf = self.connection_file
         root, ext = os.path.splitext(cf)
@@ -275,6 +277,7 @@ def init_kernel_manager(self):
                                     iopub_port=self.iopub_port,
                                     stdin_port=self.stdin_port,
                                     hb_port=self.hb_port,
+                                    control_port=self.control_port,
                                     connection_file=self.connection_file,
                                     kernel_name=self.kernel_name,
                                     parent=self,
@@ -302,6 +305,7 @@ def init_kernel_manager(self):
         self.iopub_port=km.iopub_port
         self.stdin_port=km.stdin_port
         self.hb_port=km.hb_port
+        self.control_port=km.control_port
         self.connection_file = km.connection_file
 
         atexit.register(self.kernel_manager.cleanup_connection_file)
@@ -318,6 +322,7 @@ def init_kernel_client(self):
                                 iopub_port=self.iopub_port,
                                 stdin_port=self.stdin_port,
                                 hb_port=self.hb_port,
+                                control_port=self.control_port,
                                 connection_file=self.connection_file,
                                 parent=self,
             )

jupyter_client/manager.py

@@ -297,7 +297,7 @@ def shutdown_kernel(self, now=False, restart=False):
 
         This attempts to shutdown the kernels cleanly by:
 
-        1. Sending it a shutdown message over the shell channel.
+        1. Sending it a shutdown message over the control channel.
         2. If that fails, the kernel is shutdown forcibly by sending it
            a signal.
 

jupyter_client/threaded.py

@@ -230,14 +230,14 @@ def ioloop(self):
 
     ioloop_thread = Instance(IOLoopThread, allow_none=True)
 
-    def start_channels(self, shell=True, iopub=True, stdin=True, hb=True):
+    def start_channels(self, shell=True, iopub=True, stdin=True, hb=True, control=True):
         self.ioloop_thread = IOLoopThread()
         self.ioloop_thread.start()
 
         if shell:
             self.shell_channel._inspect = self._check_kernel_info_reply
 
-        super(ThreadedKernelClient, self).start_channels(shell, iopub, stdin, hb)
+        super(ThreadedKernelClient, self).start_channels(shell, iopub, stdin, hb, control)
 
     def _check_kernel_info_reply(self, msg):
         """This is run in the ioloop thread when the kernel info reply is received
@@ -255,3 +255,4 @@ def stop_channels(self):
     shell_channel_class = Type(ThreadedZMQSocketChannel)
     stdin_channel_class = Type(ThreadedZMQSocketChannel)
     hb_channel_class = Type(HBChannel)
+    control_channel_class = Type(ThreadedZMQSocketChannel)