Repository navigation
Conversation
be87667 to
8dfbaf0
Compare
f2e5781 to
274a7cd
Compare
get_conn() resolves the hook's connection and every fallback connection with get_connection() and reads extra_dejson, whose secret masking sends to the supervisor synchronously; inside an async task with another async SDK call in flight both raise DeadlockImminentError, although the agent the hook builds is async-native. aget_conn() fetches the connection and its fallbacks with get_async_connection() (extra from Connection.aextra_dejson(), apache#71890, Airflow 3.3.2+, or the raw extra via json.loads on older versions as in apache#72130), seeds the existing connection cache and builds the model with the unchanged get_conn(); acreate_agent() does the same before create_agent(), spec_file included. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Replace the local aextra_dejson fallback with the common.compat helper added in apache#74147, which masks the extra's secrets asynchronously on Airflow 3.3.2+ and reads extra_dejson from a worker thread on older versions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Prefetching the connections only covered the base hook's own lookups: a subclass's _get_provider_kwargs (Snowflake's builds a SnowflakeHook) still made synchronous Task SDK calls on the event loop, and on Airflow 3.1 to 3.3.1 get_async_connection itself reads extra_dejson synchronously. aget_conn() and acreate_agent() now run the unchanged get_conn() / create_agent() in asyncio.to_thread: from any thread other than the loop's, a synchronous send() waits for the in-flight asend() to release the comms lock instead of raising DeadlockImminentError. This drops the prefetch, the prefetched fallback cache and the extra parameter of _seed_connection. aget_hook() is the async counterpart of get_hook(), with the same connection type check, so an async operator can dispatch to the pydanticai* subclasses. get_hook() now caches the connection it fetched on the hook it returns, so the hook's first get_conn() does not fetch it again. The tests refuse the synchronous lookups only on the event loop, drive a subclass that reads another connection as primary and as fallback, and check that every argument of acreate_agent reaches the agent. The hook docs gain a section on using it from an async task, and uv.lock is back to main's.
274a7cd to
158be99
Compare
_get_conn_and_extra set _conn before reading extra_dejson, whose secret masking is a supervisor round trip. A second acreate_agent gathered on the same hook found _conn set and read the empty default extra, so it raised "No model specified" or dropped the fallback chain; a failed extra read left the hook half-seeded the same way. Both are now published once both are in hand. get_hook's type check already guarantees a PydanticAIHook, so the isinstance guard before seeding is gone, and _seed_connection's docstring describes its contract instead of one caller.
| :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) |
There was a problem hiding this 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.
There was a problem hiding this comment.
Reproduced the mechanism in the source: asend() takes _thread_lock through run_in_executor(None, ...) and reads with asyncio.to_thread, both on the default executor, so hook lookups filling that pool left nothing to run the read. The three methods now run on a single-worker executor of their own (_run_off_loop, e72b7d6), with the caller's context variables copied so task logging keeps its context; one worker also serializes the hook's sends on 3.0/3.1. Tests: each hop completes while the default executor's only worker is parked, and two gathered aget_conn calls never send at the same time.
AirflowToolset.run_blocking, strands_storage, managed_agent and sandbox use asyncio.to_thread too and share the exposure; I can move run_blocking to the same executor in a follow-up if you agree.
| 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. |
There was a problem hiding this 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.
There was a problem hiding this comment.
Done: the section says Airflow 3.2.1 and later, and the sentence on the loop now says the call blocks, raising DeadlockImminentError from 3.3.1 (e72b7d6). The aget_hook and aget_conn docstrings use the same wording.
| ) | ||
|
|
||
| @pytest.mark.asyncio | ||
| @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent") |
There was a problem hiding this comment.
Small one: this Agent patch is missing autospec=True, unlike the other new tests in this class. It still passes with it.
asyncio.to_thread shared the default executor with asend(), which holds the comms lock while a pool thread reads its reply. Enough gathered hook calls filled the pool: every thread blocked in send() on the lock and nothing was left to run the read, so the task hung. The three async methods now run on a single-worker executor of the hook module, with the caller's context variables copied because task logging binds them. One worker also serializes the hook's sends on Airflow 3.0 and 3.1, where the channel to the supervisor has no lock of its own. The docs section states Airflow 3.2.1 and later (3.2.0 reads the reply outside the lock) and that the loop blocks before 3.3.1, where DeadlockImminentError arrived.
PydanticAIHook.get_conn()resolves its connection (and every connection of afallback_conn_idschain) withget_connection()and readsextra_dejson, whose secret masking sendsMaskSecretto the supervisor synchronously, and a subclass's_get_provider_kwargsmay read more (Snowflake's builds aSnowflakeHook). On the event loop of an async task with another async SDK call in flight, each of those raisesDeadlockImminentError, although the pydantic-ai agent the hook builds is itself async-native (await agent.run(...)).This adds the async counterparts without a second resolution path:
aget_conn()andacreate_agent()run the unchangedget_conn()/create_agent()on a single-worker executor of the hook module, never on asyncio's default executor:asend()reads its reply on that executor while it holds the comms lock, so lookups filling it would hang the task. From any thread other than the loop's, a synchronoussend()waits for the in-flightasend()to release the comms lock instead of raising, so fallback chains,spec_file, model-id qualification and subclass overrides behave exactly as in the sync path; one worker also serializes the hook's sends on 3.0/3.1.aget_hook()is the async counterpart ofget_hook(), with the same connection type check, so an async operator can dispatch to thepydanticai*subclasses.get_hook()now caches the connection it fetched on the hook it returns, so the hook's firstget_conn()doesn't fetch it a second time (sync operators included).await hook.acreate_agent(...)/await agent.run(...), for Airflow 3.2.1 and later (3.2.0 reads a reply outside the comms lock;DeadlockImminentErrorexists from 3.3.1, before that the loop blocks).This is also the first building block for running
@task.agent/AgentOperatornatively async.Tests: the sync lookups fail when they run on the event loop. The new tests resolve a model, a fallback chain, a subclass whose provider kwargs read another connection (as primary and as fallback),
acreate_agentwith every argument forwarded (with and withoutspec_file), a secondacreate_agentthat fetches nothing again, twoacreate_agentcalls gathered on one hook, a retry after a failed extra read, each method completing while the default executor is full, two hooks never sending at the same time, andaget_hookwith its type check. Existing tests pass unchanged.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Opus 5.5 following the guidelines