[MLOps] 1. KServe-Torchserve 를 위한 TorchModelArchiver 서빙 준비

서론

회사에서 kserve로 pytorch 모델을 배포하는 것을 시도해보게 되었다. 이 과정에서 배운 것들을 정리하고자 한다.
먼저, kserve에서 pytorch 를 framework로 지정하고 사용하면, TorchServe가 사용된다. 그리고 TorchServe를 사용하기 위해서는 TorchServe Model Archive(.mar) 와 Model Configuration (config.properties) 이 필요한다. 이를 지원하는 것이 바로 torch-model-archiver 이다. 이번 글에서는 이에 대해 설명해본다.

본론

1. 동작

먼저 모델이 개발이 완료되면, model.py 파일 그리고 .pt 파일이 있을 것이다. 예를 들어 아래와 같은 구조로 파일이 있을 것이다.
.

mnist
├── src
│   ├── mnist_model.py
│   ├── requirements.txt
├── ckpt
│   ├── mnist.pt

.
편의상 모델은 아래와 같이 1층 레이어 구조라고 하자.
.

# mnist_model.py
from torch import nn  

class Net(nn.Module):  
    def __init__(self):  
        super(Net, self).__init__()  
        self.linear = nn.Linear(28 * 28, 10)  

    def forward(self, x):  
        x = self.linear(x)  
        return x

.
이렇게 파일이 있으면 더이상 아무것도 준비하지 않고 모델을 서빙할 수 있으면 참 좋겠지만, 그렇지 못하다.

모델을 동작시킬 때는 일단 handler.py 파일이 필요하다. handler.py 파일의 역할은 데이터가 들어왔을 때, 전처리/추론/후처리를 어떻게 할지 정의해 놓은 파일이다. 일단 해당 파일이 어떻게 구성 되어 있는지 살펴 보자. BaseHandler 파일 을 보면 initialize, preprocess, inference, postprocess, handle, explain_handle, describe_handle 함수들이 있는 것을 볼 수 있는데 (private 함수 제외), 이 중 필요한 부분을 직접 구현해서 custom handler를 만들어 보도록 하자. 참고로, 향후 kserve에서 처음 서버를 띄울때, handler 파일의 initialize 함수가 실행되고, POST로 파일을 보냈을 때, handler 함수가 실행되어서 이 두 함수만 구현했다. (이부분은 다음 kserve 포스팅에서 더 자세히 설명 될 것이다)
.

# handler.py 
import torch  
import os
from mnist_model import Net  
from ts.torch_handler.base_handler import BaseHandler

class MNISTBasicNNHandler(BaseHandler):  
    """  
    A custom model handler implementation.    
    """  
    def __init__(self):  
        self._context = None  
        self.initialized = False  
        self.model = None  
        self.device = None  

    def initialize(self, context):  
        """  
        Invoke by torchserve for loading a model
        :param context: context contains model server system properties        
        :return:        
        """  
        #  load the model        
        self.manifest = context.manifest  
        properties = context.system_properties  
        model_dir = properties.get("model_dir")  
        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")  

        # Read model serialize/pt file  
        serialized_file = self.manifest['model']['serializedFile']  

        model_pt_path = os.path.join(model_dir, serialized_file)  
        if not os.path.isfile(model_pt_path):  
            raise RuntimeError("Missing the model.pt file")  

        self.model = Net().to(self.device)  
        self.model.load_state_dict(torch.load(model_pt_path, map_location=self.device))  
        self.model.eval()  

        self.initialized = True  

    def handle(self, data, context):  
        """  
        Invoke by TorchServe for prediction request.
        Do pre-processing of data, prediction using model and postprocessing of prediciton output
        :param data: Input data for prediction        
        :param context: Initial context contains model server system properties.        
        :return: prediction output        
        """        
        pred_out_list = []  
        for cur_data in data:  
            imgstring = cur_data['image_bytes']['b64']  # expected data structure : {'image_bytes':{'b64': '...'}}            
            cur_data = self.preprocess(cur_data) # preprocess data from b64 to torch. data should be viewed with torch.view(-1,28,28)
            pred_out = self.model.forward(cur_data)  
            pred_out_list.append(pred_out.tolist())  
        return pred_out_list

.
이렇게 파일이 준비되었다면 이제, 폴더는 아래와 같이 되었을 것이다.
.

mnist
├── src
│   ├── mnist_model.py
│   ├── requirements.txt
│   ├── handler.py
├── ckpt
│   ├── mnist.pt

.
이제 위 파일들을 torchserve에 사용할 수 있도록 TorchServe Model Archive(.mar) 와 Model Configuration (config.properties) 파일을 만들어 보자. 이는 torch-model-archiver 라는 CLI 를 사용하면 쉽게 만들 수 있다. 자세한 내용은 TorchModelArchiver 문서에서 확인할 수 있고, 여기서는 위 예시에 대한 사용법만 알아보자.

먼저 pip install torch-model-archiver 을 했다면 모든 준비는 끝났다(참 쉽죠??). 이제 명령어로 필요한 두 파일을 생성하면 끝이다. 명령어는 아래와 같다
.
torch-model-archiver --model-name mnist --version 1.0 --serialized-file ./ckpt/mnist.pt --extra-files ./src/mnist_model.py --handler ./src/handler.py --requirements-file ./src/requirements.txt
.
필자의 경우에는 --model_file 인자를 사용하지 않고 --extra_files에 모델 파일을 넘겼다. 두 경우 모두 가능함을 보여주기 위함이다(아마 model_file 에 넘겨도 되는데 차이점은 torchserve내에서 os.env 에 저장된 model_path(?) 이런 경로로 바로 모델 파일을 찾을 수 있는지, 혹은 그냥 사용자가 모델 파일의 위치를 알아서 그곳을 import 할지가 다를 것이다). 만약 모델 파일이 여러 파일들로 나뉘어져 있다면, --extra-files 에 모델 파일들이 있는 최상위 디렉토리를 넘겨주면 되고, 만약 두개이상의 파일이나 디렉토리를 넘겨주려면 ',' 로 연결해서 써주면 된다. 그리고 위의 TorchModelArchiver 가이드 링크의 예시에서는 handler를 custom 을 사용하지 않고, image_classifier를 사용하고 있는데, 본인의 모델과 추론할 이미지가 기본 제공 image_classifier handler 와 호환 가능하면 그대로 사용해도 된다. 기본 제공 handler는 image_classifier, object_detector, text_classifier, image_segmenter 가 있는데, 잘 알아보고 사용하기 바란다.
아무튼 이렇게 명령어를 실행하면 드디어 mar파일이 생성된다!! 파일명은 --model-name으로 파일이 생성된다. 최종 결과는 아래와 같다.
.

mnist
├── src
│   ├── mnist_model.py
│   ├── requirements.txt
│   ├── handler.py
├── ckpt
│   ├── mnist.pt
├── mnist.mar

.
MAR File의 내부구조는 args에 포함한 파일 + MAR_INF 폴더 내부에 있는 json 파일로 구성된다(내부구조가 궁금하시면 찾아보시길 바란다). 그래서 이 파일을 다른데로 TorchServe 에서 필요로 하는 곳으로 옮기고 실행하면 된다.

결론

torchserve를 돌리기 위해서는 mar 파일이 필요한데, 이를 만들어주는 명령어를 사용하기 위한 작업을 모델 개발 완료 시점 이후 부터 알아보았다. 아주아주 허접한 글이지만... 누군가에게 조금이나마 도움이 되었으면 한다. 다음 글에서 지금 만든 .mar 파일로 torchserve를 띄우는 법을 알아보자!