Use os.posix_spawn instead of fork+exec in the task-sdk supervisor

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

@@ -495,11 +495,13 @@ def exit(n: int) -> NoReturn:
 (e.g. via ``socket.getaddrinfo`` -> system DNS resolver -> proxy lookup), the
 runtime detects the corrupted state and crashes with SIGABRT.
 
-Calling ``os.execv`` immediately after ``os.fork`` replaces the child's address
-space, giving it clean ObjC state.  Before exec, the supervisor ``dup2``s the
-socketpairs onto fixed FDs the exec'd child reconstructs: 0 (requests/stdin),
-1 (stdout), 2 (stderr), 3 (structured logs).  ``os.set_inheritable`` clears
-``FD_CLOEXEC`` on those FDs so they survive the upcoming exec.
+Starting a fresh interpreter via ``os.posix_spawn`` gives it clean ObjC state --
+and, unlike ``fork()`` followed by ``execv()``, never runs
+``os.register_at_fork()``/``pthread_atfork()`` handlers at all.  The socketpairs
+are remapped onto fixed FDs the spawned child reconstructs, via ``posix_spawn``'s
+own ``file_actions`` (``POSIX_SPAWN_DUP2``): 0 (requests/stdin), 1 (stdout),
+2 (stderr), 3 (structured logs) -- no separate ``set_inheritable`` call is
+needed, since the remapping runs as part of the spawn itself.
 
 Task execution (``ActivitySubprocess``), the DAG processor
 (``DagFileProcessorProcess``) and the triggerer (``TriggerRunnerSupervisor``)
@@ -569,17 +571,18 @@ def _resolve_child_target(dotted: str) -> Callable[[], None]:
 
 def _child_exec_main():
     """
-    Entry point for the child process when using fork+exec.
+    Entry point for the child process when using ``os.posix_spawn``.
 
-    After exec, FDs 0/1/2/3 are the requests/stdout/stderr/log sockets the parent
-    placed there via dup2.  The target to run is named in ``_AIRFLOW_CHILD_TARGET``
-    (``module:qualname``); it is rehydrated and handed to :func:`_fork_main`, which
-    sets up the structured log channel from FD 3 exactly as the bare-fork path does.
+    FDs 0/1/2/3 are the requests/stdout/stderr/log sockets ``posix_spawn``'s own
+    ``file_actions`` (``POSIX_SPAWN_DUP2``) remapped there as part of the spawn.
+    The target to run is named in ``_AIRFLOW_CHILD_TARGET`` (``module:qualname``);
+    it is rehydrated and handed to :func:`_fork_main`, which sets up the structured
+    log channel from FD 3 exactly as the bare-fork path does.
     """
     # The bootstrap already restored PR_SET_DUMPABLE before importing Airflow; this is the
-    # logged fallback (execve had reset what supervise_task() set before the fork).
+    # logged fallback (posix_spawn's own exec had reset what supervise_task() set beforehand).
     _make_process_nondumpable()
-    # FDs 0, 1, 2 were dup2'd onto the socketpairs before exec.
+    # FDs 0, 1, 2 were remapped onto the socketpairs by posix_spawn's file_actions.
     child_requests = socket(fileno=0)
     child_stdout = socket(fileno=1)
     child_stderr = socket(fileno=2)
@@ -729,12 +732,13 @@ def start(
         """
         Fork and start a new subprocess with the specified target function.
 
-        :param use_exec: If True, immediately ``os.execv`` a fresh Python interpreter
-            after ``os.fork``: forced on platforms that need it (macOS, whose Objective-C
-            frameworks are not fork-safe) and opted into for the task process elsewhere via
-            ``[core] execute_tasks_new_python_interpreter`` (a lock a supervisor thread
-            held at fork time cannot survive into a fresh address space).
-            ``target`` is rehydrated in the exec'd child from its ``module:qualname``,
+        :param use_exec: If True, start a fresh Python interpreter via ``os.posix_spawn``
+            instead of a bare ``os.fork``: forced on platforms that need it (macOS, whose
+            Objective-C frameworks are not fork-safe) and opted into for the task process
+            elsewhere via ``[core] execute_tasks_new_python_interpreter``. Unlike
+            ``fork()`` followed by ``execv()``, ``posix_spawn`` never runs
+            ``os.register_at_fork()``/``pthread_atfork()`` handlers at all.
+            ``target`` is rehydrated in the spawned 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
@@ -746,7 +750,7 @@ def start(
         """
         if use_exec and "<" in getattr(target, "__qualname__", "<"):
             # Closures/lambdas (``<locals>`` / ``<lambda>`` in the qualname) and
-            # objects without a qualname can't be named for the exec'd child.
+            # objects without a qualname can't be named for the spawned child.
             raise ValueError(f"use_exec=True requires a top-level importable target, got {target!r}")
         # Create socketpairs/"pipes" to connect to the stdin and out from the subprocess
         child_stdout, read_stdout = socketpair()
@@ -755,80 +759,81 @@ def start(
         # Place for child to send requests/read responses, and the server side to read/respond
         child_requests, read_requests = socketpair()
 
-        # Open the socketpair before forking off the child, so that it is open when we fork.
+        # Open the socketpair before starting the child, so that it is open when we do.
         child_logs, read_logs = socketpair()
 
-        pid = os.fork()
-        if pid == 0:
+        if use_exec:
+            # file_actions run as part of the spawn itself -- no forked child to run
+            # imperative dup2 code in. All four source FDs here are guaranteed >= 3
+            # (0/1/2 are already open in every launch path), so no dup2 target below
+            # clobbers a source that hasn't been placed onto its own target FD yet.
+            file_actions = [
+                (os.POSIX_SPAWN_DUP2, child_requests.fileno(), 0),
+                (os.POSIX_SPAWN_DUP2, child_stdout.fileno(), 1),
+                (os.POSIX_SPAWN_DUP2, child_stderr.fileno(), 2),
+                (os.POSIX_SPAWN_DUP2, child_logs.fileno(), 3),
+            ]
+            spawn_kwargs: dict[str, Any] = {"file_actions": file_actions}
             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)
-
-            # Python GC should delete these for us, but lets make double sure that we don't keep anything
-            # around in the forked processes, especially things that might involve open files or sockets!
-            del constructor_kwargs
-            del logger
+                # Atomic with the spawn -- no parent-side mirror needed, unlike the
+                # bare-fork path below.
+                spawn_kwargs["setpgroup"] = 0
+            child_env = dict(os.environ, _AIRFLOW_CHILD_TARGET=f"{target.__module__}:{target.__qualname__}")
+            pid = os.posix_spawn(
+                sys.executable,
+                [sys.executable, "-c", _CHILD_EXEC_BOOTSTRAP],
+                child_env,
+                **spawn_kwargs,
+            )
+        else:
+            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)
+
+                # Python GC should delete these for us, but lets make double sure that we don't keep anything
+                # around in the forked processes, especially things that might involve open files or sockets!
+                del constructor_kwargs
+                del logger
 
-            try:
-                if use_exec:
-                    # exec a fresh Python interpreter to drop inherited state that is not
-                    # fork-safe (ObjC/CoreFoundation on macOS; a held lock elsewhere). Redirect the
-                    # socketpairs onto the fixed FDs the exec'd child reconstructs:
-                    # 0 (requests/stdin), 1 (stdout), 2 (stderr), 3 (structured logs).
-                    # The source fds are always >= 3 (0/1/2 stay open in every launch
-                    # path), so no dup2 clobbers a not-yet-placed source. set_inheritable
-                    # guarantees FD_CLOEXEC is clear on all four (dup2 leaves it set when
-                    # a source already equals its target), so they survive execv. The
-                    # entry point is passed to the child by name.
-                    os.environ["_AIRFLOW_CHILD_TARGET"] = f"{target.__module__}:{target.__qualname__}"
-                    os.dup2(child_requests.fileno(), 0)
-                    os.dup2(child_stdout.fileno(), 1)
-                    os.dup2(child_stderr.fileno(), 2)
-                    os.dup2(child_logs.fileno(), 3)
-                    for fd in (0, 1, 2, 3):
-                        os.set_inheritable(fd, True)
-                    os.execv(
-                        sys.executable,
-                        [sys.executable, "-c", _CHILD_EXEC_BOOTSTRAP],
-                    )
-                    # execv replaces the process -- unreachable on success
-                else:
+                try:
                     # Run the child entrypoint
                     _fork_main(child_requests, child_stdout, child_stderr, child_logs.fileno(), target)
-            except BaseException as e:
-                import traceback
-
-                with suppress(BaseException):
-                    # We can't use log here, as if we except out of the child something _weird_ went on.
-                    print("Exception in child process, exiting with code 124", file=sys.stderr)
-                    traceback.print_exception(type(e), e, e.__traceback__, file=sys.stderr)
-
-            # It's really super super important we never exit this block. We are in the forked child, and if we
-            # 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
+                except BaseException as e:
+                    import traceback
+
+                    with suppress(BaseException):
+                        # We can't use log here, as if we except out of the child something _weird_ went on.
+                        print("Exception in child process, exiting with code 124", file=sys.stderr)
+                        traceback.print_exception(type(e), e, e.__traceback__, file=sys.stderr)
+
+                # It's really super super important we never exit this block. We are in the forked child, and if we
+                # 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. We still have the
         # other end of the pair open
         cls._close_unused_sockets(child_stdout, child_stderr, child_logs)
 

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

@@ -4703,6 +4703,127 @@ def test_start_rejects_non_importable_target_under_exec(self):
             supervisor.WatchedSubprocess.start(target=lambda: None, use_exec=True)
 
 
+class TestStartUsesPosixSpawn:
+    """use_exec=True goes through os.posix_spawn, never os.fork -- that's the whole point."""
+
+    def _start(self, mocker, **kwargs):
+        spawn = mocker.patch("airflow.sdk.execution_time.supervisor.os.posix_spawn", return_value=4321)
+        fork = mocker.patch(
+            "airflow.sdk.execution_time.supervisor.os.fork",
+            side_effect=AssertionError("os.fork() must not be called when use_exec=True"),
+        )
+        mocker.patch("airflow.sdk.execution_time.supervisor.psutil.Process")
+        supervisor.WatchedSubprocess.start(
+            id=uuid7(), target=supervisor._subprocess_main, use_exec=True, **kwargs
+        )
+        return spawn, fork
+
+    def test_does_not_call_fork(self, mocker):
+        """The defining property of the fix: no os.fork() call exists on this path at all."""
+        spawn, fork = self._start(mocker)
+        fork.assert_not_called()
+        spawn.assert_called_once()
+
+    def test_spawns_the_bootstrap_with_the_target_env_var(self, mocker):
+        spawn, _ = self._start(mocker)
+        args, kwargs = spawn.call_args
+        path, argv, env = args
+        assert path == sys.executable
+        assert argv == [sys.executable, "-c", supervisor._CHILD_EXEC_BOOTSTRAP]
+        assert env["_AIRFLOW_CHILD_TARGET"] == "airflow.sdk.execution_time.supervisor:_subprocess_main"
+
+    def test_file_actions_dup2_the_four_fds(self, mocker):
+        spawn, _ = self._start(mocker)
+        file_actions = spawn.call_args.kwargs["file_actions"]
+        targets = {new_fd for _, _, new_fd in file_actions}
+        assert targets == {0, 1, 2, 3}
+        assert all(action == os.POSIX_SPAWN_DUP2 for action, _, _ in file_actions)
+
+    def test_setpgroup_passed_when_new_process_group(self, mocker):
+        spawn, _ = self._start(mocker, new_process_group=True)
+        assert spawn.call_args.kwargs["setpgroup"] == 0
+
+    def test_setpgroup_omitted_when_not_new_process_group(self, mocker):
+        spawn, _ = self._start(mocker, new_process_group=False)
+        assert "setpgroup" not in spawn.call_args.kwargs
+
+    @pytest.mark.skipif(sys.platform == "win32", reason="os.fork/os.register_at_fork are POSIX-only")
+    def test_hanging_after_fork_handler_wedges_bare_fork_but_not_posix_spawn(self):
+        """
+        A handler registered via os.register_at_fork(after_in_child=...) that never
+        returns wedges a bare-forked child forever, but does not affect a posix_spawn'd
+        child at all -- posix_spawn never runs it.
+
+        Runs in a disposable subprocess: os.register_at_fork() has no unregister call,
+        so registering a permanently-hanging one here would otherwise poison every later
+        fork in this pytest worker for the rest of the test run.
+        """
+        probe = """
+import os, sys, time
+
+def _hangs_forever():
+    while True:
+        time.sleep(3600)
+
+os.register_at_fork(after_in_child=_hangs_forever)
+
+r, w = os.pipe()
+pid = os.fork()
+if pid == 0:
+    os.write(w, b"unreachable")
+    os._exit(0)
+os.close(w)
+os.set_blocking(r, False)
+deadline = time.monotonic() + 2
+hung = True
+while time.monotonic() < deadline:
+    try:
+        if os.read(r, 1):
+            hung = False
+            break
+    except BlockingIOError:
+        time.sleep(0.01)
+os.close(r)
+os.kill(pid, 9)
+os.waitpid(pid, 0)
+if not hung:
+    print("FAIL: bare fork did not hang despite the handler")
+    sys.exit(1)
+
+r2, w2 = os.pipe()
+os.set_inheritable(w2, True)
+pid2 = os.posix_spawn(
+    sys.executable,
+    [sys.executable, "-c", "print('ok')"],
+    os.environ,
+    file_actions=[(os.POSIX_SPAWN_DUP2, w2, 1)],
+)
+os.close(w2)
+os.set_blocking(r2, False)
+deadline = time.monotonic() + 2
+spawned_ok = False
+while time.monotonic() < deadline:
+    try:
+        data = os.read(r2, 8)
+        if data.strip() == b"ok":
+            spawned_ok = True
+            break
+    except BlockingIOError:
+        time.sleep(0.01)
+os.close(r2)
+os.waitpid(pid2, 0)
+if not spawned_ok:
+    print("FAIL: posix_spawn hung too, despite the same handler still registered")
+    sys.exit(1)
+
+print("PASS")
+"""
+        result = subprocess.run(
+            [sys.executable, "-c", probe], capture_output=True, text=True, timeout=15, check=False
+        )
+        assert result.stdout.strip() == "PASS", f"stdout={result.stdout!r} stderr={result.stderr!r}"
+
+
 @pytest.mark.usefixtures("disable_capturing")
 def test_fork_exec_bootstrap_runs_an_importable_target_end_to_end(
     captured_logs, time_machine, monkeypatch, client_with_ti_start