Docker 기초 (9) - ML Training and Serving Images

CPU 전용 wheel로 서빙 image를 줄이는 방법부터 모델 파일 배치, GPU 학습 container 구성까지 ML에서의 Docker 활용을 정리합니다.

Posted 2024/08/27

By 강호승

2 min read

Docker 기초 시리즈의 9편입니다. 전체 목차는 0편에 있습니다.

학습 image와 서빙 image

두 image는 목표가 다릅니다.

학습 image: GPU와 무거운 프레임워크를 포함합니다. 크기보다 재현성이 우선이고, 같은 image로 어느 GPU 서버에서든 같은 학습이 돌아가는 것이 목표입니다
서빙 image: 예측 요청에 응답하는 데 필요한 것만 담습니다. 배포와 scale out 때마다 pull 되므로 크기와 시작 속도가 곧 배포 속도입니다

서빙 image: CPU 전용 wheel

반려동물 품종 분류기 7편의 FastAPI 서빙 앱을 이미지로 만듭니다. 서빙 이미지에서 가장 큰 감량 포인트는 PyTorch입니다. 기본 wheel은 CUDA 라이브러리를 포함해 GB 단위인데, CPU로 추론하는 서빙이라면 수백 MB의 CPU 전용 wheel로 충분합니다.

  
FROM python:3.12-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt \
    --index-url https://download.pytorch.org/whl/cpu \
    --extra-index-url https://pypi.org/simple

COPY app/ ./app
COPY model/pet_classifier.pt ./model/   # TorchScript로 export한 모델

EXPOSE 8000
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

--index-url을 PyTorch의 CPU 저장소로 주면 torch가 CPU 판으로 설치되고, 나머지 패키지는 --extra-index-url의 PyPI에서 받습니다.

--host 0.0.0.0은 함정 방지입니다. uvicorn 기본값인 127.0.0.1로 띄우면 container 안의 localhost에만 묶여서, -p 8000:8000을 열어도 밖에서 접속이 안 됩니다. 7편의 localhost 격리가 서버 방향으로 나타나는 사례입니다.

모델 파일은 어디에 두나?

두 방식이 있고 트레이드오프가 명확합니다.

image에 굽기: 위 Dockerfile 방식. image 하나가 코드와 모델의 완결된 배포 단위가 되고, 시작이 빠릅니다. 대신 모델이 바뀔 때마다 image를 다시 빌드하고, 모델 크기만큼 image가 커집니다
시작할 때 받기: container가 시작하면서 model registry나 S3에서 모델을 내려받습니다. image는 작고 모델 교체에 재빌드가 필요 없지만, 시작 시점에 외부 저장소 의존성이 생깁니다. MLflow 해부 4편의 model registry가 이 방식의 저장소 역할을 합니다

모델이 작고 배포 단위를 단순하게 가져가고 싶으면 굽고, 모델 교체가 잦거나 모델이 크면 받는 쪽으로 갑니다.

pip cache mount

의존성이 바뀔 때마다 전체를 다시 내려받는 것이 부담스러우면 BuildKit의 cache mount로 pip 다운로드 캐시를 빌드 사이에 재사용합니다.

  
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt

이 방식을 쓸 때는 --no-cache-dir를 빼야 캐시가 쌓입니다. layer cache가 통째로 무효화된 경우에도 다운로드는 건너뛰게 되어, 의존성이 무거운 ML 이미지에서 효과가 큽니다.

GPU 학습 image

GPU 학습 container에는 세 층이 필요합니다. 호스트의 NVIDIA driver, container에 GPU를 연결해 주는 NVIDIA Container Toolkit, 그리고 CUDA 라이브러리를 포함한 image(pytorch/pytorch 공식 image 또는 nvidia/cuda base)입니다.

  
docker run --rm --gpus all pytorch/pytorch \
  python -c "import torch; print(torch.cuda.is_available())"

--gpus all이 container에 GPU를 노출하는 옵션이고, Toolkit이 없으면 이 옵션 자체가 실패합니다
호스트 driver가 지원하는 CUDA 버전이 image의 CUDA 버전 이상이어야 합니다. driver가 낮으면 container 안에서 GPU 초기화가 실패합니다
macOS에서는 GPU 학습 container를 쓸 수 없습니다. Docker가 Linux VM 위에서 돌기 때문에 CUDA도 Apple Silicon의 MPS도 container에 연결되지 않습니다. GPU 학습은 Linux 서버나 클라우드 GPU 인스턴스에서 합니다

빌드한 image를 registry에 push 하고 서버에서 pull 받는 배포 동선은 AWS 해부 5편에, 서빙 API의 실전 구성은 NYC 택시 프로젝트 7편에 있습니다.

시리즈를 마치며

일곱 편의 요약입니다.

container는 호스트 커널을 공유하는 격리된 프로세스입니다 (2편)
run은 매번 새 container를 만들고, 수명은 메인 프로세스가 정합니다 (3편, 4편)
Dockerfile은 자주 바뀌는 것일수록 아래에 적습니다. 빌드 속도의 대부분이 이 순서에서 나옵니다 (6편)
상태는 volume에 두고, container끼리는 user-defined network의 이름으로 통신합니다 (7편)
Compose로 실행 방법을 저장소에 커밋하고, prune으로 디스크를 관리합니다 (5편, 8편)
서빙 image는 CPU 전용 wheel과 slim base로 작게, 학습 image는 driver, Toolkit, CUDA image의 세 층으로 만듭니다 (9편)

같은 글감의 다른 도구 편은 MLflow 해부, Airflow 해부, AWS 해부 시리즈입니다.

Study, Docker