[ML] Docker 환경에서 GPU 사용하기

요약 ✌

Multi-GPU 기반 dockerize

1. NVIDIA Container Toolkit 설치

# pkg repo & gpg key 설정  
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
      && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
      && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
            
            
# nvidia-docker2 설치
$ sudo apt-get update && sudo apt-get install -y nvidia-docker2     

# docker 재시작
$ sudo systemctl restart docker 

2. Dockerfile 작성

FROM nvidia/cuda:10.2-runtime-ubuntu18.04

RUN apt-get update && apt-cache madison python3-pip

RUN apt-get install -y software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get install -y python3.8 python3-pip
RUN apt-get install -y libpython3.8-dev

# Update symlink to point to latest
RUN rm /usr/bin/python3 && ln -s /usr/bin/python3.8 /usr/bin/python3
RUN python3 --version
RUN pip3 --version
RUN python3 -m pip install --upgrade pip


ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility


COPY . MY_PROJECT/
WORKDIR MY_PROJECT
RUN pip3 install -r requirements.txt


CMD ["/bin/bash"]

3. Docker image build & 실행

$ docker build -t {TAG_NAME}:{TAG_VERSION} .
$ docker run --gpus all -it --name {CONTAINER_NAME} {TAG_NAME}:{TAG_VERSION}

4. docker-compose.yaml 작성 및 container 실행

docker-compose.yaml

services:
  test:
    image: {image_name}:{image_version}
    command: nvidia-smi
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: [gpu]

Docker compose 실행

$ docker-compose up

---

GPU Dockerize step들을 하나 하나 자세하게 알아보겠습니다.

우리의 목표는 Model Training 환경을 dockerize 하는 것입니다!

다들 한번쯤은 동료의 코드를 돌릴 때 tf나 torch version이 안맞거나,
pandas 등의 package 버전이 안맞아서 고생해보신 적 있을거에요.
(대부분의 requirements.txt 에는 버전 명시가 안되어 있으니까요. ㅠㅠ)

아래 step들을 따라 모델링 코드를 docker화 하게 되면, 매번 힘들게 했던 versioning 이슈들을 쉽게 해결할 수 있습니다. !

1. NVIDIA Container Toolkit

NVIDIA-Container Toolkit 이란?

NVIDIA Container Toolkit 부터 설치하라는데 이 툴킷은 어떤 역할을 할까요?

먼저 우분투 기반으로 도커를 설치하고, 이 컨테이너에 들어가 봅시다.

$ sudo docker run -it --name test2 docker_test:latest /bin/bash
root@0ea54af1874c:/home# nvidia-smi
bash: nvidia-smi: command not found 

현재 컨테이너는 CPU 위에서 만들어진 container로 GPU를 사용할 수 없습니다.
내부에서 NVIDIA Toolkit 설치해서도 사용할 수는 있지만, 이 때 host에 설치된 버전과 일치되어야하는 제약 사항이 있어서 쉽지는 않습니다.

NVIDIA Container Toolkit은 이러한 귀찮음을 해결해주는 Toolkit으로
Container에서 GPU를 사용시 libcuda.so 와 같은 host에 있는 CUDA Toolkit들을 마운트해주는 역할을 합니다.

NVIDIA-Container-Toolkit 설치

# pkg repo & gpg key 설정  
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
      && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
      && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
            
            
# nvidia-docker2 설치
$ sudo apt-get update && sudo apt-get install -y nvidia-docker2     

# docker 재시작
$ sudo systemctl restart docker 

# (optional) 테스트 
$ sudo docker run --rm --gpus all nvidia/cuda:11.6.2-base-ubuntu20.04 nvidia-smi

2. Dockerfile

위에서 nvidia-container-toolkit을 설치했다면,
이제 docker-cli에서 docker run 시 gpus 옵션을 통해 사용할 gpu를 설정해줄 수 있습니다.

$ docker run --gpus all {이미지 이름}

이 때 사용하는 이미지는 nvidia/cuda base의 이미지여야합니다.

우리는 우리만의 새로운 이미지를 만들고 배포하기 위해 직접 한번 Dockerfile 작성해보도록 할게요.
먼저 Project root에 Dockerfile 하나 생성해줍시다.

Dockerfile

FROM nvidia/cuda:10.2-runtime-ubuntu18.04

가장 먼저 해당 Docker의 base가 될 cuda 이미지를 선택해줍니다.
https://hub.docker.com/r/nvidia/cuda

저는 cuda 10.2가 필요하여 nvidia/cuda:10.2-runtime-ubuntu18.04를 설치해줍니다.

참고로, 선택을 할 때 보면 base, runtime, devel 세종류가 있습니다.
base < runtime < devel 순으로 제공하는 lib 들이 많고, 그만큼 크기도 점점 커집니다.
왠만하면 base나 runtime을 사용하시면 됩니다.

Dockerfile

FROM nvidia/cuda:10.2-runtime-ubuntu18.04

RUN apt-get update && apt-cache madison python3-pip

RUN apt-get install -y software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get install -y python3.8 python3-pip
RUN apt-get install -y libpython3.8-dev

# Update symlink to point to latest
RUN rm /usr/bin/python3 && ln -s /usr/bin/python3.8 /usr/bin/python3
RUN python3 --version
RUN pip3 --version
RUN python3 -m pip install --upgrade pip

ubuntu 18.04의 경우, 단순 python 설치 시 기본으로 python 3.6이 설치됩니다.
3.6에서는 tf 2.5, torch 1.8 이상을 설치할 수 없기 때문에 3.8로 python을 재설치해줍시다.

Dockerfile

FROM nvidia/cuda:10.2-runtime-ubuntu18.04

RUN apt-get update && apt-cache madison python3-pip

RUN apt-get install -y software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get install -y python3.8 python3-pip
RUN apt-get install -y libpython3.8-dev

# Update symlink to point to latest
RUN rm /usr/bin/python3 && ln -s /usr/bin/python3.8 /usr/bin/python3
RUN python3 --version
RUN pip3 --version
RUN python3 -m pip install --upgrade pip


ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility


COPY . MY_PROJECT/
WORKDIR MY_PROJECT
RUN pip3 install -r requirements.txt


CMD ["/bin/bash"]

추가로 ENV로 내가 사용할 NVIDIA_VISIBLE_DEVICES를 설정해준다.
이 옵션은 위에서처럼 --gpus가 아니라 아래처럼 실행 시 영향을 준다.

$ docker run --runtime nvidia XX 

즉 아래의 run 방법은 모두 동일하게 nvidia device를 사용할 수 있다.
(참고 : https://kyumdoctor.co.kr/18)

$ docker run --runtime nvidia -e NVIDIA_VISIBLE_DEVICES=all XX
$ docker run --gpus all XX

3. Docker Image 빌드 & 실행

Dockerfile이 만들어졌으니 Docker image를 빌드하고 실행해보자.

$ docker build -t {TAG_NAME}:{TAG_VERSION} .
$ docker run --gpus all -it --name {CONTAINER_NAME} {TAG_NAME}:{TAG_VERSION} 
# nvidia-smi
-----------------------------------------------------------------------------+
| NVIDIA-SMI XXX   Driver Version: XXX   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|

자, 이제 만들어진 Docker image를 docker-compose를 통해 쉽게 사용할 수 있도록 해봅시다.

4. docker-compose.yaml

https://docs.docker.com/compose/gpu-support/

물론 docker run을 해서 사용해도 되지만, 보통은 docker compose 를 이용해 container들을 관리합니다.

특히나 Multi-GPU 사용 시 GPU OOM을 발생하는 원인이 되기도 하거든요.

docker-compose.yaml 을 작성하고 실행해봅시다.

services:
  test:
    image: {image_name}:{image_version}
    command: nvidia-smi
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: [gpu]

만약 특정 GPU를 사용하고 싶다면 devices > driver 아래에 count : {device num} 을 추가해주면 됩니다.

           - driver : nvidia
           	 count : 1
             capabilities : [gpu]

이제 이 docker-compose를 이용해 container를 띄워볼게요

$ docker compose up

Reference

https://github.com/NVIDIA/nvidia-docker
https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/user-guide.html
https://engineer-mole.tistory.com/265