Triton String Input

박국현·2023년 5월 25일
0

공부

목록 보기
9/9

Triton Inference Server의 장점 중 하나는 여러 모듈을 묶는 ensemble 기능이다. 이 전 게시글에서 언급했듯이 전처리 + 모델 inference + 후처리를 묶어 하나의 요청에서 끝낼 수 있기 때문이다.

최근 AI 영역에서 가장 핫한 분야는 언어 모델(Language Model)인데 이 모델들은 입력을 자연어, 즉 string 자료형으로 받는다.
이 상황에서 내가 풀고 싶었던 문제는 다음과 같다.

  1. numpy를 사용하지 않고 문자열을 담은 요청을 보내는 법
  2. tritonclient를 사용하지 않고 요청을 보내는 법

이 두 내용을 정리한 문서가 거의 없다...결국 ChatGPT와 Bard에게 물어가며 알아낸 방법을 여기에 정리해본다.

이 글은 tritonserver 컨테이너가 실행되고 있다고 가정한다.

서버(요청을 받는 곳)

공식 깃헙에 들어있는 코드의 일부를 사용해 간단한 모델을 만들자. 모델의 이름은 string_identity. 문자열 데이터를 담은 요청이 오면 그 문자열을 그대로 응답으로 보내는 모델이다.

다음과 같은 디렉토리를 만든다. 아무 곳이든 상관없다.

models
	- string_identity
    	- config.pbtxt
    	- 1
        	- model.py

각 파일의 내용은 다음과 같이 만든다.

# config.pbtxt
name: "string_identity"                                                  
backend: "python"                                          
                                                        
input [                                                                  
  {                                                                      
    name: "INPUT0"                                                       
    data_type: TYPE_STRING                                               
    dims: [ 1 ]                                                          
  }                                                                      
]                                                                        
output [                                                                 
  {                                                                      
    name: "OUTPUT0"                                                      
    data_type: TYPE_STRING                                               
    dims: [ 1 ]                                                          
  }                                                                      
]                                                                        
# model.py
import sys                                                                               
import json                                                                              
                                                                                         
sys.path.append('../../')                                                                
import triton_python_backend_utils as pb_utils                                           
import numpy as np                                                                       
                                                                                         
                                                                                         
class TritonPythonModel:                                                                 
    """This model always returns the input that it has received.                         
    """                                                                                  
                                                                                         
    def initialize(self, args):                                                          
        self.model_config = json.loads(args['model_config'])                             
                                                                                         
    def execute(self, requests):                                                         
        """ This function is called on inference request.                                
        """                                                                              
        responses = []                                                                   
        for request in requests:                                                         
            in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT0")                  
            out_tensor_0 = pb_utils.Tensor("OUTPUT0", in_0.as_numpy().astype(np.object_))
            print(f"INPUT VALUE: {in_0.as_numpy()[0].decode()}")  # 요청 데이터 출력
            responses.append(pb_utils.InferenceResponse([out_tensor_0]))                 
        return responses                                                                 

위 내용대로 파일이 준비됐으면 triton 서버를 실행한다. models 경로에서 아래 명령어를 실행하면 된다.

tritonserver --model-repository .

Client(요청을 보내는 곳)

다시 처음으로 돌아가서, 내가 확인하고 싶었던 것은

  1. numpy를 사용하지 않고 문자열을 담은 요청을 보내는 법
  2. tritonclient를 사용하지 않고 요청을 보내는 법

이렇게 두 가지이다.
따라서 requests 모듈만 가지고 문자열을 담은 요청을 보내보자.

import requests                                                
                                                               
                                                               
URL = "http://localhost:8000/v2/models/string_identity/infer"  
def main():                                                    
    data = {                                               
                "name": "string_identity",             
                "inputs": [                            
                    {                           
                         "name": "INPUT0",      
                         "shape": [1],          
                         "datatype": "BYTES",   
                         "data": ["한글입니다!"]
                    }                           
                ]                                              
    }                                                      
    res = requests.post(URL, json=data)                    
    print(res.json())                                      
    return                                                 
                                                               
if __name__ == "__main__":                                     
    main()                                                 

위 코드에서 주목할 점은 다음과 같다.

  • tritonclient를 사용할 때와 URL이 다르다.
  • 요청에 담을 데이터의 구조는 요청을 받을 모델의 config에서 설정한 구조와 똑같아야 한다.
  • datatypeBYTES로 설정해야 한다.

위 파일을 실행해서 문제 없이 서버에서 다음이 출력되면 된다!

INPUT VALUE: 한글입니다!
profile
공부하자!!

2개의 댓글

comment-user-thumbnail
2023년 10월 11일

감사합니다!! 트리톤 보고 있었는데 많은 도움이 되었습니다..

답글 달기
comment-user-thumbnail
2023년 10월 11일

감사합니다!! 트리톤 보고 있었는데 많은 도움이 되었습니다..

답글 달기