diff --git a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py index 55eff9df31e59..496b0179d8913 100644 --- a/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py +++ b/providers/ssh/src/airflow/providers/ssh/utils/remote_job.py @@ -184,11 +184,22 @@ def build_posix_wrapper_command( escaped_command = command.replace("'", "'\"'\"'") # Launch detached under ``setsid`` so the job is its own session/process-group - # leader. ``$!`` is then the leader PID *and* the PGID (verified: setsid does not - # fork when started as a background job), recorded synchronously just like before, - # so cancellation can signal the whole job tree instead of orphaning the user - # command. Without ``setsid`` (some macOS/BSD hosts) ``$!`` is just the wrapper PID - # and cancellation degrades to the previous single-process behaviour. + # leader, letting cancellation signal the whole job tree instead of orphaning the + # user command. We do NOT rely on the launcher's ``$!`` to identify that group: + # ``setsid(1)`` forks internally when the process about to exec into it is already + # a process-group leader (``setsid(2)`` cannot create a new session from a group + # leader), which happens whenever job control is on in the launching shell. When + # it forks, ``$!`` is the short-lived setsid parent, not the job's real PGID, and + # cancellation signals a dead group. Instead the job script reports its OWN pid: + # immediately after ``setsid(2)`` POSIX guarantees ``pid == pgid == sid`` for the + # caller, and that identity survives the following exec, so ``$$`` inside the job + # script is always the true PGID regardless of whether setsid forked to get there. + # The pid file is read only when cancelling (:func:`build_posix_kill_command`), + # which happens long after submission, so the job's asynchronous write lands well + # before any reader and the launcher does not wait for it (recording the launcher's + # ``$!`` would just reintroduce the wrong-pid bug on the fork path). Without + # ``setsid`` (some macOS/BSD hosts) ``$$`` is just the job's own PID and + # cancellation degrades to the previous single-process behaviour. wrapper = f"""set -euo pipefail job_dir='{paths.job_dir}' log_file='{paths.log_file}' @@ -202,6 +213,7 @@ def build_posix_wrapper_command( job_script=' set +e +echo -n "$$" > "'"$pid_file"'" export LOG_FILE="'"$log_file"'" export STATUS_FILE="'"$status_file"'" {env_exports}{escaped_command} >>"'"$log_file"'" 2>&1 @@ -216,7 +228,6 @@ def build_posix_wrapper_command( else nohup bash -c "$job_script" >/dev/null 2>&1 & fi -echo -n $! > "$pid_file" echo "{paths.job_id}" """ return wrapper diff --git a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py index 56fe5d7abb33a..82b285d4587fc 100644 --- a/providers/ssh/tests/unit/ssh/utils/test_remote_job.py +++ b/providers/ssh/tests/unit/ssh/utils/test_remote_job.py @@ -18,7 +18,10 @@ from __future__ import annotations import base64 +import contextlib import os +import pty +import select import shutil import subprocess import time @@ -132,7 +135,7 @@ def test_escapes_quotes(self): assert wrapper is not None def test_runs_in_own_process_group(self): - """The job launches under setsid (when available); $! is the leader PID/PGID.""" + """The job launches under setsid (when available) and self-reports its PGID.""" paths = RemoteJobPaths(job_id="test_job", remote_os="posix") wrapper = build_posix_wrapper_command("/path/to/script.sh", paths) @@ -140,8 +143,11 @@ def test_runs_in_own_process_group(self): assert "command -v setsid" in wrapper assert "setsid bash -c" in wrapper assert "nohup bash -c" in wrapper - # Leader PID recorded synchronously by the launcher ($! == PGID under setsid) - assert 'echo -n $! > "$pid_file"' in wrapper + # The job self-reports its own pid ($$ == PGID after setsid), which is correct + # even when setsid(1) forks under job control -- unlike the launcher's $!. + assert 'echo -n "$$" > "' in wrapper + # Launcher must NOT record $! (would be the short-lived setsid parent on a fork). + assert 'echo -n $! > "$pid_file"' not in wrapper class TestBuildWindowsWrapperCommand: @@ -262,46 +268,123 @@ class TestPosixKillBehaviour: Regression test for the orphaned-process bug: killing only the recorded PID left the user command (and its children) running, so the exit_code file was never written and - the trigger timed out. The job now runs in its own process group and the kill signals - the group. + the trigger timed out. The job runs in its own process group and self-reports that + group's PGID, so the kill signals the whole group even when setsid(1) forks. """ + def _marker(self, tag: str) -> str: + # Unique per (test, xdist worker): CI runs these with ``-n auto`` (default + # ``load`` distribution), so sibling tests can execute concurrently in separate + # workers against the same OS process table. A shared literal would let one + # test's ``pgrep -f`` / ``pkill -f`` match or kill another's job. os.getpid() + # differs per worker; the tag differs per test. + return f"sleep 9{tag}{os.getpid()}" + @staticmethod def _group_alive(pgid: int) -> bool: # pgrep -g matches by process-group id; rc 0 => at least one member alive. return subprocess.run(["pgrep", "-g", str(pgid)], capture_output=True, check=False).returncode == 0 - # setsid only avoids forking when the launching shell is not a process-group leader; on some - # CI runners it forks, so the recorded $! is the short-lived setsid parent rather than the job - # PGID and the pre-kill pgrep -g finds an empty group. Re-launch on a fresh draw. - @pytest.mark.flaky(reruns=5) - def test_kill_terminates_whole_job_tree(self, tmp_path): - paths = RemoteJobPaths(job_id="killtree", remote_os="posix", base_dir=str(tmp_path / "jobs")) - # `sleep 300` runs as a child of the wrapper subshell -> the tree the old kill orphaned. - # Run under bash, which is the remote login shell this operator requires (the wrapper - # uses `set -o pipefail`); the kill is run the same way below. - wrapper = build_posix_wrapper_command("sleep 300", paths) - subprocess.run(["bash", "-c", wrapper], check=True, capture_output=True, text=True) + @staticmethod + def _job_running(marker: str) -> bool: + return subprocess.run(["pgrep", "-f", marker], capture_output=True, check=False).returncode == 0 + + @staticmethod + def _pgid_of(pid: str) -> str: + return subprocess.run( + ["ps", "-o", "pgid=", "-p", pid], capture_output=True, text=True, check=False + ).stdout.strip() - # The launcher records $! synchronously, so the pid file is present on return. + def _await_recorded_pid(self, paths) -> int: + # The job writes its pid asynchronously (the launcher does not wait), so poll. pid_path = Path(paths.pid_file) - assert pid_path.exists(), "job never wrote its pid file" - pid_text = pid_path.read_text().strip() - assert pid_text, "pid file is empty" - pgid = int(pid_text) + deadline = time.monotonic() + 5 + pid_text = "" + while time.monotonic() < deadline: + pid_text = pid_path.read_text().strip() if pid_path.exists() else "" + if pid_text: + break + time.sleep(0.02) + assert pid_text, "job never wrote its pid file" + return int(pid_text) + @staticmethod + def _run_bash_mc_under_pty(script: str, marker: bytes, timeout: float = 8.0) -> None: + """Run ``bash -mc script`` under a pty we own so job control genuinely activates + (bash silently disables ``-m`` without a controlling terminal). Read until the + marker, NOT to EOF: the detached job inherits the pty slave as its stdin, so EOF + would not arrive until the job itself exits (the full sleep runtime).""" + pid, fd = pty.fork() + if pid == 0: + try: + os.execvp("bash", ["bash", "-mc", script]) + except OSError: + os._exit(127) # never fall through as a duplicate pytest process try: - assert self._group_alive(pgid), "job tree should be running before kill" + deadline = time.monotonic() + timeout + buf = b"" + while time.monotonic() < deadline: + r, _, _ = select.select([fd], [], [], 0.2) + if fd in r: + try: + chunk = os.read(fd, 4096) + except OSError: + break + if not chunk: + break + buf += chunk + if marker in buf: + break + finally: + os.close(fd) # hangs up the pty; the launcher (not the detached job) exits + with contextlib.suppress(ChildProcessError): + os.waitpid(pid, 0) # reap the launcher so it does not linger as a zombie + def _assert_kill_tears_down(self, paths, pgid: int, marker: str) -> None: + try: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and not self._group_alive(pgid): + time.sleep(0.02) + assert self._group_alive(pgid), "job group should be running before kill" subprocess.run(["bash", "-c", build_posix_kill_command(paths.pid_file)], check=True) - deadline = time.monotonic() + 5 - while time.monotonic() < deadline and self._group_alive(pgid): + while time.monotonic() < deadline and self._job_running(marker): time.sleep(0.05) - assert not self._group_alive(pgid), "kill left part of the job tree running" + assert not self._job_running(marker), "kill left the job running (orphaned)" finally: - # Belt-and-suspenders: never leave a stray `sleep 300` behind if an assert fails. subprocess.run(["bash", "-c", f"kill -9 -{pgid} 2>/dev/null || true"], check=False) + subprocess.run(["bash", "-c", f"pkill -9 -f '{marker}' 2>/dev/null || true"], check=False) + + def test_kill_terminates_whole_job_tree(self, tmp_path): + """Default path (no job control): the job self-reports its PGID and the kill + signals the whole group.""" + marker = self._marker("1") + paths = RemoteJobPaths(job_id="killtree", remote_os="posix", base_dir=str(tmp_path / "jobs")) + wrapper = build_posix_wrapper_command(marker, paths) + subprocess.run(["bash", "-c", wrapper], check=True, capture_output=True, text=True) + pgid = self._await_recorded_pid(paths) + self._assert_kill_tears_down(paths, pgid, marker) + + def test_kill_terminates_whole_job_tree_under_job_control(self, tmp_path): + """With job control on, setsid(1) forks and the launcher's ``$!`` would name the + short-lived setsid parent, not the job -- the condition the old wrapper orphaned + the job under. Force it deterministically via a real controlling terminal and + assert the recorded pid IS the job's true PGID and the kill reaches the job.""" + marker = self._marker("2") + paths = RemoteJobPaths(job_id="killtree_jc", remote_os="posix", base_dir=str(tmp_path / "jobs")) + wrapper = build_posix_wrapper_command(marker, paths) + self._run_bash_mc_under_pty(wrapper + "\necho SUBMIT_DONE\n", b"SUBMIT_DONE") + pgid = self._await_recorded_pid(paths) + + job_pids = subprocess.run( + ["pgrep", "-f", marker], capture_output=True, text=True, check=False + ).stdout.split() + assert job_pids, "job never started" + true_pgid = self._pgid_of(job_pids[0]) + # Core regression assertion: recorded pid == the job's real PGID. Under the old + # $!-based wrapper this differs (setsid forked) and on_kill orphans the job. + assert str(pgid) == true_pgid, f"recorded pid {pgid} is not the job PGID {true_pgid}" + self._assert_kill_tears_down(paths, pgid, marker) class TestCleanupCommands: