Cancel the agent run when a common.ai LLM or agent task is killed - #73495
Conversation
2af8ec8 to
112a455
Compare
No common.ai operator implemented on_kill, so killing a running task absorbed the SIGTERM grace window. The pydantic-ai run carried on issuing model requests and its toolsets never unwound, leaving a provisioned sandbox running until SIGKILL, where it leaked. Killing a task should stop the work promptly and let it tear down cleanly. The operators now hold the run's CancellationToken and cancel it from on_kill, from a separate thread so the cancel actually interrupts a run driven by run_sync on the signal-handler thread. Requires pydantic-ai 2.26.0 for first-party run cancellation.
112a455 to
6f76fb0
Compare
|
On kill logs |
|
|
||
| self._cancellation_token = CancellationToken() | ||
| try: | ||
| return agent.run_sync(user_prompt, cancellation_token=self._cancellation_token, **run_kwargs) |
There was a problem hiding this comment.
One consequence to consider: RunCancelled is a RuntimeError, so once it propagates out of execute the runner routes it through _handle_current_task_failed and the task's retry policy. With LLMRetryPolicy that means a fresh, uncancellable agent.run_sync classification call inside whatever is left of the 5 s grace window, and a rules policy could classify the cancel as retryable. Sandbox teardown has already finished by then, so the PR's goal holds either way. Is that acceptable, or should the kill map to something the runner treats as terminal without consulting the policy?
There was a problem hiding this comment.
Confirmed: RunCancelled is a RuntimeError, so it hits the final except BaseException in _run_task_and_map_outcome and flows through _handle_current_task_failed into the retry policy, exactly as you describe.
My inclination here, and I'd welcome your view, is to leave it flowing through the policy rather than mapping the kill to a terminal exception. My thinking is that mapping it terminal (AirflowFailException / AirflowTaskTerminated) would skip retries entirely, which I worry would regress the standard "evicted worker retries the task" behaviour, and retry-vs-terminal feels like the retry policy's call rather than something the mixin should force. On the grace-window cost, as far as I can tell, if the classifier outlasts the window it gets SIGKILLed mid-call and retry falls back to the default count (the pre-fix behaviour), so I don't think it regresses there.
If you'd prefer kills to skip the LLM classifier, I think the cleanest place is retry.py (short-circuiting RunCancelled to the fallback) rather than this PR, which doesn't touch retry.py. I'm happy to take that on in a follow-up PR whenever you'd like, and of course happy to go whichever way you think is best.
- Move the CancellationToken import to module scope. pydantic-ai is already imported at module top through PydanticAIHook, so the lazy import saved nothing. - Drop the cancel-path partial-transcript push. A retry clears the task instance's XCom before it starts, so nothing consumes it. A kill now just propagates RunCancelled to fail the task. - Document the Task SDK floor in the mixin docstring: on_kill is called on SIGTERM from Airflow 3.0.4 and 3.1.0 onward, below which a kill falls back to the pre-existing behaviour. - Qualify the sbx leak note in agent_security.rst to SIGKILL. - Assert the test forwards the exact held cancellation token, not just any token.
2 participants
Killing a running common.ai task (LLMOperator, AgentOperator, or an operator deriving
from them) used to absorb the SIGTERM grace window. No operator implemented on_kill, so
the pydantic-ai run kept issuing model requests and its toolsets never unwound, leaving a
provisioned sandbox running until SIGKILL, where it leaked.
The operators now hold the run's pydantic-ai CancellationToken and cancel it from
on_kill, so a killed run stops and tears down cleanly within the grace window instead of
running dead until SIGKILL. The cancel is issued from a short-lived thread, because
on_kill runs in the Task SDK's SIGTERM handler on the same thread that drives run_sync,
where an inline cancel does not interrupt the blocked run (it only takes effect once the
in-flight await returns). The run-cancellation APIs (CancellationToken and RunCancelled)
need pydantic-ai 2.26.0. The provider already requires 2.33.0 or newer, so this change
carries no dependency bump.
Cancellation relies on the Task SDK calling on_kill on SIGTERM, which is Airflow 3.0.4+
and 3.1+. On 3.0.0 to 3.0.3 a kill falls back to the pre-existing behaviour (the run
continues until SIGKILL).
This does not change the SIGKILL case. A task killed with SIGKILL runs no Python, so a
sandbox killed that way still depends on the backend's own backstop.
Test plan
Manual end-to-end (not covered by CI): on a breeze start-airflow stack, trigger a
common.ai agent/LLM task and, while it is running, mark it failed or clear it from the UI.
The task log shows on_kill firing and the run raising RunCancelled within the grace
window. Verified for AgentOperator (live UI) and LLMOperator (RunCancelled ~1ms after
on_kill, versus running dead to SIGKILL without the change).
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 4.8) following the guidelines
🤖 Generated with Claude Code