uranusjr commented on code in PR #75707:
URL: https://github.com/apache/airflow/pull/75707#discussion_r4385334568
##########
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:
Redundant, this is already covered by the two other tests above and below.
##########
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:
Redundant, this is already covered by the two other added tests.
--
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]
