pankajkoti commented on code in PR #73495: URL: https://github.com/apache/airflow/pull/73495#discussion_r4075208208
########## providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py: ########## @@ -0,0 +1,66 @@ +# 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. +"""Mixin that cancels an operator's in-flight pydantic-ai run when the task is killed.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pydantic_ai import Agent, AgentRunResult, CancellationToken + + +class CancellableAgentRunMixin: + """ + Run a pydantic-ai agent synchronously with kill-time cancellation wired in. + + The wrapper holds the in-flight run's ``CancellationToken`` so :meth:`on_kill` can + cancel it. Cancelling makes ``run_sync`` raise ``RunCancelled`` and unwind, giving the + agent's toolsets a chance to exit (tearing down a provisioned sandbox, for one) before + SIGKILL rather than leaving the run to die mid-flight. + """ + + # Set only while a run is in flight. Read by on_kill from the signal handler. + _cancellation_token: CancellationToken | None = None + + # Provided by BaseOperator at runtime. Declared here for the type checker. + log: Any + + def run_agent_sync( + self, agent: Agent[Any, Any], user_prompt: Any, **run_kwargs: Any + ) -> AgentRunResult[Any]: + """Call ``agent.run_sync`` under a fresh cancellation token held for :meth:`on_kill`.""" + from pydantic_ai import CancellationToken Review Comment: Good catch, moved `CancellationToken` to a module-level import. The hook already imports pydantic_ai at module top and both operators import the hook at top, so the run was loading it regardless. The `RunCancelled` function-body import is gone too, as part of dropping the cancel-path transcript (see the reply below). ########## providers/common/ai/src/airflow/providers/common/ai/operators/agent.py: ########## @@ -555,19 +559,29 @@ def execute(self, context: Context) -> Any: storage = self._durable_storage counter = self._durable_counter - if self.durable and storage is not None and counter is not None: - from pydantic_ai.models import infer_model - - from airflow.providers.common.ai.durable.caching_model import CachingModel + # A killed run raises RunCancelled (see run_agent_sync). Emit the partial + # transcript for a message_history session before re-raising. The durable + # cache cleanup below is skipped on the raise, preserving it for the retry. + from pydantic_ai import RunCancelled - if agent.model is None: - raise ValueError("Agent model must be set when durable=True") - resolved_model = infer_model(agent.model) - caching_model = CachingModel(resolved_model, storage=storage, counter=counter) - with agent.override(model=caching_model): - result = agent.run_sync(self.prompt, **run_kwargs) - else: - result = agent.run_sync(self.prompt, **run_kwargs) + try: + if self.durable and storage is not None and counter is not None: + from pydantic_ai.models import infer_model + + from airflow.providers.common.ai.durable.caching_model import CachingModel + + if agent.model is None: + raise ValueError("Agent model must be set when durable=True") + resolved_model = infer_model(agent.model) + caching_model = CachingModel(resolved_model, storage=storage, counter=counter) + with agent.override(model=caching_model): + result = self.run_agent_sync(agent, self.prompt, **run_kwargs) + else: + result = self.run_agent_sync(agent, self.prompt, **run_kwargs) + except RunCancelled as cancelled: + if self.message_history is not None: + self._emit_message_history(context, cancelled) Review Comment: You're right: a retry clears the TI's XCom so it can never resume from this, and the only consumer would be a downstream task with a failure-tolerant trigger rule, which isn't a wired-up flow. The "for the next turn to resume" wording was inherited from the normal-path docstring and doesn't apply on cancel. Rather than keep a push nothing consumes, I dropped the cancel-path emit. A kill now just propagates RunCancelled to fail the task, and the operator test asserts nothing is pushed. ########## providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py: ########## @@ -0,0 +1,66 @@ +# 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. +"""Mixin that cancels an operator's in-flight pydantic-ai run when the task is killed.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pydantic_ai import Agent, AgentRunResult, CancellationToken + + +class CancellableAgentRunMixin: + """ + Run a pydantic-ai agent synchronously with kill-time cancellation wired in. + + The wrapper holds the in-flight run's ``CancellationToken`` so :meth:`on_kill` can + cancel it. Cancelling makes ``run_sync`` raise ``RunCancelled`` and unwind, giving the + agent's toolsets a chance to exit (tearing down a provisioned sandbox, for one) before Review Comment: Updated agent_security.rst so the sbx leak line reads "killed outright (SIGKILL)", matching the phrasing the sandbox pages already use. ########## providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py: ########## @@ -0,0 +1,66 @@ +# 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. +"""Mixin that cancels an operator's in-flight pydantic-ai run when the task is killed.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pydantic_ai import Agent, AgentRunResult, CancellationToken + + +class CancellableAgentRunMixin: + """ + Run a pydantic-ai agent synchronously with kill-time cancellation wired in. + + The wrapper holds the in-flight run's ``CancellationToken`` so :meth:`on_kill` can + cancel it. Cancelling makes ``run_sync`` raise ``RunCancelled`` and unwind, giving the + agent's toolsets a chance to exit (tearing down a provisioned sandbox, for one) before + SIGKILL rather than leaving the run to die mid-flight. Review Comment: Added a docstring note. I checked the boundary against the tags: the SIGTERM on_kill handler is absent in 3.0.2/3.0.3 and present from 3.0.4 (and all 3.1+), so it now says this needs Airflow 3.0.4+ / 3.1+ and falls back to the pre-existing behaviour below that. ########## providers/common/ai/src/airflow/providers/common/ai/mixins/cancellable_run.py: ########## @@ -0,0 +1,66 @@ +# 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. +"""Mixin that cancels an operator's in-flight pydantic-ai run when the task is killed.""" + +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pydantic_ai import Agent, AgentRunResult, CancellationToken + + +class CancellableAgentRunMixin: + """ + Run a pydantic-ai agent synchronously with kill-time cancellation wired in. + + The wrapper holds the in-flight run's ``CancellationToken`` so :meth:`on_kill` can + cancel it. Cancelling makes ``run_sync`` raise ``RunCancelled`` and unwind, giving the + agent's toolsets a chance to exit (tearing down a provisioned sandbox, for one) before + SIGKILL rather than leaving the run to die mid-flight. + """ + + # Set only while a run is in flight. Read by on_kill from the signal handler. + _cancellation_token: CancellationToken | None = None + + # Provided by BaseOperator at runtime. Declared here for the type checker. + log: Any + + def run_agent_sync( + self, agent: Agent[Any, Any], user_prompt: Any, **run_kwargs: Any + ) -> AgentRunResult[Any]: + """Call ``agent.run_sync`` under a fresh cancellation token held for :meth:`on_kill`.""" + from pydantic_ai import CancellationToken + + self._cancellation_token = CancellationToken() + try: + return agent.run_sync(user_prompt, cancellation_token=self._cancellation_token, **run_kwargs) Review 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. ########## providers/common/ai/tests/unit/common/ai/mixins/test_cancellable_run.py: ########## @@ -0,0 +1,89 @@ +# 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 threading +import time +from unittest.mock import ANY, MagicMock + +import pytest +from pydantic_ai import CancellationToken + +from airflow.providers.common.ai.mixins.cancellable_run import CancellableAgentRunMixin +from airflow.providers.common.ai.operators.agent import AgentOperator +from airflow.providers.common.ai.operators.llm import LLMOperator + + +class TestRunAgentSync: + def test_forwards_cancellation_token_and_clears_after_success(self): + mixin = CancellableAgentRunMixin() + agent = MagicMock(spec=["run_sync"]) + + result = mixin.run_agent_sync(agent, "prompt", usage_limits=None) + + assert result is agent.run_sync.return_value + agent.run_sync.assert_called_once_with("prompt", cancellation_token=ANY, usage_limits=None) + assert isinstance(agent.run_sync.call_args.kwargs["cancellation_token"], CancellationToken) Review Comment: Good point. The test now captures the token the mixin holds during the run_sync call and asserts run_sync received that exact object, not just any CancellationToken. I mutation-checked it (make run_agent_sync pass a fresh token while storing a different one, and it goes red). I left out the FunctionModel + SIGALRM real-run test: the off-thread requirement is already covered by test_cancels_active_token_off_the_calling_thread (mutation-verified, inlining the cancel turns it red) plus the live UI-kill, so a second timing-based test would mostly add flakiness risk. -- 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]
