dabla commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4276409840
##########
providers/openai/tests/unit/openai/hooks/test_openai.py:
##########
@@ -1032,3 +1033,199 @@ 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_dejson"] + (["aextra_dejson"]
if aextra_dejson else []))
+ conn.password = password
+ conn.host = host
+ 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,
+ ):
+ hook = OpenAIHook(conn_id="openai_async")
+ client = await hook.aget_conn()
+
+ mock_get_async_connection.assert_awaited_once_with("openai_async",
hook=hook)
+ conn.aextra_dejson.assert_awaited_once_with()
+ 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()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+ build.assert_called_once_with(extra)
+ mock_client.assert_called_once_with(workload_identity={"provider":
"custom"}, base_url="https://host/v1";)
+
+
[email protected]
+@patch("airflow.providers.openai.hooks.openai.AsyncOpenAI")
+async def test_aget_conn_resolves_the_connection_through_the_hook(mock_client):
+ """A subclass overriding ``aget_connection`` supplies the connection, as
``get_connection`` does for ``get_conn``."""
+ conn = _async_conn({}, password="api_key_from_subclass")
+
+ class SubclassedOpenAIHook(OpenAIHook):
+ @classmethod
+ async def aget_connection(cls, conn_id: str):
+ assert conn_id == "openai_async"
+ return conn
+
+ with patch.object(OpenAIHook, "get_connection",
side_effect=AssertionError("sync get_connection")):
+ await SubclassedOpenAIHook(conn_id="openai_async").aget_conn()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+ mock_client.assert_called_once_with(api_key="api_key_from_subclass",
base_url=None)
+
+
[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()
+
+ conn.aextra_dejson.assert_awaited_once_with()
+
+
[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 = _async_client(response_items)
+ with patch.object(OpenAIHook, "aget_conn",
new=AsyncMock(return_value=client)) as mock_aget_conn:
+ assert (
+ await
OpenAIHook(conn_id="openai_async").acreate_embeddings(input_text,
dimensions=1024)
+ == expected
+ )
+
+ mock_aget_conn.assert_awaited_once_with()
+ client.embeddings.create.assert_awaited_once_with(
+ model="text-embedding-3-small", input=input_text, dimensions=1024
+ )
+ client.__aexit__.assert_awaited_once()
+
+
[email protected]
+async def
test_acreate_embeddings_concurrent_calls_each_close_their_own_client():
+ """Concurrent calls on one hook each build a client and close it; none is
shared or left open."""
+ clients = [_async_client([(0, [0.1, 0.2])]) for _ in range(3)]
+ hook = OpenAIHook(conn_id="openai_async")
+ with patch.object(OpenAIHook, "aget_conn",
new=AsyncMock(side_effect=clients)):
Review Comment:
Fixed in 9240abe0fa: the test now patches `get_async_connection` and
`AsyncOpenAI` (`side_effect=clients`) and asserts three constructions, so it
fails if `aget_conn` caches a client on the hook again.
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -207,23 +211,41 @@ 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.
+
+ Each call builds a new client. Close it when done, with ``async with``
or
Review Comment:
Thanks for measuring it. Added in 6fdead531a: the `aget_conn` docstring now
says the client caches the exchanged token with `workload_identity`, so for
many calls keep one client open and reuse 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]
