kaxil commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4218753376
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -201,23 +205,44 @@ def get_conn(self) -> OpenAI:
OpenAI client's workload identity support. See
:meth:`_build_workload_identity`.
"""
conn = self.get_connection(self.conn_id)
- extras = conn.extra_dejson
- openai_client_kwargs = extras.get("openai_client_kwargs", {})
+ return OpenAI(**self._client_kwargs(conn, conn.extra_dejson))
+
+ async def aget_conn(self) -> AsyncOpenAI:
+ """
+ Return an ``AsyncOpenAI`` client, for use inside an async task or
trigger.
+
+ Configured exactly as :meth:`get_conn` configures its client, but the
connection is
+ resolved without any synchronous call to the Task SDK.
``get_connection`` and the
+ secret masking of ``extra_dejson`` send to the supervisor
synchronously, which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ conn = await get_async_connection(self.conn_id)
Review Comment:
`get_conn` resolves through `self.get_connection`, so a subclass override
applies there. With no `hook=`, this lookup goes through `BaseHook` instead.
`AnthropicHook.get_async_conn` passes `hook=self`;
`get_async_connection(self.conn_id, hook=self)` would do the same here. The
`get_connection` patch in `test_aget_conn_api_key` never fires because
`get_async_connection` is mocked, so a test with a subclass that overrides
`aget_connection` is needed to cover this.
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -566,6 +591,43 @@ def create_embeddings(
return response.data[0].embedding
return [item.embedding for item in sorted(response.data, key=lambda
item: item.index)]
+ @overload
+ async def acreate_embeddings(
+ self,
+ text: str | list[int],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[float]: ...
+
+ @overload
+ async def acreate_embeddings(
+ self,
+ text: list[str] | list[list[int]],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[list[float]]: ...
+
+ async def acreate_embeddings(
+ self,
+ text: str | list[str] | list[int] | list[list[int]],
+ model: str = "text-embedding-3-small",
+ **kwargs: Any,
+ ) -> list[float] | list[list[float]]:
+ """
+ Generate embeddings for the given text using the given model,
asynchronously.
+
+ The async counterpart of :meth:`create_embeddings`, through
:meth:`aget_conn`.
+
+ :param text: The text to generate embeddings for.
+ :param model: The model to use for generating embeddings.
+ :return: One embedding for a single text or token array; one embedding
per item for a batch.
+ """
+ client = await self._aget_cached_conn()
+ response = await client.embeddings.create(model=model, input=text,
**kwargs)
+ if isinstance(text, str) or (text and isinstance(text[0], int)):
Review Comment:
This shaping is copied from `create_embeddings`. Since `_client_kwargs` was
pulled out so the two paths can't drift, a small `_embeddings_result(text,
response)` used by both would do the same for the return shape.
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -998,3 +1000,156 @@ def
test_managed_agents_reports_sdk_upgrade_without_breaking_hook():
hook.conn = SimpleNamespace(beta=SimpleNamespace())
with pytest.raises(OpenAIAgentSessionError, match="requires
openai>=3.13.0"):
hook.create_agent_session(input="Hello", environment={"type": "none"},
agent_id="agent")
+
+
+def _async_conn(
+ extra: dict, password: str | None = None, host: str | None = None, *,
aextra_dejson: bool = True
+) -> Mock:
+ """
+ A connection whose ``extra_dejson`` raises on the event loop, where its
secret masking deadlocks.
+
+ With ``aextra_dejson=False`` it lacks ``Connection.aextra_dejson()``, like
on Airflow < 3.3.2,
+ where ``get_async_extra_dejson`` reads ``extra_dejson`` from a worker
thread instead.
+ """
+ conn = Mock(
+ spec=["password", "host", "extra", "extra_dejson"] +
(["aextra_dejson"] if aextra_dejson else [])
+ )
+ conn.password = password
+ conn.host = host
+ conn.extra = json.dumps(extra)
+ if aextra_dejson:
+ conn.aextra_dejson = AsyncMock(return_value=extra)
+ type(conn).extra_dejson = PropertyMock(side_effect=lambda:
_read_off_the_event_loop(extra))
+ return conn
+
+
+def _read_off_the_event_loop(extra: dict) -> dict:
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return extra
+ raise AssertionError("extra_dejson sends mask_secret synchronously on the
event loop")
+
+
[email protected]
[email protected](
+ ("password", "extra", "expected_kwargs"),
+ [
+ pytest.param(
+ "api_key_in_password",
+ {},
+ {"api_key": "api_key_in_password", "base_url": "https://host/v1"},
+ id="api-key-in-password",
+ ),
+ pytest.param(
+ None,
+ {"openai_client_kwargs": {"api_key": "api_key_in_extra",
"base_url": "https://other/v1"}},
+ {"api_key": "api_key_in_extra", "base_url": "https://other/v1"},
+ id="api-key-and-base-url-in-extra",
+ ),
+ pytest.param(
+ "api_key",
+ {"openai_client_kwargs": {"timeout": 30}},
+ {"api_key": "api_key", "base_url": "https://host/v1";, "timeout":
30},
+ id="openai-client-kwargs",
+ ),
+ ],
+)
+async def test_aget_conn_api_key(password, extra, expected_kwargs):
+ conn = _async_conn(extra, password=password, host="https://host/v1";)
+ with (
+ patch(
+ "airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)
+ ) as mock_get_async_connection,
+ patch.object(OpenAIHook, "get_connection",
side_effect=AssertionError("sync get_connection")),
+ patch("airflow.providers.openai.hooks.openai.AsyncOpenAI") as
mock_client,
+ ):
+ client = await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ mock_get_async_connection.assert_awaited_once_with("openai_async")
+ mock_client.assert_called_once_with(**expected_kwargs)
+ assert client is mock_client.return_value
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def
test_aget_conn_without_aextra_dejson_reads_extra_dejson_off_the_event_loop(mock_client):
+ """On Airflow < 3.3.2 the connection has no ``aextra_dejson()``:
``extra_dejson`` runs in a worker thread."""
+ conn = _async_conn({"openai_client_kwargs": {"api_key":
"api_key_in_extra"}}, aextra_dejson=False)
+ with patch(
+ "airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ mock_client.assert_called_once_with(api_key="api_key_in_extra",
base_url=None)
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def test_aget_conn_workload_identity(mock_client):
+ extra = {
+ "auth_type": "workload_identity",
+ "identity_provider_id": "idp",
+ "service_account_id": "sa",
+ "workload_identity_provider": "custom",
+ "token_provider": "some.module.get_token",
+ "openai_client_kwargs": {"api_key": "ignored"},
+ }
+ conn = _async_conn(extra, host="https://host/v1";)
+ with (
+ patch("airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)),
+ patch.object(OpenAIHook, "_build_workload_identity",
return_value={"provider": "custom"}) as build,
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+ build.assert_called_once_with(extra)
+ mock_client.assert_called_once_with(workload_identity={"provider":
"custom"}, base_url="https://host/v1";)
+
+
[email protected]
+async def test_aget_conn_invalid_auth_type():
+ conn = _async_conn({"auth_type": "magic"})
+ with (
+ patch("airflow.providers.openai.hooks.openai.get_async_connection",
new=AsyncMock(return_value=conn)),
+ pytest.raises(ValueError, match="Unsupported auth_type 'magic'"),
+ ):
+ await OpenAIHook(conn_id="openai_async").aget_conn()
+
+
[email protected]
[email protected](
+ ("input_text", "response_items", "expected"),
+ [
+ pytest.param("Sample text", [(0, [0.1, 0.2])], [0.1, 0.2], id="text"),
+ pytest.param([1, 2, 3], [(0, [0.1, 0.2])], [0.1, 0.2], id="tokens"),
+ pytest.param(
+ ["First text", "Second text"],
+ [(1, [0.3, 0.4]), (0, [0.1, 0.2])],
+ [[0.1, 0.2], [0.3, 0.4]],
+ id="text-batch",
+ ),
+ ],
+)
+async def test_acreate_embeddings(input_text, response_items, expected):
+ client = MagicMock()
Review Comment:
`MagicMock(spec=AsyncOpenAI)`, like the sync tests' `spec=OpenAI`? Without
it a typo such as `client.embedding` on the async path would still pass.
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -201,23 +205,44 @@ def get_conn(self) -> OpenAI:
OpenAI client's workload identity support. See
:meth:`_build_workload_identity`.
"""
conn = self.get_connection(self.conn_id)
- extras = conn.extra_dejson
- openai_client_kwargs = extras.get("openai_client_kwargs", {})
+ return OpenAI(**self._client_kwargs(conn, conn.extra_dejson))
+
+ async def aget_conn(self) -> AsyncOpenAI:
+ """
+ Return an ``AsyncOpenAI`` client, for use inside an async task or
trigger.
+
+ Configured exactly as :meth:`get_conn` configures its client, but the
connection is
+ resolved without any synchronous call to the Task SDK.
``get_connection`` and the
+ secret masking of ``extra_dejson`` send to the supervisor
synchronously, which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ conn = await get_async_connection(self.conn_id)
+ return AsyncOpenAI(**self._client_kwargs(conn, await
get_async_extra_dejson(conn)))
+
+ async def _aget_cached_conn(self) -> AsyncOpenAI:
+ """Return the hook's ``AsyncOpenAI`` client, created on first use like
:attr:`conn`."""
+ if self._async_conn is None:
Review Comment:
With `asyncio.gather(*(hook.acreate_embeddings(b) for b in batches))` on a
fresh hook (the concurrent case from the description), every call sees `None`
here before the first `aget_conn()` returns, since it awaits twice. That's N
connection lookups and N clients, and N-1 of them are dropped without being
closed. The one that is kept is never closed either, and it stays bound to the
first event loop: reusing the hook across two `asyncio.run()` calls fails the
second one with `APIConnectionError` when `max_retries` is 0, or spends a retry
otherwise.
`AnthropicHook.get_async_conn` avoids all of this by building a client per
call and telling the caller to close it, so `async with await self.aget_conn()
as client:` here would be enough. If you'd rather keep the cache, it needs a
lock and an `aclose()`.
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -998,3 +1000,156 @@ def
test_managed_agents_reports_sdk_upgrade_without_breaking_hook():
hook.conn = SimpleNamespace(beta=SimpleNamespace())
with pytest.raises(OpenAIAgentSessionError, match="requires
openai>=3.13.0"):
hook.create_agent_session(input="Hello", environment={"type": "none"},
agent_id="agent")
+
+
+def _async_conn(
+ extra: dict, password: str | None = None, host: str | None = None, *,
aextra_dejson: bool = True
+) -> Mock:
+ """
+ A connection whose ``extra_dejson`` raises on the event loop, where its
secret masking deadlocks.
+
+ With ``aextra_dejson=False`` it lacks ``Connection.aextra_dejson()``, like
on Airflow < 3.3.2,
+ where ``get_async_extra_dejson`` reads ``extra_dejson`` from a worker
thread instead.
+ """
+ conn = Mock(
+ spec=["password", "host", "extra", "extra_dejson"] +
(["aextra_dejson"] if aextra_dejson else [])
+ )
+ conn.password = password
+ conn.host = host
+ conn.extra = json.dumps(extra)
Review Comment:
Nothing in the hook reads `conn.extra`, and no test asserts `aextra_dejson`
was awaited. So replacing `await get_async_extra_dejson(conn)` with
`json.loads(conn.extra)`, which avoids the deadlock but skips the masking,
would keep every `aget_conn` test green. Dropping `extra` from the spec and
adding `conn.aextra_dejson.assert_awaited_once_with()` would catch it.
--
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]
