diff --git a/1/model.py b/1/model.py index 5197d74..4e04d35 100644 --- a/1/model.py +++ b/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] 모델 응답 로그 출력 - self.logger.log_info(f"\n<<< [Output]: {output_text}") + # [LOGGING] Request ID 포함하여 결과 출력 + self.logger.log_info(f"\n[RID: {request_id}] <<< [Output]: {output_text}") - # 5. 응답 생성 - responses.append(self._create_response(output_text)) + # 5. 응답 생성 (ID 전달) + 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: # Chat 형식이지만 템플릿 미지원 시 문자열 변환 시도 + if is_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): - """생성된 텍스트를 Triton Response 객체로 변환""" + def _create_response(self, output_text: str, request_id: str): + """생성된 텍스트를 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 - - #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 + return self._np_decoder(tensor.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): @@ -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() \ No newline at end of file + torch.cuda.empty_cache()