diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 4c3ae38d21fab..702eda5acac3f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -19,7 +19,6 @@ from __future__ import annotations -import atexit import contextlib import functools import io @@ -442,17 +441,12 @@ def exit(n: int) -> NoReturn: os.close(requests.fileno()) os._exit(n) - if hasattr(atexit, "_clear"): - # Since we're in a fork we want to try and clear them. If we can't do it cleanly, then we won't try - # and run new atexit handlers. - with suppress(Exception): - atexit._clear() - base_exit = exit - - def exit(n: int) -> NoReturn: - # This will only run any atexit funcs registered after we've forked. - atexit._run_exitfuncs() - base_exit(n) + # We deliberately never run atexit handlers registered post-fork (used to via + # atexit._run_exitfuncs() here). A library's atexit handler that isn't fork-safe -- holds a + # native lock or thread that didn't survive the fork -- can block forever, which means this + # forked task child, and the pod hosting it, never exits (e.g. pyarrow's S3 client finalizer + # deadlocking in its AWS-CRT teardown). There's no way to tell a safe handler from an unsafe + # one in advance, so os._exit() skips all of them rather than risk running an unknown one. try: block_orm_access() 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 c9a977c9793c8..b1cccee8c6224 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -386,6 +386,41 @@ def subprocess_main(): ] ) + def test_fork_main_never_runs_post_fork_atexit_handlers(self, tmp_path, client_with_ti_start): + """A fork-unsafe library's atexit handler must never run in the forked task child. If it + did and it blocked (e.g. pyarrow's S3 client finalizer deadlocking in AWS-CRT teardown), + the child -- and the pod hosting it -- would never exit.""" + marker = tmp_path / "atexit-ran" + + def subprocess_main(): + # This is run in the subprocess! + CommsDecoder()._get_response() + + import atexit + + atexit.register(marker.write_text, "atexit ran") + + proc = ActivitySubprocess.start( + dag_rel_path=os.devnull, + bundle_info=FAKE_BUNDLE, + what=TaskInstance( + id="4d828a62-a417-4936-a7a6-2b3fabacecab", + 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, + ) + + rc = proc.wait() + + assert rc == 0 + assert not marker.exists() + @pytest.mark.flaky(reruns=3) def test_reopen_log_fd(self, captured_logs, client_with_ti_start): def subprocess_main():