FastAPI와 PyTorch 기반 실시간 비동기 AI 서빙 환경에서 동적 바칭 시 발생하는 VRAM 누수 및 OOM 장애 트러블슈팅

🌐 English Abstract

This article provides an in-depth technical guide to diagnosing and resolving GPU VRAM memory leaks and Out-Of-Memory (OOM) errors during real-time asynchronous AI model serving using FastAPI and PyTorch with Dynamic Batching. It explores the interaction between Python's asyncio event loop, PyTorch's CUDA Caching Allocator, and reference counting mechanisms. Readers will learn actionable strategies, including custom memory profiling, torch.inference_mode() application, context manager optimization, and robust background worker architectural patterns to maintain high-throughput production stability.

FastAPI와 PyTorch를 결합한 실시간(Real-time) 비동기(Asynchronous) AI 모델 서빙(Model Serving) 아키텍처는 고성능 웹 서비스 구현 시 널리 채택되는 조합입니다. 특히 트래픽 처리량(Throughput)을 극대화하기 위해 여러 요청을 하나로 묶어 GPU에 전송하는 동적 바칭(Dynamic Batching) 기법은 필수적인 최적화 요소로 자리 잡았습니다.

그러나 서비스 운영 중 특정 시점에 VRAM(Video Random Access Memory) 사용량이 꾸준히 증가하다 결국 OOM(Out Of Memory) 장애를 일으키는 현상을 자주 목격하게 됩니다. 재시작 시점에는 정상 동작하다가 고하중 상태에서 점진적으로 메모리가 증발하는 이 문제는 원인 추적이 매우 까다롭습니다. 본 글에서는 이 문제의 근본적인 원인을 PyTorch 내부 메커니즘 및 Python 비동기 이벤트 루프(Event Loop) 수준에서 분석하고, 완벽한 해결책과 실무 트러블슈팅 전략을 공유합니다.

1. 동적 바칭과 비동기 이벤트 루프에서의 문제 발생 원인

FastAPI 애플리케이션 내에서 `asyncio.Queue` 기반으로 동적 바칭을 구현할 때, 메모리 누수가 발생하는 원인은 크게 세 가지 메커니즘이 복합적으로 작용하기 때문입니다.

1.1 연산 그래프(Computation Graph)의 미세한 참조 유지

PyTorch는 기본적으로 연산을 수행할 때 미분 계수 계산을 위한 연산 그래프(Computation Graph)를 생성합니다. 추론(Inference) 단계에서 `@torch.no_grad()` 또는 `@torch.inference_mode()`를 엄격하게 적용하지 않으면, 텐서(Tensor)가 그래프 참조를 계속 유지합니다. 동적 바칭 구현부에서 비동기 퓨처(Future) 객체에 텐서 출력을 직접 전달할 때 이 참조가 제거되지 않고 이벤트 루프 메모리에 상주하게 됩니다.

1.2 PyTorch CUDA Caching Allocator의 메모리 파편화(Memory Fragmentation)

PyTorch는 GPU 메모리 할당 속도를 높이기 위해 자체적인 **CUDA 메모리 캐싱 할당자(CUDA Caching Allocator)**를 사용합니다. `torch.cuda.memory_allocated()`로 측정되는 실제 텐서 점유량은 줄어들더라도, `torch.cuda.memory_reserved()`로 표시되는 PyTorch 예약 메모리는 운영체제(NVIDIA Driver)로 반환되지 않습니다. 동적 바칭으로 인해 가변적인 크기의 텐서가 연속하여 할당 및 해제될 때 메모리 파편화(Memory Fragmentation)가 심화되어 실제 사용량보다 더 일찍 OOM이 발생합니다.

1.3 Python Garbage Collection과 비동기 Task 간 순환 참조

비동기 백그라운드 태스크(Background Task)가 바치 단위로 배치 데이터를 팝(Pop)하여 처리할 때, 예외(Exception) 처리 구문이나 로깅 라이브러리 내부에서 텐서 참조를 `sys.exc_info()` 또는 캡처된 스택 프레임에 남겨두는 경우가 있습니다. Python의 세대별 가비지 컬렉터(Garbage Collector)가 이를 즉시 정률 해제하지 못하면 GPU VRAM 참조 역시 함께 묶이게 됩니다.

2. VRAM 누수 진단 및 트러블슈팅 단계별 가이드

장애 원인을 규명하기 위해서는 단순 모니터링 기법을 넘어 PyTorch 내부 메모리 상태를 정밀하게 추적해야 합니다.

2.1 Memory Snapshot을 활용한 정밀 분석

PyTorch 2.0 이상에서는 VRAM 할당 스냅샷을 HTML 형태로 시각화할 수 있는 강력한 프로파일링 기능을 제공합니다. 누수가 의심되는 시점에 아래 코드를 실행하여 텐서 할당 궤적을 캡처할 수 있습니다.

import torch

# 스냅샷 기록 시작
torch.cuda.memory._record_memory_history(max_entries=100000)

# 바칭 및 추론 로직 실행 후 스냅샷 저장
try:
    # 예: OOM 발생 시점 또는 특정 주기
    torch.cuda.memory._dump_snapshot("vram_leak_snapshot.pickle")
except Exception as e:
    print(f"Failed to dump snapshot: {e}")

생성된 `.pickle` 파일은 PyTorch 공식 Visualizer 웹페이지(`pytorch.org/memory_viz`)에 업로드하여 어떤 레이어와 어떤 변수가 메모리를 반환하지 않고 있는지 정확히 추적할 수 있습니다.

2.2 Allocated Memory vs Reserved Memory 실시간 로깅

FastAPI 미들웨어 또는 미들웨어 헬스체크 엔드포인트에 메모리 상태 로깅을 추가하여 파편화 여부를 판단합니다.

from fastapi import FastAPI
import torch

app = FastAPI()

@app.get("/health/memory")
async def get_cuda_memory_status():
    if not torch.cuda.is_available():
        return {"cuda": False}
    
    allocated = torch.cuda.memory_allocated() / (1024 ** 2)  # MB 단위
    reserved = torch.cuda.memory_reserved() / (1024 ** 2)    # MB 단위
    max_allocated = torch.cuda.max_memory_allocated() / (1024 ** 2)
    
    return {
        "allocated_mb": round(allocated, 2),
        "reserved_mb": round(reserved, 2),
        "max_allocated_mb": round(max_allocated, 2),
        "fragmentation_ratio": round((reserved - allocated) / reserved, 2) if reserved > 0 else 0
    }

3. 동적 바칭 아키텍처 결함 수정 및 최적화 솔루션

문제를 근본적으로 해결하기 위해 FastAPI 비동기 큐 처리 로직과 PyTorch 추론 루프를 아래 표준 패턴으로 개편해야 합니다.

3.1 완벽한 Inference Mode 및 Tensor Detach 처리

추론 코드 전반에 `torch.inference_mode()`를 적용하고, 텐서 결과를 CPU로 이관 시 반드시 `.detach().cpu()`를 적용하여 연산 그래프와의 연결을 완전히 끊어야 합니다.

import asyncio
import torch
from typing import List, Dict, Any

class AsyncBatchPredictor:
    def __init__(self, model: torch.nn.Module, max_batch_size: int = 16, max_delay: float = 0.05):
        self.model = model
        self.max_batch_size = max_batch_size
        self.max_delay = max_delay
        self.queue = asyncio.Queue()
        asyncio.create_task(self._batch_worker())

    async def predict(self, input_tensor: torch.Tensor) -> Any:
        loop = asyncio.get_running_loop()
        future = loop.create_future()
        await self.queue.put((input_tensor, future))
        return await future

    async def _batch_worker(self):
        while True:
            batch = []
            futures = []
            
            # 첫 번째 요청 대기
            tensor, fut = await self.queue.get()
            batch.append(tensor)
            futures.append(fut)
            
            # Max Delay 동안 추가 요청 수집 (Dynamic Batching)
            start_time = asyncio.get_event_loop().time()
            while len(batch) < self.max_batch_size:
                timeout = self.max_delay - (asyncio.get_event_loop().time() - start_time)
                if timeout <= 0:
                    break
                try:
                    tensor, fut = await asyncio.wait_for(self.queue.get(), timeout=timeout)
                    batch.append(tensor)
                    futures.append(fut)
                except asyncio.TimeoutError:
                    break

            # 추론 수행
            await self._run_inference(batch, futures)

    async def _run_inference(self, batch: List[torch.Tensor], futures: List[asyncio.Future]):
        try:
            # 1. Tensor Stack 및 GPU 전송
            input_batch = torch.stack(batch).cuda(non_blocking=True)
            
            # 2. Inference Mode 적용 (no_grad 대비 연산 overhead 낮고 메모리 안전)
            with torch.inference_mode():
                outputs = self.model(input_batch)
                # CPU 이관 및 numpy/python 객체 변환을 통해 GPU 참조 완전 제거
                outputs_cpu = outputs.detach().cpu()

            # 3. 개별 Future 결과 할당
            for i, fut in enumerate(futures):
                if not fut.cancelled():
                    fut.set_result(outputs_cpu[i])
                    
        except Exception as e:
            for fut in futures:
                if not fut.cancelled():
                    fut.set_exception(e)
        finally:
            # 4. 명시적 변수 정리 (GC 유도 및 CUDA 캐시 파편화 방지)
            del batch, input_batch, outputs, outputs_cpu
            # 필요시 주기적인 gc.collect() 수행

3.2 PyTorch Memory Allocator 환경 변수 최적화

동적 바칭처럼 가변적인 크기의 메모리가 반복 할당될 때는 PyTorch의 캐시 할당 정책을 조정하여 파편화를 줄일 수 있습니다. 애플리케이션 실행 전 환경 변수를 설정합니다.

• `PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128`: 큰 메모리 블록이 파편화되는 것을 방지하기 위해 분할 가능한 최대 블록 크기를 제한합니다. 이 설정을 적용하면 메모리 재사용 효율이 급격히 증가합니다.
• `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (PyTorch 2.1 이상): CUDA Virtual Memory Management를 활용해 메모리 세그먼트를 가변적으로 확장함으로써 메모리 파편화를 거의 완벽하게 방지합니다.

3.3 주기적인 CUDA Cache Clear의 올바른 사용법

`torch.cuda.empty_cache()`는 VRAM을 강제로 OS에 반환하지만, 호출 자체가 GPU 동기화(Synchronization)를 유발하여 커널 처리 속도를 대폭 저하시킵니다. 따라서 모든 추론 요청마다 실행하면 안 되며, 특정 조건(예: 캐시 파편화율이 일정 수준 이상이거나 VRAM 사용량이 85%를 초과할 때) 백그라운드 스레드에서 제한적으로 호출해야 합니다.

# 매 요청마다 실행 금지 (Overhead 극심)
# 파편화 지표가 특정 임계치를 넘을 때만 선택적 호출
def smart_empty_cache():
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    if reserved > 0 and (reserved - allocated) / reserved > 0.35: # 파편화율 35% 초과 시
        torch.cuda.empty_cache()

4. 프로덕션 멀티 프로세스 서빙 아키텍처 고려사항

FastAPI 메인 단일 프로세스 내부에서 GPU 연산을 직접 수행하는 아키텍처는 Python의 **GIL(Global Interpreter Lock)** 및 이벤트 루프 블로킹 위험이 항상 존재합니다. 따라서 운영 환경(Production)에서는 더욱 강력한 아키텍처 분리가 권장됩니다.

• Process Isolation (IPC 통신): FastAPI 웹 서버 프로세스와 PyTorch 추론 전용 워커 프로세스를 완전히 분리하고 `SharedMemory` 또는 `gRPC`를 통해 통신합니다. 이 경우 추론 워커 프로세스에 메모리 누수가 발생하더라도 특정 시점에 워커만 Graceful Restart할 수 있습니다.
• Triton Inference Server 도입: 동적 바칭 및 GPU 메모리 관리를 Python 레이어에서 처리하는 대신, NVIDIA Triton Inference Server나 vLLM 등의 C++ 기반 전용 서빙 엔진으로 위임합니다. FastAPI는 순수 게이트웨이(API Gateway) 역할만 담당함으로써 VRAM 메모리 관리 문제를 근본적으로 격리시킬 수 있습니다.

5. 결론 및 요약

FastAPI와 PyTorch 기반 동적 바칭 시스템에서 VRAM 누수 및 OOM 장애를 예방하기 위한 핵심 요약은 다음과 같습니다.

1. `torch.inference_mode()` 사용 필수: 연산 그래프 생성을 차단하여 텐서 참조가 유지되는 것을 방지합니다.
2. 텐서 반환 시 명시적 이관 및 변수 삭제: `detach().cpu()`를 적용해 GPU 점유를 즉시 해제하고 `del` 구문으로 참조를 정리합니다.
3. `PYTORCH_CUDA_ALLOC_CONF` 설정: `expandable_segments:True` 또는 `max_split_size_mb` 설정을 통해 메모리 파편화를 최소화합니다.
4. 프로파일링 관측성(Observability) 확보: Memory Snapshot 기능을 활용하여 실제 누수 텐서의 할당 지점을 주기적으로 모니터링합니다.

위 트러블슈팅 패턴과 최적화 설정을 적용하면 실시간 비동기 AI 서빙 환경에서도 OOM 장애 없이 안정적이고 높은 트래픽 처리량을 유지할 수 있습니다.