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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
