kaxil commented on code in PR #73495:
URL: https://github.com/apache/airflow/pull/73495#discussion_r4073976026
##########
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:
Who reads this partial transcript? The next try of the same task clears
every XCom key on the TI before it starts (`xcom_keys_to_clear` in
`execution_api/routes/task_instances.py`), so a retry never sees it, and a
downstream turn only gets it if it runs after a failed upstream (`all_done` or
`one_failed`). If that salvage case is the intended consumer this is fine, but
the docstring's "for the next turn to resume" reads like the retry case, and
then this push is extra work on the kill path with nothing consuming it.
##########
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:
Does this need to be a function-body import? `hooks/pydantic_ai.py` imports
`pydantic_ai` at module top, and both `agent.py` and `llm.py` import that hook
at the top, so the module is already loaded by the time this runs. Same for the
`RunCancelled` import inside `AgentOperator.execute`, where `agent.py` already
has `from pydantic_ai.capabilities import Toolset` at the top.
##########
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:
Worth updating `docs/agent_security.rst` (line 118) alongside this? It still
says the `sbx` backend leaks orphaned microVMs "if the worker is killed", which
after this change only holds for SIGKILL. The sandbox pages already say "killed
outright", so that line is the outlier.
##########
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:
This checks the type of the token passed to `run_sync` but not that it is
the same object `on_kill` will cancel. A version of `run_agent_sync` that
passes a fresh `CancellationToken()` to `run_sync` while storing a different
one on `self._cancellation_token` passes the whole suite, and that is the
feature silently off (the kill cancels a token nothing listens to). A
`side_effect` that captures `mixin._cancellation_token` at call time and
asserts `held is passed` closes it. The thread-versus-inline premise is also
only tested on a mocked token; a real `Agent(FunctionModel(...))` whose model
awaits a few seconds, with a SIGALRM handler calling `on_kill()` on the main
thread, unwinds with `RunCancelled` in about half a second at HEAD and only
after the model returns with the cancel inlined, so it makes a cheap
deterministic regression test for the claim the design rests on.
##########
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:
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?
##########
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:
Worth stating the Airflow floor here? The Task SDK only calls `on_kill` from
its SIGTERM handler starting with 3.0.4 and 3.1.0 (`_on_term` in
`task_runner.py`). On 3.0.0 to 3.0.3 the handler does not exist, so a kill goes
straight to the old behaviour and this cancellation never fires. That is not a
regression and does not need a `min-airflow-version` bump, but a sentence like
"requires a Task SDK that calls `on_kill` on SIGTERM (Airflow 3.0.4+ / 3.1+)"
would save someone on 3.0.3 a confusing debugging session.
--
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]
