kaxil commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4362809730
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -206,7 +208,27 @@ def get_hook(cls, conn_id: str, hook_params: dict | None =
None):
f"one of {expected_types}. Change the connection type to
'pydanticai', put the API key "
"in the API Key field (stored as the password) and the model
(e.g. 'openai:gpt-5') in Model."
)
- return conn.get_hook(hook_params=hook_params)
+ # The check above guarantees a ``PydanticAIHook``. Seeding saves the
second fetch of
+ # this same connection its first ``get_conn`` would make.
+ hook = conn.get_hook(hook_params=hook_params)
+ hook._seed_connection(conn)
+ return hook
+
+ @classmethod
+ async def aget_hook(cls, conn_id: str, hook_params: dict | None = None):
+ """
+ Return the hook for ``conn_id``, built with ``hook_params``, for use
inside an async task.
+
+ The async counterpart of :meth:`get_hook`, with the same connection
type check and
+ the connection already cached on the hook. The lookup runs in a worker
thread:
+ ``get_connection`` sends to the supervisor synchronously, which raises
+ ``DeadlockImminentError`` on the event loop while another async call
is in flight,
+ but only waits its turn from any other thread.
+
+ :raises ValueError: if the connection's type is registered to a hook
that is not a
+ ``PydanticAIHook``.
+ """
+ return await asyncio.to_thread(cls.get_hook, conn_id, hook_params)
Review Comment:
These `to_thread` calls share the default executor with `asend()`, which
holds `_thread_lock` while it waits for a pool thread to read its reply
(`comms.py:329`, `:346`). Gather enough hook calls to fill that pool (`min(32,
cpu_count + 4)`, so 8 on a 4-CPU worker) while some other async SDK call is in
flight, and they all block in `send()` on that lock. Nothing is left to run the
read, so the task hangs. I reproduced it with the real `CommsDecoder` over a
socketpair: on an 18-thread pool, 18 gathered `aget_conn()` calls hang and 17
finish.
Could `aget_hook`, `aget_conn` and `acreate_agent` use their own
single-worker executor instead? That would also serialize the hook's sends on
3.0/3.1.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -2207,3 +2216,249 @@ def test_hook_resolves_from_environment_variable(self,
hook_class, serializer, m
assert resolved.conn_type == hook_class.conn_type
assert type(resolved.get_hook()) is hook_class
+
+
+def _off_the_event_loop(lookup):
+ """
+ Wrap a synchronous Task SDK lookup so it fails the test when it runs on
the event loop.
+
+ ``get_connection`` and the secret masking of ``extra_dejson`` send to the
supervisor
+ synchronously. On the event loop, with another async call in flight, that
raises
+ ``DeadlockImminentError``; from a worker thread it only waits its turn.
+ """
+
+ def guarded(*args, **kwargs):
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return lookup(*args, **kwargs)
+ raise AssertionError("synchronous Task SDK call on the event loop")
+
+ return guarded
+
+
[email protected]
+def loop_guarded_registry():
+ """The ``registry`` fixture, with every synchronous lookup refused on the
event loop."""
+ reg = _ConnRegistry()
+ with (
+ patch.object(
+ PydanticAIHook, "get_connection",
side_effect=_off_the_event_loop(reg.get_connection)
+ ) as mock_get_connection,
+ patch.object(
+ Connection,
+ "extra_dejson",
+ new=property(_off_the_event_loop(lambda conn:
json.loads(conn.extra) if conn.extra else {})),
+ ),
+ patch.object(Connection, "get_hook", side_effect=reg.get_hook,
autospec=True),
+ patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai._get_registered_hooks",
+ autospec=True,
+ return_value=reg.hooks,
+ ),
+ ):
+ reg.get_connection_mock = mock_get_connection
+ yield reg
+
+
+class _ReadsAnotherConnectionHook(PydanticAIHook):
+ """A subclass whose provider kwargs come from another connection, as
Snowflake's do."""
+
+ def _get_provider_kwargs(self, api_key, base_url, extra):
+ return {"api_key":
self.get_connection("credentials").extra_dejson["token"]}
+
+
+class TestPydanticAIHookAsync:
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_model_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert model is infer_model_stub.models["openai:gpt-5.6-sol"]
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("primary")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_fallback_chain_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second", "third"]}
+ )
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ loop_guarded_registry.add("third", extra={"model": "groq:llama-4"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ infer_model_stub.models["anthropic:claude-opus-4-6"],
+ infer_model_stub.models["groq:llama-4"],
+ ]
+ assert loop_guarded_registry.get_connection_mock.call_args_list == [
+ call("primary"),
+ call("second"),
+ call("third"),
+ ]
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("subclass_conn_id", ["primary", "second"],
ids=["primary", "fallback"])
+ async def
test_aget_conn_runs_a_subclass_provider_kwargs_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub, subclass_conn_id
+ ):
+ loop_guarded_registry.add("credentials", extra={"token":
"from-another-connection"})
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]}
+ )
+ loop_guarded_registry.add(
+ "second", hook_class=_ReadsAnotherConnectionHook, extra={"model":
"anthropic:claude-opus-4-6"}
+ )
+ hook_class = _ReadsAnotherConnectionHook if subclass_conn_id ==
"primary" else PydanticAIHook
+
+ await hook_class(llm_conn_id="primary").aget_conn()
+
+
loop_guarded_registry.get_connection_mock.assert_any_call("credentials")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_is_cached(self, loop_guarded_registry,
infer_model_stub):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ assert await hook.aget_conn() is await hook.aget_conn()
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-5.6-sol"]
+ loop_guarded_registry.get_connection_mock.assert_called_once()
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_forwards_every_argument(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ agent = await PydanticAIHook(llm_conn_id="primary").acreate_agent(
+ int, instructions="Count.", retries=3
+ )
+
+ assert agent is mock_agent_cls.return_value
+ mock_agent_cls.assert_called_once_with(
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ output_type=int,
+ instructions="Count.",
+ retries=3,
+ capabilities=[MaskingCapability()],
+ )
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent")
Review Comment:
Small one: this `Agent` patch is missing `autospec=True`, unlike the other
new tests in this class. It still passes with it.
##########
providers/common/ai/docs/hooks/pydantic_ai.rst:
##########
@@ -97,3 +97,22 @@ The spec file's ``model`` is the last resort: ``model_id``
and the connection's
``model`` extra both take precedence over it, in the
:ref:`order the connection page lists <pydanticai-model-resolution>`. Passing
``instructions`` to ``create_agent`` when a ``spec_file`` is also given
appends additional instructions to the file value.
+
+Using the hook in an async task
+-------------------------------
+
+In an ``async def`` task, use the async counterparts ``aget_hook``,
``aget_conn``
+and ``acreate_agent``, and run the agent with ``await agent.run(...)``. They
take the
+same parameters as ``get_hook``, ``get_conn`` and ``create_agent`` and return
the same
+objects, fallback chain included. Connection lookups make synchronous calls to
the Task
+SDK, which raise ``DeadlockImminentError`` on the event loop while another
async call
+is in flight, so the async methods run them in a worker thread.
Review Comment:
Two version caveats here. `DeadlockImminentError` only exists from 3.3.1; on
3.2.0 to 3.3.0 the same call on the loop just blocks. And on 3.2.0 `send()`
holds `_thread_lock` for the write only and reads the reply outside it (fixed
in 3.2.1 by #64894), so a worker-thread `send()` can still read another call's
reply there, including one from core's own `asend()` that nothing in the
provider can serialize. I'd say Airflow 3.2.1+ in this section.
--
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]
