Skip to content

Commit 9f41fc8

Browse files
committed
Warm-shutdown supervisor on SIGTERM instead of killing the running task
PR #61627 installed a SIGTERM/SIGINT handler in `ActivitySubprocess.wait()` that forwards the signal to the task subprocess via `os.kill(self.pid, ...)`, so the operator's `on_kill()` runs when the supervisor is signalled. On KubernetesExecutor this kills the task on every routine pod shutdown. `dumb-init` runs as PID 1 with DUMB_INIT_SETSID=1 (the image default), so it broadcasts SIGTERM to the whole process group — the task subprocess already receives it directly. After #61627 the supervisor forwards it a second time, so the task gets SIGTERM twice; either delivery makes the child handler call `on_kill()` and abort the task the moment the pod starts terminating. That is wrong for a graceful pod shutdown (spot interruption, scale-down, rolling update). For operators that release external resources in `on_kill()` (Databricks tearing down a job cluster, ADF cancelling a pipeline, etc.) this destroys in-flight work we want to let finish. This change makes the task warm-shut-down instead, in two parts: 1. supervisor: replace the forward-and-kill handler with a *warm-shutdown* handler that does NOT signal the task. It logs the received signal and keeps the supervisor in its monitoring loop so the running task continues to completion. The handler is installed in `_PythonCoordinator.execute_task` around both `start()` (the RUNNING transition) and `wait()` (run + terminal-state report + log upload), so no window is left where the default SIGTERM disposition could tear the supervisor down. Previous dispositions are restored on exit so a reused Celery prefork worker doesn't leak the handler. 2. KubernetesExecutor pod template: set DUMB_INIT_SETSID=0, so dumb-init delivers pod-termination signals only to its direct child (the supervisor), not to the whole process group. Without this the task subprocess would still receive the group-broadcast SIGTERM directly and call `on_kill()`, defeating the warm shutdown. This mirrors what Celery worker pods already do. `on_kill()` is still reached through the legitimate kill paths (heartbeat-failure / server-terminated kills, overtime termination, or an explicit success/failure transition) — all of which already drive `self.kill(...)` and are unchanged.
1 parent 473e332 commit 9f41fc8

9 files changed

Lines changed: 231 additions & 125 deletions

File tree

‎chart/files/pod-template-file.kubernetes-helm-yaml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,16 @@ spec:
106106
env:
107107
- name: AIRFLOW__CORE__EXECUTOR
108108
value: {{ .Values.executor | quote }}
109+
# Deliver pod-termination signals only to the task supervisor (dumb-init's
110+
# direct child) instead of broadcasting them to the whole process group.
111+
# On graceful pod shutdown the supervisor's warm-shutdown handler then lets
112+
# the running task finish -- the same mechanism Celery workers use -- rather
113+
# than the task subprocess being killed directly by dumb-init's group-wide
114+
# SIGTERM (which could also reach the subprocess before it installs its own
115+
# signal handler). Hard kills (heartbeat loss / overtime / the post-grace
116+
# SIGKILL) are unaffected.
117+
- name: DUMB_INIT_SETSID
118+
value: "0"
109119
{{- if or .Values.workers.kubernetes.kerberosSidecar.enabled .Values.workers.kubernetes.kerberosInitContainer.enabled }}
110120
- name: KRB5_CONFIG
111121
value: {{ .Values.kerberos.configPath | quote }}
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Default ``DUMB_INIT_SETSID`` changed to ``"0"`` for KubernetesExecutor task pods.
2+
3+
Pod-termination signals (e.g. SIGTERM on graceful shutdown) are now delivered only to the
4+
task supervisor (``dumb-init``'s direct child) instead of being broadcast to the whole
5+
process group. This lets a running task finish via the supervisor's warm-shutdown handler --
6+
the same behaviour Celery worker pods already had -- rather than the task subprocess being
7+
killed directly. Hard kills (heartbeat loss / overtime / the post-grace SIGKILL) are unaffected.

‎chart/tests/helm_tests/airflow_aux/test_pod_template_file.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1111,6 +1111,17 @@ def test_should_add_extraEnvs(self):
11111111
"valueFrom": {"configMapKeyRef": {"name": "my-config-map", "key": "my-key"}},
11121112
} in jmespath.search("spec.containers[0].env", docs[0])
11131113

1114+
def test_should_set_dumb_init_setsid_for_warm_shutdown(self):
1115+
"""Pod-termination signals must reach only the supervisor so a running task can warm-shut-down."""
1116+
docs = render_chart(
1117+
show_only=["templates/pod-template-file.yaml"],
1118+
chart_dir=self.temp_chart_dir,
1119+
)
1120+
1121+
assert {"name": "DUMB_INIT_SETSID", "value": "0"} in jmespath.search(
1122+
"spec.containers[0].env", docs[0]
1123+
)
1124+
11141125
def test_should_add_component_specific_labels(self):
11151126
docs = render_chart(
11161127
values={

‎docker-stack-docs/entrypoint.rst‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,15 @@ The table below summarizes ``DUMB_INIT_SETSID`` possible values and their use ca
300300
| | If you are running it through ``["bash", "-c"]`` command, |
301301
| | you need to start the worker via ``exec airflow celery worker`` |
302302
| | as the last command executed. |
303+
| | |
304+
| | The same applies to KubernetesExecutor task pods. Here ``dumb-init``|
305+
| | runs as the init process and its direct child is the task |
306+
| | *supervisor*, which supervises a single task subprocess. Setting the |
307+
| | variable to 0 propagates a graceful SIGTERM only to the supervisor, |
308+
| | which then performs a warm shutdown and waits for the running task |
309+
| | to finish, instead of the signal being broadcast to the whole |
310+
| | process group and killing the task subprocess directly. The Airflow |
311+
| | Helm chart sets this on the KubernetesExecutor pod template for you. |
303312
+----------------+----------------------------------------------------------------------+
304313

305314
Additional quick test options

‎task-sdk/src/airflow/sdk/execution_time/coordinator.py‎

Lines changed: 56 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040

4141
import contextlib
4242
import functools
43+
import os
44+
import signal
4345
from typing import TYPE_CHECKING, Any
4446

4547
import attrs
@@ -50,7 +52,7 @@
5052
from airflow.sdk.configuration import conf
5153

5254
if TYPE_CHECKING:
53-
from collections.abc import Mapping
55+
from collections.abc import Generator, Mapping
5456
from os import PathLike
5557

5658
from structlog.typing import FilteringBoundLogger
@@ -114,6 +116,43 @@ class _CoordinatorSpec(pydantic.BaseModel):
114116
extra: dict[str, Any] | None = None
115117

116118

119+
@contextlib.contextmanager
120+
def _warm_shutdown_signals() -> Generator[None, None, None]:
121+
"""
122+
Install SIGTERM/SIGINT warm-shutdown handlers for the duration of task supervision.
123+
124+
While supervising a task the supervisor must not be torn down by a
125+
termination signal; instead it keeps running so the task can finish (or be
126+
shut down gracefully) and its terminal state and logs are reported. The
127+
handlers are installed around BOTH ``start()`` (which transitions the TI to
128+
RUNNING) and ``wait()`` (which runs the task and then reports the terminal
129+
state / uploads logs), so there is no window where Python's default SIGTERM
130+
disposition could kill the supervisor and tear the just-started task down
131+
with it.
132+
133+
The previous dispositions are restored on exit so a long-lived supervisor
134+
process (e.g. a reused Celery prefork worker) does not leak the handler into
135+
later tasks or clobber the worker's own signal handling.
136+
"""
137+
138+
def _warm_shutdown(signum, frame):
139+
log.info(
140+
"Received signal; warm shutdown in progress, waiting for the running task to complete.",
141+
signal=signal.Signals(signum).name,
142+
pid=os.getpid(),
143+
)
144+
145+
prev_sigterm = signal.getsignal(signal.SIGTERM)
146+
prev_sigint = signal.getsignal(signal.SIGINT)
147+
signal.signal(signal.SIGTERM, _warm_shutdown)
148+
signal.signal(signal.SIGINT, _warm_shutdown)
149+
try:
150+
yield
151+
finally:
152+
signal.signal(signal.SIGTERM, prev_sigterm)
153+
signal.signal(signal.SIGINT, prev_sigint)
154+
155+
117156
class _PythonCoordinator(BaseCoordinator):
118157
"""
119158
Coordinator implementation to execute Python tasks.
@@ -140,17 +179,22 @@ def execute_task(
140179
# process handling.
141180
from airflow.sdk.execution_time.supervisor import ActivitySubprocess
142181

143-
process = ActivitySubprocess.start(
144-
dag_rel_path=dag_rel_path,
145-
what=what,
146-
client=client,
147-
logger=logger,
148-
bundle_info=bundle_info,
149-
subprocess_logs_to_stdout=subprocess_logs_to_stdout,
150-
sentry_integration=sentry_integration,
151-
)
152-
exit_code = process.wait()
153-
return self.ExecutionResult(exit_code, process.final_state)
182+
# Keep the warm-shutdown handlers installed across both start() (which
183+
# transitions the TI to RUNNING) and wait() (which runs the task and
184+
# reports its terminal state / uploads logs) so a SIGTERM at any point
185+
# in this window can't kill the supervisor and tear the task down.
186+
with _warm_shutdown_signals():
187+
process = ActivitySubprocess.start(
188+
dag_rel_path=dag_rel_path,
189+
what=what,
190+
client=client,
191+
logger=logger,
192+
bundle_info=bundle_info,
193+
subprocess_logs_to_stdout=subprocess_logs_to_stdout,
194+
sentry_integration=sentry_integration,
195+
)
196+
exit_code = process.wait()
197+
return self.ExecutionResult(exit_code, process.final_state)
154198

155199

156200
@functools.cache

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

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1392,31 +1392,10 @@ def wait(self) -> int:
13921392
if self._exit_code is not None:
13931393
return self._exit_code
13941394

1395-
# Forward termination signals to the task subprocess so the operator's
1396-
# on_kill() hook runs on graceful shutdown (e.g. K8s pod SIGTERM).
1397-
# Without this the supervisor exits on SIGTERM without notifying the
1398-
# child, leaving spawned resources (pods, subprocesses, etc.) running.
1399-
prev_sigterm = signal.getsignal(signal.SIGTERM)
1400-
prev_sigint = signal.getsignal(signal.SIGINT)
1401-
1402-
def _forward_signal(signum, frame):
1403-
log.info(
1404-
"Received signal, forwarding to task subprocess",
1405-
signal=signal.Signals(signum).name,
1406-
pid=self.pid,
1407-
)
1408-
with suppress(ProcessLookupError):
1409-
os.kill(self.pid, signum)
1410-
1411-
signal.signal(signal.SIGTERM, _forward_signal)
1412-
signal.signal(signal.SIGINT, _forward_signal)
1413-
14141395
try:
14151396
self._monitor_subprocess()
14161397
finally:
14171398
self.selector.close()
1418-
signal.signal(signal.SIGTERM, prev_sigterm)
1419-
signal.signal(signal.SIGINT, prev_sigint)
14201399

14211400
# self._monitor_subprocess() will set the exit code when the process has finished
14221401
# If it hasn't, assume it's failed

‎task-sdk/tests/task_sdk/dags/signal_forward_test.py‎

Lines changed: 0 additions & 44 deletions
This file was deleted.

‎task-sdk/tests/task_sdk/execution_time/test_coordinator.py‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
from __future__ import annotations
2020

2121
import json
22+
import os
23+
import signal
24+
from unittest import mock
2225

2326
import pytest
2427

@@ -28,6 +31,7 @@
2831
BaseCoordinator,
2932
CoordinatorManager,
3033
_PythonCoordinator,
34+
_warm_shutdown_signals,
3135
get_coordinator_manager,
3236
reset_coordinator_manager,
3337
)
@@ -229,3 +233,136 @@ def test_every_example_coordinator_constructs(self, sdk_config):
229233
for queue, key in queue_to_coordinator.items():
230234
coordinator = manager.for_queue(queue)
231235
assert isinstance(coordinator, import_string(specs[key]["classpath"]))
236+
237+
238+
class TestWarmShutdownSignals:
239+
"""Tests for the warm-shutdown signal handling that wraps task supervision."""
240+
241+
@pytest.fixture(autouse=True)
242+
def _restore_disposition(self):
243+
"""Guarantee SIGTERM/SIGINT dispositions are restored even if a test leaks one."""
244+
original_term = signal.getsignal(signal.SIGTERM)
245+
original_int = signal.getsignal(signal.SIGINT)
246+
yield
247+
signal.signal(signal.SIGTERM, original_term)
248+
signal.signal(signal.SIGINT, original_int)
249+
250+
def test_installs_handlers_inside_context(self):
251+
"""While the context is active a warm-shutdown handler is installed for both signals."""
252+
sentinel_term = signal.getsignal(signal.SIGTERM)
253+
sentinel_int = signal.getsignal(signal.SIGINT)
254+
255+
with _warm_shutdown_signals():
256+
inside_term = signal.getsignal(signal.SIGTERM)
257+
inside_int = signal.getsignal(signal.SIGINT)
258+
259+
assert callable(inside_term)
260+
assert callable(inside_int)
261+
# The installed handler is the warm-shutdown closure, not the previous disposition.
262+
assert inside_term is not sentinel_term
263+
assert inside_int is not sentinel_int
264+
# Both signals share the same warm-shutdown closure.
265+
assert inside_term is inside_int
266+
267+
def test_restores_previous_dispositions_on_exit(self):
268+
"""The exact previous dispositions are restored when the context exits normally."""
269+
270+
def _prev_term(signum, frame): # pragma: no cover - never invoked
271+
pass
272+
273+
def _prev_int(signum, frame): # pragma: no cover - never invoked
274+
pass
275+
276+
signal.signal(signal.SIGTERM, _prev_term)
277+
signal.signal(signal.SIGINT, _prev_int)
278+
279+
with _warm_shutdown_signals():
280+
pass
281+
282+
assert signal.getsignal(signal.SIGTERM) is _prev_term
283+
assert signal.getsignal(signal.SIGINT) is _prev_int
284+
285+
def test_restores_previous_dispositions_on_exception(self):
286+
"""Dispositions are restored even if the wrapped body raises."""
287+
288+
def _prev_term(signum, frame): # pragma: no cover - never invoked
289+
pass
290+
291+
signal.signal(signal.SIGTERM, _prev_term)
292+
293+
with pytest.raises(RuntimeError, match="boom"), _warm_shutdown_signals():
294+
raise RuntimeError("boom")
295+
296+
assert signal.getsignal(signal.SIGTERM) is _prev_term
297+
298+
def test_sigterm_inside_context_does_not_kill(self):
299+
"""
300+
A SIGTERM delivered while supervising must be swallowed, not kill the process.
301+
302+
This is the regression guard: with the default SIGTERM disposition (SIG_DFL)
303+
in place as the *previous* handler, sending SIGTERM to ourselves would
304+
terminate the process. The warm-shutdown handler installed by the context
305+
manager must absorb it so the running task is allowed to finish.
306+
"""
307+
# Make the pre-context disposition the default so a missing warm-shutdown
308+
# handler would actually kill this process (and fail the test by dying).
309+
signal.signal(signal.SIGTERM, signal.SIG_DFL)
310+
311+
reached_after_signal = False
312+
with _warm_shutdown_signals():
313+
os.kill(os.getpid(), signal.SIGTERM)
314+
# If the handler did not absorb the signal, we never get here.
315+
reached_after_signal = True
316+
317+
assert reached_after_signal
318+
# And the default disposition is put back afterwards.
319+
assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL
320+
321+
322+
class TestPythonCoordinatorWarmShutdown:
323+
"""The Python coordinator must wrap start() and wait() in the warm-shutdown handlers."""
324+
325+
def test_execute_task_wraps_start_and_wait(self, monkeypatch):
326+
"""
327+
Handlers are installed for the whole start()+wait() window and restored after.
328+
329+
Capturing the SIGTERM disposition at the moment ``start()`` and ``wait()``
330+
run proves the handler spans the RUNNING transition (start) and the
331+
terminal-state report / log upload (wait), with no window left uncovered.
332+
"""
333+
original_term = signal.getsignal(signal.SIGTERM)
334+
captured: dict[str, object] = {}
335+
336+
class _FakeProcess:
337+
final_state = "success"
338+
339+
def wait(self_inner):
340+
captured["wait"] = signal.getsignal(signal.SIGTERM)
341+
return 0
342+
343+
def _fake_start(*args, **kwargs):
344+
captured["start"] = signal.getsignal(signal.SIGTERM)
345+
return _FakeProcess()
346+
347+
import airflow.sdk.execution_time.supervisor as supervisor_mod
348+
349+
monkeypatch.setattr(supervisor_mod.ActivitySubprocess, "start", staticmethod(_fake_start))
350+
351+
coordinator = _PythonCoordinator()
352+
result = coordinator.execute_task(
353+
what=mock.MagicMock(),
354+
dag_rel_path="some_dag.py",
355+
bundle_info=mock.MagicMock(),
356+
client=mock.MagicMock(),
357+
subprocess_logs_to_stdout=False,
358+
)
359+
360+
assert result.exit_code == 0
361+
assert result.final_state == "success"
362+
# During both start() and wait() a warm-shutdown handler was installed...
363+
assert callable(captured["start"])
364+
assert callable(captured["wait"])
365+
assert captured["start"] is not original_term
366+
assert captured["start"] is captured["wait"]
367+
# ...and the original disposition is restored once execute_task returns.
368+
assert signal.getsignal(signal.SIGTERM) is original_term

0 commit comments

Comments
 (0)