Only kill children in process group at shutdown

- not the whole process group, which includes self and possibly parents
- not children which have started their own process groups

Author: minrk

ipykernel/kernelbase.py

@@ -808,7 +808,8 @@ def _send_interupt_children(self):
             pid = os.getpid()
             pgid = os.getpgid(pid)
             # Prefer process-group over process
-            if pgid and hasattr(os, "killpg"):
+            # but only if the kernel is the leader of the process group
+            if pgid and pgid == pid and hasattr(os, "killpg"):
                 try:
                     os.killpg(pgid, SIGINT)
                     return
@@ -1136,67 +1137,73 @@ def _input_request(self, prompt, ident, parent, password=False):
             raise EOFError
         return value
 
-    def _killpg(self, signal):
+    def _signal_children(self, signum):
         """
-        similar to killpg but use psutil if it can on windows
-        or if pgid is none
+        Send a signal to all our children
 
+        Like `killpg`, but does not include the current process
+        (or possible parents).
         """
-        pgid = os.getpgid(os.getpid())
-        if pgid and hasattr(os, "killpg"):
-            try:
-                os.killpg(pgid, signal)
-            except (OSError) as e:
-                self.log.exception(f"OSError running killpg, not killing children.")
+        if psutil is None:
+            self.log.info("Need psutil to signal children")
             return
-        elif psutil is not None:
-            children = parent.children(recursive=True)
-            for p in children:
-                try:
-                    if signal == SIGTERM:
-                        p.terminate()
-                    elif signal == SIGKILL:
-                        p.kill()
-                except psutil.NoSuchProcess:
-                    pass
 
-    async def _progressively_terminate_all_children(self):
+        for p in self._process_children():
+            self.log.debug(f"Sending {Signals(signum)!r} to subprocess {p}")
+            try:
+                if signum == SIGTERM:
+                    p.terminate()
+                elif signum == SIGKILL:
+                    p.kill()
+                else:
+                    p.send_signal(signum)
+            except psutil.NoSuchProcess:
+                pass
+
+    def _process_children(self):
+        """Retrieve child processes in the kernel's process group
+
+        Avoids:
+        - including parents and self with killpg
+        - including all children that may have forked-off a new group
+        """
+        if psutil is None:
+            return []
+        kernel_process = psutil.Process()
+        all_children = kernel_process.children(recursive=True)
+        if os.name == "nt":
+            return all_children
+        kernel_pgid = os.getpgrp()
+        process_group_children = []
+        for child in all_children:
+            try:
+                child_pgid = os.getpgid(child.pid)
+            except OSError:
+                pass
+            else:
+                if child_pgid == kernel_pgid:
+                    process_group_children.append(child)
+        return process_group_children
 
-        pgid = os.getpgid(os.getpid())
+    async def _progressively_terminate_all_children(self):
         if psutil is None:
-            # blindly send quickly sigterm/sigkill to processes if psutil not there.
+            # we need psutil to safely clean up children
             self.log.info("Please install psutil for a cleaner subprocess shutdown.")
-            self._send_interupt_children()
-            await asyncio.sleep(0.05)
-            self.log.debug("Sending SIGTERM to {pgid}")
-            self._killpg(SIGTERM)
-            await asyncio.sleep(0.05)
-            self.log.debug("Sending SIGKILL to {pgid}")
-            self._killpg(pgid, SIGKILL)
+            return
 
         sleeps = (0.01, 0.03, 0.1, 0.3, 1, 3, 10)
-        children = psutil.Process().children(recursive=True)
-        if not children:
+        if not self._process_children():
             self.log.debug("Kernel has no children.")
             return
-        self.log.debug(f"Trying to interrupt then kill subprocesses : {children}")
-        self._send_interupt_children()
 
         for signum in (SIGTERM, SIGKILL):
-            self.log.debug(
-                f"Will try to send {signum} ({Signals(signum)!r}) to subprocesses :{children}"
-            )
             for delay in sleeps:
-                children = psutil.Process().children(recursive=True)
-                try:
-                    if not children:
-                        self.log.warning(
-                            "No more children, continuing shutdown routine."
-                        )
-                        return
-                except psutil.NoSuchProcess:
-                    pass
-                self._killpg(15)
+                children = self._process_children()
+                if not children:
+                    self.log.debug("No more children, continuing shutdown routine.")
+                    return
+                # signals only children, not current process
+                self._signal_children(signum)
                 self.log.debug(
                     f"Will sleep {delay}s before checking for children and retrying. {children}"
                 )

ipykernel/tests/test_kernel.py

@@ -6,14 +6,20 @@
 import ast
 import os.path
 import platform
+import signal
 import subprocess
 import sys
 import time
+from subprocess import Popen
 from tempfile import TemporaryDirectory
 
 from flaky import flaky
 import pytest
-from packaging import version
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
 
 import IPython
 from IPython.paths import locate_profile
@@ -496,3 +502,79 @@ def test_control_thread_priority():
     # comparing first to last ought to be enough, since queues preserve order
     # use <= in case of very-fast handling and/or low resolution timers
     assert control_dates[-1] <= shell_dates[0]
+
+
+def _child():
+    print("in child", os.getpid())
+
+    def _print_and_exit(sig, frame):
+        print(f"Received signal {sig}")
+        # take some time so retries are triggered
+        time.sleep(0.5)
+        sys.exit(-sig)
+
+    signal.signal(signal.SIGTERM, _print_and_exit)
+    time.sleep(30)
+
+
+def _start_children():
+    ip = IPython.get_ipython()
+    ns = ip.user_ns
+
+    cmd = [sys.executable, "-c", f"from {__name__} import _child; _child()"]
+    child_pg = Popen(cmd, start_new_session=False)
+    child_newpg = Popen(cmd, start_new_session=True)
+    ns["pid"] = os.getpid()
+    ns["child_pg"] = child_pg.pid
+    ns["child_newpg"] = child_newpg.pid
+    # give them time to start up and register signal handlers
+    time.sleep(1)
+
+
+@pytest.mark.skipif(
+    platform.python_implementation() == "PyPy",
+    reason="does not work on PyPy",
+)
+@pytest.mark.skipif(
+    psutil is None,
+    reason="requires psutil",
+)
+def test_shutdown_subprocesses():
+    """Kernel exits after polite shutdown_request"""
+    with new_kernel() as kc:
+        km = kc.parent
+        msg_id, reply = execute(
+            f"from {__name__} import _start_children\n_start_children()",
+            kc=kc,
+            user_expressions={
+                "pid": "pid",
+                "child_pg": "child_pg",
+                "child_newpg": "child_newpg",
+            },
+        )
+        print(reply)
+        expressions = reply["user_expressions"]
+        kernel_process = psutil.Process(int(expressions["pid"]["data"]["text/plain"]))
+        child_pg = psutil.Process(int(expressions["child_pg"]["data"]["text/plain"]))
+        child_newpg = psutil.Process(
+            int(expressions["child_newpg"]["data"]["text/plain"])
+        )
+        wait_for_idle(kc)
+
+        kc.shutdown()
+        for i in range(300):  # 30s timeout
+            if km.is_alive():
+                time.sleep(0.1)
+            else:
+                break
+        assert not km.is_alive()
+        assert not kernel_process.is_running()
+        # child in the process group shut down
+        assert not child_pg.is_running()
+        # child outside the process group was not shut down (unix only)
+        if os.name != 'nt':
+            assert child_newpg.is_running()
+        try:
+            child_newpg.terminate()
+        except psutil.NoSuchProcess:
+            pass