스트리밍 기반 임베딩으로 메모리 최적화하기

Code-Study 프로젝트에서의 경험을 다룬 글입니다.

RAG 챗 봇을 만들기 위해 PDF 파일을 받아 백터DB에 임베딩하는 기능이 필요했습니다.

이번 글에서는 제가 적용한 PDF 스트리밍 기반 임베딩 방식을 소개하고, 기존 방식 대비 메모리 사용량을 어떻게 줄였는지 공유해보려 합니다

Spring AI -> PDFBox 변경

처음에는 Spring AI가 제공하는 PagePdfDocumentReader를 사용했습니다.
서비스가 이미 Spring AI에 강하게 의존하고 있고, 잘 추상화된 기능이 있다면 굳이 새로 구현할 이유가 없다고 생각했기 때문입니다.
그래서 다음과 같이 임시 디렉토리에 파일을 저장하고, PagePdfDocumentReader를 이용해 문서를 읽어 처리했습니다.

로컬에서 간단히 테스트했을 때는 정상적으로 데이터가 잘 적재되는 것을 확인할 수 있었습니다.
(테스트 데이터로는 제 이력서를 사용했습니다.)

이력서를 임베딩한 모습

이번에는 실제 서비스 상황을 가정해 아틸러리(Artillery)로 부하 테스트를 진행했습니다.

결과는 예상 밖이었습니다.
효율적이진 못해도 전체가 성공할줄은 알았는데 에러 응답이 절반정도 발생했고, 안정적인 처리가 불가능했습니다.

파일을 다루는 로직이다 보니 메모리와 관련한 문제일것이라 생각했고 JFR로 메모리 부분을 모니터링 해보니 다음과 같았다. 

메모리 할당을 GB 단위로 하고있는 걸 확인함, 아마 한번에 너무 많은 메모리를 할당해 써서 에러 응답이 나오는 것 같음 

원인을 분석해보니 PagePdfDocumentReader는 내부적으로 Apache PDFBox의 PDFParser를 사용해 전체 PDF 파일을 한 번에 로딩하고 있었습니다.

또 내부 구현을 살펴본 결과 저 방법외에 다른 방법을 지원하지 않음을 확인했다 

 

개선 시도: Apache PDFBox 기반 직접 구현

 

이 문제를 해결하기 위해 PagePdfDocumentReader 대신 Apache PDFBox를 직접 활용한 스트리밍 방식을 구현했습니다.

개선 아이디어는 단순합니다.

• PDF 파일을 페이지 단위로 읽고
• 읽은 내용을 바로 청크 분할
• 일정 배치 단위로 모아 벡터DB에 flush

즉, “읽은 건 바로 처리하고 버린다”는 원칙입니다.

이렇게 하면 메모리에 전체 PDF를 쌓아두지 않고, 데이터 양을 제한할 수 있으며
테스트 결과 또한 확연히 달랐습니다.