Skip atexit in forked task children rather than running post-fork handlers

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()

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():