This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new e2edfff0418 [v3-4-test] Run a language SDK runtime in its own process
group (#75707) (#75731)
e2edfff0418 is described below
commit e2edfff0418bb4e38e8287b278067d537845c8d3
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Tue Oct 27 17:09:30 2026 +0800
[v3-4-test] Run a language SDK runtime in its own process group (#75707)
(#75731)
Stopping a task that a language SDK runs signalled the runtime's pid alone,
so
anything the task had started was reparented to PID 1 and kept running.
While
such a child is alive it also holds the runtime open, so the runtime itself
survived until the SIGKILL escalation `killed_task_cleanup_time` later.
The runtime now starts with `process_group=0` and the supervisor is told the
child has its own group, so the existing `os.killpg` path reaches the whole
tree, as it has for Python tasks since #65505. A startup that fails signals
the group too, and a runtime that exits before connecting is already reaped,
so its group is reached by the leader's pid.
This covers every way the supervisor stops a task -- a cleared task,
`execution_timeout`,
repeated heartbeat failures and task-success overtime -- for the
TypeScript, Java and
executable runtimes alike.
(cherry picked from commit 750f849cc2372ee42dc6910541d4b6a866102ac9)
closes: #75533
Co-authored-by: Claude <[email protected]>
---
.../src/airflow/sdk/coordinators/_subprocess.py | 34 ++++-
.../src/airflow/sdk/execution_time/supervisor.py | 43 +++---
.../tests/task_sdk/coordinators/test_subprocess.py | 170 ++++++++++++++++++++-
3 files changed, 225 insertions(+), 22 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
index 0bd84e147be..5d5223072fe 100644
--- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
+++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py
@@ -55,7 +55,12 @@ from airflow.sdk.execution_time.bundles import
initialize_ti_bundle
from airflow.sdk.execution_time.comms import RetryTask, TaskState
from airflow.sdk.execution_time.coordinator import BaseCoordinator,
TaskLaunchError
from airflow.sdk.execution_time.schema import get_schema_version_migrator
-from airflow.sdk.execution_time.supervisor import ActivitySubprocess,
NeverRaised, ProcessTracker
+from airflow.sdk.execution_time.supervisor import (
+ ActivitySubprocess,
+ NeverRaised,
+ ProcessTracker,
+ signal_process_tree,
+)
if TYPE_CHECKING:
import pathlib
@@ -206,6 +211,19 @@ def _build_runtime_env() -> dict[str, str]:
}
+def _kill_leftover_group(pid: int) -> None:
+ """
+ Kill what an exited runtime left in its process group, which ``getpgid``
can no longer resolve.
+
+ A group outlives its leader, so it is reached by the leader's pid. Not
when another process has
+ reused that pid, since the group could then be that process's.
+ """
+ if psutil.pid_exists(pid):
+ return
+ with contextlib.suppress(ProcessLookupError, PermissionError):
+ os.killpg(pid, signal.SIGKILL)
+
+
def _accept_connections(
servers: dict[str, socket.socket],
drains: dict[str, socket.socket],
@@ -305,11 +323,14 @@ class _ResourceTracker:
case socket.socket():
o.close()
case subprocess.Popen():
- o.terminate()
+ signal_process_tree(o.pid, signal.SIGTERM,
fallback=o.send_signal)
try:
o.wait(self.timeout)
except subprocess.TimeoutExpired:
- o.kill()
+ signal_process_tree(o.pid, signal.SIGKILL,
fallback=o.send_signal)
+ # A runtime that exited before connecting was already
reaped, so its group
+ # is out of reach above while what it spawned is still
running.
+ _kill_leftover_group(o.pid)
def track(self, *objects: Tracked) -> tuple[Tracked, ...]:
self.tracked.update((id(o), o) for o in objects)
@@ -365,6 +386,11 @@ class _PopenActivitySubprocess(ActivitySubprocess):
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 after exec, so the group exists before kill().
+ # Must agree with new_process_group=True below.
+ process_group=0,
)
tracker.track(proc)
for soc in tracker.untrack(stdout_w, stderr_w):
@@ -387,6 +413,8 @@ class _PopenActivitySubprocess(ActivitySubprocess):
start_time=time.monotonic(),
stdin=socks[comm_server],
subprocess_schema_version=subprocess_schema_version,
+ # The supervisor signals the group; must agree with
process_group=0 above.
+ new_process_group=True,
comm_server=comm_server,
logs_server=logs_server,
**kwargs,
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 388075858bd..0c198ab637a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -783,6 +783,27 @@ class PsutilTracker(ProcessTracker):
return self._impl.wait(timeout)
+def signal_process_tree(pid: int, sig: signal.Signals, *, fallback:
Callable[[signal.Signals], None]) -> None:
+ """
+ Deliver ``sig`` to *pid*'s process group, so every process it spawned is
reached too.
+
+ Calls *fallback* instead when the group cannot be resolved or signalled,
and -- critically --
+ when *pid* is still in the caller's own process group, because ``killpg``
on our own group
+ would signal the caller itself and its siblings.
+ """
+ try:
+ pgid = os.getpgid(pid)
+ except (ProcessLookupError, PermissionError):
+ pgid = None
+ if pgid is not None and pgid != os.getpgid(0):
+ try:
+ os.killpg(pgid, sig)
+ return
+ except (ProcessLookupError, PermissionError):
+ pass
+ fallback(sig)
+
+
def _validate_schema_version(instance, _, value) -> str | None:
if value is None:
return None
@@ -1264,25 +1285,13 @@ class WatchedSubprocess:
"""
Deliver ``sig`` to the child process, or to its whole process group
when it has its own.
- When ``new_process_group`` was set at ``start()`` time, the signal is
sent with
- ``os.killpg`` so subprocesses spawned by the child (venv children,
bash shells, etc.)
- are reached too (see issue #65505). Falls back to signalling the child
PID alone when
- the group cannot be resolved or signalled -- and, critically, when the
child still
- shares the supervisor's own process group (``setpgid`` failed),
because ``killpg`` on
- our own group would signal the supervisor itself and its siblings.
+ When ``new_process_group`` was set at ``start()`` time, the signal
reaches the subprocesses
+ the child spawned too (venv children, bash shells, etc.); see issue
#65505.
"""
if self._new_process_group:
- try:
- pgid = os.getpgid(self._process.pid)
- except (ProcessLookupError, PermissionError):
- pgid = None
- if pgid is not None and pgid != os.getpgid(0):
- try:
- os.killpg(pgid, sig)
- return
- except (ProcessLookupError, PermissionError):
- pass
- self._process.send_signal(sig)
+ signal_process_tree(self._process.pid, sig,
fallback=self._process.send_signal)
+ else:
+ self._process.send_signal(sig)
def cleanup_sockets_after_kill(self) -> None:
"""Drain log-bearing sockets, then close every remaining socket after
a forced kill."""
diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
index 63eafc6d1ab..779eac3711c 100644
--- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
+++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py
@@ -524,6 +524,25 @@ class TestConnectionFromProcess:
mock_owned.assert_called_once_with(("127.0.0.1", 5000), ("127.0.0.1",
6000), 999)
+_RUNTIME_PID = 6767
+
+
[email protected]
+def runtime_process_group(monkeypatch):
+ """
+ Resolve ``_RUNTIME_PID`` to its own group and any other pid to the
supervisor's real group.
+
+ ``os.killpg`` is mocked and returned so a test can assert what was
signalled, and the runtime's
+ pid is reported as still in use so the leftover sweep stays out of the way.
+ """
+ own_group = os.getpgid(0)
+ killpg = MagicMock()
+ monkeypatch.setattr(os, "getpgid", lambda pid: _RUNTIME_PID if pid ==
_RUNTIME_PID else own_group)
+ monkeypatch.setattr(os, "killpg", killpg)
+ monkeypatch.setattr(psutil, "pid_exists", lambda pid: True)
+ return killpg
+
+
class TestResourceTracker:
"""
Unit tests for the _ResourceTracker context manager.
@@ -559,11 +578,26 @@ class TestResourceTracker:
tracker.track(sock)
sock.close.assert_called_once()
- def test_context_manager_terminates_tracked_popen_on_exit(self):
+ def test_context_manager_terminates_tracked_popen_on_exit(self,
runtime_process_group):
+ """A failed startup signals the runtime's whole process group, not
just its pid."""
+ proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = _RUNTIME_PID
+ with _ResourceTracker(timeout=0.1) as tracker:
+ tracker.track(proc)
+ runtime_process_group.assert_called_once_with(_RUNTIME_PID,
signal.SIGTERM)
+ proc.send_signal.assert_not_called()
+
+ def test_context_manager_kills_tracked_popen_group_on_timeout(self,
runtime_process_group):
+ """A runtime that ignores SIGTERM during startup has its group
SIGKILLed."""
proc = MagicMock(spec=subprocess.Popen)
+ proc.pid = _RUNTIME_PID
+ proc.wait.side_effect = subprocess.TimeoutExpired(cmd="runtime",
timeout=0.1)
with _ResourceTracker(timeout=0.1) as tracker:
tracker.track(proc)
- proc.terminate.assert_called_once()
+ assert runtime_process_group.call_args_list == [
+ call(_RUNTIME_PID, signal.SIGTERM),
+ call(_RUNTIME_PID, signal.SIGKILL),
+ ]
def test_untracked_socket_not_closed_on_exit(self):
sock = MagicMock(spec=socket.socket)
@@ -1150,6 +1184,138 @@ class TestPopenActivitySubprocessClearXComs:
} 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)
+"""
+
+
+_RUNTIME_SPAWNING_A_CHILD_THEN_DYING = """
+import subprocess
+import sys
+
+child = subprocess.Popen([sys.executable, "-c", "import time;
time.sleep(300)"])
+with open(sys.argv[1], "w") as handle:
+ handle.write(str(child.pid))
+sys.exit(1)
+"""
+
+
+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 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
+ assert proc._new_process_group is True
+
+ def
test_a_runtime_that_dies_before_connecting_does_not_leave_its_child(self,
tmp_path):
+ """The runtime is reaped before start() fails, so its group is swept
by the leader's pid."""
+ pid_file = tmp_path / "child.pid"
+ with pytest.raises(RuntimeError, match="before connecting"):
+ _PopenActivitySubprocess.start(
+ what=_make_ti(),
+ dag_rel_path="bundle",
+ bundle_info=FAKE_BUNDLE,
+ client=MagicMock(),
+ command=[
+ sys.executable,
+ "-c",
+ _RUNTIME_SPAWNING_A_CHILD_THEN_DYING,
+ str(pid_file),
+ ],
+ subprocess_logs_to_stdout=False,
+ )
+
+ spawned_pid = int(pid_file.read_text())
+ try:
+ assert _wait_until_gone(spawned_pid), "the process the runtime
spawned outlived the failure"
+ finally:
+ with contextlib.suppress(OSError):
+ os.kill(spawned_pid, signal.SIGKILL)
+
+ def test_stopping_the_task_also_stops_what_the_runtime_spawned(self,
tmp_path):
+ """Regression for #75533: SIGTERM reaches the CLI the task started,
not only the runtime."""
+ 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,
+ )
+
+ deadline = time.monotonic() + 10.0
+ while not (reported := pid_file.read_text().strip() if
pid_file.exists() else ""):
+ assert time.monotonic() < deadline, "the runtime did not report
the pid it spawned"
+ time.sleep(0.05)
+ spawned_pid = int(reported)
+
+ try:
+ proc.kill(signal.SIGTERM, force=True, escalation_delay=5.0)
+ assert _wait_until_gone(proc.pid), "the runtime outlived the stop"
+ assert _wait_until_gone(spawned_pid), "the process the task
spawned outlived the stop"
+ finally:
+ with contextlib.suppress(OSError):
+ os.kill(spawned_pid, signal.SIGKILL)
+
+
class TestResolveArtifactBundle:
"""Execute-time resolution of the Dag bundle holding the artifacts."""
