jason810496 commented on code in PR #75707:
URL: https://github.com/apache/airflow/pull/75707#discussion_r4385551218
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -387,6 +406,7 @@ def start( # type: ignore[override]
start_time=time.monotonic(),
stdin=socks[comm_server],
subprocess_schema_version=subprocess_schema_version,
+ new_process_group=True,
Review Comment:
6428e774187: each site now names the other.
##########
task-sdk/tests/task_sdk/coordinators/test_subprocess.py:
##########
@@ -1150,6 +1171,164 @@ def handle_request(request: httpx2.Request) ->
httpx2.Response:
} in captured_logs
+_RUNTIME_SPAWNING_A_CHILD = """
+import socket
+import subprocess
+import sys
+
+def connect(flag):
+ host, _, port = next(arg for arg in sys.argv if
arg.startswith(flag)).split("=", 1)[1].rpartition(":")
+ return socket.create_connection((host, int(port)))
+
+comm = connect("--comm=")
+logs = connect("--logs=")
+child = subprocess.Popen([sys.executable, "-c", "import time;
time.sleep(300)"])
+with open(sys.argv[1], "w") as handle:
+ handle.write(str(child.pid))
+comm.settimeout(30)
+try:
+ comm.recv(1)
+except TimeoutError:
+ sys.exit(3)
+"""
+
+
+def _wait_until_gone(pid: int, timeout: float = 10.0) -> bool:
+ """Return whether *pid* has stopped running, treating a not-yet-reaped
zombie as gone."""
+ deadline = time.monotonic() + timeout
+ while True:
+ try:
+ if psutil.Process(pid).status() == psutil.STATUS_ZOMBIE:
+ return True
+ except psutil.NoSuchProcess:
+ return True
+ if time.monotonic() >= deadline:
+ return False
+ time.sleep(0.05)
+
+
+class TestSignalProcessTree:
+ """``_signal_process_tree`` reaches the runtime's group, or the pid when
it cannot."""
+
+ def test_signals_the_group(self):
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ with (
+ patch("os.getpgid", side_effect=lambda pid: 6767 if pid == 6767
else os.getpid()),
+ patch("os.killpg") as mock_killpg,
+ ):
+ _signal_process_tree(proc, signal.SIGTERM)
+ mock_killpg.assert_called_once_with(6767, signal.SIGTERM)
+ proc.send_signal.assert_not_called()
+
+ def test_signals_the_pid_when_the_runtime_shares_our_group(self):
+ """Signalling our own group would hit the supervisor itself and its
siblings."""
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ with (
+ patch("os.getpgid", return_value=os.getpid()),
+ patch("os.killpg") as mock_killpg,
+ ):
+ _signal_process_tree(proc, signal.SIGTERM)
+ mock_killpg.assert_not_called()
+ proc.send_signal.assert_called_once_with(signal.SIGTERM)
+
+ @pytest.mark.parametrize("failing", ["getpgid", "killpg"])
+ @pytest.mark.parametrize("error", [ProcessLookupError, PermissionError])
+ def test_signals_the_pid_when_the_group_cannot_be_signalled(self, failing,
error):
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ own_group = os.getpgid(0)
+
+ def resolve(pid):
+ return 6767 if pid == 6767 else own_group
+
+ getpgid = {"side_effect": error if failing == "getpgid" else resolve}
+ killpg = {"side_effect": error} if failing == "killpg" else {}
+ with patch("os.getpgid", **getpgid), patch("os.killpg", **killpg) as
mock_killpg:
+ _signal_process_tree(proc, signal.SIGTERM)
+ assert mock_killpg.called is (failing == "killpg")
+ proc.send_signal.assert_called_once_with(signal.SIGTERM)
+
+
+class TestPopenActivitySubprocessProcessGroup:
+ """A language SDK runtime runs in its own process group so a stop reaches
what it spawned."""
+
+ def test_popen_asks_for_a_new_process_group(self, mock_client):
+ mock_proc = MagicMock(spec=subprocess.Popen)
+ mock_proc.pid = 12345
+ with (
+ patch(
+ "airflow.sdk.coordinators._subprocess.subprocess.Popen",
return_value=mock_proc
+ ) as popen_mock,
+ patch(
+ "airflow.sdk.coordinators._subprocess._accept_connections",
+ side_effect=lambda servers, drains, proc, **kw: (
+ {soc: MagicMock(spec=socket.socket) for soc in
servers.values()},
+ {soc: b"" for soc in drains.values()},
+ ),
+ ),
+ patch.object(ActivitySubprocess, "_register_pipe_readers"),
+ patch.object(ActivitySubprocess, "_on_child_started"),
+ ):
+ proc = _PopenActivitySubprocess.start(
+ what=_make_ti(),
+ dag_rel_path="bundle",
+ bundle_info=MagicMock(),
+ client=mock_client,
+ command=["/bin/true"],
+ subprocess_logs_to_stdout=False,
+ )
+
+ assert popen_mock.call_args.kwargs["process_group"] == 0
+ # The supervisor only uses killpg when it knows the child has its own
group.
+ assert proc._new_process_group is True
+
+ def test_runtime_runs_in_its_own_process_group(self, tmp_path):
+ """Regression for #75533: a real runtime is a group leader, so killpg
reaches its children."""
+ pid_file = tmp_path / "child.pid"
+ with patch.object(ActivitySubprocess, "_on_child_started"):
+ proc = _PopenActivitySubprocess.start(
+ what=_make_ti(),
+ dag_rel_path="bundle",
+ bundle_info=FAKE_BUNDLE,
+ client=MagicMock(),
+ command=[sys.executable, "-c", _RUNTIME_SPAWNING_A_CHILD,
str(pid_file)],
+ subprocess_logs_to_stdout=False,
+ )
+ try:
+ assert os.getpgid(proc.pid) == proc.pid
+ finally:
+ proc.kill(signal.SIGKILL)
Review Comment:
6428e774187: deleted.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -206,6 +206,20 @@ def _build_runtime_env() -> dict[str, str]:
}
+def _signal_process_tree(proc: subprocess.Popen, sig: signal.Signals) -> None:
+ """
+ Deliver ``sig`` to *proc*'s process group, so everything the task spawned
is reached.
+
+ Falls back to the pid when the group cannot be signalled, and when *proc*
is in the
+ supervisor's own group, which ``killpg`` would signal too.
+ """
+ with contextlib.suppress(OSError):
+ if (pgid := os.getpgid(proc.pid)) != os.getpgid(0):
+ os.killpg(pgid, sig)
+ return
+ proc.send_signal(sig)
Review Comment:
6428e774187: pulled it out as `signal_process_tree` in `supervisor.py` and
called it in subprocess coordiantor.
The missing test found a real hole, thanks.
A runtime that exits before connecting is reaped by `proc.poll()` in
`_accept_connections`, so `getpgid` raises and the `Popen.send_signal` fallback
no-ops on an exited process: the child leaked.
Added `_kill_leftover_group`, which reaches the group by the leader's pid
and bails when that pid was reused.
`test_a_runtime_that_dies_before_connecting_does_not_leave_its_child` fails
without it.
##########
task-sdk/tests/task_sdk/coordinators/test_subprocess.py:
##########
@@ -560,10 +561,30 @@ def
test_context_manager_closes_tracked_socket_on_exit(self):
sock.close.assert_called_once()
def test_context_manager_terminates_tracked_popen_on_exit(self):
+ """A failed startup signals the runtime's whole process group, not
just its pid."""
proc = MagicMock(spec=subprocess.Popen)
- with _ResourceTracker(timeout=0.1) as tracker:
+ proc.pid = 6767
+ with (
+ patch("os.getpgid", side_effect=lambda pid: 6767 if pid == 6767
else os.getpid()),
+ patch("os.killpg") as mock_killpg,
+ _ResourceTracker(timeout=0.1) as tracker,
+ ):
Review Comment:
6428e774187: consolidated them into `runtime_process_group` fixture.
##########
task-sdk/tests/task_sdk/coordinators/test_subprocess.py:
##########
@@ -1150,6 +1171,164 @@ def handle_request(request: httpx2.Request) ->
httpx2.Response:
} in captured_logs
+_RUNTIME_SPAWNING_A_CHILD = """
+import socket
+import subprocess
+import sys
+
+def connect(flag):
+ host, _, port = next(arg for arg in sys.argv if
arg.startswith(flag)).split("=", 1)[1].rpartition(":")
+ return socket.create_connection((host, int(port)))
+
+comm = connect("--comm=")
+logs = connect("--logs=")
+child = subprocess.Popen([sys.executable, "-c", "import time;
time.sleep(300)"])
+with open(sys.argv[1], "w") as handle:
+ handle.write(str(child.pid))
+comm.settimeout(30)
+try:
+ comm.recv(1)
+except TimeoutError:
+ sys.exit(3)
+"""
+
+
+def _wait_until_gone(pid: int, timeout: float = 10.0) -> bool:
+ """Return whether *pid* has stopped running, treating a not-yet-reaped
zombie as gone."""
+ deadline = time.monotonic() + timeout
+ while True:
+ try:
+ if psutil.Process(pid).status() == psutil.STATUS_ZOMBIE:
+ return True
+ except psutil.NoSuchProcess:
+ return True
+ if time.monotonic() >= deadline:
+ return False
+ time.sleep(0.05)
+
+
+class TestSignalProcessTree:
+ """``_signal_process_tree`` reaches the runtime's group, or the pid when
it cannot."""
+
+ def test_signals_the_group(self):
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ with (
+ patch("os.getpgid", side_effect=lambda pid: 6767 if pid == 6767
else os.getpid()),
+ patch("os.killpg") as mock_killpg,
+ ):
+ _signal_process_tree(proc, signal.SIGTERM)
+ mock_killpg.assert_called_once_with(6767, signal.SIGTERM)
+ proc.send_signal.assert_not_called()
+
+ def test_signals_the_pid_when_the_runtime_shares_our_group(self):
+ """Signalling our own group would hit the supervisor itself and its
siblings."""
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ with (
+ patch("os.getpgid", return_value=os.getpid()),
+ patch("os.killpg") as mock_killpg,
+ ):
+ _signal_process_tree(proc, signal.SIGTERM)
+ mock_killpg.assert_not_called()
+ proc.send_signal.assert_called_once_with(signal.SIGTERM)
+
+ @pytest.mark.parametrize("failing", ["getpgid", "killpg"])
+ @pytest.mark.parametrize("error", [ProcessLookupError, PermissionError])
+ def test_signals_the_pid_when_the_group_cannot_be_signalled(self, failing,
error):
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = 6767
+ own_group = os.getpgid(0)
+
+ def resolve(pid):
+ return 6767 if pid == 6767 else own_group
+
+ getpgid = {"side_effect": error if failing == "getpgid" else resolve}
+ killpg = {"side_effect": error} if failing == "killpg" else {}
+ with patch("os.getpgid", **getpgid), patch("os.killpg", **killpg) as
mock_killpg:
+ _signal_process_tree(proc, signal.SIGTERM)
+ assert mock_killpg.called is (failing == "killpg")
+ proc.send_signal.assert_called_once_with(signal.SIGTERM)
+
+
+class TestPopenActivitySubprocessProcessGroup:
+ """A language SDK runtime runs in its own process group so a stop reaches
what it spawned."""
+
+ def test_popen_asks_for_a_new_process_group(self, mock_client):
+ mock_proc = MagicMock(spec=subprocess.Popen)
+ mock_proc.pid = 12345
+ with (
+ patch(
+ "airflow.sdk.coordinators._subprocess.subprocess.Popen",
return_value=mock_proc
+ ) as popen_mock,
+ patch(
+ "airflow.sdk.coordinators._subprocess._accept_connections",
+ side_effect=lambda servers, drains, proc, **kw: (
+ {soc: MagicMock(spec=socket.socket) for soc in
servers.values()},
+ {soc: b"" for soc in drains.values()},
+ ),
+ ),
+ patch.object(ActivitySubprocess, "_register_pipe_readers"),
+ patch.object(ActivitySubprocess, "_on_child_started"),
+ ):
+ proc = _PopenActivitySubprocess.start(
+ what=_make_ti(),
+ dag_rel_path="bundle",
+ bundle_info=MagicMock(),
+ client=mock_client,
+ command=["/bin/true"],
+ subprocess_logs_to_stdout=False,
+ )
+
+ assert popen_mock.call_args.kwargs["process_group"] == 0
+ # The supervisor only uses killpg when it knows the child has its own
group.
+ assert proc._new_process_group is True
Review Comment:
6428e774187: taken.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -365,6 +379,11 @@ def start( # type: ignore[override]
stdout=stdout_w.fileno(),
stderr=stderr_w.fileno(),
env=_build_runtime_env(),
+ # Own process group, so the supervisor can signal everything
the runtime
+ # spawned with os.killpg, as the Python task runner has done
since #65505.
+ # Popen returns only once the child has exec'd, so the group
always exists
+ # before kill() can run.
Review Comment:
6428e774187: taken.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
