Triton은 Python 백엔드를 제공하여 Python 코드로 작성된 모델 로직을 실행할 수 있도록 한다. 이는 전처리/후처리 작업이나 Scikit-Learn 모델을 서빙하는데 유용하다. Python 백엔드로 모델 서빙을 하기 위해서는 환경 세팅과 모델 코드 작성에 있어서 신경쓸 점들이 좀 많다. Python 모델로 구성된 model repository 구조를 다시 한번 복습해보자.

<model-repository-path>/
    <model-1-name>/
    	config.pbtxt
        1/
        	model.py
        2/
        	model.py
        3/
        	model.py
        ...
        
    <model-2-name>/
    	config.pbtxt
    	1/
        	model.py
        ...
        
    ...

마찬가지로 모델 configuration을 담은 config.pbtxt가 필요하고 모델 추론 Python 코드인 model.py가 추가로 요구된다.

config.pbtxt

backend : "python"
max_batch_size: 8
input [
  {
    name: "input0"
    data_type: TYPE_FP32
    dims: [ 224, 224, 3 ]
  },
  {
    name: "input1"
    data_type: TYPE_FP32
    dims: [ 224, 224, 3 ]
  } 
]
output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 1000 ]
  },
  {
    name: "output1"
    data_type: TYPE_FP32
    dims: [ 2 ]
  }
]

parameters: {
  key: "model_path",
  value: {string_value: "/home/ubuntu/model"}
}

backend 항목을 "python"으로 지정해줘야한다. 추가로 parameters 항목에 사용자 지정 파라미터를 정의하여 model.py에서 사용할 수 있다.

model.py

모델 추론 코드인 model.py는 반드시 다음과 같은 구조를 지켜줘야한다.

import triton_python_backend_utils as pb_utils


class TritonPythonModel:
    def initialize(self, args):
        """
        'initialize' 메소드는 모델이 Triton 서버에 load될 때 한 번 호출된다.
        모델의 state를 초기화하는데 사용할 수 있다. 반드시 구현하지 않아도 된다.

        Parameters
        ----------
        args : dict
          - "model_config" : config.pbtxt의 내용을 담은 JSON string
          - "model_instance_kind"
          - "model_instance_device_id"
          - "model_repository" : model repository 경로
          - "model_version"
          - "model_name"
        """
        print('Initialized...')
    
    def execute(self, requests):
        """
        'execute' 메소드는 필수적으로 구현되어야한다. 매 추론 요청마다 호출된다.

        Parameters
        ----------
        requests : list of pb_utils.InferenceRequest
          배치를 구성하는 input data의 list다.
          길이는 클라이언트 혹은 dynamic batching에 의해 구성된 배치 사이즈다.

        Returns
        -------
        list of pb_utils.InferenceResponse. requests와 길이가 동일해야된다.
        """
        
        responses = []
        for request in requests:
            ...
            responses.append(pb_utils.InferenceResponse(output_tensors=output_tensors))

        return responses
    
    def finalize(self):
    	"""
        모델이 Triton 서버에서 unload될 때 한 번 호출된다. 반드시 구현하지 않아도 된다.
        """
        print("Cleaning up...")

사용 방법이 좀더 와닿게 하기 위해 initialize 함수에서 config.pbtxt를 읽고 로컬 경로에 저장된 tensorflow 모델을 로드한 다음 execute 함수에서 해당 모델을 배치 추론하는 간단한 예시를 작성해보았다.

import json
import numpy as np
import triton_python_backend_utils as pb_utils
import tensorflow as tf


class TritonPythonModel:
    def initialize(self, args):
    	self.model_config = model_config = json.loads(args['model_config'])
        params = model_config['parameters']
        model_path = params['model_path']
        
        self.input_list = []
        for config_input in model_config['input']:
        	self.input_list.append(
            	(config_input['name'],
                 config_input['data_type'],
                 config_input['dims'])
            )
        
        self.output_list = []
        for config_output in model_config['output']:
        	self.output_list.append(
            	(config_output['name'],
                 config_output['data_type'],
                 config_output['dims'])
            )
        self.model = tf.saved_model.load(model_path)
        
        
    def execute(self, requests):
    	num_requests = len(requests)
        input_array_list = []
        responses = []
        
        for input_idx in range(len(self.input_list)):
            name, dt, shape = self.input_list[input_idx]
            first_tensor = pb_utils.get_input_tensor_by_name(requests[0], name).as_numpy()
            batched_tensor = first_tensor
            for req_idx in range(1, num_requests):
                tensor = pb_utils.get_input_tensor_by_name(requests[req_idx], name).as_numpy()
                batched_tensor = np.concatenate((batched_tensor, tensor), axis=0)
            input_array_list.append(batched_tensor)
            
        output_array_list = self.model(input_array_list)
        
        for req_idx in range(num_requests):
            output_tensors = []
            for output_idx in range(len(self.output_list)):
                name, dt, shape = self.output_list[output_idx]
                tensor = output_array_list[output_idx][req_idx]
                output_tensor = pb_utils.Tensor(
                    name,
                    tensor.astype(pb_utils.triton_string_to_numpy(dt))
                )
                output_tensors.append(output_tensor)
                
            inference_response = pb_utils.InferenceResponse(output_tensors=output_tensors)
            responses.append(inference_response)
        return responses

Multi-Threading

execute 함수에서 여러 thread에 요청들을 나누어 모델 추론을 하는 case들이 있다. Python은 GIL로 인해 여러 thread를 사용해도 한 번에 한 thread만 실행된다. 그래서 필자가 multi threading으로 실험해봤을 때 성능 향상은 없었다. 

환경 세팅

Python 백엔드는 다른 Triton 백엔드에 비해 환경 설정이 까다롭다. 서버(ex. AWS ec2 instance)세팅과 Nvidia Triton 서버 이미지 마련하기까지 끝냈다면, triton-python-backend-stub 빌드와 Python 가상환경 및 라이브러리 설치를 추가로 해야한다.

Python Backend Stub

Triton Python 백엔드는 stub process를 통해 우리의 model.py를 Triton C++ core와 연결한다. Triton 이미지에 python backend stub까지 같이 제공되지만 만약 제공되는 것과 다른 Python 버전을 사용하고 싶다면 이 stub를 새로 빌드해야된다.

필자가 사용한 Triton 22.12 릴리즈는 Python 3.8을 지원한다. Python 3.7을 사용하기 위하여 stub를 새로 빌드하는 과정을 거쳤다.

Python 가상환경과 라이브러리

필자는 Conda 가상환경을 생성하고 model.py에서 필요한 라이브러리를 pip install했다. Conda 외의 다른 가상 환경을 사용하거나 가상 환경 없이 Python을 설치해서 사용해도 된다. 단, python backend stub과 Python 버전을 반드시 맞출 것!

Code

필자가 구축한 서버는 폐쇄망에 위치했기 때문에 모든 설치 작업을 로컬에서 한 후 도커 이미지로 말아서 ec2에 배포했다. 다음과 같이 setup.sh에 python backend stub를 빌드하고 Miniconda를 설치한 다음 생성한 가상환경에서 pip install하는 코드를 작성했다.

#!/bin/bash

export CONDA_PATH = "/home/ubuntu/miniconda"
export PYTHON_BACKEND_PATH = "/home/ubuntu/python_backend"

rm -rf ${CONDA_PATH}
wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-py37_4.10.3-Linux-x86_64.sh \
    -O ${PWD}/miniconda.sh --no-check-certificate && \
    /bin/bash ${PWD}/miniconda.sh -b -p ${CONDA_PATH} && \
    rm ${PWD}/miniconda.sh && \
    $(eval echo "${CONDA_PATH}/bin/conda clean -ya")
export PATH=${CONDA_PATH}/bin:${PATH}
conda info

conda create -q -y -n venv python=3.7
source ${CONDA_PATH}/bin/activate venv

git config --global http.sslVerify false
git clone -b r22.12 https://github.com/triton-inference-server/python_backend.git ${PYTHON_BACKEND_PATH}
cd ${PYTHON_BACKEND_PATH}
rm -rf build && mkdir build && cd build
cmake -DTRITON_ENABLE_GPU=OFF -DCMAKE_INSTALL_PREFIX:PATH='pwd'/install \
      -D TRITON_BACKEND_REPO_TAG=r22.12 \
      -D TRITON_COMMON_REPO_TAG=r22.12 \
      -D TRITON_CORE_REPO_TAG=r22.12 \
      ..
make triton-python-backend-stub -j16

pip3 install tensorflow

cp ${PYTHON_BACKEND_PATH}/build/triton_python_backend_stub \
	/opt/tritonserver/backends/python/triton_python_backend_stub
export LD_LIBRARY_PATH=${CONDA_PATH}/envs/venv/lib:$LD_LIBRARY_PATH

그 다음 같은 디렉토리에 Dockerfile을 작성했다.

FROM nvcr.io/nvidia/tritonserver:22.12-py3
SHELL ["/bin/bash", "-c"]
WORKDIR /home/ubuntu
COPY . .

RUN apt-get update \
    && apt-get install -y \
                       zlib1g-dev \
                       wget \
                       libarchive-dev \
                       rapidjson-dev \
                       vim \
                       build-essential
                       
RUN wget --no-check-certificate https://github.com/Kitware/CMake/releases/download/v3.24.1/cmake-3.24.1-Linux-x86_64.sh -q -O /tmp/cmake-install.sh \
    && chmod u+x /tmp/cmake-install.sh
    && mkdir /opt/cmake-3.24.1
    && /tmp/cmake-install.sh --skip-license --prefix=/opt/cmake-3.24.1 \
    && rm /tmp/cmake-install.sh
    && ln -s /opt/cmake-3.24.1/bin/* /usr/local/bin \
    && cmake --version
    
RUN source setup.sh

이미지 빌드 -> 서버에서 docker run -> venv activate -> triton 서버 실행까지 하면 모델에 추론 요청을 보낼 수 있다.

참고로 model repository에 있는 각 모델마다 다른 Python 버전, 다른 가상환경으로 실행할 수도 있다고 하는데 참고자료에 있는 Triton python backend 공식 문서를 참고하면 좋을 것 같다.

참고자료

https://github.com/triton-inference-server/python_backend