[Python] FastAPI 백그라운드 학습에서 메모리 사용량 급증을 줄이는 리팩터링 정리

FastAPI의 백그라운드 작업에서 대용량 학습 로직을 실행하면, 요청이 누적될수록 프로세스 메모리가 지속 증가하는 현상이 발생할 수 있습니다. 본 글에서는 실제 코드 흐름에서 메모리 사용량이 커지는 원인을 짚고, DataFrame/NumPy 처리 및 finally 정리 전략을 통해 메모리 부담을 줄이는 방법을 정리합니다.

문제 정의

다음과 같은 상황에서 메모리 사용량 급증 또는 누수처럼 보이는 현상이 발생할 수 있습니다.

• API 요청 처리 프로세스 내부에서 장시간 학습/전처리 로직을 실행한다.
• 시나리오(또는 그룹) 루프마다 대형 DataFrame을 copy()하여 반복 생성한다.
• NumPy 배열 변환 과정에서 astype() 등으로 대형 배열 복사가 여러 번 일어난다.
• 동일한 크기의 마스크(np.isnan)를 반복 계산하여 임시 객체가 계속 생성된다.
• 결과 딕셔너리에 큰 모델/설명기(직렬화 결과 등)를 중복 저장하여 참조가 오래 유지된다.
• 예외 발생 시 정리(cleanup)가 보장되지 않는다.

핵심은 “루프마다 새로 만들 필요가 있는 것은 데이터 전체 복사본이 아니라, 시나리오별로 달라지는 라벨/행 인덱스/선택된 피처 및 최종 입력(X, y)만”이라는 점입니다.

코드/방법

1) data_copy = data는 원본에 영향을 준다

아래는 복사가 아니라 참조입니다.

data_view = base_table

따라서 data_view에서 컬럼을 추가/수정하거나 .loc[...] = ... 같은 in-place 변경을 하면 base_table도 함께 변경됩니다. 원본을 불변으로 유지하려면 “데이터 전체 복사” 대신 “라벨/마스크만 분리”하는 접근이 필요합니다.

2) 시나리오 루프에서 DataFrame 전체 복사 대신 라벨 배열 분리

기존 패턴(메모리 부담 큼):

• 시나리오마다 대형 테이블을 copy()로 복제
• 라벨 컬럼을 추가하고 구간별로 .loc로 수정
• 마지막에 numpy 변환 및 astype로 재복사

개선 패턴(권장):

• 원본 테이블은 불변(복사하지 않음)
• 시나리오별로 label_array(NumPy)를 생성/갱신
• 시나리오별로 필요한 행/열만 선택해 X를 한 번만 생성

예시(개념 코드):

import numpy as np

def train_by_group(base_table, time_index, group_ids):
    n = len(base_table)

    for group_id in group_ids:
        # 시나리오별 라벨만 생성(가벼움)
        label_array = np.full(n, 2, dtype=np.int8)

        # 구간 규칙에 따라 label_array만 업데이트
        # label_array[mask_a] = 0
        # label_array[mask_b] = 1

        # 시나리오별로 사용할 컬럼/행 인덱스 계산
        feature_cols = compute_feature_columns(base_table, group_id)
        row_idx = compute_valid_rows(base_table, feature_cols)

        # 최종 입력 X, y만 생성(필요한 부분만 복사)
        X = base_table.iloc[row_idx, feature_cols].to_numpy(dtype=np.float32, copy=True)
        y = label_array[row_idx]

        run_model_training(X, y)

이 방식의 장점은 다음과 같습니다.

• 시나리오 개수만큼 대형 DataFrame 복사본을 만들지 않는다.
• 원본 데이터 오염 위험이 없다.
• 메모리 피크가 “필요한 X 생성 시점”으로 제한된다.

3) 전처리 함수에서 deepcopy, 반복 astype, 반복 isnan 제거

전처리/피처 선택 함수에서 흔히 메모리를 크게 쓰는 패턴은 다음입니다.

• deepcopy(input_array)로 입력 전체를 복사
• astype(float32/float64)를 중간에 여러 번 수행
• np.isnan(data)를 여러 번 호출하여 같은 크기의 boolean 배열을 반복 생성
• np.unique(..., return_counts=True)로 불필요한 count 배열까지 생성

개선 목표는 “float 변환은 1회, isnan 마스크는 1회, unique는 필요한 컬럼에만”입니다.

아래는 동작을 유지하면서 메모리 생성을 줄인 최적화 예시입니다. (변수/함수명은 일반화했습니다.)

import numpy as np

def select_features_optimized(
    input_array: np.ndarray,
    categorical_threshold: int = 15,
    nan_ratio_threshold: float = 0.3,
):
    """
    반환:
      nan_ratio_cols, categorical_cols, valid_row_index, drop_cols, use_cols

    핵심 최적화:
    - 입력 deepcopy 제거
    - float 변환 1회
    - np.isnan 마스크 1회 생성 후 재사용
    - unique는 후보 컬럼에만 수행, return_counts 제거
    """

    if input_array is None:
        return [], [], np.array([], dtype=int), set(), []

    arr = np.asarray(input_array)
    if arr.size == 0 or arr.shape[0] < 1:
        return [], [], np.array([], dtype=int), set(), []

    n_rows, n_cols = arr.shape

    # 1) 문자열/객체 컬럼 탐지(필요 최소)
    text_cols = []
    if arr.dtype == object or arr.dtype.kind in ("U", "S"):
        for c in range(n_cols):
            col = arr[:, c]
            try:
                _ = col.astype(np.float32, copy=False)
            except Exception:
                text_cols.append(c)

    # 2) float 배열 1회 생성
    if text_cols:
        text_set = set(text_cols)
        arr_f = np.empty((n_rows, n_cols), dtype=np.float32)
        for c in range(n_cols):
            if c in text_set:
                arr_f[:, c] = np.nan
            else:
                arr_f[:, c] = arr[:, c].astype(np.float32, copy=False)
    else:
        arr_f = arr.astype(np.float32, copy=False)

    # 3) isnan 마스크 1회 생성
    nan_mask = np.isnan(arr_f)

    # 4) 컬럼별 NaN 포함/비율 계산
    nan_included = nan_mask.any(axis=0)
    nan_ratio = nan_mask.mean(axis=0)  # axis=0이면 분모는 행 수가 되어야 함

    nan_ratio_cols = np.flatnonzero(nan_ratio > nan_ratio_threshold).tolist()

    # 5) 범주형 컬럼 탐색(삭제될 컬럼은 제외)
    drop_by_nan = set(nan_ratio_cols)
    categorical_cols = []

    for c in range(n_cols):
        if c in drop_by_nan:
            continue
        if c in set(text_cols):
            continue

        col = arr_f[:, c]
        col = col[~np.isnan(col)]
        if col.size == 0:
            continue

        if np.unique(col).size < categorical_threshold:
            categorical_cols.append(c)

    categorical_cols.extend(text_cols)
    categorical_cols = sorted(set(categorical_cols))

    # 6) 삭제/사용 컬럼
    drop_cols = set(nan_ratio_cols) | set(categorical_cols)
    use_cols = [c for c in range(n_cols) if c not in drop_cols]

    # 7) 사용할 컬럼 기준으로 NaN 없는 행 인덱스 계산
    if use_cols:
        row_has_nan = nan_mask[:, use_cols].any(axis=1)
        valid_row_index = np.flatnonzero(~row_has_nan)
    else:
        valid_row_index = np.array([], dtype=int)

    return nan_ratio_cols, categorical_cols, valid_row_index, drop_cols, use_cols

이 방식으로 바꾸면 다음이 개선됩니다.

• 입력 전체 deepcopy 제거로 메모리 피크 감소
• np.isnan 중복 호출 제거로 임시 배열 생성 감소
• astype 다중 호출 제거로 대형 복사 횟수 감소

4) finally에서 try 블록의 객체 정리 가능 여부와 안전한 작성

finally에서 try 블록 안에서 생성된 객체를 del로 제거하는 것은 가능합니다. 다만 예외가 중간에 발생하면 변수가 아직 생성되지 않았을 수 있으므로, 안전한 패턴을 권장합니다.

권장 패턴(사전 초기화):

import gc

def run_job(...):
    big_table = None
    work_array = None
    model_blob = None

    try:
        big_table = load_table(...)
        work_array = make_array(big_table)
        model_blob = train_model(work_array)
        return {"status": "ok"}
    finally:
        # 생성 여부와 관계없이 안전
        del big_table
        del work_array
        del model_blob
        gc.collect()

주의 사항:

• 외부에서 주입된 커넥션/프로듀서(예: 메시지 브로커 클라이언트)는 finally에서 무조건 close()하면 다음 요청에서 재사용 불가 문제가 생길 수 있습니다.
• 전송 완료 보장은 flush()를 적절한 위치(루프 밖 1회)에서 수행하는 방식이 더 안전합니다.

결과/출력

위와 같은 변경을 적용하면 다음 효과를 기대할 수 있습니다.

• 시나리오 수가 늘어도 DataFrame 전체 복사본 생성이 없어져 메모리 피크가 크게 감소한다.
• 같은 크기의 마스크/배열을 반복 생성하지 않아 GC 부담이 줄어든다.
• 큰 결과물(모델/설명기)을 딕셔너리에 중복 저장하지 않도록 하면 참조 유지로 인한 메모리 고착이 감소한다.
• 예외가 발생해도 finally에서 정리가 보장되어 “누수처럼 보이는 현상”을 완화한다.

응용/팁

1) 큰 결과물은 반환/로그/메시지에 중복 저장하지 않는다

큰 모델 직렬화 결과(문자열/바이트/딕셔너리)를 다음과 같이 여러 곳에 동시에 담으면 메모리가 오래 유지됩니다.

• 상태 반환 객체
• 내부 누적 결과 딕셔너리
• 메시지 전송 payload

권장 방식:

• 저장소(오브젝트 스토리지/DB)에 저장 후 키만 반환/전송
• 혹은 반환에는 길이/해시/버전 등 메타 정보만 담기

2) 전송 클라이언트의 flush는 루프 밖에서 1회만

시나리오마다 flush()를 호출하면 불필요한 대기와 내부 버퍼 관리 부담이 커질 수 있습니다. 가능하면 루프 끝에서 1회 호출로 합칩니다.

3) 데이터 타입을 float64로 올리는 순간 메모리는 즉시 커진다

가능하면 float32로 충분한지 먼저 검토합니다. 특히 입력 행/열이 큰 경우 효과가 큽니다.

4) 근본적으로는 학습 작업을 API 프로세스 밖으로 분리한다

API는 “요청 접수 및 작업 큐 발행”만, 학습은 별도 워커 프로세스에서 수행하는 구조가 가장 안정적입니다. 이는 단순한 코드 최적화보다 장애 가능성을 근본적으로 낮춥니다.

결론

본 정리의 요지는 다음과 같습니다.

• 시나리오별로 데이터가 달라진다고 해서 매번 전체 DataFrame을 복사할 필요는 없습니다. 원본을 불변으로 두고, 시나리오별로 달라지는 라벨/인덱스/X,y만 생성하는 방식이 메모리 효율이 가장 좋습니다.
• 전처리 함수에서는 deepcopy와 반복적인 astype, 반복적인 np.isnan 계산을 제거하면 메모리 피크와 임시 객체 생성을 크게 줄일 수 있습니다.
• finally에서는 try 내 객체 정리가 가능하며, 사전 초기화 패턴으로 예외 시에도 안전하게 정리할 수 있습니다.
• 최종적으로는 장시간 학습 로직은 API 프로세스에서 분리하는 것이 가장 안정적인 방향입니다.

배운 점과 시사점은 “복사는 안전하지만 비싸며, 불변 원본 + 필요한 부분만 생성하는 설계가 대규모 데이터 처리에서 가장 중요하다”는 것입니다. 또한 예외 처리와 정리 로직이 보장되지 않으면 메모리 문제는 재현이 어렵고 누수처럼 보이기 쉬워, 구조적 분리와 cleanup 보장이 함께 필요합니다.