YAshhh29 commented on code in PR #74149:
URL: https://github.com/apache/airflow/pull/74149#discussion_r4275483442
##########
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:
One thing I noticed with the per-call client and
`auth_type="workload_identity"`: the SDK keeps the exchanged token on the
client's `WorkloadIdentityAuth`, so a new client per call also means a new
token exchange per call.
I measured it on openai 3.26.1 with the exchange stubbed: 5 concurrent calls
on per-call clients did 5 exchanges, and 5 calls on one shared client did 1.
`api_key` auth isn't affected.
Not a blocker, since closing each client is the right call after the earlier
review. Might be worth a line in the `aget_conn` docstring, so someone making
many calls with workload identity knows they can hold one client 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]
