[PYTHON] GPU 메모리 누수를 방지하는 2가지 컨텍스트 매니저 설계 방법과 해결책

CUDA Out of Memory (OOM)

 

딥러닝 모델을 개발하고 서비스에 배포할 때 개발자를 가장 괴롭히는 문제 중 하나는 바로 'CUDA Out of Memory (OOM)' 에러입니다. 모델의 파라미터가 커지고 배치 사이즈가 늘어남에 따라 제한된 GPU 자원을 관리하는 것은 선택이 아닌 필수입니다. 특히 파이썬의 예외 처리 과정에서 GPU 메모리가 제대로 해제되지 않고 고착되는 현상은 서비스의 안정성을 크게 해칩니다.

본 포스팅에서는 파이썬의 표준 라이브러리인 contextlib를 활용하여, 복잡한 try...finally 구문 없이도 GPU 리소스를 선언적으로 관리할 수 있는 **고급 컨텍스트 매니저(Context Manager)** 설계 기법을 심층 분석합니다. 이를 통해 메모리 할당과 해제의 자동화를 구현하는 실전적인 해결책을 제시하겠습니다.

---

1. 일반적인 리소스 관리 vs Context Manager 기반 관리의 차이

GPU 메모리 관리에서 컨텍스트 매니저를 도입했을 때 얻을 수 있는 구조적 이점을 비교 요약합니다.

비교 항목	전통적인 방식 (Manual)	컨텍스트 매니저 방식 (Automated)
가독성	try-finally 블록으로 인한 코드 비대화	with 문을 활용한 간결하고 직관적인 구조
안전성	개발자의 실수로 해제 누락 가능성 높음	블록 탈출 시 예외 발생 여부와 상관없이 자동 해제
재사용성	동일한 해제 로직을 매번 수동 작성	클래스나 제네레이터 형태로 캡슐화하여 재사용
메모리 최적화	누적된 캐시 수동 정리 필요	특정 작업 후 자동으로 empty_cache() 수행 가능
해결 과제	변수 스코프 관리 및 수동 del 호출	contextlib을 통한 정교한 스코프 제어

---

2. contextlib을 활용한 GPU 메모리 관리 실전 Example (7+)

현업에서 즉시 적용 가능한 7가지 이상의 고급 예제 코드를 통해 GPU 메모리 관리의 정수를 살펴봅니다. (PyTorch 기준)

Ex 1. @contextmanager를 활용한 자동 캐시 정리 (Basic)

import torch
from contextlib import contextmanager

@contextmanager
def gpu_memory_config():
    """작업 후 남아있는 미사용 캐시 메모리를 강제로 비워주는 매니저"""
    try:
        yield
    finally:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            print("GPU Cache Cleared Successfully.")

# 사용 예시
with gpu_memory_config():
    # 대규모 행렬 연산 수행
    data = torch.randn(10000, 10000).cuda()
    result = data @ data.T

Ex 2. 추론 모드 전용 하이브리드 컨텍스트 매니저

from contextlib import ExitStack

@contextmanager
def inference_mode(model):
    """모델을 평가 모드로 전환하고 그래디언트 계산을 완전히 차단"""
    model.eval()
    with torch.no_grad():
        try:
            yield
        finally:
            model.train() # 필요 시 다시 학습 모드로 복구

Ex 3. 예외 발생 시 특정 장치 메모리 강제 덤프 및 해제

import sys

@contextmanager
def monitor_gpu_oom():
    try:
        yield
    except RuntimeError as e:
        if 'out of memory' in str(e):
            print("OOM Detected! Emergency memory flush initiated.")
            for obj in gc.get_objects():
                try:
                    if torch.is_tensor(obj) and obj.is_cuda:
                        del obj
                except:
                    pass
            torch.cuda.empty_cache()
        raise e

Ex 4. 특정 GPU 장치 임시 할당 및 복구 매니저

@contextmanager
def use_device(device_id: int):
    """특정 작업에만 임시로 다른 GPU를 사용하도록 설정"""
    if not torch.cuda.is_available():
        yield
        return
    
    prev_device = torch.cuda.current_device()
    torch.cuda.set_device(device_id)
    try:
        yield
    finally:
        torch.cuda.set_device(prev_device)

Ex 5. 중첩된 리소스를 관리하는 ExitStack 활용법

from contextlib import ExitStack

def process_multi_gpu(model_a, model_b):
    """여러 컨텍스트를 동적으로 관리해야 할 때의 해결책"""
    with ExitStack() as stack:
        stack.enter_context(gpu_memory_config())
        stack.enter_context(torch.no_grad())
        
        # 복잡한 멀티 모델 추론 로직 수행
        out_a = model_a(input_data)
        out_b = model_b(out_a)
        return out_b

Ex 6. 텐서 프로파일링을 위한 메모리 스냅샷 매니저

@contextmanager
def trace_gpu_usage(label: str):
    """특정 구간의 메모리 증감량을 측정하여 로깅"""
    torch.cuda.synchronize()
    before = torch.cuda.memory_allocated()
    try:
        yield
    finally:
        torch.cuda.synchronize()
        after = torch.cuda.memory_allocated()
        print(f"[{label}] Memory Delta: {(after - before) / 1024**2:.2f} MB")

Ex 7. contextlib.closing을 응용한 커스텀 리소스 종료

from contextlib import closing

class RemoteGPULoader:
    def __init__(self, url): self.url = url
    def load(self): print("Loading model from remote..."); return self
    def close(self): print("Remote connection closed and GPU buffer freed.")

# 객체에 close 메서드가 있을 때 안전하게 종료 보장
with closing(RemoteGPULoader("http://ai-server").load()) as loader:
    # 모델 로드 및 작업 수행
    pass

---

3. GPU 메모리 관리 설계 시 핵심 해결 과제 (제약 사항)

컨텍스트 매니저를 설계할 때 반드시 고려해야 할 3가지 기술적 해결책입니다.

1. Python Garbage Collector와의 동기화: 파이썬 객체가 삭제되어도 실제 CUDA 메모리는 비워지지 않을 수 있습니다. finally 블록에서 torch.cuda.empty_cache()를 명시적으로 호출하는 설계가 필요합니다.
2. 비동기 커널 실행: GPU 연산은 비동기적으로 이루어지므로, 정확한 메모리 측정을 위해서는 torch.cuda.synchronize()를 적절히 배치하여 컨텍스트 경계를 명확히 해야 합니다.
3. 순환 참조 방지: 컨텍스트 매니저 내부에서 큰 텐서를 지역 변수로 잡고 있을 경우, 클로저나 예외 객체에 의해 참조가 유지되어 메모리가 해제되지 않을 수 있으므로 주의가 필요합니다.

---

4. 결론: 왜 contextlib인가?

단순히 try...finally를 사용하는 것보다 contextlib을 통해 컨텍스트 매니저를 설계하는 것은 '의도의 명확성' 때문입니다. 다른 팀원이 코드를 보았을 때, with gpu_cleanup(): 구문 하나만으로 해당 블록이 자원 집약적이며 안전하게 관리되고 있음을 즉시 파악할 수 있습니다.

결국 고도화된 AI 파이프라인의 안정성은 알고리즘만큼이나 이러한 리소스 거버넌스(Resource Governance) 설계에 달려 있습니다. 오늘 소개한 7가지 패턴을 실무에 적용하여 더 이상 OOM 에러로 고통받지 않는 견고한 서빙 환경을 구축하시기 바랍니다.

---

참고 출처

• Python Standard Library Documentation - contextlib
• PyTorch Documentation - CUDA Memory Management Best Practices
• Effective Python (2nd Edition) - Brett Slatkin (Item 64: Consider contextlib and with statements for reusable try/finally behavior)
• NVIDIA Developer Blog - Optimizing PyTorch Memory Usage