Triton Inference Server랑 친해지기 (1)

업무 중 NVIDIA Triton Infernece Server를 사용하여 모델을 서빙해야하는 일이 생겼는데요, 사실 저에게는 굉장히 생소했습니다. (일단 해볼 일이 없었음..) 굉장히 많이 헤메고 시간을 많이 쏟았어서, Triton Server를 알아보고, 직접 이미지를 빌드해서 AWS의 SageMaker의 endpoint로 말아올리는 과정까지를 정리해보고자 합니다.

Triton Inference Server, 왜 사용하나요? 🤔

먼저 PyTorch, Tensorflow 등으로 만들어진 모델들을 ONNX나 TensorRT등을 이용하여 경량화하는데 이러한 프레임워크들을 모두 지원하고 있습니다.
또한 모델을 model_repository를 통해 관리할 수 있는데요, 아래에 추가로 언급할 예정이지만 해당 레포지토리를 통해 모델의 버전을 쉽게 관리할 수 있습니다.
추가적으로는 여러개의 모델을 하나의 파이프라인으로 연결할 수 있는 앙상블 모델도 지원합니다.

Triton Inference Server 버전 찾기

Triton inference server 컨테이너의 릴리즈 및 각 프레임워크의 버전들은 아래의 링크에서 확인할 수 있습니다.

https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/index.html

model_repository 구성하기

model_repository는 단순하게 모델이 저장되는 레포지토리라고 생각하면 되는데요. 해당 디렉토리 아래에는 다양한 프레임워크를 사용하여 생성된 모델들, 모델의 버전 및 모델들의 메타데이터(config.pbtxt)가 저장됩니다.
공식 문서에 따르면 다음과 같은 Layout을 맞추어 경로를 지정해주어야 합니다.

  <model-repository-path>/
    <model-name>/
      [config.pbtxt]
      [<output-labels-file> ...]
      [configs]/
        [<custom-config-file> ...]
      <version>/
        <model-definition-file>
      <version>/
        <model-definition-file>

또한 Server를 실행할 때 model_repository 경로를 인자로 넣어주어야 합니다.

config.pbtxt

공식 문서에서 소개하는대로, 모든 모델들은 각각의 config.pbtxt 파일을 하나씩 가지고 있어야 하며 해당 파일의 이름 역시 고정입니다.

  platform: "tensorrt_plan"
  max_batch_size: 8
  input [
    {
      name: "input0"
      data_type: TYPE_FP32
      dims: [ 16 ]
    },
    {
      name: "input1"
      data_type: TYPE_FP32
      dims: [ 16 ]
    }
  ]
  output [
    {
      name: "output0"
      data_type: TYPE_FP32
      dims: [ 16 ]
    }
  ]

위는 가장 기본적으로 config.pbtxt 파일에 포함되어야 하는 내용입니다.
input, output에는 데이터의 타입과 dimension을 명시해야하며, platform은 모델의 framework로 생성되었는지를 적어주는 부분입니다. max_batch_size의 경우 해당 모델이 지원하는 최대 배치의 사이즈를 의미합니다.
위의 config.pbtxt 파일에 의하면 해당 모델은 입력이 2개이며 (input1, input2), 출력은 output0 하나이며 모든 입출력은 FLOAT32 텐서인 TensorRT로 작성된 모델입니다.

Ensemble 모델의 경우 일반적인 모델과 달라지는 부분이 있는데요, 보통 preprocess > inference > postprocess 의 구조를 가지고 있기 때문에 이 과정을 추가로 명시해주어야 합니다.

name: "ensemble_dali_inception"
platform: "ensemble"
max_batch_size: 256
input [
  {
    name: "INPUT"
    data_type: TYPE_UINT8
    dims: [ -1 ]
  }
]
output [
  {
    name: "OUTPUT"
    data_type: TYPE_FP32
    dims: [ 1001 ]
  }
]
ensemble_scheduling {
  step [
    {
      model_name: "dali"
      model_version: -1
      input_map {
        key: "DALI_INPUT_0"
        value: "INPUT"
      }
      output_map {
        key: "DALI_OUTPUT_0"
        value: "preprocessed_image"
      }
    },
    {
      model_name: "inception_graphdef"
      model_version: -1
      input_map {
        key: "input"
        value: "preprocessed_image"
      }
      output_map {
        key: "InceptionV3/Predictions/Softmax"
        value: "OUTPUT"
      }
    }
  ]
}

위와 같이 ensemble_scheduling에 의해 모델이 어떻게 작동해야하는지를 지정할 수 있습니다. 위 config.pbtxt에 따르면 Ensemble 모델 입력 > dali 모델 입력 > dali 모델 출력 > inception_graphdef 모델 입력 > inception_graphdef 모델 출력 > Ensemble 모델 출력 의 순서로 모델이 동작합니다.

다음 글에서는 무엇을 할거냐면요

Triton github에서 제공하고 있는 build.py를 이용하면 triton server image를 쉽게 빌드할 수 있는데요, 빌드 후 AWS의 ECR 레포지토리로 푸시하여 SageMaker의 endpoint로 말아올리는 과정에 대하여 소개할 예정입니다. 추가적으로 빌드를 하면서 겪었던 이슈들과 빌드 시 지정할 수 있는 옵션들과 옵션의 의미, serve 파일에서 일어나는 일 등을 글로 정리해보고자 합니다.

읽어주셔서 감사합니다 (_ _ )