[Pytorch] Docker를 이용해 Jupyter 서버 만들기

Pytorch with Jupyter Lab을 Docker로 올려보자

서론

그냥 파이썬 가상환경 만들어서 돌려도 되지만 Docker로 만들어서 돌리면 얼마나 편할까 라는 생각을 하게 되었습니다. 그래서 참을 수 없는 삽질이 시작되었습니다.

Pytorch를 사용할 수 있는 Jupyter Lab 환경을 도커 상에서 구축하기 위해 다음과 같은 과정을 거쳤습니다.

본론

1. Prerequisites

먼저 시작하기에 앞서 필요한 소프트웨어를 모두 설치해야합니다. 다음 표에 나와 있는 기본 소프트웨어들이 모두 설치되어 있는지 확인합니다.

소프트웨어	설명	설치확인
nvidia driver	NVIDIA GPU DRIVER	$ nvidia-smi
CUDA	GPU CALCULATION API	$ nvcc --version
cuDNN	Library for DNN	CUDA 디렉토리에서 직접 찾아야 함
Docker	Linux Kernel의 가상환경 제공	$ docker -v

CUDA 설치과정에서 .bashrc 파일에 해당 경로를 export 하는 코드를 꼭 추가해야합니다. 전체 과정에 대해 설명하고 있는 사이트를 같이 제공합니다.

이미 설치가 되어있다면 다음 단계로 바로 진행해도 됩니다.

2. Installation

진행과정은 아래와 같습니다.

1. nvidia-container-toolkit 설치
2. pytorch Dockerfile 작성
3. docker-compose.yaml 작성
4. 여러가지 응용

a. nvidia-container-toolkit 설치

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Docker 컨테이너에서는 원래 GPU를 사용할 수 없습니다. 하지만 nvidia-container-toolkit을 이용하여 이를 가능하게 할 수 있습니다.

위 명령어를 그대로 입력하여 Docker 컨테이너에서도 GPU를 사용할 수 있도록 설정해줍니다.

b. Dockerfile 작성

FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel #설치환경에 따라 cuda 및 cudnn 버전 변경 필요!

RUN pip install jupyter

EXPOSE 8888

docker hub/pytorch에서 설치한 CUDA와 cuDNN의 버전에 부합하는 pytorch devel이미지를 찾습니다.

그리고 위와 같이 도커파일을 생성합니다. 현재 머신은 cuda11.6에 cudnn8을 사용하고 있기에 해당 버전에 부합하는 버전의 이미지를 사용하였습니다.

또한 Jupyter를 통해 접속해야 하므로 Image를 Build할 때 jupyter를 설치하도록 합니다. 만약 다른 패키지가 많이 필요하다면 requirements.txt를 작성하여 설치하면 좋을 것 같습니다.

c . docker-compose.yaml 작성

version: "3"
services:
        torch:
                build: .
                container_name: jupter-torch
                ports:
                        - "your_port:8888" #외부에서 접속할 포트를 설정해야 합니다!
                deploy:
                        resources:
                                reservations:
                                        devices:
                                                - capabilities: [gpu]
                                                  driver: nvidia
                environment:
                        - NVIDIA_VISIBLE_DEVICES=all
                        - NVIDIA_DRIVER_CAPABILITIES=compute,utility
                command: jupyter lab --ip='*' --allow-root --no-browser --NotebookApp.token='your_token' #접속 인증을 위해 사용할 토큰을 설정해야 합니다!
                restart: always
                volumes:
                        - /home/your_id/any_dir:/workspace #workspace로 사용할 경로를 지정합니다.

Dockerfile을 생성한 디렉토리에 docker-compose.yaml파일을 생성합니다. 이때 외부에서 서버에 접속할 포트와 인증에 사용할 토큰을 미리 설정합니다.

3. Deployment

jupyter
ㄴ Dockerfile
ㄴ docker-compose.yaml

Installation 과정을 진행한 디렉토리(Dockerfile과 docker-compose.yaml가 존재하는)에서 컨테이너를 적재하여 GPU 연산이 가능한 pytorch를 사용할 수 있는 Jupyter Lab을 구동할 수 있습니다.

docker compose up

4. Test

http://your_machine_address:your_port/lab?token=your_token

위 URL 주소에 Jupyter Lab으로 사용할 머신의 주소와 포트번호 그리고 토큰을 넣어 접속을 확인해 봅니다. 접근이 성공했다면 새로운 ipynb파일을 만들고 아래 명령어를 통해 GPU가 정상적으로 불러와지는지 확인합니다.

import torch

if torch.cuda.is_available():
    print("CUDA is available. Number of GPUs:", torch.cuda.device_count())
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
    print("CUDA is not available. Using CPU.")

"""
OUTPUT
CUDA is available. Number of GPUs: 2
GPU 0: NVIDIA GeForce RTX 3080
GPU 1: NVIDIA GeForce RTX 3080
"""

위와 같이 성공적으로 GPU를 불러왔다면 앞으로 원하는 모델을 만들어 사용해볼 수 있습니다.

결론

설정에 하루종일 매달려 있는 것보다는 만들어진 이미지 등을 사용해서 빠르게 만드는 게 좋아보입니다. Docker를 사용하는 것만으로도 설정에 필요한 많은 시간이 줄어드는 것을 알 수 있었습니다.

추가

Colab같이 다양한 라이브러리를 사용해보고 싶다면 아래와 같이 Dockerfile을 수정해 여러 파이썬 패키지를 설치하고 사용하면 좋을 것 같습니다.

FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel

RUN pip install jupyter

ADD requirements.txt /

RUN pip install -r /requirements.txt

EXPOSE 8888

dockerfile이 있는 위치에 requirements.txt 파일을 만들고 내용을 아래와 같이 입력합니다.

matplotlib
transformers
scikit-learn
pandas
scipy
seaborn
dash
nltk
lightgbm
xgboost
catboost
statsmodels
optuna
pycaret
joblib
stopit
tqdm
update_checker
deap
dask[delayed]
dask[dataframe]
dask-ml
fsspec>=0.3.3
distributed>=2.10.0
scikit-mdr
skrebate
tpot
flaml
fastai
lightning
setuptools
wheel
spacy
gensim

위 패키지들은 웹 페이지를 참고하였습니다.