82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
# model.py
|
|
import triton_python_backend_utils as pb_utils
|
|
import numpy as np
|
|
import json
|
|
import random
|
|
import string
|
|
|
|
class TritonPythonModel:
|
|
def initialize(self, args):
|
|
"""
|
|
모델이 로드될 때 딱 한 번만 호출됩니다.
|
|
"""
|
|
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:
|
|
try:
|
|
# 입력 텐서 직접 가져오기
|
|
self.logger.log_info("입력 텐서 직접 가져오기")
|
|
input_tensor = pb_utils.get_input_tensor_by_name(request, "INPUT")
|
|
self.logger.log_info(f"input_tensor: '{input_tensor}'")
|
|
|
|
if input_tensor is None:
|
|
# 입력이 없는 경우 기본값 처리
|
|
input_text = "default_input"
|
|
else:
|
|
# numpy 배열로 변환 후 문자열로 디코딩
|
|
input_array = input_tensor.as_numpy()
|
|
if len(input_array) > 0:
|
|
# bytes인 경우 decode, 아니면 str로 변환
|
|
first_element = input_array[0]
|
|
if isinstance(first_element, bytes):
|
|
input_text = first_element.decode('utf-8')
|
|
elif isinstance(first_element, np.bytes_):
|
|
input_text = first_element.decode('utf-8')
|
|
else:
|
|
input_text = str(first_element)
|
|
else:
|
|
input_text = "empty_input"
|
|
|
|
self.logger.log_info(f"INPUT 처리됨: {input_text}")
|
|
|
|
# 랜덤 문자열 생성
|
|
random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))
|
|
output_text = f"Processed: {input_text} -> {random_string}"
|
|
|
|
self.logger.log_info(f"OUTPUT 생성됨: {output_text}")
|
|
|
|
# 출력 텐서 생성 - bytes 형태로 직접 생성
|
|
output_bytes = output_text.encode('utf-8')
|
|
output_tensor = pb_utils.Tensor("OUTPUT", np.array([output_bytes], dtype=np.object_))
|
|
|
|
# 응답 생성
|
|
response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
|
|
responses.append(response)
|
|
|
|
except Exception as e:
|
|
# 오류 발생시 오류 응답 생성
|
|
self.logger.log_error(f"실행 중 오류 발생: {str(e)}")
|
|
error_msg = f"Error: {str(e)}"
|
|
output_bytes = error_msg.encode('utf-8')
|
|
output_tensor = pb_utils.Tensor("OUTPUT", np.array([output_bytes], dtype=np.object_))
|
|
response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
|
|
responses.append(response)
|
|
|
|
return responses
|
|
|
|
def finalize(self):
|
|
"""
|
|
모델 종료시 정리 작업
|
|
"""
|
|
self.logger.log_info(f"'{self.model_name}' 모델 종료")
|
|
pass |