PyTorch 와 PyTorch Lightning을 위한 간단 도커파일 (Dockerfile) 작성 & 사용

도커 (docker) 사용 이유

학습 코드를 작성하다 보면 기본적으로 cuda version, python version, pytorch version, Nvidia GPU와 cuda compatibility같은 것을 신경 써야 합니다. 코드가 복잡해지면 이것저것 신경 써야 하는 것이 많아져서 그냥 처음부터 도커 환경을 구축해서 작업하는 것이 속 편하더라고요. 이번 포스트에서는 pytorch와 pytorch lightning을 사용하기 위한 간단한 도커파일 작성하는 법을 소개하려고 합니다 (ubuntu/wsl2 사용).

 

도커(docker)란?

이번 포스트에서 사용되는 도커의 용도는 특정 머신 (예를 들어 제가 사용하는 pc)에서 원하는 개발 환경을 구축하는 것입니다. 도커를 사용하면 머신에 이미 있는 환경과 분리된 환경을 구축할 수 있기 때문에 편리합니다 (로컬 pc에서 이미 설치되어 있는 우분투, cuda, 등등과 분리된 환경을 구축해서 원하는 우분투, cuda 설치 가능). 더 자세하고 정확한 정보는 공식 문서 또는 잘 정리되어있는 글을 참고해주세요. 특히 docker image와 docker container 개념을 먼저 이해하는 것이 좋습니다. 대략적으로 도커 이미지는 원하는 개발 환경이 세팅되어 있는 패키지라고 생각하면 되고, 도커 컨테이너는 이 패키지를 실행해서 생기는 인스턴스라고 생각하면 됩니다. 

도커파일(Dockerfile)이란?

도커 이미지를 custom하게 만들 때 사용되는 파일입니다. 이번 포스트에서 다루지는 않겠지만 도커 이미지를 만드는 다른 방법도 있습니다.

방법

1. 도커 설치

• 공식 설치 패이지에서 os에 맞는 버전을 설치한다 (https://docs.docker.com/get-docker/)

2. Nvidia driver 설치

• Nvidia driver이 설치되어있어야 한다
• os 마다 설치 방법이 다르다. 우분투는 여기 참고. 공식 페이지에서 컴퓨터 사양에 맞게 골라서 다운로드해도 된다.

3. Nvidia-docker 설치 (3,4 참고자료)

• github에서 nvidia container toolkit 설치 후 nvidia-docer2 설치

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
&& curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
apt-get update
apt-get install -y nvidia-docker2

• docker을 다시 시작시켜야 설치가 완료된다

sudo systemctl restart docker

4. 잘 설치되었는지 테스트

docker run -it --rm --gpus all nvidia/cuda:11.6.2-base-ubuntu20.04 nvidia-smi

• 이 코드를 실행하면 우선 nvidia/cuda:11.6.2-base-ubuntu20.04 이미지가 머신에 존재하는지 확인을 하고, 없으면 다운로드를 합니다.
• 머신에 도커 이미지가 있으면 도커 이미지를 기반으로 도커 컨테이너를 만들고 실행합니다. 머신에 있는 도커 이미지 리스트는 docker images 커맨드로 확인 가능합니다. 다음과 같이 결과가 나오는데, 제 컴퓨터에는 현제 2개의 도커 이미지가 있는 것을 확인할 수 있습니다. nvidia/cuda는 위에 코드를 실행하면서 설치되었습니다.

• -it는 interactive하게 도커 컨테이너를 실행하겠다는 뜻입니다 (나중에 추가 설명)
• --rm은 도커 컨테이너를 exit할 때 container을 제거하겠다는 뜻입니다
• --gpus all은 도커 컨태이너 안에서 모든 gpu를 접근할 수 있게 하겠다는 뜻입니다
• nvidia-smi는 도커 컨테이너에서 nvidia-smi를 실행하겠다는 뜻입니다.  위에 있는 docker run 코드를 실행하면 nvidia-smi를 실행하고 도커 컨테이너를 나가게 되는데, 만약 --rm flag를 사용하지 않으면 도커 컨테이너가 exit 한 상태로 남습니다. 모든 도커 컨테이너는 docker ps -a를 통해서 확인할 수 있습니다. 참고로 -a flag가 있어야 exit 된 컨태이너도 확인 가능합니다 (docker ps는 돌아가고 있는 컨태이너만 보여줍니다). --rm flag 없이 docker run을 실행하면 다음과 같이 container이 남아있는 것을 확인할 수 있습니다. 불필요한 컨태이너이기 때문에 docker rm 36669da14934 (36669da14934는 container ID) 로 제거하면 됩니다. 참고로 exit된 container을 다시 살릴 수 있는데, docker start와 docker attach 커맨드를 찾아보시면 됩니다.

• 참고로 종종 도커 이미지가 업데이트되면서 예전에 있었던 도커 이미지가 없어서 다운로드 가 안될 수도 있습니다. 없는 도커 이미지를 다운로드하려고 하면 다음과 같은 오류가 뜹니다:

• 이런 경우 docker hub (https://hub.docker.com/)에서 적절한 도커 이미지를 찾아서 사용하면 됩니다. 예를 들어 nvidia/cuda를 검색해서 Tags를 보면 어떤 docker image들을 다운로드할 수 있는지 확인 가능합니다. 만약 캡처 화면에 보이는 12.2.0-devel-ubuntu20.04를 사용하고 싶으면 nvidia/cuda:11.6.2-base-ubuntu20.04. 대신 nvidia/cuda/12.2.0-devel-ubuntu20.04을 사용하면 됩니다.

5. Dockerfile 작성 & 이미지 만들기 (참고자료: https://docs.docker.com/engine/reference/builder/)

• Dockerfile은 원하는 도커 이미지를 만들기 위한 파일입니다.  Docker hub에 다양한 도커 이미지가 있지만, 원하는 환경이 완벽하게 세팅되어 있는 도커 이미지는 보통 없습니다. 이런 경우  Docker hub에 있는 도커 이미지를 기반으로 원하는 도커 이미지를 구축하면 되는데, 이때 Dockerfile이 유용합니다.
• 이번 포스트에서는 다음과 같은 도커 파일을 소개하려고 합니다. 참고로 파일 명은 꼭 Dockerfile로 해야 합니다. 우선 도커파일 내용부터 설명을 하겠습니다.

FROM nvidia/cuda:11.6.2-devel-ubuntu20.04

# Install some basic utilities
RUN apt-get update && apt-get install -y \
    curl

# Create a working directory
RUN mkdir /project
WORKDIR /project

# Install Miniconda and Python 3.9
ENV CONDA_AUTO_UPDATE_CONDA=false
ENV PATH=/root/miniconda/bin:$PATH
RUN curl -sLo ~/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-py39_23.5.2-0-Linux-x86_64.sh \
 && chmod +x ~/miniconda.sh \
 && ~/miniconda.sh -b -p ~/miniconda \
 && rm ~/miniconda.sh \
 && conda install -y python==3.9.12 \
 && conda clean -ya
 
COPY requirements.txt ./

RUN python -m pip install --no-cache-dir --upgrade pip \
    && python -m pip install --no-cache-dir -r requirements.txt

• 첫 줄은 항상 FROM으로 시작해야 합니다. FROM은 도커파일을 통해 만들 도커 이미지에서 base로 사용할 도커 이미지를 뜻합니다. 이 도커 이미지에 이것저것 설치해서 최종 이미지를 만들게 됩니다 (FROM을 여러 번 사용할 수도 있습니다)
• 특정 커맨드를 실행하고 싶으면  앞에 RUN을 붙이면 됩니다. 리눅스에서는 default로 /bin/sh -c 쉘로 실행해 줍니다.
• 저는 curl을 설치했는데, 도커 이미지를 만들 때 최소한의 소프트웨어/코드/등등을 담는 것을 원칙으로 하기 때문에 많은 것들이 설치가 안되어있습니다. nvidia/cuda 이미지는 파이썬도 없기 때문에 위 도커파일에서 미니콘다를 설치합니다 (miniconda를 설치하기 위해서 curl을 설치했습니다). 물론 miniconda 대신 더 슬림하게 python만 설치해도 되지만 종종 pip으로 설치가 잘 안 되지만 conda로 설치가 잘 되는 패키지들도 있어서 저는 conda도 설치했습니다.
• 그다음 /project라는 폴더를 만들고 working directory로 설정했는데, RUN, COPY 등등을 실행하면 working directory 안에서 실행하게 됩니다 (여러모로 편한데 저는 작업 코드를 /project에 넣는 것을 선호합니다)
• 그 다음 miniconda를 설치합니다. 참고로 ENV는 환경 변수를 설정할 때 사용합니다.
• 그 다음 COPY를 통해서 requirements.txt를 도커 이미지에 넣고 (/project 폴더에 copy 됨) pip을 통해서 설치합니다. 여기서 사용하는 requirements.txt는 다음과 같습니다. 

--extra-index-url https://download.pytorch.org/whl/cu116
torch==1.12.1+cu116
--extra-index-url https://download.pytorch.org/whl/cu116
torchvision==0.13.1+cu116
--extra-index-url https://download.pytorch.org/whl/cu116
torchaudio==0.12.1

rich==12.6.0
pytorch-lightning==1.8.1

• 도커 이미지를 만들기 위해서 같은 폴더 안에 Dockerfile과 requirements.txt가 있는 상태에서 docker build . -t test를 실행하면 됩니다. Dockerfile이 있는지 감지한 후 docker image를 만들어줍니다. 여기서 -t는 도커 이미지의 tag을 의미하는데 name:tag 꼴을 갖습니다. name만 지정하면 (test로) tag은 latest로 됩니다. 자세한 내용은 공식 문서 참고. 
• 보다시피 pytorch와 pytorch lightning을 둘 다 설치합니다. rich 같은 경우 pytorch lightning에서 progress bar을 사용할 때 필요한 패키지입니다. 

6. 도커 컨테이너 실행

• 이제 이미지가 있으니 docker run을 통해서 컨테이너를 만들면 끝납니다:

docker run \
--gpus all \ # 모든 GPU 사용
--shm-size 32g \ # shared memory
-it \ # interactive하게 실행
--name test_container \ # docker container 이름
-v /local/path/to/project/:/project \ # 프로젝트 폴더 마운트
test \ # 도커 이미지 이름
/bin/bash # /bin/bash 실행

• 몇 가지 새로운 옵션들이 추가되었습니다.
• --shm-size는 shared memory 사이즈를 뜻하는데 default 값이 64MB입니다 (작음). 학습하면서 프로세스가 shared memory를 많이 잡아먹으면 오류가 뜨기 때문에 넉넉하게 잡는 것이 좋습니다.
• --it: interactive 하게 container 실행. 생략하면 /bin/bash를 실행하고 exit 합니다.
• --name: docker container 이름입니다. 
• -v: volume을 mount 할 때 사용합니다. 컨테이너 밖에서 코드를 수정했을 때 컨태이너 안에 반영이 되게 하고싶거나, 컨테이너에서 파일을 만들었는데 (예를 들어 모델 checkpoint를 저장했는데) 이것이 컨태이너 밖에 반영이 되게 하고 싶으면 폴더를 mount하는 것이 편합니다. 프로젝트 작업 폴더 (/local/path/to/project)를 컨태이너 내부 프로젝트 폴더 (/project)에 매핑해주면 됩니다. 
•  test는 도커 이미지 이름입니다
• /bin/bash는 도커 컨태이너에서 /bin/bash를 사용하겠다는 뜻입니다.
• 실행하면 다음과 같이 /bin/bash로 도커 컨테이너와 소통? 이 가능합니다

• 보다시피 working directory (/project)로 들어갑니다 (entry point)
• python train.py 등 코드를 실행하면 됩니다
• 컨테이너에서 나가고 싶으면  exit을 실행하면 됩니다.