본문 바로가기

AI

[AI Agent] ASGI 에서 LangGraph 스트림 버퍼링 문제 (iterator 와 서버의 조합)

반응형

AI agent 서비스를 개발하며 겪은 문제들을 예제 프로젝트로 재구성해 정리하는 시리즈이다. 예제 코드는 parcel-bot repo 에 있고, docker compose 로 전부 재현할 수 있다.

 

이번 글은 이 시리즈를 시작하게 만든 문제를 다룬다. Django 와 LangGraph 로 만든 서비스에 성능을 높이기 위해 async 를 도입하는 과정에서, LangGraph 스트림이 통째로 버퍼링되는 문제를 만났다. 이전 글에서 정리한 iterator 와 서버의 관계를 정리했던 것처럼 이번 글에서는 문제를 재현하고 원인과 해결 방법을 확인해보겠다.

문제 상황

Django 와 LangGraph 로 구성된 AI agent 서비스가 있다. agent 의 응답은 SSE 로 스트리밍한다.

 

def recommend(state):
    grades = SecurityGrade.objects.filter(user=...)   # Django ORM (sync)
    result = llm_client.messages.create(...)          # sync LLM SDK 호출
    return {"recommendation": ...}

 

이 node 안의 두 호출은 모두 결과가 올 때까지 스레드를 멈추고 기다리는 블로킹 호출이다. Django ORM 쿼리는 DB 응답이 올 때까지 기다리고, LLM 호출도 sync 클라이언트를 쓰면 응답까지 수 초를 그대로 기다린다. 이런 블로킹 호출을 담는 node 는 자연스럽게 sync 함수로 작성되고, node 가 sync 이므로 graph 실행도 sync API 인 graph.stream() 을 쓰게 된다. 이 함수가 반환하는 것이 sync generator 이다.

 

여기까지는 문제없이 동작하고 있었다. 그러다 성능을 높이기 위해 async 를 도입하기로 했다. SSE 같은 스트리밍 연결은 오래 유지되는데, sync 서버는 연결 하나가 워커 하나를 스트림이 끝날 때까지 점유하기 때문이다. 서버를 ASGI (uvicorn) 로 전환하고 view 를 async 로 재작성했다.

 

그런데 배포 후 확인해보니 스트리밍이 정상적으로 동작하지 않았다. 토큰이 한참 동안 출력되지 않다가 답변 전체가 한 번에 출력됐다. 스트림을 만드는 코드는 바꾸지 않았고 서빙 방식만 바뀌었을 뿐인데 동작이 달라진 것이다.

 

증상은 이전 글과 같은데 방향이 반대다. 이전 글에서는 async iterator 를 WSGI 에서 서빙해서 발생한 버퍼링이었다면, 이번에는 sync generator 인 graph.stream() 을 ASGI 에서 서빙해서 발생한 버퍼링이다.

예제 구성

parcel-bot 에 최소 구성의 LangGraph graph 를 추가해서 재현한다. node 는 하나이고, 실제 서비스처럼 sync 함수이다.

 

# parcel_bot/graph.py
import time
from typing import TypedDict

from langgraph.config import get_stream_writer
from langgraph.graph import END, START, StateGraph


class ChatState(TypedDict):
    message: str


def reply(state: ChatState) -> ChatState:
    # sync node: 실제 서비스에서 ORM/sync LLM 호출이 섞이는 상황을 재현한다.
    writer = get_stream_writer()
    text = (
        f"'{state['message']}' 문의 확인했습니다. 배송 접수를 도와드릴게요. "
        "보내시는 분 성함과 받으시는 분 주소를 알려주세요."
    )
    for token in text.split(" "):
        time.sleep(0.3)
        writer(token + " ")
    return state


chat_graph = (
    StateGraph(ChatState)
    .add_node("reply", reply)
    .add_edge(START, "reply")
    .add_edge("reply", END)
    .compile()
)

 

두 가지 장치를 설명하고 넘어가겠다.

 

time.sleep(0.3) 은 node 안의 블로킹 대기 (DB 왕복, LLM 응답 대기) 를 한 줄로 축약한 것이다. await asyncio.sleep() 이 아니라 time.sleep() 인 이유는 실제 node 의 블로킹 호출들과 같은 방식으로 스레드를 붙잡기 위함이다.

 

get_stream_writer() 는 node 실행 도중에 임의의 값을 graph 스트림으로 밀어 넣는 통로이다. LangGraph 의 기본 스트리밍 모드 (updates/values) 는 node 가 끝난 시점의 결과만 내놓기 때문에, node 중간의 토큰을 흘리려면 stream_mode="custom" 과 writer 를 사용한다. 실제 LLM 을 붙일 때는 stream_mode="messages" 를 사용하여 node 안에서 일어나는 LangChain LLM 호출의 토큰을 자동으로 스트림에 실어주는데, 여기서는 LLM 호출이 없는 mock 이라 custom 모드로 같은 동작을 재현했다.

 

실제 서비스 이 예제
node 안 LLM 호출의 토큰 콜백 writer(token) 호출
stream_mode="messages" 로 수신 stream_mode="custom" 으로 수신

 

이 graph 를 세 가지 방식으로 서빙하는 endpoint 를 만든다.

 

# parcel_bot/views.py (발췌)
_STREAM_END = object()


def graph_sync(request):
    # sync view + graph.stream() (sync generator)
    message = _message(request)

    def stream():
        for token in chat_graph.stream({"message": message}, stream_mode="custom"):
            yield _sse(token)
        yield "data: [DONE]\n\n"

    return StreamingHttpResponse(stream(), content_type="text/event-stream")


async def graph_async(request):
    # async view + graph.astream() (async generator)
    message = _message(request)

    async def stream():
        async for token in chat_graph.astream({"message": message}, stream_mode="custom"):
            yield _sse(token)
        yield "data: [DONE]\n\n"

    return StreamingHttpResponse(stream(), content_type="text/event-stream")


async def graph_bridge(request):
    # sync graph.stream() 을 async generator 로 감싸 한 개씩 꺼내는 브리지
    message = _message(request)

    async def stream():
        events = chat_graph.stream({"message": message}, stream_mode="custom")
        while True:
            token = await sync_to_async(next, thread_sensitive=True)(events, _STREAM_END)
            if token is _STREAM_END:
                break
            yield _sse(token)
        yield "data: [DONE]\n\n"

    return StreamingHttpResponse(stream(), content_type="text/event-stream")

 

graph_sync 가 async 도입 전의 구성이고, graph_asyncgraph_bridge 는 해결 섹션에서 다룰 구성이다.

실험: iterator 종류 × 서버 종류 매트릭스

이전 글과 같은 구성인 WSGI (runserver, :8000) 와 ASGI (uvicorn, :8001) 컨테이너에, 같은 타임스탬프 curl 로 도착 시각을 측정한다.

 

먼저 WSGI 에 sync 스트림 (/graph/sync/) 이다. async 도입 전의 구성이고, 300ms 간격으로 실시간 스트리밍된다.

 

42.030 data: {"token": "'책' "}
42.330 data: {"token": "문의 "}
42.631 data: {"token": "확인했습니다. "}

 

같은 WSGI 에 async 스트림 (/graph/async/) 을 요청하면 이전 글에서 본 그 버퍼링이다. 전부 같은 시각에 도착한다.

 

49.573 data: {"token": "'책' "}
49.576 data: {"token": "문의 "}
49.578 data: {"token": "확인했습니다. "}

 

이번에는 ASGI 에 sync 스트림 (/graph/sync/) 이다. async 도입 과정에서 만난 문제 상황이 바로 이 조합이다. 전부 같은 시각에 도착한다.

 

06.200 data: {"token": "'책' "}
06.201 data: {"token": "문의 "}
06.203 data: {"token": "확인했습니다. "}

 

ASGI 에 async 스트림 (/graph/async/) 은 정상 스트리밍된다.

 

06.534 data: {"token": "'책' "}
06.834 data: {"token": "문의 "}
07.134 data: {"token": "확인했습니다. "}

 

결과를 정리하면 다음과 같다.

 

  sync iterator (graph.stream()) async iterator (graph.astream())
WSGI (runserver) 실시간 스트리밍 통째 버퍼링 (이전 글)
ASGI (uvicorn) 통째 버퍼링 (이번 글) 실시간 스트리밍

 

iterator 종류와 서버 종류가 어긋나는 조합에서만 버퍼링이 발생한다. 이전 글의 문제는 이 매트릭스의 한 칸 (WSGI / async iterator) 이었던 것이다.

원인: 어긋난 조합은 Django 가 통째로 소비한다

ASGI 에서 sync 스트림을 서빙할 때의 로그를 보면, 이전 글의 경고와 정확히 대칭인 경고가 출력된다.

 

Warning: StreamingHttpResponse must consume synchronous iterators
in order to serve them asynchronously. Use an asynchronous iterator instead.

 

Django 소스에서 이유를 확인해보겠다. 먼저 streaming_content 속성부터 짚고 넘어간다.

streaming_content: 응답 본문 iterator 를 담는 속성

StreamingHttpResponse 생성자에 넘긴 iterator 는 streaming_content 라는 속성으로 보관된다. 저장할 때 sync/async 판별이 일어난다.

 

# django/http/response.py의 StreamingHttpResponse (발췌)
def _set_streaming_content(self, value):
    try:
        self._iterator = iter(value)     # sync iterator 면 성공
        self.is_async = False
    except TypeError:
        self._iterator = aiter(value)    # async iterator 면 이쪽
        self.is_async = True

 

iter() 가 성공하면 sync 로, TypeError 가 나면 aiter() 로 다시 받아 async 로 표시한다. 조회할 때는 is_async 에 따라 다른 형태로 감싸서 반환한다.

 

@property
def streaming_content(self):
    if self.is_async:
        _iterator = self._iterator

        async def awrapper():
            async for part in _iterator:
                yield self.make_bytes(part)

        return awrapper()                          # async generator
    else:
        return map(self.make_bytes, self._iterator)  # sync iterator (map 객체)

 

정리하면 streaming_content 는 응답 본문 iterator 를 bytes 변환까지 포함해 감싼 것이고, 원본이 sync 였다면 sync iterator (map 객체) 로, async 였다면 async generator 로 나온다. 원본의 sync/async 성질이 그대로 유지된다는 점이 핵심이다.

ASGI 서버가 소비하는 경로

ASGI 서버는 응답을 async for 로 소비한다. 그 진입점인 __aiter__ 를 보면 다음과 같다.

 

# django/http/response.py의 StreamingHttpResponse.__aiter__ (발췌)
async def __aiter__(self):
    try:
        async for part in self.streaming_content:
            yield part
    except TypeError:
        warnings.warn("StreamingHttpResponse must consume synchronous iterators ...")
        # sync iterator. Consume via sync_to_async and yield via async generator.
        for part in await sync_to_async(list)(self.streaming_content):
            yield part

 

본문이 async iterator 이면 첫 번째 경로의 async for 가 그대로 동작해서 한 조각씩 나간다. 그런데 본문이 sync 이면 streaming_contentasync for 를 지원하지 않는 map 객체라서 TypeError 가 나고, fallback 경로로 떨어진다. 여기서 sync_to_async(list) 로 전체를 리스트로 만든 뒤에야 yield 를 시작한다. list() 가 끝나려면 sync generator 가 끝까지 실행되어야 하므로, 모든 토큰이 생성 완료된 후 한 번에 나가는 것이다.

 

이전 글에서 본 WSGI 쪽 (__iter__) 도 같은 구조의 대칭이다. async iterator 를 받으면 async_to_sync 로 전체를 소비한 뒤 넘긴다.

 

event loop 관점에서 보면 어쩔 수 없는 선택이다. sync iterator 의 next() 는 블로킹 호출이라 event loop 위에서 한 개씩 부르면 loop 전체가 멈춘다. Django 는 안전하게 스레드로 보내서 통째로 소비하는 쪽을 택했고, 그 대가로 스트리밍이 사라진다.

sync node 인데 astream() 은 왜 실시간인가

해결로 넘어가기 전에 확인할 것이 있다. 위 매트릭스에서 graph.astream() 은 ASGI 에서 실시간 스트리밍이 됐다. 그런데 node 는 time.sleep 으로 블로킹하는 sync 함수 그대로였다. 어떻게 실시간 스트리밍이 가능한 것일까.

 

LangGraph 는 async API (astream()) 로 실행할 때 sync node 를 event loop 에서 직접 부르지 않고 워커 스레드에 위임한다 (run_in_executor). 이때 contextvars 를 스레드로 복사하기 때문에 executor 스레드 안에서도 get_stream_writer() 가 동작한다. node 가 스레드에서 블로킹하며 writer(token) 을 호출하면 값이 내부 큐에 들어가고, event loop 쪽의 astream() 소비자가 큐를 await 하다가 값이 들어오는 즉시 yield 한다.

 

여기서 핵심은 생산과 소비가 동시에 진행된다는 점이다. node 를 실행하는 워커 스레드와 event loop 는 서로 다른 스레드라서, node 가 다음 토큰을 만들기 위해 블로킹하는 동안에도 event loop 는 멈추지 않고 이미 큐에 들어온 토큰을 클라이언트로 전송한다.

 

워커 스레드:  sleep(0.3) → writer(t1) → sleep(0.3) → writer(t2) → ...
event loop:  큐 대기 ......... t1 전송 ... 큐 대기 ....... t2 전송 ...

 

앞의 원인 섹션에서 본 버퍼링 (Django 의 sync_to_async(list)) 도 스레드에 위임하는 것은 같다. 차이는 전달 단위다. list() 는 generator 가 전부 끝난 뒤 한 번에 넘기고, 큐는 토큰이 생길 때마다 한 개씩 넘긴다. 실시간 스트리밍의 조건은 동시 실행에 더해 한 개씩 전달이라는 것이다. 그래서 sync node 그대로도 astream() 이면 토큰 단위 스트리밍이 되고, node 를 전부 async 로 다시 짤 필요가 없다.

해결

1. graph.astream() 으로 전환

graph 호출부를 async API 로 바꾸는 것이 정석이다. 위에서 본 것처럼 node 는 sync 인 채로도 된다. LangGraph 가 스레드 위임을 해주기 때문이다. view 의 async generator 가 async for 로 소비하므로 ASGI 와 매칭되어 버퍼링이 사라진다. 성능을 위해 async 를 도입한다는 원래 목적과도 맞는 방향이다. 서버 전환 (ASGI) 에서 멈추지 않고 graph 호출부까지 async 로 맞춰야 전환이 완성되는 것이다.

 

이후 여유가 될 때 node 를 하나씩 async 화 (async def + async LLM 클라이언트) 하면 스레드 홉이 줄고 병렬 node 의 I/O 대기가 겹쳐지는 이득이 있다. 단 두 가지 주의할 점이 있다.

  • async node 를 하나라도 넣으면 sync API (invoke()/stream()) 로는 그 graph 를 실행할 수 없다. 호출부 전환이 선행되어야 하는 이유이다.
  • async node 안에서 sync 블로킹 호출 (sync ORM, time.sleep, sync LLM 클라이언트) 을 하면 스레드가 아니라 event loop 전체가 멈춘다. 같은 프로세스의 모든 요청이 함께 정지하므로, async node 안의 sync 코드는 반드시 sync_to_async 로 감싸야 한다.

2. 호출부를 못 바꿀 때: per-item 브리지

graph.stream() 을 유지해야 하는 상황이라면 (예: 스트림을 소비하는 서비스 계층이 sync generator 를 전제로 짜여 있는 경우), 소비 지점에서 한 개씩 스레드로 꺼내는 브리지를 만들 수 있다. 위 예제의 graph_bridge 가 이 패턴이다.

 

token = await sync_to_async(next, thread_sensitive=True)(events, _STREAM_END)

 

핵심은 sync_to_async(list) 처럼 전체를 꺼내는 것이 아니라 sync_to_async(next) 로 한 개씩 꺼낸다는 것이다. async generator 안에서 sync generator 의 next() 를 스레드로 한 번씩 호출하므로, 겉모습은 ASGI 와 매칭되는 async iterator 이고 내용물은 한 개씩 실시간으로 넘어온다. 실측 결과도 300ms 간격 스트리밍이다.

 

10.459 data: {"token": "'책' "}
10.759 data: {"token": "문의 "}
11.060 data: {"token": "확인했습니다. "}

 

Django 가 fallback 에서 하는 일 (스레드 위임) 을 통째가 아니라 한 개씩 하도록 직접 구현한 것이라고 보면 된다. LangGraph 의 astream() + sync node 조합이 내부적으로 하는 일과도 같은 구조이다.

직접 재현하기

위 예제를 그대로 재현해보고 싶다면 아래와 같이 코드를 받아서 docker 서비스를 실행하면 된다.

 

git clone https://github.com/rmk1075/parcel-bot && cd parcel-bot
docker compose up -d --build
# :8000 = WSGI, :8001 = ASGI
# /graph/sync/ /graph/async/ /graph/bridge/ 를 이전 글의 타임스탬프 curl 로 비교
반응형