Merge pull request #500 from cangtianhuang/develop

Fix `API Tracer` about exceeding length

Author: cangtianhuang

tools/api_tracer/config_serializer.py

@@ -2,7 +2,6 @@
 import time
 from collections import defaultdict, deque
 from threading import Event, Thread
-from types import EllipsisType
 from typing import Any, Dict, List, TextIO
 
 import yaml
@@ -27,29 +26,33 @@ def __init__(
         self.merge_output = merge_output
 
         self.file_handlers: Dict[int, Dict[str, TextIO]] = {}
-
         self.buffer_limit = 20000
         self.buffers: Dict[int, List[Dict]] = defaultdict(list)
 
+        self.max_args_count = 100
+        self.max_item_count = 100
+        self.max_line_length = 1024
+        self.max_nest_depth = 5
+
         # asyncio
         self.log_queue = deque()
         self._stop_event = Event()
         self.writer_thread = Thread(target=self._writer_loop)
         self.total_calls_processed = 0
 
         self._serialize_handlers = {
-            type(None): lambda x: x,
-            bool: lambda x: x,
-            int: lambda x: x,
-            float: lambda x: x,
-            str: lambda x: x,
+            type(None): lambda x, depth: x,
+            bool: lambda x, depth: x,
+            int: lambda x, depth: x,
+            float: lambda x, depth: x,
+            str: lambda x, depth: x,
             list: self._serialize_list,
             tuple: self._serialize_tuple,
             set: self._serialize_set,
             dict: self._serialize_dict,
             type: self._serialize_type,
             slice: self._serialize_slice,
-            EllipsisType: self._serialize_ellipsis,
+            type(Ellipsis): self._serialize_ellipsis,
         }
 
     def open(self):
@@ -157,60 +160,86 @@ def dump_call(
     ):
         """记录一次API调用"""
         try:
+            total_args = len(args) + len(kwargs)
+            if total_args > self.max_args_count:
+                if len(args) < self.max_args_count:
+                    kwargs = dict(
+                        list(kwargs.items())[: self.max_args_count - len(args) - 1]
+                    )
+                    kwargs["__truncated__"] = "<Truncated: max args exceeded>"
+                else:
+                    args = tuple(
+                        list(args)[: self.max_args_count - 1]
+                        + ["<Truncated: max args exceeded>"]
+                    )
+                    kwargs = {}
             call_record = {
                 "level": level,
                 "api": api_name,
-                "args": [self._serialize_item(arg) for arg in args],
+                "args": [self._serialize_item(arg, depth=0) for arg in args],
                 "kwargs": {
-                    key: self._serialize_item(value) for key, value in kwargs.items()
+                    key: self._serialize_item(value, depth=0)
+                    for key, value in kwargs.items()
                 },
-                # "output_summary": self._serialize_item(output)
+                # "output_summary": self._serialize_item(output, depth=0)
             }
             self.log_queue.append(call_record)
         except Exception as e:
             print(f"[ConfigSerializer] Error serializing call for '{api_name}': {e}")
 
-    def _serialize_list(self, item: list) -> Dict:
+    def _serialize_list(self, item: list, depth: int) -> Dict:
+        if len(item) > self.max_item_count:
+            item = item[: self.max_item_count - 1] + ["<Truncated: max item count>"]
         return {
             "type": "list",
-            "value": [self._serialize_item(sub_item) for sub_item in item],
+            "value": [self._serialize_item(sub_item, depth) for sub_item in item],
         }
 
-    def _serialize_tuple(self, item: tuple) -> Dict:
+    def _serialize_tuple(self, item: tuple, depth: int) -> Dict:
+        if len(item) > self.max_item_count:
+            item = item[: self.max_item_count - 1] + ("<Truncated: max item count>",)
         return {
             "type": "tuple",
-            "value": [self._serialize_item(sub_item) for sub_item in item],
+            "value": [self._serialize_item(sub_item, depth) for sub_item in item],
         }
 
-    def _serialize_set(self, item: set) -> Dict:
+    def _serialize_set(self, item: set, depth: int) -> Dict:
+        if len(item) > self.max_item_count:
+            item = set(list(item)[: self.max_item_count - 1])
         return {
             "type": "set",
-            "value": [self._serialize_item(sub_item) for sub_item in item],
+            "value": [self._serialize_item(sub_item, depth) for sub_item in item],
         }
 
-    def _serialize_dict(self, item: dict) -> Dict:
+    def _serialize_dict(self, item: dict, depth: int) -> Dict:
+        if len(item) > self.max_item_count:
+            item = dict(list(item.keys())[: self.max_item_count - 1])
+            item["__truncated__"] = "<Truncated: max item count>"
         return {
             "type": "dict",
-            "value": {str(k): self._serialize_item(v) for k, v in item.items()},
+            "value": {str(k): self._serialize_item(v, depth) for k, v in item.items()},
         }
 
-    def _serialize_type(self, item: type) -> Dict:
+    def _serialize_type(self, item: type, depth: int) -> Dict:
         return {"type": "type", "value": f"{item.__module__}.{item.__name__}"}
 
-    def _serialize_slice(self, item: slice) -> Dict:
+    def _serialize_slice(self, item: slice, depth: int) -> Dict:
         return {
             "type": "slice",
             "value": {"start": item.start, "stop": item.stop, "step": item.step},
         }
 
-    def _serialize_ellipsis(self, item: Any) -> Dict:
+    def _serialize_ellipsis(self, item: Any, depth: int) -> Dict:
         return {"type": "ellipsis", "value": "..."}
 
-    def _serialize_item(self, item: Any) -> Any:
+    def _serialize_item(self, item: Any, depth=0) -> Any:
         """递归序列化对象"""
+        if depth > self.max_nest_depth:
+            return "<Truncated: max depth exceeded>"
+
         handler = self._serialize_handlers.get(type(item))
         if handler:
-            return handler(item)
+            return handler(item, depth=depth + 1)
 
         special_serialization = self.dialect.serialize_special_type(item)
         if special_serialization is not None:
@@ -228,6 +257,8 @@ def format_arg(arg: Any) -> str:
             if arg is None or isinstance(arg, (bool, int, float)):
                 return str(arg)
             if isinstance(arg, str):
+                if len(arg) > 100:
+                    return f'"{arg[:97]}..."'
                 return f'"{arg}"'
 
             if isinstance(arg, dict) and "type" in arg:
@@ -252,7 +283,11 @@ def format_arg(arg: Any) -> str:
 
         args_str = ", ".join(format_arg(arg) for arg in args)
         kwargs_str = ", ".join(f"{k}={format_arg(v)}" for k, v in kwargs.items())
-        return f"{api_name}({args_str + (', ' + kwargs_str if kwargs_str else '')})"
+        result = f"{api_name}({args_str + (', ' + kwargs_str if kwargs_str else '')})"
+
+        if len(result) > self.max_line_length:
+            result = result[: self.max_line_length - 4] + "...)"
+        return result
 
     def get_apis_and_configs(self):
         if self.merge_output:

tools/api_tracer/test_infer.py

@@ -14,7 +14,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODELS = [
-    # "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     # "Qwen/Qwen2-0.5B",
     # "Qwen/Qwen3-0.6B",
     # "Qwen/Qwen3-30B-A3B",
@@ -27,8 +26,11 @@
 
 def run_inference_test(model_name: str):
     print(f"🚀 Running inference test for: {model_name}")
-    output_path = f"tools/api_tracer/trace_output_test_infer/{model_name}"
-    tracer = APITracer("torch", output_path=output_path, levels=[0, 1])
+    true_model_name = "/".join(model_name.rsplit("/", 2)[-2:])
+    output_path = f"tools/api_tracer/trace_output_test_infer/{true_model_name}"
+    tracer = APITracer(
+        "torch", output_path=output_path, levels=[0, 1], merge_output=True
+    )
 
     try:
         model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
@@ -59,10 +61,10 @@ def run_inference_test(model_name: str):
         print("\n--- Generated Response ---")
         print(response)
         print("--------------------------\n")
-        print(f"✅ Test for {model_name} finished.")
+        print(f"✅ Test for {true_model_name} finished.")
     except Exception as e:
         traceback.print_exc()
-        print(f"❌ An error occurred during inference for {model_name}: {e}")
+        print(f"❌ An error occurred during inference for {true_model_name}: {e}")
 
 
 def main():

tools/api_tracer/test_train.py

@@ -1,5 +1,4 @@
 import os
-import time
 import traceback
 
 os.environ["HF_HOME"] = "tools/api_tracer/.huggingface"
@@ -17,9 +16,9 @@
 from tools.api_tracer import APITracer
 
 MODELS = [
-    # "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     # "Qwen/Qwen2-0.5B",
     # "Qwen/Qwen3-0.6B",
+    # "Qwen/Qwen3-30B-A3B",
     # "Qwen/Qwen2.5-VL-3B-Instruct",
     # "deepseek-ai/DeepSeek-V2-Lite",
     # "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
@@ -29,20 +28,41 @@
 
 def run_training_test(model_name: str):
     print(f"🚀 Running training test for: {model_name})")
-    output_path = f"tools/api_tracer/trace_output_test_train/{model_name}"
-    tracer = APITracer("torch", output_path=output_path, levels=[0, 1])
+    true_model_name = "/".join(model_name.rsplit("/", 2)[-2:])
+    output_path = f"tools/api_tracer/trace_output_test_train/{true_model_name}"
+    tracer = APITracer(
+        "torch", output_path=output_path, levels=[0, 1], merge_output=True
+    )
 
     try:
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             torch_dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
+            use_cache=False,
         )
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
+        if "Llama" in true_model_name:
+            llama_chat_template = (
+                "{% for message in messages %}"
+                "{% if message['role'] == 'system' %}"
+                "{{'<|start_header_id|>system<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>'}}"
+                "{% elif message['role'] == 'user' %}"
+                "{{'<|start_header_id|>user<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>'}}"
+                "{% elif message['role'] == 'assistant' %}"
+                "{{'<|start_header_id|>assistant<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>'}}"
+                "{% endif %}"
+                "{% endfor %}"
+                "{% if add_generation_prompt %}"
+                "{{'<|start_header_id|>assistant<|end_header_id|>\n\n'}}"
+                "{% endif %}"
+            )
+            tokenizer.chat_template = llama_chat_template
+
         print(f"Model Class: {model.__class__}")
         print(f"Tokenizer Class: {tokenizer.__class__}")
 
@@ -82,7 +102,7 @@ def preprocess_function(examples):
             save_strategy="no",
             bf16=True,
             report_to="none",
-            max_steps=5,
+            max_steps=1,
             gradient_checkpointing=True,
         )
 
@@ -98,10 +118,10 @@ def preprocess_function(examples):
         with tracer:
             trainer.train()
 
-        print(f"✅ Test for {model_name} finished.")
+        print(f"✅ Test for {true_model_name} finished.")
     except Exception as e:
         traceback.print_exc()
-        print(f"❌ An error occurred during training for {model_name}: {e}")
+        print(f"❌ An error occurred during training for {true_model_name}: {e}")
 
 
 def run_training_test_vision(model_name: str):