diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 87311f02da7a1..9758d2b18d7c8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/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 (```` / ```` 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( diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index f777b2d5a8a90..0aa307edbf232 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/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,8 +1337,11 @@ 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) @@ -1279,16 +1349,99 @@ def test_kill_process_already_exited(self, watched_subprocess, mock_process): 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"), [