본문으로 건너뛰기

[Open WebUI] 채팅 메시지 마이그레이션을 스트리밍+배치 처리로 전환하여 메모리 폭발 방지

PR 링크: open-webui/open-webui#21542 상태: Merged | 변경: +78 / -38

들어가며

Open WebUI의 채팅 메시지 테이블 마이그레이션(8452d01d26d7)은 기존 채팅 데이터를 새 chat_message 테이블로 백필하는 작업을 수행합니다. 기존 구현은 .fetchall()로 모든 채팅 행을 메모리에 로드한 후, 각 메시지마다 개별 INSERT + SAVEPOINT를 실행했습니다. 대규모 인스턴스에서는 수백만 행이 메모리에 올라가고, 메시지당 SAVEPOINT 오버헤드가 누적되어 마이그레이션이 극도로 느려졌습니다.

핵심 코드 분석

Before: 전체 로드 + 건별 INSERT

# 모든 채팅을 메모리에 로드
chats = conn.execute(
    sa.select(chat_table.c.id, chat_table.c.user_id, chat_table.c.chat)
    .where(~chat_table.c.user_id.like("shared-%"))
).fetchall()

for chat_row in chats:
    for message in messages:
        savepoint = conn.begin_nested()  # 메시지마다 SAVEPOINT
        try:
            conn.execute(sa.insert(chat_message_table).values(...))
            savepoint.commit()
        except Exception:
            savepoint.rollback()

After: 스트리밍 + 배치 INSERT

BATCH_SIZE = 5000

# 서버 사이드 커서로 스트리밍
result = conn.execute(
    sa.select(...)
    .execution_options(yield_per=1000, stream_results=True)
)

messages_batch = []
for chat_row in result:
    for message in messages:
        messages_batch.append({...})
        if len(messages_batch) >= BATCH_SIZE:
            inserted, failed = _flush_batch(conn, chat_message_table, messages_batch)
            messages_batch.clear()

def _flush_batch(conn, table, batch):
    """벌크 INSERT 시도, 실패 시 건별 fallback"""
    savepoint = conn.begin_nested()
    try:
        conn.execute(sa.insert(table), batch)  # 벌크 INSERT
        savepoint.commit()
        return len(batch), 0
    except Exception:
        savepoint.rollback()
        # 건별 fallback으로 문제 행만 건너뜀
        for msg in batch:
            sp = conn.begin_nested()
            try:
                conn.execute(sa.insert(table).values(**msg))
                sp.commit()
            except Exception:
                sp.rollback()

왜 이게 좋은가

  1. 메모리 효율: yield_per=1000과 stream_results=True로 서버 사이드 커서를 사용하여 전체 데이터를 메모리에 올리지 않는다.
  2. SAVEPOINT 오버헤드 감소: 5000건 단위 벌크 INSERT로 SAVEPOINT 횟수를 1/5000로 줄인다.
  3. graceful fallback: 벌크 INSERT 실패 시 건별 fallback으로 문제 행만 건너뛰어 마이그레이션 전체가 중단되지 않는다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글