91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
import triton_python_backend_utils as pb_utils
|
|
import numpy as np
|
|
import json
|
|
import random
|
|
|
|
class TritonPythonModel:
|
|
def initialize(self, args):
|
|
"""
|
|
모델이 로드될 때 딱 한 번만 호출됩니다.
|
|
`initialize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
|
|
이 모델과 관련된 모든 상태를 초기화할 수 있습니다.
|
|
"""
|
|
self.logger = pb_utils.Logger
|
|
|
|
self.model_name = args["model_name"]
|
|
self.model_config = json.loads(args["model_config"])
|
|
|
|
self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료")
|
|
|
|
|
|
def execute(self, requests):
|
|
"""
|
|
Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다.
|
|
"""
|
|
responses = []
|
|
|
|
# 각 추론 요청을 순회하며 처리합니다.
|
|
for request in requests:
|
|
# Triton 입력 파싱
|
|
input_text = self._get_input_value(request, "INPUT")
|
|
self.logger.log_info(f"INPUT 출력:\n{input_text}")
|
|
|
|
random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))
|
|
|
|
output = random_string
|
|
self.logger.log_info(f"OUTPUT 출력:\n{output}")
|
|
|
|
# 생성된 텍스트를 Triton 출력 텐서로 변환합니다.
|
|
output_tensor = pb_utils.Tensor("OUTPUT", np.array(output.encode('utf-8'), dtype=np.bytes_))
|
|
|
|
# 응답 객체를 생성하고 출력 텐서를 추가합니다.
|
|
responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))
|
|
|
|
return responses
|
|
|
|
def _get_input_value(self, request, input_name: str, default=None):
|
|
"""
|
|
Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다.
|
|
|
|
Args:
|
|
request (pb_utils.InferenceRequest): Triton 추론 요청 객체.
|
|
input_name (str): 가져올 입력 텐서의 이름.
|
|
default (any, optional): 입력 텐서가 없을 경우 반환할 기본값. Defaults to None.
|
|
|
|
Returns:
|
|
any: 디코딩된 입력 텐서의 값. 텐서가 없으면 기본값을 반환합니다.
|
|
"""
|
|
tensor_value = pb_utils.get_input_tensor_by_name(request, input_name)
|
|
|
|
if tensor_value is None:
|
|
return default
|
|
|
|
return self._np_decoder(tensor_value.as_numpy()[0])
|
|
|
|
def _np_decoder(self, obj):
|
|
"""
|
|
NumPy 객체의 데이터 타입을 확인하고 Python 기본 타입으로 변환합니다.
|
|
|
|
Args:
|
|
obj (numpy.ndarray element): 변환할 NumPy 배열의 요소.
|
|
|
|
Returns:
|
|
any: 해당 NumPy 요소에 대응하는 Python 기본 타입 (str, int, float, bool).
|
|
bytes 타입인 경우 UTF-8로 디코딩합니다.
|
|
"""
|
|
if isinstance(obj, bytes):
|
|
return obj.decode('utf-8')
|
|
if np.issubdtype(obj, np.integer):
|
|
return int(obj)
|
|
if np.issubdtype(obj, np.floating):
|
|
return round(float(obj), 3)
|
|
if isinstance(obj, np.bool_):
|
|
return bool(obj)
|
|
|
|
def finalize(self):
|
|
"""
|
|
모델 실행이 완료된 후 Triton 서버가 종료될 때 한 번 호출되는 함수입니다.
|
|
`finalize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
|
|
종료 전에 필요한 모든 정리 작업을 수행할 수 있습니다.
|
|
"""
|
|
pass |