kaxil commented on code in PR #74201:
URL: https://github.com/apache/airflow/pull/74201#discussion_r4556546177
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -4351,3 +4362,342 @@ def
test_a_toolset_that_masks_its_own_output_is_not_wrapped_again(self):
op._build_agent()
assert hook.create_agent.call_args.kwargs["toolsets"] == [sql]
+
+
+def _refused_on_the_event_loop(call):
+ """Wrap a blocking call so it fails the test when it runs on the event
loop."""
+
+ def guarded(*args, **kwargs):
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return call(*args, **kwargs)
+ raise AssertionError("blocking call to the supervisor on the event
loop")
+
+ return guarded
+
+
+def _make_async_task_state_store_accessor(*, async_methods: bool = True):
+ """
+ A task state store backed by ``accessor.data``, whose blocking methods
fail on the event loop.
+
+ Every write and delete, async or blocking, is recorded in
``accessor.ops``. With
+ ``async_methods=False`` the store has no async methods, as before Airflow
3.4, and the operator
+ has to call the blocking ones from a worker thread.
+ """
+ from airflow.sdk.execution_time.context import TaskStateStoreAccessor
+
+ accessor = MagicMock(spec=TaskStateStoreAccessor)
+ data: dict[str, Any] = {}
+ ops: list[tuple[str, str]] = []
+ accessor.data = data
+ accessor.ops = ops
+
+ def get(key, default=None):
+ return data.get(key, default)
+
+ def set_(key, value, retention=None):
+ ops.append(("set", key))
+ data[key] = value
+
+ def delete(key):
+ ops.append(("delete", key))
+ data.pop(key, None)
+
+ accessor.get.side_effect = _refused_on_the_event_loop(get)
+ accessor.set.side_effect = _refused_on_the_event_loop(set_)
+ accessor.delete.side_effect = _refused_on_the_event_loop(delete)
+ if async_methods:
+ accessor.aget = AsyncMock(side_effect=get)
+ accessor.aset = AsyncMock(side_effect=set_)
+ accessor.adelete = AsyncMock(side_effect=delete)
+ else:
+ del accessor.aget, accessor.aset, accessor.adelete
+ return accessor
+
+
+def _make_async_context(task_state_store=None):
+ """A context for ``aexecute``: XComs go through ``axcom_push`` on Airflow
3.3+, ``xcom_push`` on 3.2."""
+ ti = _make_ti()
+ ti.axcom_push = AsyncMock()
+ context = {"task_instance": ti}
+ if task_state_store is not None:
+ context["task_state_store"] = task_state_store
+ return context
+
+
+def _pushed_xcoms(ti) -> dict[str, Any]:
+ """The XComs ``aexecute`` pushed, in order: through ``axcom_push`` on
Airflow 3.3+, ``xcom_push`` on 3.2."""
+ calls = ti.axcom_push.await_args_list if AIRFLOW_V_3_3_PLUS else
ti.xcom_push.call_args_list
+ return {c.kwargs["key"]: c.kwargs["value"] for c in calls}
+
+
+def _make_async_mock_agent(mock_hook_cls, **run):
+ """Wire a mock agent whose ``run`` is awaitable into the hook's async
path, and return it."""
+ mock_agent = MagicMock(spec=["run", "instrument"])
+ mock_agent.run = AsyncMock(**run)
+ mock_hook_cls.aget_hook = AsyncMock()
+ mock_hook_cls.aget_hook.return_value.acreate_agent =
AsyncMock(return_value=mock_agent)
+ return mock_agent
+
+
[email protected](not AIRFLOW_V_3_2_PLUS, reason="async tasks need Airflow
>= 3.2")
+class TestAgentOperatorAsync:
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_aexecute_awaits_the_run_through_the_async_hook_path(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ mock_agent = _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result("done"))
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm")
+ result = await op.aexecute(_make_async_context())
+
+ assert result == "done"
+ mock_hook_cls.aget_hook.assert_awaited_once_with(
+ "my_llm", hook_params={"model_id": None, "fallback_conn_ids": None}
+ )
+ mock_hook_cls.get_hook.assert_not_called()
+ mock_agent.run.assert_awaited_once_with(
+ "run", usage_limits=None, run_id="ti-1", cancellation_token=ANY,
usage=ANY
+ )
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.parametrize("async_store", [True, False], ids=["async-store",
"store-before-3.4"])
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def
test_successful_run_makes_no_blocking_call_to_the_supervisor(self,
mock_hook_cls, async_store):
+ """A real run with a usage budget and a message history: no blocking
call reaches the event
+ loop, whether the task state store has async methods (Airflow 3.4+) or
not."""
+ mock_hook_cls.aget_hook = AsyncMock()
+ mock_hook_cls.aget_hook.return_value.acreate_agent = AsyncMock(
+ side_effect=lambda **kw:
Agent(FunctionModel(_build_priced_response), **kw)
+ )
+ store =
_make_async_task_state_store_accessor(async_methods=async_store)
+ context = _make_async_context(task_state_store=store)
+ context["task_instance"].xcom_push.side_effect =
AssertionError("blocking xcom_push")
+
+ op = AgentOperator(
+ task_id="test",
+ prompt="run",
+ llm_conn_id="my_llm",
+ usage_limits={"cost_limit": str(PRICED_COST * 2)},
+ message_history=[],
+ )
+ with (
+
patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_4_PLUS",
async_store),
+
patch("airflow.providers.common.ai.utils.usage_budget.AIRFLOW_V_3_4_PLUS",
async_store),
+ ):
+ result = await op.aexecute(context)
+
+ assert result == "the answer"
+ assert list(_pushed_xcoms(context["task_instance"])) == [
+ "run_id",
+ "usage",
+ MODEL_NAME_XCOM_KEY,
+ "message_history",
+ ]
+ assert ("delete", _TOOL_APPROVAL_TRANSCRIPT_KEY) in store.ops
+ assert ("set", USAGE_BUDGET_KEY) in store.ops
+ assert store.ops[-1] == ("delete", USAGE_BUDGET_KEY)
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_failed_run_saves_the_budget_and_reports_its_usage(self,
mock_hook_cls):
+ _make_async_mock_agent(mock_hook_cls, side_effect=RuntimeError("boom"))
+ store = _make_async_task_state_store_accessor()
+ context = _make_async_context(task_state_store=store)
+
+ op = AgentOperator(
+ task_id="test", prompt="run", llm_conn_id="my_llm",
usage_limits={"request_limit": 5}
+ )
+ with pytest.raises(RuntimeError, match="boom"):
+ await op.aexecute(context)
+
+ assert ("set", USAGE_BUDGET_KEY) in store.ops
+ # The failure report reuses the synchronous code, from a worker thread.
+ pushed = {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list}
+ assert pushed == {"run_id", "usage"}
+
+ @pytest.mark.parametrize("feature", ["durable", "enable_hitl_review"])
+ @pytest.mark.asyncio
+ async def test_aexecute_rejects_features_that_block_the_event_loop(self,
feature):
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm",
**{feature: True})
+
+ with pytest.raises(ValueError, match="not supported with an async
callable"):
+ await op.aexecute(_make_async_context())
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def
test_run_waiting_on_a_tool_approval_pauses_through_the_synchronous_path(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ class Paused(Exception):
+ pass
+
+ _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result(DeferredToolRequests()))
+ context = _make_async_context()
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm")
+ with patch.object(
+ AgentOperator,
+ "_pause_for_tool_approval",
+ side_effect=_refused_on_the_event_loop(Mock(side_effect=Paused)),
+ ) as pause:
+ with pytest.raises(Paused):
+ await op.aexecute(context)
+
+ pause.assert_called_once()
+ assert _pushed_xcoms(context["task_instance"]) == {}
Review Comment:
This is the test failing the Compat job now. `Paused` is not what
`_is_awaiting_input` recognises, so `_acomplete_run` takes the failure branch
and reports `run_id`/`usage` through `xcom_push`, which `_pushed_xcoms` reads
when `axcom_push` is unavailable. Where `axcom_push` exists it only passes
because those pushes are sync and the helper reads the async one. Could this
raise the real `TaskAwaitingInput` and assert `_report_failed_run` is not
called? That would also cover the `not _is_awaiting_input(e)` guard, which no
test hits now.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -139,6 +140,14 @@
_RUN_USAGE_ADAPTER: TypeAdapter[RunUsage] = TypeAdapter(RunUsage)
+async def _axcom_push(ti: Any, *, key: str, value: Any) -> None:
+ if AIRFLOW_V_3_3_PLUS:
+ await ti.axcom_push(key=key, value=value)
+ else:
+ # Airflow 3.2 runs async tasks but has no ``axcom_push``: push from a
worker thread.
+ await asyncio.to_thread(ti.xcom_push, key=key, value=value)
Review Comment:
`asyncio.to_thread` uses the default executor, and `CommsDecoder.asend()`
reads its reply on that same executor while it holds `_thread_lock`. A sync
`send()` from a worker thread just blocks on the lock. So if enough of these
calls pile up, they take every slot and the read never gets a thread. Same goes
for `_report_failed_run` (1161, 1177), the transcript delete and the prompt
render in the decorator.
`hooks/pydantic_ai.py` already hit this and runs its lookups on a one-worker
executor of its own. Route these through that, or `AirflowToolset.run_blocking`?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,16 +1117,76 @@ def execute(self, context: Context) -> Any:
# Without a history to inherit one from, pydantic-ai starts a
new conversation
# on every attempt; keep one for the durable run, for the same
reason.
run_kwargs["conversation_id"] = run_kwargs["run_id"]
+ return run_kwargs
+
+ def _reject_unsupported_on_async_path(self) -> None:
+ """Refuse the features that block the event loop: both do synchronous
I/O during the run."""
+ if self.durable or self.enable_hitl_review:
+ raise ValueError(
+ f"{type(self).__name__}: durable=True and
enable_hitl_review=True are not supported "
+ "with an async callable. Use a synchronous callable for them."
+ )
+
+ async def aexecute(self, context: Context) -> Any:
+ """
+ Run the agent on the worker's event loop.
+
+ ``@task.agent`` calls this for an ``async def`` function; the operator
itself runs through
+ :meth:`execute`. The agent run is awaited on the event loop, and the
bookkeeping of a successful run (usage budget, run metadata, message history)
makes no
Review Comment:
The edit joined two lines here, so this one runs to 160 columns and the next
sentence starts mid-paragraph. Could it be re-wrapped?
--
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]
