Change logic of wait_for_ready to wait for the IOPub to be subscribed

jupyter · Dec 9, 2020 · 6d41ead · 6d41ead
1 parent 778fa85
commit 6d41ead
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/jupyter_client/asynchronous/client.py b/jupyter_client/asynchronous/client.py
@@ -126,14 +126,21 @@ async def wait_for_ready(self, timeout=None):
 
         # Wait for kernel info reply on shell channel
         while True:
+            self.kernel_info()
             try:
                 msg = await self.shell_channel.get_msg(timeout=1)
             except Empty:
                 pass
             else:
                 if msg['msg_type'] == 'kernel_info_reply':
-                    self._handle_kernel_info_reply(msg)
-                    break
+                    # Checking that IOPub is connected. If it is not connected, start over.
+                    try:
+                        iopub_msg = await self.iopub_channel.get_msg(block=True, timeout=0.2)
+                    except Empty:
+                        pass
+                    else:
+                        self._handle_kernel_info_reply(msg)
+                        break
 
             if not await self.is_alive():
                 raise RuntimeError('Kernel died before replying to kernel_info')

diff --git a/jupyter_client/blocking/client.py b/jupyter_client/blocking/client.py
@@ -91,14 +91,21 @@ def wait_for_ready(self, timeout=None):
 
         # Wait for kernel info reply on shell channel
         while True:
+            self.kernel_info()
             try:
                 msg = self.shell_channel.get_msg(block=True, timeout=1)
             except Empty:
                 pass
             else:
                 if msg['msg_type'] == 'kernel_info_reply':
-                    self._handle_kernel_info_reply(msg)
-                    break
+                    # Checking that IOPub is connected. If it is not connected, start over.
+                    try:
+                        iopub_msg = self.iopub_channel.get_msg(timeout=0.2)
+                    except Empty:
+                        pass
+                    else:
+                        self._handle_kernel_info_reply(msg)
+                        break
 
             if not self.is_alive():
                 raise RuntimeError('Kernel died before replying to kernel_info')

diff --git a/jupyter_client/client.py b/jupyter_client/client.py
@@ -102,11 +102,10 @@ def start_channels(self, shell=True, iopub=True, stdin=True, hb=True, control=Tr
         :meth:`start_kernel`. If the channels have been stopped and you
         call this, :class:`RuntimeError` will be raised.
         """
-        if shell:
-            self.shell_channel.start()
-            self.kernel_info()
         if iopub:
             self.iopub_channel.start()
+        if shell:
+            self.shell_channel.start()
         if stdin:
             self.stdin_channel.start()
             self.allow_stdin = True