Published on

SageMaker MME 콜드스타트 5분을 10초로 줄이는 법

Authors

SageMaker Multi-Model Endpoint(MME)는 "여러 모델을 한 엔드포인트에 올려 비용을 절감"하는 데 강력하지만, 운영에서 가장 먼저 부딪히는 벽이 콜드스타트입니다. 특정 모델이 트래픽을 처음 받는 순간, 컨테이너는 S3에서 모델 아티팩트를 내려받고 압축을 풀고, 프레임워크가 그래프를 로딩하고, (GPU라면) 메모리에 올린 뒤 첫 추론을 수행합니다. 이 경로가 길어지면 1~5분까지도 쉽게 늘어납니다.

이 글은 MME 콜드스타트를 **"5분에서 10초대"**로 줄이기 위해, 지연을 구성 요소로 쪼개고 각각을 줄이는 방법을 실전 관점에서 정리합니다. 핵심은 한 가지 요령이 아니라, 다운로드 시간 + 압축 해제 + 로딩/초기화 + 첫 추론 워밍업을 각각 최적화하는 것입니다.

MME 콜드스타트가 느려지는 지점부터 분해하기

MME의 첫 요청 지연은 대개 아래 합으로 설명됩니다.

  1. S3 다운로드 시간: 모델 아티팩트 크기, S3 리전/엔드포인트 네트워크, 동시성
  2. 압축 해제/파일 I/O: tar.gz 해제, 작은 파일 다량(메타데이터/샤드)로 인한 IOPS 병목
  3. 프레임워크 로딩: TorchScript/TF SavedModel/ONNX 런타임 초기화
  4. 모델 초기화: 토크나이저 로딩, vocab/merges 파싱, 커널 컴파일, CUDA 컨텍스트
  5. 첫 추론 워밍업: JIT, 캐시 미스, 커널 선택(autotune)

MME는 모델을 온디맨드로 로딩하고, 로딩된 모델을 인스턴스 로컬 디스크/메모리에 캐시합니다. 따라서 "첫 요청"만 빠르게 만들면 끝이 아니라,

  • 자주 호출되는 모델은 항상 뜨겁게 유지
  • 드물게 호출되는 모델도 최초 1회 지연을 최소화

이 두 축을 동시에 잡아야 합니다.

목표를 10초대로 만들려면: 기준선(측정)부터 잡기

최적화는 측정 없이는 감(感)입니다. MME에서는 아래 로그/메트릭을 먼저 확보하세요.

  • CloudWatch Logs: 모델 로딩 시작/완료, 다운로드/압축해제 시간(커스텀 로깅 권장)
  • CloudWatch Metrics: ModelLoadingTime, Invocation4XXErrors, Invocation5XXErrors, CPUUtilization, DiskUtilization
  • 컨테이너 내부 타이밍: download_s3_ms, untar_ms, load_model_ms, warmup_ms

(예시) inference 코드에 로딩 타이밍 로깅 넣기

아래는 Python inference toolkit 스타일의 단순 예시입니다. 본문에 < > 가 노출되면 MDX에서 문제가 될 수 있으므로, 타입 표기나 화살표는 모두 인라인 코드로 처리합니다.

import os
import time
import tarfile
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

_model_cache = {}


def load_model(model_dir: str):
    t0 = time.time()

    # 예시: 압축 파일이 model_dir 아래에 있다고 가정
    tar_path = os.path.join(model_dir, "model.tar.gz")
    if os.path.exists(tar_path):
        t1 = time.time()
        with tarfile.open(tar_path, "r:gz") as tar:
            tar.extractall(path=model_dir)
        logger.info({
            "event": "untar_done",
            "untar_ms": int((time.time() - t1) * 1000)
        })

    # 예시: 실제 프레임워크 로딩
    t2 = time.time()
    model = "loaded_model_object"  # torch.load / onnxruntime.InferenceSession 등으로 교체
    logger.info({
        "event": "model_load_done",
        "load_model_ms": int((time.time() - t2) * 1000)
    })

    logger.info({
        "event": "total_load_done",
        "total_load_ms": int((time.time() - t0) * 1000)
    })
    return model


def get_model(model_id: str, model_dir: str):
    if model_id in _model_cache:
        return _model_cache[model_id]

    model = load_model(model_dir)

    # 선택: 워밍업(더미 입력 1회)
    t3 = time.time()
    _ = "warmup"  # model(dummy_input)
    logger.info({
        "event": "warmup_done",
        "warmup_ms": int((time.time() - t3) * 1000)
    })

    _model_cache[model_id] = model
    return model

이렇게만 해도 "5분"이 어디서 발생하는지(다운로드인지, 압축해제인지, 로딩인지) 바로 드러납니다.

1) 모델 아티팩트 크기 줄이기: 10초 최적화의 80%

MME 콜드스타트 최적화에서 가장 큰 레버는 아티팩트 크기입니다. S3 다운로드와 압축해제는 크기에 거의 선형으로 비례합니다.

체크리스트

  • 불필요 파일 제거: 학습 체크포인트, optimizer state, 중간 산출물
  • 샤드 수 줄이기: 작은 파일이 너무 많으면 압축 해제/파일 I/O가 느려집니다
  • 포맷 변경: SavedModel 대비 TorchScript 또는 ONNX가 로딩이 빠른 경우가 많습니다(모델에 따라 다름)
  • 양자화: FP16, INT8(특히 CPU 추론)

(예시) PyTorch 모델을 TorchScript로 내보내기

import torch

model = ...
model.eval()

example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)
traced.save("model.pt")

TorchScript는 로딩 경로가 단순해지고, Python 코드 의존성이 줄어들어 MME에서 "모델별 커스텀 로딩"이 단순해지는 장점이 있습니다.

2) 압축 해제 병목 줄이기: tar.gz가 만능은 아니다

SageMaker 배포에서 흔히 model.tar.gz를 쓰지만, 콜드스타트 관점에서는 비용이 큽니다.

  • gzip 해제는 CPU를 사용하고
  • 파일이 많으면 메타데이터 처리로 느려지고
  • EBS I/O가 함께 병목이 됩니다

개선 전략

  • 파일 개수 줄이기(가장 효과 큼)
  • 가능한 경우 압축률보다 해제 속도 우선: gzip 대신 zstd를 쓰고 싶겠지만, SageMaker 기본 경로는 제약이 있으니 컨테이너에서 커스텀 다운로드/해제를 고려합니다
  • 모델을 단일 파일로 만들기(ONNX 단일 파일, TorchScript 단일 파일 등)

3) 컨테이너 이미지 최적화: "이미지 풀"도 콜드스타트의 일부

MME의 "모델 콜드스타트"만 보다가 놓치기 쉬운 게 컨테이너 이미지 풀 시간입니다. 특히 새로 스케일아웃되거나, 새 배포 직후에는 이미지 다운로드가 추가됩니다.

체크리스트

  • 베이스 이미지 슬림화: 불필요한 빌드 툴/패키지 제거
  • Python 패키지 고정 및 최소화: 거대한 pip 의존성은 이미지 레이어를 키웁니다
  • ECR 같은 리전 내 레지스트리 사용

(예시) 멀티스테이지 Dockerfile로 런타임만 남기기

FROM python:3.11-slim AS base
WORKDIR /app

# 필요한 런타임 패키지만 설치
RUN pip install --no-cache-dir fastapi uvicorn

COPY inference.py /app/inference.py

CMD ["python", "-c", "import inference; print('ready')"]

실제 SageMaker inference toolkit을 쓰는 경우에도 원리는 같습니다. "빌드에 필요한 것"과 "실행에 필요한 것"을 분리하면 이미지가 작아지고, 풀 시간도 줄어듭니다.

4) MME 캐시를 적극 활용: 자주 쓰는 모델은 "뜨겁게" 유지

MME는 로딩된 모델을 캐시하지만, 인스턴스의 메모리/디스크 한계로 인해 LRU 방식으로 언로드될 수 있습니다. 즉, "한 번 로딩했으니 계속 빠르겠지"가 성립하지 않습니다.

운영 전략

  • 상위 N개 모델은 주기적으로 ping하여 캐시에서 밀려나지 않게 유지
  • 트래픽 패턴이 시간대별로 바뀌면, 그 시간대 직전에 워밍업
  • 모델 크기별로 엔드포인트를 분리(초대형 모델과 소형 모델을 한 MME에 섞으면 캐시 효율이 급락)

(예시) 워밍업 Lambda/크론에서 호출하는 간단한 스크립트

import json
import boto3

smr = boto3.client("sagemaker-runtime")

endpoint = "my-mme-endpoint"
model_ids = ["model-a", "model-b", "model-c"]

for mid in model_ids:
    payload = json.dumps({"text": "warmup", "model_id": mid})
    # TargetModel 헤더로 특정 모델을 지정하는 패턴을 사용(구현/서빙 방식에 따라 다름)
    resp = smr.invoke_endpoint(
        EndpointName=endpoint,
        ContentType="application/json",
        TargetModel=f"{mid}.tar.gz",
        Body=payload,
    )
    _ = resp["Body"].read()

여기서 TargetModel 사용 방식은 "S3에 저장된 모델 아티팩트 키" 또는 "MME가 기대하는 모델 식별자"에 맞춰 조정해야 합니다.

5) 동시 요청 폭주 대비: 로딩 중 스탬피드(thundering herd) 막기

콜드스타트 순간에 동일 모델로 요청이 몰리면, 로딩이 끝나기 전에 같은 모델을 여러 번 로딩하려다 병목이 생길 수 있습니다(서빙 코드 구현에 따라 다름). 이때는 애플리케이션 레벨에서 단일 플라이트(single-flight) 를 구현해 로딩을 1회로 제한합니다.

(예시) Python에서 모델 로딩 락으로 단일화

import threading

_model_cache = {}
_model_locks = {}
_global_lock = threading.Lock()


def _get_lock(model_id: str):
    with _global_lock:
        if model_id not in _model_locks:
            _model_locks[model_id] = threading.Lock()
        return _model_locks[model_id]


def get_or_load(model_id: str, model_dir: str):
    if model_id in _model_cache:
        return _model_cache[model_id]

    lock = _get_lock(model_id)
    with lock:
        # 더블 체크
        if model_id in _model_cache:
            return _model_cache[model_id]
        model = load_model(model_dir)
        _model_cache[model_id] = model
        return model

이 한 가지로도 콜드스타트 시점의 tail latency가 크게 안정화됩니다.

6) 스토리지/인스턴스 선택: EBS와 CPU가 의외로 지배적

"모델이 GPU에서 돌아가니까 GPU 인스턴스만 보면 되겠지"라고 생각하기 쉽지만, 콜드스타트는 종종 CPU(압축 해제)디스크 I/O(EBS) 가 지배합니다.

실전 팁

  • 모델이 크고 압축 해제가 무거우면 vCPU가 많은 인스턴스가 유리
  • EBS 처리량/IOPS가 낮으면 압축 해제와 파일 로딩이 늘어짐
  • 가능하면 파일 수를 줄여 IOPS 민감도를 낮추기

인스턴스/볼륨 튜닝은 비용이 바로 늘어나는 영역이라, 먼저 "아티팩트 크기"와 "파일 수"부터 줄인 뒤 마지막에 적용하는 것이 ROI가 좋습니다.

7) "10초"를 만드는 전형적인 조합(레시피)

현장에서 1~5분 콜드스타트를 10초대로 낮출 때 자주 쓰는 조합은 아래와 같습니다.

  1. 모델을 단일 파일 포맷으로 변환(예: model.onnx 또는 model.pt)
  2. 아티팩트에서 불필요 파일 제거, 샤드/파일 수 최소화
  3. 로딩 경로 단순화(토크나이저/설정 파일도 가능하면 단일화)
  4. 워밍업 호출로 상위 모델을 캐시에 유지
  5. 로딩 단일화 락으로 스탬피드 방지
  6. (필요 시) 이미지 슬림화로 스케일아웃 시 이미지 풀 지연 감소

이 조합은 "특정 AWS 기능 하나"가 아니라, 콜드스타트의 각 단계를 1~2초씩 깎아 합산을 줄이는 방식입니다.

장애/지연이 같이 보일 때: 콜드스타트와 504를 분리해서 보기

콜드스타트가 길어지면 API Gateway/ALB/클라이언트 타임아웃이 먼저 터지면서 504로 관측되기도 합니다. 이때는 "타임아웃을 늘려서" 해결하기보다, 콜드스타트 원인을 줄이는 게 정공법입니다. 콜드스타트 진단 관점은 Cloud Run 사례와도 유사한데, 지연을 단계별로 쪼개는 접근이 동일하게 통합니다.

또한 MME를 EKS에서 호출하거나, 워밍업 잡을 쿠버네티스로 돌리는 경우 IAM 권한/IRSA 설정이 발목을 잡아 "모델 다운로드 실패"가 콜드스타트처럼 보이기도 합니다.

배포 전 체크리스트(운영 기준)

아래 항목을 배포 게이트로 두면, 콜드스타트가 다시 1~5분으로 회귀하는 것을 막을 수 있습니다.

  • 모델 아티팩트 크기(예: 500MB 이하 목표)와 파일 수(예: 1000개 이하) 측정
  • 로딩 타이밍 로그 필수 수집(download/untar/load/warmup 분리)
  • 상위 모델 워밍업 스케줄 보유(트래픽 패턴 기반)
  • 로딩 단일화(락) 적용 여부 확인
  • 엔드포인트 분리 기준 수립(초대형 모델은 별도 엔드포인트)
  • 타임아웃은 "원인 제거" 후 최소한으로 조정

결론: 5분을 10초로 줄이는 핵심은 "모델을 가볍게" 그리고 "자주 쓰는 건 뜨겁게"

SageMaker MME 콜드스타트는 대체로 S3 다운로드와 압축 해제, 그리고 프레임워크 로딩이 합쳐져 길어집니다. 10초대를 만들려면

  • 아티팩트 크기/파일 수를 줄여 물리적인 이동량을 감소시키고
  • 로딩 경로를 단순화하며
  • 워밍업과 캐시 유지로 첫 요청을 회피하고
  • 동시 로딩 스탬피드를 막아 tail latency를 안정화

이 네 가지를 동시에 적용해야 합니다.

다음 단계로는, 현재 엔드포인트의 콜드스타트를 위 로깅 방식으로 분해한 뒤(다운로드/압축해제/로딩/워밍업), 가장 큰 구간부터 1개씩 줄이는 방식으로 진행하면 됩니다. 이렇게 하면 "막연히 빠르게"가 아니라, 숫자로 5분을 10초대로 수렴시키는 로드맵을 만들 수 있습니다.