May 20, 2025

[Project-ttrade] #1 Torchserve - k8s 다중 모델 Serve

TorchServe

k8s에서 MLOps하면 Kubeflow가 대표적이죠. 그리고 Model Serving은 Kubeflow 안의 KServe 모듈이 유명하죠. 근데 무겁습니다. 다양한 종류의 모델과 서비스를 한번에 할 것이 아니라면 구성할 이유가 없습니다.

그래서 저에게 맞게 TorchServe를 사용해서 Docker Image를 만들고 진행합니다.
한번에 여러 모델을 사용할 수 있고, 바뀌는 모델을 적용할 수 있는 이미지입니다.

Dockerfile

이미지의 구조를 알기 위해선 Dockerfile부터 봐야죠.
구조는 단순합니다. TorchServe 이미지를 받아다가 환경변수 설정해주고,

FROM pytorch/torchserve:latest

ENV MODEL_FOLDERS="./MODEL1|./MODEL2"
ENV MODEL_NAMES="model1|model2"
ENV MODEL_VERSIONS="1.0|1.1"
ENV TS_CONFIG_FILE=/home/model-server/config.properties

WORKDIR /home/model-server
COPY app.py .
COPY --chmod=755 create_mar_files.sh .

CMD ["/home/model-server/create_mar_files.sh"]

파일 복사해서 "create_mar_files.sh" 스크립트를 실행해 주는 것이 전부입니다. 다음으로 스크립트를 보시죠.

Create Mar File

TorchServe는 mar 파일만 있으면 모델을 배포할 수 있습니다.
mar파일에 모델의 가중치,구조 등 모든 정보를 담고 있습니다.

mar 파일을 외부에서 만들고 폴더를 연결시켜줄까... 하다가 실시간으로 생성하는 것으로 변경했습니다.

환경변수

#!/bin/bash
IFS="|" read -ra FOLDERS <<< "$MODEL_FOLDERS"         # [./MODEL1, ./MODEL2]
IFS="|" read -ra NAMES <<< "$MODEL_NAMES"             # [model1,model2]
IFS="|" read -ra VERSIONS <<< "$MODEL_VERSIONS"       # [1.0,1.0]
IFS="|" read -ra PARAMS <<< "$PARAM_FILES"            # [model1.pth,model2.pth]

환경변수를 각각 배열로 바꿔줍니다. 모델을 여러개 할 수 있게 배열로 설정해 주었습니다. 모델의 폴더, 이름, 버전, 파라미터를 설정합니다.

모델 Mar 파일 생성

모델 폴더들의 Model.py(메인 모델이 작성된 파일)과 그 외 파일들을 가지고 mar 파일을 만드는 과정입니다.

for i in "${!FOLDERS[@]}"; do
  FOLDER="${FOLDERS[i]}"
  MODEL_FILE="${FOLDER}/Model.py"
  EXTRA_FILES=$(find "${FOLDER}" -maxdepth 1 -type f ! -name "Model.py" | tr '\n' ',' | sed 's/,$//')
  
  torch-model-archiver --model-name "${NAMES[i]}" \
                       --version "${VERSIONS[i]}" \
                       --model-file "${MODEL_FILE}" \
                       --serialized-file "${PARAMS[i]}" \
                       --handler app.py \
                       --extra-files "${EXTRA_FILES}" \
                       --export-path model-store
done

모델 serving

TorchServe 의 config에 맞게 모델을 전부 serving 합니다.

torchserve --start --ts-config $TS_CONFIG_FILE --models all

TorchServe Handler

mar 파일을 만들기 위한 Handler를 작성해 줘야 합니다. 기본적인 구성은 __init__, initialize(), preprocess(), inference(), hanlde() 로 되어 있습니다.

init

import torch
import numpy as np
from ts.torch_handler.base_handler import BaseHandler
import json
from Model import Model,Configs
import os

class ModelHandler(BaseHandler):
    def __init__(self):
        super().__init__()
        self.initialized = False

initialize

__init__ 은 Handler 객체가 생성될 때 실행되고, initialize는 모델이 로드될 때 실행됩니다. 둘 다 mar 파일 생성 시에는 호출되지 않고, torchserve를 시작하면 실행됩니다. initialize를 나눠서 재로드 될 때, initialize만 호출되게 합니다.

    def initialize(self,context):
        self.manifest = context.manifest
        properties = context.system_properties
        model_dir = properties.get("model_dir")
        serialized_file = self.manifest['model']['serializedFile']
        model_pt_path = os.path.join(model_dir, serialized_file)
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        with open(f"{model_dir}/model.json",'r') as f:
            self.config= json.load(f)
        self.model = Model(Configs(self.config))
        state_dict = torch.load(model_pt_path,map_location=torch.device('cpu'),weights_only=True)
        self.model.load_state_dict(state_dict)
        self.model.to(self.device).eval()
        self.initialized = True

Torchserve를 시작하면서 받은 파라미터를 이용해서 모델의 경로를 설정하고 로드합니다.

preprocess

들어오는 data를 미리 처리하는 과정입니다.
보통 http를 이용해서 data를 body에 넣어서 보내기에 numpy로 받고 tensor로 변경합니다.

    def preprocess(self,data):
        input_array=np.frombuffer(data[0].get("body"), dtype=np.float64).reshape(1, -1, self.config['enc_in'])
        return torch.from_numpy(input_array).float()

inference

데이터를 연산하는 과정입니다. 그냥 모델이 집어 넣어 줍시다.

    def inference(self,input_tensor):
        with torch.no_grad():
            output = self.model(input_tensor)
        return output

handle

핸들러의 handle 부분입니다. 데이터를 받아서 모델에 넣고 출력해 줍니다.

    def handle(self,data,context):
        model_input = self.preprocess(data)
        model_output=self.inference(model_input)
        return model_output.cpu().numpy().tolist()

TorchServe Config

torchserve 를 시작할 때에 $TS_CONFIG_FILE 라는 환경변수로 config값을 넣어 주었습니다. torchserve를 시작할 때의 config 값입니다.

inference_address=http://0.0.0.0:8080
management_address=http://0.0.0.0:8081
metrics_address=http://0.0.0.0:8082
model_store=/home/model-server/model-store
disable_token_authorization=true
default_workers_per_model=2
enable_envvars_config=true
install_py_dep_per_model=true
default_response_timeout=20
unregister_model_timeout=20

config.properties

inference_address: 모델을 사용하는 아이피와 포트
management_address: 모델을 확인하는 아이피와 포트
metrics_address: torchserve의 상태를 확인하는 아이피와 포트
model_store: mar파일이 저장되어 있는 폴더

이제 Dockerfile을 build하고, 적절한 위치에 파일을 넣어 실행시키면 됩니다.