Fix orphaned subprocesses and supervisor crash on heartbeat 409

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -418,8 +418,6 @@ def _fork_main(
     - Catch un-handled exceptions and attempt to show _something_ in case of error
     - Finally, run the actual task runner code (``target`` argument, defaults to ``.task_runner:main`)
     """
-    # TODO: Make this process a session leader
-
     # Store original stderr for last-chance exception handling
     last_chance_stderr = _get_last_chance_stderr()
 
@@ -673,6 +671,9 @@ class WatchedSubprocess:
     subprocess_logs_to_stdout: bool = False
     """Duplicate log messages to stdout, or only send them to ``self.process_log``."""
 
+    _new_process_group: bool = False
+    """Whether the child was placed in its own process group at fork time (see ``start``)."""
+
     start_time: float = attrs.field(factory=time.monotonic)
     """The start time of the child process."""
 
@@ -683,6 +684,7 @@ def start(
         target: Callable[[], None] = _subprocess_main,
         logger: FilteringBoundLogger | None = None,
         use_exec: bool = False,
+        new_process_group: bool = False,
         **constructor_kwargs,
     ) -> Self:
         """
@@ -694,6 +696,12 @@ def start(
             ``target`` is rehydrated in the exec'd child from its ``module:qualname``,
             so any importable entry point (task execution, DAG processor, triggerer)
             is supported.
+        :param new_process_group: If True, place the child in its own process
+            group (PGID == its PID, like
+            ``airflow.utils.process_utils.set_new_process_group``) so signals
+            can be delivered to the child's whole process tree via
+            ``os.killpg``. Task execution opts in; DAG processor and triggerer
+            keep the supervisor's process group and are signalled directly.
         """
         if use_exec and "<" in getattr(target, "__qualname__", "<"):
             # Closures/lambdas (``<locals>`` / ``<lambda>`` in the qualname) and
@@ -711,6 +719,19 @@ def start(
 
         pid = os.fork()
         if pid == 0:
+            if new_process_group:
+                # Put the task-runner into its own process group so its PGID
+                # equals its own PID. The supervisor can then deliver signals
+                # to the whole tree via os.killpg(), reaching every subprocess
+                # the task-runner spawned (e.g. venv children from
+                # PythonVirtualenvOperator). Without this, a SIGTERM from
+                # kill() only hits the task-runner and any Popen children are
+                # reparented to PID 1 and leak as orphans. Also set from the
+                # parent below so the group exists no matter which side of the
+                # fork runs first. See issue #65505.
+                with suppress(OSError):
+                    os.setpgid(0, 0)
+
             # Close and delete of the parent end of the sockets.
             cls._close_unused_sockets(read_requests, read_stdout, read_stderr, read_logs)
 
@@ -762,6 +783,15 @@ def start(
             # do then _THINGS GET WEIRD_.. (Normally `_fork_main` itself will `_exit()` so we never get here)
             os._exit(124)
 
+        if new_process_group:
+            # Mirror of the child-side setpgid, so the group is guaranteed to
+            # exist once start() returns. Without this, kill() invoked before
+            # the child is first scheduled (e.g. task_instances.start()
+            # failing synchronously in _on_child_started) would resolve the
+            # child's PGID to the supervisor's own group and killpg it.
+            with suppress(OSError):
+                os.setpgid(pid, pid)
+
         # Close the remaining parent-end of the sockets we've passed to the child via fork. We still have the
         # other end of the pair open
         cls._close_unused_sockets(child_stdout, child_stderr, child_logs)
@@ -773,6 +803,7 @@ def start(
             process=PsutilTracker(psutil.Process(pid)),
             process_log=logger,
             start_time=time.monotonic(),
+            new_process_group=new_process_group,
             **constructor_kwargs,
         )
 
@@ -1007,6 +1038,30 @@ def _cleanup_open_sockets(self):
         self.selector.close()
         self.stdin.close()
 
+    def _signal_subprocess(self, sig: signal.Signals) -> None:
+        """
+        Deliver ``sig`` to the child process, or to its whole process group when it has its own.
+
+        When ``new_process_group`` was set at ``start()`` time, the signal is sent with
+        ``os.killpg`` so subprocesses spawned by the child (venv children, bash shells, etc.)
+        are reached too (see issue #65505). Falls back to signalling the child PID alone when
+        the group cannot be resolved or signalled -- and, critically, when the child still
+        shares the supervisor's own process group (``setpgid`` failed), because ``killpg`` on
+        our own group would signal the supervisor itself and its siblings.
+        """
+        if self._new_process_group:
+            try:
+                pgid = os.getpgid(self._process.pid)
+            except (ProcessLookupError, PermissionError):
+                pgid = None
+            if pgid is not None and pgid != os.getpgid(0):
+                try:
+                    os.killpg(pgid, sig)
+                    return
+                except (ProcessLookupError, PermissionError):
+                    pass
+        self._process.send_signal(sig)
+
     def kill(
         self,
         signal_to_send: signal.Signals = signal.SIGINT,
@@ -1036,7 +1091,7 @@ def kill(
 
         for sig in escalation_path:
             try:
-                self._process.send_signal(sig)
+                self._signal_subprocess(sig)
 
                 start = time.monotonic()
                 end = start + escalation_delay
@@ -1360,7 +1415,13 @@ def start(  # type: ignore[override]
         # infrastructure; keep bare fork for those.
         use_exec = target is _subprocess_main and _should_use_exec()
         proc: Self = super().start(
-            id=what.id, client=client, target=target, logger=logger, use_exec=use_exec, **kwargs
+            id=what.id,
+            client=client,
+            target=target,
+            logger=logger,
+            use_exec=use_exec,
+            new_process_group=True,
+            **kwargs,
         )
         # Tell the task process what it needs to do!
         proc._on_child_started(

task-sdk/tests/task_sdk/execution_time/test_supervisor.py

@@ -164,6 +164,7 @@
     InProcessSupervisorComms,
     InProcessTestSupervisor,
     ProcessTracker,
+    WatchedSubprocess,
     _make_process_nondumpable,
     _remote_logging_conn,
     in_process_api_server,
@@ -1241,6 +1242,71 @@ def test_cleanup_sockets_after_delay(self, monkeypatch, mocker):
         proc.selector.close.assert_called_once()
         proc.stdin.close.assert_called_once()
 
+    def test_task_runner_starts_in_new_process_group(self, client_with_ti_start):
+        """Regression test for #65505: the task-runner child must be placed in
+        its own process group (PGID == its PID) so kill() can reach
+        subprocesses the task-runner spawns via os.killpg(); without it, a
+        venv/Popen child of the task-runner inherits the supervisor's process
+        group and killpg would signal the supervisor too (or miss the
+        grandchild entirely).
+
+        The group must already exist when start() returns: the parent sets it
+        too (double setpgid), closing the race where kill() runs before the
+        child is first scheduled (e.g. task_instances.start() failing
+        synchronously in _on_child_started).
+        """
+
+        def subprocess_main():
+            CommsDecoder()._get_response()
+            sleep(10)
+
+        proc = ActivitySubprocess.start(
+            dag_rel_path=os.devnull,
+            bundle_info=FAKE_BUNDLE,
+            what=TaskInstance(
+                id=uuid7(),
+                task_id="b",
+                dag_id="c",
+                run_id="d",
+                try_number=1,
+                dag_version_id=uuid7(),
+                queue="default",
+            ),
+            client=client_with_ti_start,
+            target=subprocess_main,
+        )
+        try:
+            child_pgid = os.getpgid(proc.pid)
+            assert child_pgid == proc.pid, (
+                "Task-runner child must be its own process-group leader as soon "
+                f"as start() returns. Got pgid={child_pgid}, pid={proc.pid}."
+            )
+            assert child_pgid != os.getpgid(0), (
+                "Child's process group must differ from the supervisor's so "
+                "os.killpg() from kill() does not signal the supervisor itself."
+            )
+        finally:
+            proc.kill(signal.SIGKILL, force=True)
+            proc.wait()
+
+    def test_child_keeps_supervisor_process_group_by_default(self):
+        """Subprocess types that don't opt in to new_process_group (DAG
+        processor, triggerer, callbacks) must keep the supervisor's process
+        group: they install their own signal handlers and expect direct,
+        graceful signalling rather than group-wide delivery.
+        """
+
+        def subprocess_main():
+            sleep(30)
+
+        proc = WatchedSubprocess.start(id=uuid7(), target=subprocess_main)
+        try:
+            assert os.getpgid(proc.pid) == os.getpgid(0), (
+                "Without new_process_group=True the child must stay in the supervisor's process group."
+            )
+        finally:
+            proc.kill(signal.SIGKILL, force=True)
+
 
 class TestWatchedSubprocessKill:
     @pytest.fixture
@@ -1261,6 +1327,7 @@ def watched_subprocess(self, mocker, mock_process):
             stdin=mocker.Mock(),
             client=mocker.Mock(),
             process=mock_process,
+            new_process_group=True,
         )
         # Mock the selector
         mock_selector = mocker.Mock(spec=selectors.DefaultSelector)
@@ -1270,25 +1337,111 @@ def watched_subprocess(self, mocker, mock_process):
         proc.selector = mock_selector
         return proc
 
-    def test_kill_process_already_exited(self, watched_subprocess, mock_process):
+    def test_kill_process_already_exited(self, watched_subprocess, mock_process, mocker):
         """Test behavior when the process has already exited."""
+        # When the process is gone, getpgid raises ProcessLookupError and the
+        # kill() path falls back to send_signal on the dead psutil.Process.
+        mocker.patch("os.getpgid", side_effect=ProcessLookupError)
         mock_process.wait.side_effect = psutil.NoSuchProcess(pid=1234)
         watched_subprocess.kill(signal.SIGINT, force=True)
 
         mock_process.send_signal.assert_called_once_with(signal.SIGINT)
         mock_process.wait.assert_called_once()
         assert watched_subprocess._exit_code == -1
 
-    def test_kill_process_custom_signal(self, watched_subprocess, mock_process):
-        """Test that the process is killed with the correct signal."""
+    def test_kill_process_custom_signal(self, watched_subprocess, mock_process, mocker):
+        """Test that the process is killed with the correct signal via killpg."""
+        mock_getpgid = mocker.patch("os.getpgid", side_effect=lambda pid: 12345 if pid else 54321)
+        mock_killpg = mocker.patch("os.killpg")
         mock_process.wait.return_value = 0
 
         signal_to_send = signal.SIGUSR1
         watched_subprocess.kill(signal_to_send, force=False)
 
-        mock_process.send_signal.assert_called_once_with(signal_to_send)
+        assert mock_getpgid.call_args_list == [mocker.call(12345), mocker.call(0)]
+        mock_killpg.assert_called_once_with(12345, signal_to_send)
+        mock_process.send_signal.assert_not_called()
         mock_process.wait.assert_called_once_with(timeout=0)
 
+    def test_kill_signals_process_group(self, watched_subprocess, mock_process, mocker):
+        """Regression test for #65505: kill() must signal the whole process
+        group so subprocesses spawned by the task-runner (venv children,
+        Docker exec, bash shells) are also reached.
+        """
+        mock_getpgid = mocker.patch("os.getpgid", side_effect=lambda pid: 12345 if pid else 54321)
+        mock_killpg = mocker.patch("os.killpg")
+        mock_process.wait.return_value = 0
+
+        watched_subprocess.kill(signal.SIGTERM, force=False)
+
+        assert mock_getpgid.call_args_list == [mocker.call(12345), mocker.call(0)]
+        mock_killpg.assert_called_once_with(12345, signal.SIGTERM)
+        mock_process.send_signal.assert_not_called()
+
+    def test_kill_does_not_signal_supervisors_own_process_group(
+        self, watched_subprocess, mock_process, mocker
+    ):
+        """If the child never made it into its own process group (setpgid
+        failed, or the child died and its PID's group resolves to ours),
+        os.killpg would signal the supervisor itself and every sibling in its
+        group -- and no exception would be raised for the fallback to catch.
+        kill() must detect the shared group and signal the child PID alone.
+        """
+        mocker.patch("os.getpgid", return_value=54321)
+        mock_killpg = mocker.patch("os.killpg")
+        mock_process.wait.return_value = 0
+
+        watched_subprocess.kill(signal.SIGTERM, force=False)
+
+        mock_killpg.assert_not_called()
+        mock_process.send_signal.assert_called_once_with(signal.SIGTERM)
+
+    def test_kill_signals_pid_only_without_new_process_group(self, mocker, mock_process):
+        """Subprocess types that don't opt in to new_process_group (DAG
+        processor, triggerer, callbacks) must be signalled directly, never
+        via killpg.
+        """
+        proc = ActivitySubprocess(
+            process_log=mocker.MagicMock(),
+            id=TI_ID,
+            pid=12345,
+            stdin=mocker.Mock(),
+            client=mocker.Mock(),
+            process=mock_process,
+        )
+        mock_getpgid = mocker.patch("os.getpgid")
+        mock_killpg = mocker.patch("os.killpg")
+        mock_process.wait.return_value = 0
+
+        proc.kill(signal.SIGTERM, force=False)
+
+        mock_getpgid.assert_not_called()
+        mock_killpg.assert_not_called()
+        mock_process.send_signal.assert_called_once_with(signal.SIGTERM)
+
+    @pytest.mark.parametrize("failing_call", ["getpgid", "killpg"])
+    @pytest.mark.parametrize("exc", [ProcessLookupError, PermissionError])
+    def test_kill_falls_back_to_send_signal_when_group_signal_fails(
+        self, watched_subprocess, mock_process, mocker, failing_call, exc
+    ):
+        """If os.killpg or os.getpgid raises ProcessLookupError (group
+        vanished, e.g. task already reaped) or PermissionError, fall back to
+        signalling the task-runner PID directly via send_signal.
+        """
+        if failing_call == "getpgid":
+            mocker.patch("os.getpgid", side_effect=exc)
+            mock_killpg = mocker.patch("os.killpg")
+        else:
+            mocker.patch("os.getpgid", side_effect=lambda pid: 12345 if pid else 54321)
+            mock_killpg = mocker.patch("os.killpg", side_effect=exc)
+        mock_process.wait.return_value = 0
+
+        watched_subprocess.kill(signal.SIGTERM, force=False)
+
+        if failing_call == "getpgid":
+            mock_killpg.assert_not_called()
+        mock_process.send_signal.assert_called_once_with(signal.SIGTERM)
+
     @pytest.mark.parametrize(
         ("signal_to_send", "exit_after"),
         [