본문으로 건너뛰기

[pydantic-ai] xAI 프로바이더에서 gRPC 이벤트 루프 불일치 버그 수정

PR 링크: pydantic/pydantic-ai#4316 상태: Merged | 변경: +68 / -2

들어가며

gRPC의 async 채널은 생성 시점의 이벤트 루프에 바인딩됩니다. xAI 프로바이더에서 AsyncClient가 모듈 레벨이나 async 컨텍스트 외부에서 생성된 후, asyncio.run() 내부의 다른 이벤트 루프에서 사용되면 RuntimeError가 발생했습니다. 이 PR은 이벤트 루프 변경을 감지하여 클라이언트를 자동으로 재생성하는 _LazyAsyncClient 래퍼를 도입합니다.

핵심 코드 분석

_LazyAsyncClient 구현

Before:

class XaiProvider(Provider[AsyncClient]):
    def __init__(self, *, api_key=None, xai_client=None):
        if xai_client is not None:
            self._client = xai_client
        else:
            self._client = AsyncClient(api_key=api_key)

클라이언트가 __init__ 시점에 즉시 생성되어 이벤트 루프에 바인딩되었습니다.

After:

class _LazyAsyncClient:
    """gRPC async channels bind to the event loop at creation time.
    This wrapper defers client creation and recreates it when the loop changes."""

    def __init__(self, **kwargs):
        self._kwargs = kwargs
        self._client: AsyncClient | None = None
        self._event_loop: asyncio.AbstractEventLoop | None = None

    def get_client(self) -> AsyncClient:
        running_loop = None
        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            pass

        if self._client is None or (running_loop is not None and running_loop is not self._event_loop):
            self._client = AsyncClient(**self._kwargs)
            self._event_loop = running_loop
        return self._client

get_client() 호출 시 현재 이벤트 루프를 확인하고, 이전과 다르면 새 클라이언트를 생성합니다. is 비교로 루프 객체의 동일성을 검사하여 정확한 판별을 합니다.

class XaiProvider(Provider[AsyncClient]):
    @property
    def client(self) -> AsyncClient:
        if self._lazy_client is not None:
            return self._lazy_client.get_client()
        return self._client

사용자가 직접 xai_client를 전달한 경우는 기존 방식 유지, API 키만 전달한 경우에만 lazy 패턴을 적용합니다.

테스트로 검증

def test_xai_provider_recreates_client_on_new_loop():
    provider = XaiProvider(api_key='api-key')
    clients = []
    clients.append(asyncio.run(get_client()))
    clients.append(asyncio.run(get_client()))
    assert clients[0] is not clients[1]  # 다른 루프 -> 다른 클라이언트

def test_xai_provider_reuses_client_on_same_loop():
    provider = XaiProvider(api_key='api-key')
    c1, c2 = asyncio.run(get_clients_same_loop())
    assert c1 is c2  # 같은 루프 -> 같은 클라이언트

왜 이게 좋은가

이 패턴은 gRPC 기반 SDK에서 흔히 발생하는 이벤트 루프 불일치 문제를 우아하게 해결합니다. 사용자가 프로바이더를 어디서 생성하든(asyncio.run() 내부/외부, Jupyter 노트북, 테스트 등) 안전하게 동작합니다. 동시에 같은 루프 내에서는 클라이언트를 재사용하여 불필요한 연결 생성을 방지합니다.

정리

항목 내용
문제 gRPC AsyncClient가 다른 이벤트 루프에서 사용 시 RuntimeError
해결 _LazyAsyncClient로 루프 변경 시 자동 재생성
호환성 사용자 제공 클라이언트는 기존 동작 유지

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글