You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "andreahlert (via GitHub)" <gi...@apache.org> on 2026/02/08 09:51:06 UTC
[PR] Forward termination signals from supervisor to task subprocess [airflow]
andreahlert opened a new pull request, #61627:
URL: https://github.com/apache/airflow/pull/61627
Fixes: #58936
## Summary
When a Kubernetes worker pod receives SIGTERM (e.g. spot interruption, scaling down, rolling update), the signal is delivered to the supervisor process (PID 1 in the container). Previously, the supervisor had no signal handler installed and would exit with default behavior, leaving the task subprocess orphaned without ever calling the operator's `on_kill()` hook. This meant spawned resources (pods, subprocesses, etc.) were never cleaned up.
**Root cause**: The `supervise()` function starts the task subprocess and calls `process.wait()`, but never installs signal handlers for SIGTERM/SIGINT. The task subprocess *does* have a SIGTERM handler (registered in `task_runner.py`) that calls `on_kill()`, but the signal never reaches it because the supervisor process terminates first.
**Fix**: Install SIGTERM/SIGINT signal handlers in `supervise()` that forward the received signal to the task subprocess via `os.kill()`. The child's existing handler then calls `on_kill()` as expected, restoring the Airflow 2 behavior.
**Signal flow after fix**:
1. K8s sends SIGTERM to supervisor (PID 1)
2. Supervisor's new handler forwards SIGTERM to task subprocess
3. Task subprocess's existing `_on_term` handler calls `operator.on_kill()`
4. Operator cleans up resources (pods, subprocesses, etc.)
5. Subprocess exits, supervisor's `wait()` returns normally
## Changes
- **`task-sdk/src/airflow/sdk/execution_time/supervisor.py`**: Added signal forwarding in `supervise()` function. Signal handlers are saved, installed before `process.wait()`, and restored in a `finally` block.
- **`task-sdk/tests/task_sdk/execution_time/test_supervisor.py`**: Added test that verifies SIGTERM forwarding from supervisor to subprocess triggers the operator's `on_kill()` hook.
## Test plan
- [ ] New test `test_on_kill_hook_called_when_supervisor_receives_sigterm` verifies the signal forwarding chain
- [ ] Existing `test_on_kill_hook_called_when_sigkilled` still passes (no regression)
- [ ] Existing signal-related tests (`test_kill_escalation_path`, `test_exit_by_signal`) still pass
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2794543459
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -482,6 +482,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(1000):
+ print(f"Iteration {i}")
+ sleep(1)
Review Comment:
Hadn't noticed how overkill that was, good catch. đź‘˝
I just copied the range(1000) from test_on_kill_hook_called_when_sigkilled next to it without thinking. The subprocess only needs to survive ~4s (2s before SIGTERM + 2s before SIGKILL), so I'll drop it to ~30.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3295434802
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2232,7 +2232,34 @@ def supervise_task(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
Re-flagging from prior review (also same suggestion from @ashb in id 2938993077, not taken): `signal.signal()` returns the previous handler, so the separate `signal.getsignal()` calls are redundant.
```python
prev_sigterm = signal.signal(signal.SIGTERM, _forward_signal)
prev_sigint = signal.signal(signal.SIGINT, _forward_signal)
```
Four lines collapse to two and you avoid the small race between the `getsignal()` lookup and the `signal()` install.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
Re-flagging from prior review (also @ashb in id 2939013171): this test still installs its own `_forward_signal` handler in the test process and SIGTERMs itself. The production handler registered inside `supervise_task()` is never exercised. If the forwarding logic in `supervisor.py` regressed to a no-op, this test would still pass because the test's own handler does the forwarding.
To actually cover the production path, the SIGTERM needs to arrive while `supervise_task()` is the one with `signal.signal(...)` installed. One way: run `supervise_task()` in a subprocess and SIGTERM that subprocess, then assert on its captured stdout. The current shape just re-implements the feature and asserts the re-implementation works.
##########
task-sdk/tests/task_sdk/dags/signal_forward_test.py:
##########
@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import os
+import signal
+import sys
+
+from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.definitions.dag import dag
+
+
+class SignalForwardOperator(BaseOperator):
+ """Operator that sends SIGTERM to its parent (the supervisor) to test signal forwarding."""
+
+ def execute(self, context):
+ # Print sentinel so the test knows execute() is running
+ print("EXECUTE_STARTED", flush=True)
+ # Send SIGTERM to the supervisor (parent process)
+ os.kill(os.getppid(), signal.SIGTERM)
+ # Keep running long enough for the signal to be forwarded back and on_kill() to fire
+ import time
+
+ time.sleep(30)
+
+ def on_kill(self) -> None:
+ print("ON_KILL_CALLED_VIA_SIGNAL_FORWARDING", flush=True)
+ sys.exit(0)
+
+
+@dag()
+def signal_forward_test():
+ SignalForwardOperator(task_id="signal_task")
+
+
+signal_forward_test()
Review Comment:
This new DAG file is orphaned: nothing in this PR imports it, loads it via `DagBag`, or asserts on its `EXECUTE_STARTED` / `ON_KILL_CALLED_VIA_SIGNAL_FORWARDING` sentinels (`grep -r signal_forward_test task-sdk/` returns only this file).
It looks like the intended consumer is the new `test_on_kill_hook_called_when_supervisor_receives_sigterm` test, but that test inlines a `CustomOperator` instead. Either wire this DAG into a test (which would also let you replace the `sleep(2)` synchronization with a `EXECUTE_STARTED` poll) or drop the file before merge -- shipping dead code makes future cleanup harder.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
+
+ # Send SIGTERM to ourselves (the supervisor) from a background thread,
+ # giving the subprocess time to start executing first. Then forcefully
+ # terminate the subprocess so the test does not hang.
+ def send_signals():
+ sleep(2)
Review Comment:
Re-flagging from prior review: the two `sleep(2)` calls are fixed-delay synchronization. On a slow CI runner the first sleep can elapse before the child's `execute()` has actually started printing, and on a fast runner the second sleep can elapse before `on_kill()` flushes. Either way this is a flake source.
It looks like the new `signal_forward_test.py` DAG (with the `EXECUTE_STARTED` sentinel) was intended to address this, but nothing in this PR consumes that DAG (see the separate comment on that file). The actual test still relies on these sleeps. Please either wire the DAG-based approach up, or replace the sleeps here with a poll-until-condition on the child's stdout.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2779759901
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -482,6 +482,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(1000):
+ print(f"Iteration {i}")
+ sleep(1)
Review Comment:
I get why the loop needs to run “long enough” so the subprocess is alive when the signal is delivered, but 1000 iterations at 1s each feels a bit overkill. If for some reason the subprocess doesn’t get terminated as expected, this could run for ~15 minutes and materially stall CI. Is there any reason this can’t be reduced to a smaller number (e.g. 30–60 iterations) while still leaving plenty of headroom for signal delivery?
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2078,7 +2078,33 @@ def supervise(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+ prev_sigint = signal.getsignal(signal.SIGINT)
+
+ def _forward_signal(signum, frame):
+ log.info(
+ "Received signal, forwarding to task subprocess",
+ signal=signal.Signals(signum).name,
+ pid=process.pid,
+ )
+ try:
+ os.kill(process.pid, signum)
+ except ProcessLookupError:
+ pass
Review Comment:
Are you swallowing the exception here because you anticipate that the child worker has been killed before invoking `os.kill`? If that is the case, I would add a comment here explaining that. Like this:
`# Child process may have already exited during shutdown races.`
This is more a nit but silent exception swallowing tends to raise eyebrows for readers who might have partial context.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -482,6 +482,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(1000):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ try:
+ os.kill(proc.pid, signum)
+ except ProcessLookupError:
+ pass
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
Your new implementation handles both SIGTERM and SIGINT, but you appear to be testing only SIGTERM here. Is this because you only anticipate K8s to send SIGTERM? I would suggest explaining that here in a comment so that a casual reader does not assume this is a gap. Not a blocking suggestion.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-3910446088
> I would really like @ashb @kaxil and @amoghrajesh to review this and figure out what is the desired signal handling in this case.
>
> It might have quite profound impact on how our signal handling is implemented in Airlfow.
>
> So far, we assumed that when we run any component except `celery worker` we relied on signal propagation to all processes in process group via `dumb-init` entrypoint process - and all our images had entrypoint (thus iniit process) set as `dumb-init`.
>
> Our images have `dumb-init` set at entry point and default behaviour is that signals received by dumb-init are forwarded to all the processes in the main proces group. This means that if the proces does not create a new process group when it forks it's children, all such created processes will receive the signal.
>
> This was a deliberate decision, because a lot of potential tasks and parsing process libraries might not properly propagate signals - typical example is when we run bash tasks, bash will silently swallow signals and will not propagate it to it's children processes unless specific handlers are configured in the bash script. A lot of people does not know that it happens and they are starting bash subprocesses, not knowing that signals will not be propagated - and since airflow is pretty versatile and you can generally write any code, we decided to send signals to all the precesses in the group (which is defailt dumb-init behaviour).
>
> The only exception is celery worker, because celery worker has a deliberate handling of the signals - the first signal they receive puts the worker in offline mode, where it does not accept any new tasks - and waits for all tasks to complete. This is then used by KEDA autoscaler for example to automatically scale celery workers down. Celery handles propagating of signals to it's subprocesses, including signal escalation in case the signal is not propagated - with timeuts, so even if the bash subprocess swallows the signal, celery will escalate eventually to SIGKILL and kill such tasks (but if tasks are handling signals, it will nicely and cleanly stop the before escalation).
>
> So instead of sending signals to process group, celery is configured to only receive signals by the master worker and it will take care of the propagation - DUMB_INIT_SETSID should be set to 0 for celery pods. And this is how it is implemented in our helm chart as well, and documentation of Airflow Docker image describes this behaviour:
>
> https://airflow.apache.org/docs/docker-stack/entrypoint.html#signal-propagation
>
> I am not sure what was the thought about supervisor's behaviour when it was implemented, whether it handles propagation of signals to child processes or whether it relies on process group receiving the signals, but I guess the latter. It would however be great to document the behaviour and expectations for signal handling in this case - because
>
> I see two assumptions here that are wrong:
>
> 1. The supervisor is assumed to be init process - which I think is pretty wrong. Supervisor in python should never be an init process - unless it implements all that the init process should do. Init process should do more than signal propagation, it should also reap orphan processes by other subprocesses (so called zombies), it should wait for all of the processes it sent signals to before exiting - otherwise it wil not cleanly by the kernel - which might mean filesystem corruption for example. This is why all docker containers (that includes airflow reference image) should have a "real" init process as entrypoint - in our case this is `dumb-init` and it's absolutely not recommended by any production image to have `init` process that is not a `real` init process (some people set bash interpreter as entrypoint and it's absolutely wrong).
> 2. the assumption is that the init process does not forward the signal to the whole process group of forked command (which happens in Airflow reference image) .
>
> I think before we continue with this PR we should at the very least make some assumptions here set and agreed.
Hey @potiuk, thanks for adding here.
Thanks, this makes sense and I agree we should align on assumptions first.
My intent with this PR is only to guarantee `on_kill()` is reached when supervisor receives SIGTERM in Kubernetes worker scenarios. I did not intend to make supervisor an init replacement.
I propose we explicitly document the expected model for this path:
1) `dumb-init` remains responsible for process-group signal behavior
2) supervisor behavior is documented for task subprocess signal forwarding expectations
3) we confirm whether this PR should rely only on process-group propagation, or keep explicit forwarding as a safety mechanism
If you all agree, I can update this PR with docs/tests aligned to the chosen model.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-3905278479
I would really like @ashb @kaxil and @amoghrajesh to review this and figure out what is the desired signal handling in this case.
It might have quite profound impact on how our signal handling is implemented in Airlfow.
So far, we assumed that when we run any component except `celery worker` we relied on signal propagation to all processes in process group via `dumb-init` entrypoint process - and all our images had entrypoint (thus iniit process) set as `dumb-init`.
Our images have `dumb-init` set at entry point and default behaviour is that signals received by dumb-init are forwarded to all the processes in the main proces group. This means that if the proces does not create a new process group when it forks it's children, all such created processes will receive the signal.
This was a deliberate decision, because a lot of potential tasks and parsing process libraries might not properly propagate signals - typical example is when we run bash tasks, bash will silently swallow signals and will not propagate it to it's children processes unless specific handlers are configured in the bash script. A lot of people does not know that it happens and they are starting bash subprocesses, not knowing that signals will not be propagated - and since airflow is pretty versatile and you can generally write any code, we decided to send signals to all the precesses in the group (which is defailt dumb-init behaviour).
The only exception is celery worker, because celery worker has a deliberate handling of the signals - the first signal they receive puts the worker in offline mode, where it does not accept any new tasks - and waits for all tasks to complete. This is then used by KEDA autoscaler for example to automatically scale celery workers down. Celery handles propagating of signals to it's subprocesses, including signal escalation in case the signal is not propagated - with timeuts, so even if the bash subprocess swallows the signal, celery will escalate eventually to SIGKILL and kill such tasks (but if tasks are handling signals, it will nicely and cleanly stop the before escalation).
So instead of sending signals to process group, celery is configured to only receive signals by the master worker and it will take care of the propagation - DUMB_INIT_SETSID should be set to 0 for celery pods. And this is how it is implemented in our helm chart as well, and documentation of Airflow Docker image describes this behaviour:
https://airflow.apache.org/docs/docker-stack/entrypoint.html#signal-propagation
I am not sure what was the thought about supervisor's behaviour when it was implemented, whether it handles propagation of signals to child processes or whether it relies on process group receiving the signals, but I guess the latter. It would however be great to document the behaviour and expectations for signal handling in this case - because
I see two assumptions here that are wrong:
1) The supervisor is assumed to be init process - which I think is pretty wrong. Supervisor in python should never be an init process - unless it implements all that the init process should do. Init process should do more than signal propagation, it should also reap orphan processes by other subprocesses (so called zombies), it should wait for all of the processes it sent signals to before exiting - otherwise it wil not cleanly by the kernel - which might mean filesystem corruption for example. This is why all docker containers (that includes airflow reference image) should have a "real" init process as entrypoint - in our case this is `dumb-init` and it's absolutely not recommended by any production image to have `init` process that is not a `real` init process (some people set bash interpreter as entrypoint and it's absolutely wrong).
2) the assumption is that the init process does not forward the signal to the whole process group of forked command (which happens in Airflow reference image) .
I think before we continue with this PR we should at the very least make some assumptions here set and agreed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3295434806
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
Re-flagging from prior review: this test still installs its own `_forward_signal` handler in the test process and SIGTERMs itself. The production handler registered inside `supervise_task()` is never exercised. If the forwarding logic in `supervisor.py` regressed to a no-op, this test would still pass because the test's own handler does the forwarding.
To actually cover the production path, the SIGTERM needs to arrive while `supervise_task()` is the one with `signal.signal(...)` installed. One way: run `supervise_task()` in a subprocess and SIGTERM that subprocess, then assert on its captured stdout. The current shape just re-implements the feature and asserts the re-implementation works.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-4065982487
@ashb @kaxil @amoghrajesh Could we take this topic back to discussion? Or should we close this PR at least for now and get it back a few months in the future?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-4562340332
Rebased onto current `main` after the Coordinator Layer refactor ([#65958](https://github.com/apache/airflow/pull/65958)) moved `process.wait()` out of `supervise_task()` and into `_PythonCoordinator.execute_task()`. Squashed the branch into a single commit and relocated the signal-forwarding handler from `supervise_task()` to `ActivitySubprocess.wait()`, which is now where the subprocess is waited on and where `self.pid` is available.
Confirmed the underlying bug still reproduces on plain `main` (test hangs and the supervisor process dies on the operator-issued SIGTERM without forwarding it). With the new placement the test passes in ~3s, and the 19 existing signal-related tests in `test_supervisor.py` still pass.
@kaxil PTAL when you have a moment.
---
Drafted-by: Claude Code (Opus 4.7); reviewed by @andreahlert before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-3911250293
I am ok but I think we need to hear about it from @ashb @kaxil @amoghrajesh -> they have more insights on supervisor and friends.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3316321616
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2232,7 +2232,34 @@ def supervise_task(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
Done
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
Done
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
+
+ # Send SIGTERM to ourselves (the supervisor) from a background thread,
+ # giving the subprocess time to start executing first. Then forcefully
+ # terminate the subprocess so the test does not hang.
+ def send_signals():
+ sleep(2)
Review Comment:
Done
##########
task-sdk/tests/task_sdk/dags/signal_forward_test.py:
##########
@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import os
+import signal
+import sys
+
+from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.definitions.dag import dag
+
+
+class SignalForwardOperator(BaseOperator):
+ """Operator that sends SIGTERM to its parent (the supervisor) to test signal forwarding."""
+
+ def execute(self, context):
+ # Print sentinel so the test knows execute() is running
+ print("EXECUTE_STARTED", flush=True)
+ # Send SIGTERM to the supervisor (parent process)
+ os.kill(os.getppid(), signal.SIGTERM)
+ # Keep running long enough for the signal to be forwarded back and on_kill() to fire
+ import time
+
+ time.sleep(30)
+
+ def on_kill(self) -> None:
+ print("ON_KILL_CALLED_VIA_SIGNAL_FORWARDING", flush=True)
+ sys.exit(0)
+
+
+@dag()
+def signal_forward_test():
+ SignalForwardOperator(task_id="signal_task")
+
+
+signal_forward_test()
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3324916181
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2232,7 +2232,34 @@ def supervise_task(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
This still isn't applied at HEAD (`0880eb03`). Lines 1398-1399 keep the two `signal.getsignal()` lookups and then lines 1410-1411 call `signal.signal()` separately. Since `signal.signal()` returns the previous handler, move the `_forward_signal` definition above the install and drop both `getsignal()` lines:
```python
def _forward_signal(signum, frame):
...
prev_sigterm = signal.signal(signal.SIGTERM, _forward_signal)
prev_sigint = signal.signal(signal.SIGINT, _forward_signal)
```
Minor, not blocking -- just flagging since it was marked done.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #61627:
URL: https://github.com/apache/airflow/pull/61627
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3295434802
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2232,7 +2232,34 @@ def supervise_task(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
Re-flagging from prior review: `signal.signal()` returns the previous handler, so the separate `signal.getsignal()` calls are redundant.
```python
prev_sigterm = signal.signal(signal.SIGTERM, _forward_signal)
prev_sigint = signal.signal(signal.SIGINT, _forward_signal)
```
Four lines collapse to two and you avoid the small race between the `getsignal()` lookup and the `signal()` install.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2940846223
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2097,7 +2097,34 @@ def supervise(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
`signal.signal()` returns the previous handler, so you can skip the separate `getsignal()` calls. Something like `prev_sigterm = signal.signal(signal.SIGTERM, _forward_signal)` on line 2119 (and same for SIGINT). That way you save 2 lines and 2 syscalls.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -485,6 +485,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
Review Comment:
This test re-implements the signal forwarding logic instead of exercising the production code path in `supervise()`. If someone breaks the handler registration in `supervise()`, this test still passes. One approach: have the subprocess send SIGTERM to its parent (`os.kill(os.getppid(), signal.SIGTERM)`) while `supervise()` is running, similar to how the existing `test_on_kill_hook_called_when_sigkilled` test works from the child side. That would test the actual wiring.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -485,6 +485,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ try:
+ os.kill(proc.pid, signum)
+ except ProcessLookupError:
+ pass
+
+ signal.signal(signal.SIGTERM, _forward_signal)
+
+ # Send SIGTERM to ourselves (the supervisor) from a background thread,
+ # giving the subprocess time to start executing first. Then forcefully
+ # terminate the subprocess so the test does not hang.
+ def send_signals():
+ sleep(2)
Review Comment:
The `sleep(2)` calls for synchronization are likely to cause flakes on slow CI. If the subprocess hasn't started executing within 2s, or if `on_kill()` takes longer than 2s, the test fails or hangs. A sentinel output line that the parent waits for before sending SIGTERM would be more reliable than a fixed sleep.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-4066012896
Sorry i missed the original ping. I think there's something we can do here
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2794476249
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -482,6 +482,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(1000):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ try:
+ os.kill(proc.pid, signum)
+ except ProcessLookupError:
+ pass
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
yeah, both use the same handler so I only tested the real-world signal. The test name makes the scope explicit. Happy to add a note in the docstring if you think it'd help.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2938993077
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2097,7 +2097,34 @@ def supervise(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+ prev_sigint = signal.getsignal(signal.SIGINT)
+
+ def _forward_signal(signum, frame):
+ log.info(
+ "Received signal, forwarding to task subprocess",
+ signal=signal.Signals(signum).name,
+ pid=process.pid,
+ )
+ try:
+ os.kill(process.pid, signum)
+ except ProcessLookupError:
+ # Child process may have already exited during shutdown races.
+ log.debug("Child already exited, signal not forwarded", pid=process.pid)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
+ signal.signal(signal.SIGINT, _forward_signal)
Review Comment:
```suggestion
prev_sigterm = signal.signal(signal.SIGTERM, _forward_signal)
prev_sigint = signal.signal(signal.SIGINT, _forward_signal)
```
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -485,6 +485,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ try:
+ os.kill(proc.pid, signum)
+ except ProcessLookupError:
+ pass
+
+ signal.signal(signal.SIGTERM, _forward_signal)
+
+ # Send SIGTERM to ourselves (the supervisor) from a background thread,
+ # giving the subprocess time to start executing first. Then forcefully
+ # terminate the subprocess so the test does not hang.
+ def send_signals():
+ sleep(2)
+ os.kill(os.getpid(), signal.SIGTERM)
+ # Give the child time to handle the signal and call on_kill()
+ sleep(2)
+ # Terminate the subprocess so proc.wait() returns
+ try:
+ os.kill(proc.pid, signal.SIGKILL)
+ except ProcessLookupError:
+ pass
+
+ t = threading.Thread(target=send_signals, daemon=True)
+ t.start()
Review Comment:
Nit if we want it: instead of a thread, we could do this kill from within CustomOperator.execute we already have (where it could send the signal to it's parent).
Nothing wrong with this, just it could be a code tidy up
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2097,7 +2097,34 @@ def supervise(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+ prev_sigint = signal.getsignal(signal.SIGINT)
+
Review Comment:
```suggestion
```
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -485,6 +485,102 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ try:
+ os.kill(proc.pid, signum)
+ except ProcessLookupError:
+ pass
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
This feels like we are testing this, not the main/runtime code.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r2794373113
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2078,7 +2078,33 @@ def supervise(
sentry_integration=sentry_integration,
)
- exit_code = process.wait()
+ # Forward termination signals to the task subprocess so that the operator's
+ # on_kill() hook is invoked on graceful shutdown (e.g. K8s pod SIGTERM).
+ # Without this, the supervisor exits on SIGTERM without notifying the child,
+ # leaving spawned resources (pods, subprocesses, etc.) running.
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+ prev_sigint = signal.getsignal(signal.SIGINT)
+
+ def _forward_signal(signum, frame):
+ log.info(
+ "Received signal, forwarding to task subprocess",
+ signal=signal.Signals(signum).name,
+ pid=process.pid,
+ )
+ try:
+ os.kill(process.pid, signum)
+ except ProcessLookupError:
+ pass
Review Comment:
LGTM. I'll add the comment and also throw in a log.debug there, same pattern as the psutil.NoSuchProcess handling in kill() a few lines above. to keeps the shutdown flow observable.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-4530177213
@ashb — You endorsed the approach on this PR back in March ("strict improvement on behavior... just being a Good (Process) Parent"). The PR has been sitting at REVIEW_REQUIRED since then with green CI and no unresolved threads. Could you convert your endorsement to a formal APPROVE so it can land? Or @kaxil if you want to chime in.
---
_Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes._
---
Drafted-by: Claude Code (Opus 4.7); reviewed by @potiuk before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3307553393
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,61 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ test_dags_dir,
+ captured_logs,
+ client_with_ti_start,
+ ):
+ """SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ Exercises the real ``supervise_task()`` code path: the DAG operator sends SIGTERM
+ to its parent (the supervisor, which here is this test process running
+ ``supervise_task``). The handler installed inside ``supervise_task`` must forward
+ the signal to the task subprocess so the operator's ``on_kill()`` hook runs.
+ """
+ import threading
Review Comment:
Nit: hoist `import threading` to the module-top stdlib imports (around line 28) rather than importing it inline inside the test method.
##########
task-sdk/tests/task_sdk/dags/signal_forward_test.py:
##########
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import os
+import signal
+import time
+
+from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.definitions.dag import dag
+
+
+class SignalForwardOperator(BaseOperator):
+ """Send SIGTERM to the supervisor parent process to exercise signal forwarding."""
+
+ def execute(self, context):
+ print("EXECUTE_STARTED", flush=True)
+ os.kill(os.getppid(), signal.SIGTERM)
+ time.sleep(30)
Review Comment:
`time.sleep(30)` makes this test take roughly the full supervisor watchdog window every run. Per PEP 475, `time.sleep()` resumes after a signal handler returns, so once `on_kill()` prints its sentinel the call still blocks for the remainder of the 30s and `execute()` only exits when the watchdog `SIGKILL`s the child.
Drop this to something short, e.g. `time.sleep(2)`, so the test wall time matches its real work. The SIGTERM-induced `KeyboardInterrupt` path will exit `execute()` promptly once the supervisor forwards the signal back.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,61 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
Review Comment:
Nit: this test drives `supervise_task()` end-to-end (real DAG, real `parsing_context_manager`, real signal handler), so it organizationally fits `TestSupervisor` (line 187) better than `TestWatchedSubprocess` where it currently lives. Not blocking, but worth grouping with the other `supervise_task` 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on PR #61627:
URL: https://github.com/apache/airflow/pull/61627#issuecomment-4546695209
@kaxil sorry for missing the original points on the first pass. Thanks for gently reflag. PTAL when you have a moment.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Forward termination signals from supervisor to task subprocess [airflow]
Posted by "andreahlert (via GitHub)" <gi...@apache.org>.
andreahlert commented on code in PR #61627:
URL: https://github.com/apache/airflow/pull/61627#discussion_r3316321877
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -488,6 +488,98 @@ def on_kill(self) -> None:
captured = capfd.readouterr()
assert "On kill hook called!" in captured.out
+ def test_on_kill_hook_called_when_supervisor_receives_sigterm(
+ self,
+ client_with_ti_start,
+ mocked_parse,
+ make_ti_context,
+ mock_supervisor_comms,
+ create_runtime_ti,
+ make_ti_context_dict,
+ capfd,
+ ):
+ """Test that SIGTERM to the supervisor process is forwarded to the task subprocess.
+
+ This simulates what happens when Kubernetes sends SIGTERM to the worker pod:
+ the supervisor should forward the signal to the child process so that the
+ operator's on_kill() hook is triggered for resource cleanup.
+ """
+ import threading
+
+ ti_id = "4d828a62-a417-4936-a7a6-2b3fabacecab"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ return httpx.Response(200, json=make_ti_context_dict())
+ return httpx.Response(status_code=204)
+
+ def subprocess_main():
+ CommsDecoder()._get_response()
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ for i in range(30):
+ print(f"Iteration {i}")
+ sleep(1)
+
+ def on_kill(self) -> None:
+ print("On kill hook called via signal forwarding!")
+
+ task = CustomOperator(task_id="test-signal-forward")
+ runtime_ti = create_runtime_ti(
+ dag_id="c",
+ task=task,
+ conf={},
+ )
+ run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=ti_id,
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=make_client(transport=httpx.MockTransport(handle_request)),
+ target=subprocess_main,
+ )
+
+ # Install signal forwarding handler (same mechanism as supervise() does)
+ prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+ def _forward_signal(signum, frame):
+ with suppress(ProcessLookupError):
+ os.kill(proc.pid, signum)
+
+ signal.signal(signal.SIGTERM, _forward_signal)
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org