본문으로 건너뛰기

[Open WebUI] 외부 임베딩 API 호출을 병렬화하여 50배 성능 향상

PR 링크: open-webui/open-webui#19296 상태: Merged | 변경: +289 / -118

들어가며

RAG(Retrieval-Augmented Generation) 파이프라인에서 문서를 청크로 분할한 후 임베딩을 생성하는 과정은 전체 처리 시간의 상당 부분을 차지한다. Open WebUI는 OpenAI, Azure OpenAI, Ollama 등 외부 임베딩 API를 지원하는데, 기존 구현에서는 배치를 순차적으로 처리했다. 6,000개 청크를 batch_size=1로 처리하면 6,000번의 순차 HTTP 요청이 발생하여 약 5분이 소요되었다.

핵심 코드 분석

Before: 순차 배치 처리 (동기)

# ThreadPoolExecutor로 하이브리드 검색만 병렬화
with ThreadPoolExecutor() as executor:
    future_results = [executor.submit(process_query, cn, q) for cn, q in tasks]
    task_results = [future.result() for future in future_results]

After: asyncio.gather로 전면 병렬화

# 모든 배치를 asyncio.gather로 동시 실행
async def generate_multiple_async(query, prefix, user, func_async, embedding_batch_size):
    if isinstance(query, list):
        batches = [
            query[i : i + embedding_batch_size]
            for i in range(0, len(query), embedding_batch_size)
        ]
        tasks = [
            func_async(batch, prefix=prefix, user=user)
            for batch in batches
        ]
        batch_results = await asyncio.gather(*tasks)
        
        embeddings = []
        for batch_embeddings in batch_results:
            if isinstance(batch_embeddings, list):
                embeddings.extend(batch_embeddings)
        return embeddings
    else:
        return await func_async([query], prefix=prefix, user=user)

핵심 변경 사항:

  1. generate_openai_batch_embeddings_async, generate_azure_openai_batch_embeddings_async, generate_ollama_batch_embeddings_async 등 엔진별 비동기 함수 추가
  2. aiohttp.ClientSession을 사용한 비동기 HTTP 요청
  3. query_collection, query_collection_with_hybrid_search 등 상위 함수를 async로 전환
  4. CPU 바운드인 Sentence Transformers는 asyncio.to_thread로 스레드 풀에서 실행

왜 이게 좋은가

  1. 30~50배 속도 향상: 6,000개 청크 기준, 순차 처리 시 약 5분이 걸리던 것이 10초 이내로 단축된다. 모든 배치가 동시에 API로 전송되기 때문이다.
  2. async 네이티브 전환: requests 라이브러리 대신 aiohttp를 사용하고, 전체 파이프라인을 async로 전환하여 이벤트 루프를 블로킹하지 않는다.
  3. CPU 바운드 분리: 로컬 Sentence Transformers 모델은 asyncio.to_thread로 스레드 풀에서 실행하여 이벤트 루프 차단을 방지한다. IO 바운드(외부 API)와 CPU 바운드(로컬 모델)를 적절히 분리한 설계다.

대규모 문서를 RAG으로 처리할 때 병목이 되는 임베딩 생성 단계를 근본적으로 개선한 PR이다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글