From 2d91ae7928728993be8ac194598c8be6baa77f9b Mon Sep 17 00:00:00 2001 From: Christoph <116812500+cmettler@users.noreply.github.com> Date: Thu, 23 Apr 2026 17:25:10 +0000 Subject: [PATCH 1/2] Fix orphaned subprocesses and supervisor crash on heartbeat 409 When a running TaskInstance is forcibly transitioned out of `running` (e.g. the scheduler resets a stale heartbeat, or an operator PATCHes the state to `failed`), the task-runner's next heartbeat returns HTTP 409 and the supervisor kills the task. Before this change two things went wrong on Linux: 1. Subprocesses the task-runner had spawned (`@task.virtualenv` / `PythonVirtualenvOperator` children, `DockerOperator` exec, Bash shells) were reparented to PID 1 and kept running as orphans until they finished on their own - wasting CPU, RAM and third-party API quota. 2. About 60s later, `_cleanup_open_sockets()` closed the selector while `_service_subprocess()` was still using it, so the supervisor crashed with `ValueError: I/O operation on closed epoll object` (regression from PR #51180). The task-runner is now placed in its own session via `os.setsid()` immediately after fork, so its process group ID equals its PID. The supervisor's `kill()` signals the whole group via `os.killpg(os.getpgid(pid), sig)`, which reaches every subprocess the task-runner spawned. Grandchildren without a SIGTERM handler exit promptly, close their inherited pipes, and the supervisor drains `_open_sockets` normally - so `_cleanup_open_sockets()` is never triggered and the selector is never closed mid-loop. `os.killpg`/`os.getpgid` fall back to `self._process.send_signal(sig)` on `ProcessLookupError` or `PermissionError`, preserving prior behaviour when the group has vanished (e.g. the task was already reaped) or permissions are lacking. closes: #65505 --- .../airflow/sdk/execution_time/supervisor.py | 23 +++- .../execution_time/test_supervisor.py | 105 +++++++++++++++++- 2 files changed, 123 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 87311f02da7a1..02ad361a995fd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -711,6 +711,16 @@ def start( pid = os.fork() if pid == 0: + # Put the task-runner into its own session so its PGID == its own + # PID. The supervisor can then deliver signals to the whole tree + # via os.killpg() in kill(), 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. See issue #65505. + with suppress(OSError): + os.setsid() + # Close and delete of the parent end of the sockets. cls._close_unused_sockets(read_requests, read_stdout, read_stderr, read_logs) @@ -1036,7 +1046,18 @@ def kill( for sig in escalation_path: try: - self._process.send_signal(sig) + # Signal the whole process group so subprocesses the + # task-runner spawned (venv children, Docker exec, bash + # shells, etc.) are also reached. Requires the task-runner to + # have been placed in its own session via os.setsid() at fork + # time (see start()). See issue #65505. + try: + os.killpg(os.getpgid(self._process.pid), sig) + except (ProcessLookupError, PermissionError): + # Group vanished or we lack permission (e.g. task already + # reaped, or the child never reached setsid). Fall back + # to signalling the task-runner alone. + self._process.send_signal(sig) start = time.monotonic() end = start + escalation_delay 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..c2fa69fd61095 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1241,6 +1241,60 @@ def test_cleanup_sockets_after_delay(self, monkeypatch, mocker): proc.selector.close.assert_called_once() proc.stdin.close.assert_called_once() + def test_child_is_session_leader(self, client_with_ti_start): + """Regression test for #65505: after fork, the task-runner child must + call os.setsid() so its PGID equals its own PID. This allows kill() + to reach subprocesses the task-runner spawns via os.killpg(); without + setsid, 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). + """ + + 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: + # Give the child a moment to run setsid() after fork. + deadline = time.monotonic() + 2.0 + child_pgid = None + while time.monotonic() < deadline: + try: + child_pgid = os.getpgid(proc.pid) + except ProcessLookupError: + sleep(0.05) + continue + if child_pgid == proc.pid: + break + sleep(0.05) + + assert child_pgid == proc.pid, ( + "Task-runner child must be its own session/process-group leader " + f"(os.setsid called after fork). Got pgid={child_pgid}, pid={proc.pid}." + ) + assert child_pgid != os.getpgid(os.getpid()), ( + "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() + class TestWatchedSubprocessKill: @pytest.fixture @@ -1270,8 +1324,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 +1336,56 @@ 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", return_value=12345) + 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) + mock_getpgid.assert_called_once_with(12345) + 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", return_value=12345) + mock_killpg = mocker.patch("os.killpg") + mock_process.wait.return_value = 0 + + watched_subprocess.kill(signal.SIGTERM, force=False) + + mock_getpgid.assert_called_once_with(12345) + mock_killpg.assert_called_once_with(12345, signal.SIGTERM) + mock_process.send_signal.assert_not_called() + + @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 + or child never reached setsid) 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) + mocker.patch("os.killpg") + else: + mocker.patch("os.getpgid", return_value=12345) + mocker.patch("os.killpg", side_effect=exc) + mock_process.wait.return_value = 0 + + watched_subprocess.kill(signal.SIGTERM, force=False) + + mock_process.send_signal.assert_called_once_with(signal.SIGTERM) + @pytest.mark.parametrize( ("signal_to_send", "exit_after"), [ From 3da9b1dc55e35b49d03b6de8fe9fcccb68ec9b5e Mon Sep 17 00:00:00 2001 From: Christoph <116812500+cmettler@users.noreply.github.com> Date: Sat, 4 Jul 2026 18:32:31 +0000 Subject: [PATCH 2/2] Scope process-group handling to the task runner and guard kill() against self-signalling Review feedback on #65738: os.killpg(os.getpgid(child)) trusted that the child had already run setsid() -- if setpgid failed or kill() ran before the child was first scheduled (task_instances.start() raising synchronously), getpgid resolved to the supervisor's own group and killpg would have signalled the supervisor and all its siblings, with no exception for the fallback to catch. Use a plain process group (setpgid, matching airflow.utils.process_utils.set_new_process_group) instead of a new session, set it from both sides of the fork so the group exists as soon as start() returns, refuse to killpg our own group, and make the whole behaviour opt-in per subclass (like use_exec) so the DAG processor, triggerer and callback subprocesses keep direct signalling. The graceful SIGTERM-forwarding path now also signals the group, closing the same orphan leak on e.g. K8s pod termination. --- .../airflow/sdk/execution_time/supervisor.py | 88 +++++++++---- .../execution_time/test_supervisor.py | 120 +++++++++++++----- 2 files changed, 152 insertions(+), 56 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 02ad361a995fd..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,15 +719,18 @@ def start( pid = os.fork() if pid == 0: - # Put the task-runner into its own session so its PGID == its own - # PID. The supervisor can then deliver signals to the whole tree - # via os.killpg() in kill(), 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. See issue #65505. - with suppress(OSError): - os.setsid() + 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) @@ -772,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) @@ -783,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, ) @@ -1017,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, @@ -1046,18 +1091,7 @@ def kill( for sig in escalation_path: try: - # Signal the whole process group so subprocesses the - # task-runner spawned (venv children, Docker exec, bash - # shells, etc.) are also reached. Requires the task-runner to - # have been placed in its own session via os.setsid() at fork - # time (see start()). See issue #65505. - try: - os.killpg(os.getpgid(self._process.pid), sig) - except (ProcessLookupError, PermissionError): - # Group vanished or we lack permission (e.g. task already - # reaped, or the child never reached setsid). Fall back - # to signalling the task-runner alone. - self._process.send_signal(sig) + self._signal_subprocess(sig) start = time.monotonic() end = start + escalation_delay @@ -1381,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 c2fa69fd61095..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,13 +1242,18 @@ def test_cleanup_sockets_after_delay(self, monkeypatch, mocker): proc.selector.close.assert_called_once() proc.stdin.close.assert_called_once() - def test_child_is_session_leader(self, client_with_ti_start): - """Regression test for #65505: after fork, the task-runner child must - call os.setsid() so its PGID equals its own PID. This allows kill() - to reach subprocesses the task-runner spawns via os.killpg(); without - setsid, 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). + 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(): @@ -1270,24 +1276,12 @@ def subprocess_main(): target=subprocess_main, ) try: - # Give the child a moment to run setsid() after fork. - deadline = time.monotonic() + 2.0 - child_pgid = None - while time.monotonic() < deadline: - try: - child_pgid = os.getpgid(proc.pid) - except ProcessLookupError: - sleep(0.05) - continue - if child_pgid == proc.pid: - break - sleep(0.05) - + child_pgid = os.getpgid(proc.pid) assert child_pgid == proc.pid, ( - "Task-runner child must be its own session/process-group leader " - f"(os.setsid called after fork). Got pgid={child_pgid}, pid={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(os.getpid()), ( + 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." ) @@ -1295,6 +1289,24 @@ def subprocess_main(): 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 @@ -1315,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) @@ -1338,14 +1351,14 @@ def test_kill_process_already_exited(self, watched_subprocess, mock_process, moc 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", return_value=12345) + 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_getpgid.assert_called_once_with(12345) + 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) @@ -1355,35 +1368,78 @@ def test_kill_signals_process_group(self, watched_subprocess, mock_process, mock group so subprocesses spawned by the task-runner (venv children, Docker exec, bash shells) are also reached. """ - mock_getpgid = mocker.patch("os.getpgid", return_value=12345) + 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) - mock_getpgid.assert_called_once_with(12345) + 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 - or child never reached setsid) or PermissionError, fall back to + """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) - mocker.patch("os.killpg") + mock_killpg = mocker.patch("os.killpg") else: - mocker.patch("os.getpgid", return_value=12345) - mocker.patch("os.killpg", side_effect=exc) + 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(