Update 1/model.py

1/model.py

@@ -2,6 +2,7 @@ import json
 import torch
 import numpy as np
 import triton_python_backend_utils as pb_utils
+import uuid  # [추가] 랜덤 ID 생성을 위해 필요
 from typing import List, Dict, Any, Union, Tuple
 
 from transformers import (
@@ -44,7 +45,6 @@ class TritonPythonModel:
         bnb_config = self._get_bnb_config()
 
         # 2. Base Model 로드
-        # Adapter 모델인 경우 Config에서 Base 경로를 덮어쓸 수 있음
         load_path = self.base_model_path
         if self.is_adapter_model:
             peft_config = PeftConfig.from_pretrained(self.adapter_model_path)
@@ -67,26 +67,22 @@ class TritonPythonModel:
         if self.is_adapter_model:
             self.model = PeftModel.from_pretrained(self.model, self.adapter_model_path)
         
-        # 추론 모드 설정
         self.model.eval()
 
         # 4. Tokenizer 로드
         self.tokenizer = AutoTokenizer.from_pretrained(load_path, trust_remote_code=True)
         
-        # Pad Token 설정 (없을 경우 EOS로 대체)
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
             self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
             self.logger.log_info("Pad token was None. Set to EOS token.")
 
-        # Chat Template 지원 여부 확인
         self.supports_chat_template = (
             hasattr(self.tokenizer, "chat_template") and 
             self.tokenizer.chat_template is not None
         )
 
     def _get_bnb_config(self) -> Union[BitsAndBytesConfig, None]:
-        """양자화 설정 객체를 반환합니다."""
         if self.quantization == "int4":
             return BitsAndBytesConfig(
                 load_in_4bit=True,
@@ -107,10 +103,19 @@ class TritonPythonModel:
         responses = []
 
         for request in requests:
+            # [ID 생성 로직]
+            # 클라이언트가 보낸 ID가 있으면 사용하고, 없으면 랜덤 UUID 생성
+            request_id = request.request_id()
+            if not request_id:
+                request_id = str(uuid.uuid4())
+
             try:
                 # 1. 입력 데이터 파싱
                 input_data, is_chat = self._parse_input(request)
-                self.logger.log_info(f"\n>>> [Input]: {input_data}")
+                
+                # [LOGGING] Request ID 포함하여 로그 출력
+                log_input_str = json.dumps(input_data, ensure_ascii=False) if isinstance(input_data, (list, dict)) else str(input_data)
+                self.logger.log_info(f"\n[RID: {request_id}] >>> [Input]: {log_input_str}")
                 
                 # 2. Generation Config 생성
                 gen_config = self._create_generation_config(request)
@@ -121,37 +126,31 @@ class TritonPythonModel:
                 # 4. 모델 추론 (Generate)
                 output_text = self._generate(inputs, gen_config)
 
-                # [LOGGING] 모델 응답 로그 출력
+                # [LOGGING] Request ID 포함하여 결과 출력
-                self.logger.log_info(f"\n<<< [Output]: {output_text}")
+                self.logger.log_info(f"\n[RID: {request_id}] <<< [Output]: {output_text}")
 
-                # 5. 응답 생성
+                # 5. 응답 생성 (ID 전달)
-                responses.append(self._create_response(output_text))
+                responses.append(self._create_response(output_text, request_id))
 
             except Exception as e:
-                self.logger.log_error(f"Error during execution: {e}")
+                self.logger.log_error(f"[RID: {request_id}] Error during execution: {e}")
-                # 에러 발생 시 빈 문자열 또는 에러 메시지 반환 (클라이언트 처리에 따라 변경 가능)
                 err_tensor = pb_utils.Tensor("text_output", np.array([str(e).encode('utf-8')], dtype=np.bytes_))
-                responses.append(pb_utils.InferenceResponse(output_tensors=[err_tensor]))
+                # 에러 응답에도 ID 포함
+                responses.append(pb_utils.InferenceResponse(output_tensors=[err_tensor], id=request_id))
 
         return responses
 
     def _parse_input(self, request) -> Tuple[Union[str, List[Dict]], bool]:
-        """입력 텐서를 파싱하여 텍스트 또는 대화 목록과 타입(채팅 여부)을 반환합니다."""
         input_text = self._get_input_scalar(request, "text_input")
-        
         try:
-            # JSON 형식의 대화 기록인지 시도
             conversation = json.loads(input_text)
-            # 리스트 형태여야 채팅 히스토리로 간주
             if isinstance(conversation, list):
                 return conversation, True
         except (json.JSONDecodeError, TypeError):
             pass
-        
         return input_text, False
 
     def _tokenize(self, input_data, is_chat: bool):
-        """입력 데이터를 토큰화하여 모델 입력 텐서를 반환합니다."""
         if self.supports_chat_template and is_chat:
             return self.tokenizer.apply_chat_template(
                 input_data,
@@ -161,17 +160,11 @@ class TritonPythonModel:
                 return_dict=True
             ).to(self.device)
         else:
-            # 일반 텍스트인 경우
+            if is_chat:
-            if is_chat: # Chat 형식이지만 템플릿 미지원 시 문자열 변환 시도
                  input_data = str(input_data)
-            
+            return self.tokenizer(input_data, return_tensors="pt").to(self.device)
-            return self.tokenizer(
-                input_data,
-                return_tensors="pt"
-            ).to(self.device)
 
     def _generate(self, inputs, gen_config: GenerationConfig) -> str:
-        """모델 생성 로직 수행 및 디코딩"""
         input_ids = inputs["input_ids"]
         input_len = input_ids.shape[-1]
 
@@ -183,15 +176,11 @@ class TritonPythonModel:
                 eos_token_id=self.tokenizer.eos_token_id
             )
 
-        # 입력 토큰을 제외하고 생성된 토큰만 슬라이싱
         generated_tokens = outputs[0][input_len:]
         decoded_output = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-        
         return decoded_output.strip()
 
     def _create_generation_config(self, request) -> GenerationConfig:
-        """Request에서 파라미터를 추출하여 GenerationConfig 객체 생성"""
-        # 기본값 설정 및 입력값 추출 Helper
         def get_param(name, default=None, cast_type=None):
             val = self._get_input_scalar(request, name, default)
             if val is not None and cast_type:
@@ -199,56 +188,37 @@ class TritonPythonModel:
             return val
 
         return GenerationConfig(
-            max_length=get_param("max_length", 1024, int), # max_length 기본값은 넉넉하게
+            max_length=get_param("max_length", 1024, int),
             max_new_tokens=get_param("max_new_tokens", 256, int),
             temperature=get_param("temperature", 1.0, float),
             do_sample=get_param("do_sample", False, bool),
             top_k=get_param("top_k", 50, int),
             top_p=get_param("top_p", 1.0, float),
             repetition_penalty=get_param("repetition_penalty", 1.0, float),
-            # stream=get_param("stream", False, bool) # Python Backend에서 Stream은 별도 구현 필요
         )
 
-    def _create_response(self, output_text: str):
+    def _create_response(self, output_text: str, request_id: str):
-        """생성된 텍스트를 Triton Response 객체로 변환"""
+        """생성된 텍스트를 Triton Response 객체로 변환 (ID 포함)"""
         output_tensor = pb_utils.Tensor(
             "text_output", 
             np.array([output_text.encode('utf-8')], dtype=np.bytes_)
         )
-        return pb_utils.InferenceResponse(output_tensors=[output_tensor])
+        # id 파라미터를 추가하여 클라이언트에게도 ID를 반환
+        return pb_utils.InferenceResponse(output_tensors=[output_tensor], id=request_id)
 
     def _get_config_param(self, key: str, default: str = None) -> str:
-        """config.pbtxt 파라미터 조회 Helper"""
         params = self.model_config.get('parameters', {})
         if key in params:
             return params[key].get('string_value', default)
         return default
 
     def _get_input_scalar(self, request, name: str, default=None):
-        """입력 텐서에서 스칼라 값을 추출하는 Helper"""
         tensor = pb_utils.get_input_tensor_by_name(request, name)
         if tensor is None:
             return default
-        
+        return self._np_decoder(tensor.as_numpy()[0])
-        #value = tensor.as_numpy().item() # item()을 사용하여 스칼라 값 추출
-        value = self._np_decoder(tensor.as_numpy()[0])
-        
-        # 바이트 타입 디코딩
-        # if isinstance(value, bytes):
-        #     return value.decode('utf-8')
-        return value
 
     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):
@@ -259,8 +229,7 @@ class TritonPythonModel:
            return bool(obj)
 
     def finalize(self):
-        """리소스 정리"""
         self.logger.log_info(f"Finalizing model {self.model_name}")
         self.model = None
         self.tokenizer = None
-        torch.cuda.empty_cache()
+        torch.cuda.empty_cache()