ThreadPoolExecutor와 asyncio중 무엇을 사용해야할까?

https://pypi.org/project/easy-rag-llm/

https://pypi.org/project/easy-rag-llm/

 

위 라이브러리를 개발하던 중 병렬처리에 대해 고민할 지점이 생겨 이를 기록하고자 한다. 

 

1. 상황

개발중인건 LLM RAG를 돕는 라이브러리다. 입력된 자료를 조각내고 각 조각을 임베딩하는 속도를 개선하기 위해 멀티스레딩을 적용하고자 생각했고, 파이썬의 ThreadPoolExecutor를 사용했다. v1.1.1 까지 해당 방식으로 배포했고, 작은 분량의 소스만을 입력해 사용하느라 인지하지 못했으나 500장 이상 대용량의 pdf가 첨부되면 기대한 바와 달리 3분 이상의 임베딩 시간이 소요되는 것을 체감하였다. 이에 기존에 채택한 멀티스레딩이 적합한 방식인지 의구심을 가졌고, 조사하면서 기록한다.

 

2. ThreadPool과 asyncio의 차이

 둘다 병렬처리에 이용된다는 것은 인지하고 있었으나 병렬처리하면 멀티스레드, asyncio는 비동기처리 정도만 생각하고 있었던 것 같다.

- ThreadPool, 

1) 작업을 여러 스레드로 나누어 동시 작업한다.

2) IO Bound 작업에 적합하다. 

3) Python에서는 GIL(Global Interpreter Lock)때문에 CPU Bound task에선 멀티스레딩이 제한된다.

우선 인지하고 있는 내용들이다.

 

- asyncio,

1) 단일 스레드에서 동작하지만 비동기 IO와 코루틴을 제공한다.

2) 이벤트 루프 기반으로 작업이 관리된다. 즉, 작업 대기시간동안 이벤트루프가 다른 작업을 실행한다. 

3) IO 작업에서 유리하다.

 

별로 도움이 안됐다. 인지하고 있는 내용들이었으며, 문제를 해결하는데는 도움이 되지 않았다.

 

3. 코드 용도 다시 검토

속도 개선이 필요한 부분은 1) pdf에서 텍스트를 추출하고 청크를 나누는 작업, 2) 임베딩 작업 두가지다. 하지만  1번 작업은 만족스러운 체감 시간을 보여줬고, 2번에서 문제를 느꼈다. 차이를 고민해보자.

 

1) 확실히 IO Bound다. asyncio로 처리한다고 생각해보자. 비동기 래퍼로 한번 더 싸야하는 수고스러움이 필요. 우선 패스.

2) 임베딩 생성. 청크 받고 -> pre_embedded model 호출. 이게 키일 것 같다. 호출은 네트워크 작업이다. 멀티 스레드 나눠봤자 네트워크 기다리는 건 변하지 않는다. asyncio 도입하면? 네트워크 기다리는 동안도 작업할 수 있으니 병렬처리를 더 많이 할 수 있다.

 

가설: 네트워크 대기 시간이 있는 부분에서 멀티스레딩 대신 이벤트 루프 기반 비동기처리를 하면 속도가 개선 될 것이다. 검증해보자.

 

4. 검증

454페이지 분량의 AWS batch 안내서를 임베딩했다.

1) 멀티스레드 적용: v.1.1.1

Generating Embeddings: 100%|█████████████████████████████████████████████| 1341/1341 [01:15<00:00, 17.82it/s]

 

==> 1분 15초 걸렸다.

 

2) 비동기적용: v1.1.2

-> semaphore 20으로 비동기 20개 제한

Generating Embeddings: 100%|█████████████████████████████████████████████| 1341/1341 [00:58<00:00, 22.85it/s]

 

==> 58초. 약 22.7% 속도 개선

 

세마포어 20개부터는 더 이상 개선되지 않았다.  여기선 ratelimit에 걸리기전 적절한 수준이 20으로 판단된다. 

 

 

5. 추가 개선

현재까지 구현 방식은 api 요청한번에 하나의 텍스트를 보내 임베딩하는 방식이었다. 레이트 리밋은 api 호출 횟수에 대한 제한이니 굳이 하나씩 보낼 필요없다. 배치로 묶어서 처리하였다. 배치사이즈는 20으로 고정했고,

Generating Embeddings (Batched): 100%|██████████████████████████████████████████████████| 68/68 [00:08<00:00,  8.43it/s]

 

==> 8초 까지 개선되었다.