dabla commented on code in PR #74201:
URL: https://github.com/apache/airflow/pull/74201#discussion_r4556861074
##########
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:
Done in 36005df14d: the pause now raises a real `TaskAwaitingInput` and the
test asserts `_report_failed_run` is not called, so it covers the `not
_is_awaiting_input(e)` guard and no longer depends on which push mock the
helper reads. It is skipped below 3.3, where the exception does not exist.
##########
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:
You're right, and `asend()` takes `_thread_lock` through that executor too.
In 5669cfddcf the hook's one-worker executor moved to `utils/coroutines.py` as
`run_off_loop`, and all of these go through it: the 3.2 XCom push, both
`_report_failed_run` calls, the approval pause, the transcript delete, the
prompt render, the usage budget calls before 3.4, and the skills XCom push from
the last round. The hook uses the same helper, so the provider has one worker
for its blocking supervisor calls. Each path has a test that runs with the
default executor's only worker taken.
I did not use `AirflowToolset.run_blocking`: it goes through
`asyncio.to_thread` itself, so calls waiting on its lock sit in
default-executor threads as well. That is on main and applies to every toolset.
Putting it on the same single worker would let one slow tool call hold up the
bookkeeping of every other run, so I'd look at it in a separate PR.
#76545 removes the cause in the Task SDK: there `asend()` does the whole
exchange on one worker thread that holds the lock, so waiting `send()` calls
can no longer keep the read from running. That only helps on the Airflow
version that ships it, so the provider still needs its own worker for the older
ones it supports.
##########
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:
Re-wrapped in 4980616c05, as two paragraphs.
--
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]
