딥러닝 API, GPU+도커로 서빙하기 (feat. on premise)

COG?

COG(Github)는 PyTorch와 같은 딥러닝 모델을 서빙하기 위해 나온 🐳도커 이미지 빌드 툴로, 간단한 YAML 파일과 Python 패키지 설정만으로 FastAPI에 기반한 웹 API를 제작해준다.

🤗 Transformers 라이브러리로 사용해보기

COG Yaml 파일

cog yaml 파일은 아래와 같은 형식을 가진다. 가장 주요한 특징은 gpu: true로 옵션을 제공할 경우, 아래 python_packages에 torch나 tensorflow와 같은 DL 패키지가 적혀 있는 경우 해당 패키지에 적합한 CUDA, cuDNN이 사전에 설정된 Ubuntu기반 이미지(NVIDIA가 제공)를 불러와서 빌드를 진행해준다는 점이다.

# cog.yaml
build:
  gpu: true
  python_version: "3.9"
  python_packages:
    - "torch==1.10.0"
    - "transformers==4.18"
    - "sentencepiece==0.1.96"
    - "protobuf==3.20.1"
predict: "predict.py:Predictor"

주의: 위 python_packages은 정확히 == 으로 버전을 pin해줘야 동작한다. <등의 뭉뚱그려진 Requirements는 빌드를 거절하더라.

위와 같은 Yaml 파일로 설정을 할 경우, Ubuntu 18.04 LTS에 CUDA 11.1이 설치된 이미지를 Base 이미지로 사용하고, 그 위에 PyTorch를 1.10(CUDA 11.1)을 설치한다.

다만, python_version을 지정하는 경우 해당 파이썬 버전을 .tar.xz 소스로 받아 설치하는 과정이 있어, 첫 빌드 시간이 꽤나 걸린다는 단점이 있다. 따로 파이썬 버전을 지정하지 않는 경우에는 시스템에 내장된 버전을 사용하는 듯 하다. (Ubuntu 18.04는 Python 3.7 등)

COG 파이썬 파일

아래 파이썬 파일을 통해 실제 모델을 서빙하게 된다. Huggingface Transformers 라이브러리를 사용한다고 별다른 차이가 발생하지는 않는다.

# predict.py
from cog import BasePredictor, Input, Path
import torch
from transformers import pipeline
import json


class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.pipe = pipeline(
            "text-generation",
            model="beomi/kcgpt2",
            device=0 if torch.cuda.is_available() else -1,
        )

    # The arguments and types the model takes as input
    def predict(self, x: str = Input(description="문장을 입력해 주세요.")) -> str:
        """Run a single prediction on the model"""
        output = self.pipe(
            x,
            do_sample=True,
            top_p=0.9,
            no_repeat_ngram_size=2,
            # early_stopping=True,
            # max_new_tokens=150,
            num_return_sequences=20,
            max_length=150,
        )
        return json.dumps(output, ensure_ascii=False, indent=2)

다만, 위와 같이 Transformers Hub에서 모델을 다운받아야 하는 경우, 해당 모델 파일을 도커 이미지에 함께 포함하지 않아서, 매 실행시마다 Hub에서 다운받는 문제가 있다.

이 문제는 두 가지 방법이 있는데, 상황에 따라 적절히 선택하면 될 것 같다.

1. 도커 실행시 Huggingface 라이브러리 캐싱 폴더를 Volume으로 마운트 (도커 레벨에서 재활용)
2. 모델을 애초에 다운받아서, 도커 빌드시에 포함 (도커 이미지에 포함시키기)

일반적으로는 PLM을 Finetune한 사내 모델을 사용하기 때문에, 2번과 같이 로컬 폴더에서 도커 이미지 빌드 시점에 통합해서 사용하는 것이 일반적일 것 같다.

이미지 굽기

위와 같이 준비가 완료되면, 아래와 같이 빌드를 진행할 수 있다.

cog build -t cog-test

이미지 실행하기

이미지가 준비되면, 해당 도커 이미지는 내부 5000포트로 FastAPI가 서빙되고 있다.

해당 FastAPI는 Uvicorn을 통해 ASGI로 서빙된다.

이때, --gpus 1와 같이 GPU 옵션을 제공해 넘겨주면 (기존 predict.py 코드가 GPU를 지원한다면) 웹 API 서비스 역시 GPU 가속을 지원해준다.

docker run -it -p 5000:5000 --gpus 1 cog-test

참고: --gpus옵션은 '갯수'라서, 1=1개, 2=2개,.. 이다. 만약 특정 idx의 GPU를 지정하려면 --gpus '"device=0,2"'같이 특정 idx(0,2번)의 GPU를 지정해줄 수 있다. device 번호로 지정은 상당히 귀찮다..

COG UI

COG는 웹 Swagger 형식으로 UI를 지원해준다. 간단한 POST 테스트를 웹 상에서도 진행가능!

성능

성능은 차이가 아예 없다. 모델 컴파일 작업이나 경량화 작업 혹은 최적화 작업이 없이, 기존의 모델 forward 코드를 그대로 가져다 쓸 수 있다는 것이 장점이자 단점.

속도 향상 라이브러리

성능 향상을 위해서는 여러가지 옵션을 생각해볼 수 있을 듯 하다.

nebullvm와 같은, ONNX 등을 통한 CPU/GPU상에서 서빙을 가속하는 라이브러리와 함께 사용한다면, 보다 라이브에 태울 수 있는 수준의 속도가 나오지 않을까.

Reference

• COG: https://github.com/replicate/cog