mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.0__py3-none-any.whl

msprobe/core/compare/multiprocessing_compute.py

@@ -1,9 +1,22 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import multiprocessing
 from dataclasses import dataclass
-from functools import partial
-import numpy as np
 import pandas as pd
+from tqdm import tqdm
 from msprobe.core.common.log import logger
 from msprobe.core.common.utils import CompareException
 from msprobe.core.common.const import CompareConst
@@ -29,11 +42,19 @@ def _handle_multi_process(func, input_parma, result_df, lock):
         except OSError as e:
             logger.error("pool terminate failed")
 
+    progress_bar = tqdm(total=len(result_df), desc="API/Module Item Compare Process", unit="row", ncols=100)
+
+    def update_progress(size, progress_lock):
+        with progress_lock:
+            progress_bar.update(size)
+
     for process_idx, df_chunk in enumerate(df_chunks):
         idx = df_chunk_size * process_idx
+        chunk_size = len(df_chunk)
         result = pool.apply_async(func,
                                   args=(idx, op_name_mapping_dict, df_chunk, lock, input_parma),
-                                  error_callback=err_call)
+                                  error_callback=err_call,
+                                  callback=update_progress(chunk_size, lock))
         results.append(result)
     final_results = [r.get() for r in results]
     pool.close()
@@ -42,7 +63,7 @@ def _handle_multi_process(func, input_parma, result_df, lock):
 
 
 def _ms_graph_handle_multi_process(func, result_df, mode):
-    process_num = int((multiprocessing.cpu_count() + 1) // 2)
+    process_num = int((multiprocessing.cpu_count() + 1) // 4)
     df_chunk_size = len(result_df) // process_num
     if df_chunk_size > 0:
         df_chunks = [result_df.iloc[i:i + df_chunk_size] for i in range(0, len(result_df), df_chunk_size)]
@@ -84,7 +105,8 @@ def read_dump_data(result_df):
     except IndexError as e:
         logger.error('result dataframe elements can not be access.')
         raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR) from e
-    
+
+
 @dataclass
 class ComparisonResult:
     cos_result: list
@@ -116,9 +138,12 @@ def _save_cmp_result(offset, result: ComparisonResult, result_df, lock):
             result_df.loc[process_index, CompareConst.MAX_ABS_ERR] = result.max_err_result[i]
             result_df.loc[process_index, CompareConst.MAX_RELATIVE_ERR] = result.max_relative_err_result[i]
             result_df.loc[process_index, CompareConst.ERROR_MESSAGE] = result.err_msgs[i]
-            result_df.loc[process_index, CompareConst.ACCURACY] = check_accuracy(result.cos_result[i], result.max_err_result[i])
-            result_df.loc[process_index, CompareConst.ONE_THOUSANDTH_ERR_RATIO] = result.one_thousand_err_ratio_result[i]
-            result_df.loc[process_index, CompareConst.FIVE_THOUSANDTHS_ERR_RATIO] = result.five_thousand_err_ratio_result[i]
+            result_df.loc[process_index, CompareConst.ACCURACY] = (
+                check_accuracy(result.cos_result[i], result.max_err_result[i]))
+            result_df.loc[process_index, CompareConst.ONE_THOUSANDTH_ERR_RATIO] = (
+                result.one_thousand_err_ratio_result)[i]
+            result_df.loc[process_index, CompareConst.FIVE_THOUSANDTHS_ERR_RATIO] = (
+                result.five_thousand_err_ratio_result)[i]
         return result_df
     except ValueError as e:
         logger.error('result dataframe is not found.')

msprobe/core/compare/npy_compare.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import abc
 import numpy as np
 from msprobe.core.common.utils import format_value
@@ -78,10 +93,8 @@ def get_error_message(n_value, b_value, npu_op_name, error_flag, error_file=None
 
 def npy_data_check(n_value, b_value):
     error_message = ""
-    if n_value is None or b_value is None:
-        error_message += "Dump file not found.\n"
-    if n_value == "" or b_value == "":
-        error_message += "Dump file not found.\n"
+    if not isinstance(n_value, np.ndarray) or not isinstance(b_value, np.ndarray):
+        error_message += "Dump file is not ndarray.\n"
 
     # 检查 n_value 和 b_value 是否为空
     if not error_message and (n_value.size == 0 or b_value.size == 0):
@@ -97,7 +110,8 @@ def npy_data_check(n_value, b_value):
 
     if not error_message:
         n_value, b_value = handle_inf_nan(n_value, b_value)  # 判断是否有 nan/inf 数据
-        if CompareConst.NAN in (n_value, b_value):
+        # handle_inf_nan 会返回'Nan'或ndarray类型，使用类型判断是否存在无法处理的nan/inf数据
+        if not isinstance(n_value, np.ndarray) or not isinstance(b_value, np.ndarray):
             error_message += "The position of inf or nan in NPU and bench Tensor do not match.\n"
     if error_message == "":
         error_flag = False
@@ -273,7 +287,8 @@ class GetFiveThousandErrRatio(TensorComparisonBasic):
             relative_err = get_relative_err(n_value, b_value)
         if not np.size(relative_err):
             return CompareConst.NAN, ""
-        return format_value(np.sum(relative_err < CompareConst.FIVE_THOUSAND_RATIO_THRESHOLD) / np.size(relative_err)), ""
+        return format_value(
+            np.sum(relative_err < CompareConst.FIVE_THOUSAND_RATIO_THRESHOLD) / np.size(relative_err)), ""
 
 
 class CompareOps:

msprobe/core/compare/utils.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
 import os
 import re
@@ -59,14 +73,18 @@ def check_and_return_dir_contents(dump_dir, prefix):
 
 def rename_api(npu_name, process):
     npu_split = npu_name.split(process)
-    torch_func_index, in_out = npu_split[0], npu_split[1]
+    try:
+        torch_func_index, in_out = npu_split[0], npu_split[1]
+    except IndexError as error:
+        logger.error(f'{npu_name} can not be split with {process}, please check!')
+        raise CompareException(CompareException.INDEX_OUT_OF_BOUNDS_ERROR) from error
     torch_func_split = torch_func_index.rsplit(Const.SEP, 2)
     torch_func = str(torch_func_split[0]) + str(in_out)
     return torch_func
 
 
 def read_op(op_data, op_name):
-    op_parsed_list = Const.DEFAULT_LIST
+    op_parsed_list = []
     if Const.FORWARD in op_name:
         if Const.INPUT_ARGS in op_data:
             input_item = op_data[Const.INPUT_ARGS]
@@ -103,16 +121,23 @@ def read_op(op_data, op_name):
     return op_parsed_list
 
 
-def op_item_parse(item, op_name, index, item_list=None, top_bool=True):
+def op_item_parse(item, op_name, index, item_list=None, top_bool=True, depth=0):
+    if depth > Const.MAX_DEPTH:
+        logger.error(f"parse of api/module of {op_name} exceeds the recursion limit.")
+        raise CompareException(CompareException.RECURSION_LIMIT_ERROR)
     if item_list is None:
         item_list = []
     if item is None or (isinstance(item, dict) and not item):
         if not top_bool:
-            tmp = {'full_op_name': op_name + '.' + str(index), 'Max': None, 'Min': None, 'Mean': None, 'Norm': None,
-                   'dtype': None, 'shape': None, 'md5': None, 'data_name': '-1'}
+            tmp = {
+                'full_op_name': op_name + '.' + str(index), 'Max': None, 'Min': None, 'Mean': None, 'Norm': None,
+                'dtype': None, 'shape': None, 'md5': None, 'data_name': '-1'
+            }
         else:
-            tmp = {'full_op_name': op_name + '.0', 'Max': None, 'Min': None, 'Mean': None, 'Norm': None, 'dtype': None,
-                   'shape': None, 'md5': None, 'data_name': '-1'}
+            tmp = {
+                'full_op_name': op_name + '.0', 'Max': None, 'Min': None, 'Mean': None, 'Norm': None, 'dtype': None,
+                'shape': None, 'md5': None, 'data_name': '-1'
+            }
         item_list.append(tmp)
         return item_list
     if index is None:
@@ -125,7 +150,7 @@ def op_item_parse(item, op_name, index, item_list=None, top_bool=True):
     if isinstance(item, dict):
         if 'type' not in item:
             for kwarg in item:
-                kwarg_parsed_list = op_item_parse(item[kwarg], op_name + Const.SEP + kwarg, None)
+                kwarg_parsed_list = op_item_parse(item[kwarg], op_name + Const.SEP + kwarg, None, depth=depth+1)
                 item_list += kwarg_parsed_list
                 kwarg_parsed_list.clear()
         elif 'dtype' in item:
@@ -171,7 +196,7 @@ def op_item_parse(item, op_name, index, item_list=None, top_bool=True):
             resolve_api_special_parameters(item, full_op_name, item_list)
     else:
         for j, item_spec in enumerate(item):
-            op_item_parse(item_spec, full_op_name, j, item_list=item_list, top_bool=False)
+            op_item_parse(item_spec, full_op_name, j, item_list=item_list, top_bool=False, depth=depth+1)
     return item_list
 
 
@@ -226,9 +251,10 @@ def get_accuracy(result, n_dict, b_dict, summary_compare=False, md5_compare=Fals
             b_struct = b_dict[key][index]
             err_msg = ""
             if md5_compare:
-                result_item = [n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1],
-                               n_struct[2], b_struct[2],
-                               CompareConst.PASS if n_struct[2] == b_struct[2] else CompareConst.DIFF]
+                result_item = [
+                    n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1], n_struct[2], b_struct[2],
+                    CompareConst.PASS if n_struct[2] == b_struct[2] else CompareConst.DIFF
+                ]
                 if has_stack and index == 0 and key == "input_struct":
                     result_item.extend(npu_stack_info)
                 else:
@@ -237,15 +263,19 @@ def get_accuracy(result, n_dict, b_dict, summary_compare=False, md5_compare=Fals
                 continue
 
             if summary_compare:
-                result_item = [n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1],
-                               " ", " ", " ", " ", " ", " ", " ", " "]
+                result_item = [
+                    n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1],
+                    " ", " ", " ", " ", " ", " ", " ", " "
+                ]
             else:
-                result_item = [n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1],
-                               " ", " ", " ", " ", " "]
+                result_item = [
+                    n_name, b_name, n_struct[0], b_struct[0], n_struct[1], b_struct[1],
+                    " ", " ", " ", " ", " "
+                ]
 
-            npu_summary_data = n_dict.get("summary")[n_start + index]
+            npu_summary_data = n_dict.get(CompareConst.SUMMARY)[n_start + index]
             result_item.extend(npu_summary_data)
-            bench_summary_data = b_dict.get("summary")[b_start + index]
+            bench_summary_data = b_dict.get(CompareConst.SUMMARY)[b_start + index]
             result_item.extend(bench_summary_data)
 
             if summary_compare:
@@ -257,7 +287,7 @@ def get_accuracy(result, n_dict, b_dict, summary_compare=False, md5_compare=Fals
                         if bench_val != 0:
                             relative = str(abs((diff / bench_val) * 100)) + '%'
                         else:
-                            relative = "N/A"
+                            relative = CompareConst.N_A
                         result_item[start_idx + i] = diff
                         result_item[start_idx + i + 4] = relative
                         magnitude_diff = abs(diff) / (max(abs(npu_val), abs(bench_val)) + 1e-10)
@@ -287,15 +317,19 @@ def get_accuracy(result, n_dict, b_dict, summary_compare=False, md5_compare=Fals
                 n_name = n_dict['op_name'][n_start + index]
                 n_struct = n_dict[key][index]
                 if md5_compare:
-                    result_item = [n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN,
-                                   n_struct[1], CompareConst.NAN, n_struct[2], CompareConst.NAN, CompareConst.NAN]
+                    result_item = [
+                        n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN, n_struct[1], CompareConst.NAN,
+                        n_struct[2], CompareConst.NAN, CompareConst.NAN
+                    ]
                     result.append(result_item)
                     continue
-                result_item = [n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN,
-                               n_struct[1], CompareConst.NAN, " ", " ", " ", " ", " "]
-                summary_data = n_dict.get("summary")[n_start + index]
+                result_item = [
+                    n_name, CompareConst.NAN, n_struct[0], CompareConst.NAN, n_struct[1], CompareConst.NAN,
+                    " ", " ", " ", " ", " "
+                ]
+                summary_data = n_dict.get(CompareConst.SUMMARY)[n_start + index]
                 result_item.extend(summary_data)
-                summary_data = [CompareConst.NAN for _ in range(len(n_dict.get("summary")[0]))]
+                summary_data = [CompareConst.NAN for _ in range(len(n_dict.get(CompareConst.SUMMARY)[0]))]
                 result_item.extend(summary_data)
 
                 err_msg = ""
@@ -313,15 +347,12 @@ def get_accuracy(result, n_dict, b_dict, summary_compare=False, md5_compare=Fals
 
     n_num = len(n_dict['op_name'])
     b_num = len(b_dict['op_name'])
-    n_num_input = len([name for name in n_dict['op_name'] if Const.INPUT in name])
-    b_num_input = len([name for name in b_dict['op_name'] if Const.INPUT in name])
-    n_num_kwarg = len([name for name in n_dict['op_name'] if 'kwarg' in name])
-    b_num_kwarg = len([name for name in b_dict['op_name'] if 'kwarg' in name])
-    n_num_output = n_num - n_num_input - n_num_kwarg
-    b_num_output = b_num - b_num_input - b_num_kwarg
+    n_num_input = len([name for name in n_dict['op_name'] if Const.INPUT in name.split(Const.SEP) or Const.KWARGS in name.split(Const.SEP)])
+    b_num_input = len([name for name in b_dict['op_name'] if Const.INPUT in name.split(Const.SEP) or Const.KWARGS in name.split(Const.SEP)])
+    n_num_output = n_num - n_num_input
+    b_num_output = b_num - b_num_input
     get_accuracy_core(0, n_num_input, 0, b_num_input, 'input_struct')
-    get_accuracy_core(n_num_input, n_num_kwarg, b_num_input, b_num_kwarg, "kwargs_struct")
-    get_accuracy_core(n_num_input + n_num_kwarg, n_num_output, b_num_input + b_num_kwarg, b_num_output, 'output_struct')
+    get_accuracy_core(n_num_input, n_num_output, b_num_input, b_num_output, 'output_struct')
 
 
 def get_un_match_accuracy(result, n_dict, md5_compare, summary_compare):
@@ -331,7 +362,8 @@ def get_un_match_accuracy(result, n_dict, md5_compare, summary_compare):
     err_msg = CompareConst.NO_BENCH
     accuracy_check_res = CompareConst.N_A
     for index, n_name in enumerate(n_dict["op_name"]):
-        if n_name.find("input") != -1:
+        name_ele_list = n_name.split(Const.SEP)
+        if "input" in name_ele_list:
             n_struct = n_dict["input_struct"][index]
         else:
             n_struct = n_dict["output_struct"][index_out]
@@ -383,25 +415,28 @@ def merge_tensor(tensor_list, summary_compare, md5_compare):
             op_dict['stack_info'].append(tensor['full_info'])
             break
         op_dict["op_name"].append(tensor['full_op_name'])
+        name_ele_list = tensor['full_op_name'].split(Const.SEP)
         if not md5_compare:
-            if tensor['full_op_name'].find("input") != -1:
+            if "input" in name_ele_list:
                 op_dict["input_struct"].append((tensor['dtype'], tensor['shape']))
-            elif tensor['full_op_name'].find("kwarg") != -1:
+            elif "kwarg" in name_ele_list:
                 op_dict["kwargs_struct"].append((tensor['dtype'], tensor['shape']))
-            elif tensor['full_op_name'].find("output") != -1:
+            elif "output" in name_ele_list:    
                 op_dict["output_struct"].append((tensor['dtype'], tensor['shape']))
         else:
-            if tensor['full_op_name'].find("input") != -1:
+            if "input" in name_ele_list:
                 op_dict["input_struct"].append((tensor['dtype'], tensor['shape'], tensor['md5']))
-            elif tensor['full_op_name'].find("kwarg") != -1:
+            if "kwarg" in name_ele_list:    
                 op_dict["kwargs_struct"].append((tensor['dtype'], tensor['shape'], tensor['md5']))
-            elif tensor['full_op_name'].find("output") != -1:
+            elif "output" in name_ele_list:
                 op_dict["output_struct"].append((tensor['dtype'], tensor['shape'], tensor['md5']))
-
         op_dict["summary"].append([tensor['Max'], tensor['Min'], tensor['Mean'], tensor['Norm']])
 
         if all_mode_bool:
             op_dict["data_name"].append(tensor['data_name'])
+            data_name = op_dict["data_name"][-1].rsplit(Const.SEP, 1)[0]
+            if data_name != "-1":
+                op_dict["op_name"][-1] = data_name
 
     if not op_dict["kwargs_struct"]:
         del op_dict["kwargs_struct"]
@@ -410,7 +445,7 @@ def merge_tensor(tensor_list, summary_compare, md5_compare):
 
 def _compare_parser(parser):
     parser.add_argument("-i", "--input_path", dest="input_path", type=str,
-                        help="<Required> The compare input path, a dict json.",  required=True)
+                        help="<Required> The compare input path, a dict json.", required=True)
     parser.add_argument("-o", "--output_path", dest="output_path", type=str,
                         help="<Required> The compare task result out path.", required=True)
     parser.add_argument("-s", "--stack_mode", dest="stack_mode", action="store_true",
@@ -422,9 +457,8 @@ def _compare_parser(parser):
     parser.add_argument("-cm", "--cell_mapping", dest="cell_mapping", type=str, nargs='?', const=True,
                         help="<optional> The cell mapping file path.", required=False)
     parser.add_argument("-am", "--api_mapping", dest="api_mapping", type=str, nargs='?', const=True,
-                        help="<optional> The api mapping file path.", required=False)
-    
-
-
-    
-    
+                        help="<optional> The api mapping file path.", required=False)    
+    parser.add_argument("-dm", "--data_mapping", dest="data_mapping", type=str,
+                        help="<optional> The data mapping file path.", required=False)
+    parser.add_argument("-lm", "--layer_mapping", dest="layer_mapping", type=str,
+                        help="<optional> The layer mapping file path.", required=False)

msprobe/core/data_dump/data_collector.py

@@ -1,9 +1,24 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
 
 from msprobe.core.data_dump.scope import build_scope, ListScope
 from msprobe.core.data_dump.json_writer import DataWriter
 from msprobe.core.common.log import logger
-from msprobe.core.common.const import Const, MsgConst
+from msprobe.core.common.const import Const
 from msprobe.core.data_dump.data_processor.factory import DataProcessorFactory
 
 
@@ -14,14 +29,13 @@ def build_data_collector(config):
 class DataCollector:
     multi_output_apis = ["_sort_", "npu_flash_attention"]
     tasks_need_tensor_data = [Const.OVERFLOW_CHECK, Const.TENSOR, Const.FREE_BENCHMARK]
-    level_without_construct = ["L1", "L2"]
+    level_without_construct = [Const.LEVEL_L1, Const.LEVEL_L2]
 
     def __init__(self, config):
         self.config = config
         self.data_writer = DataWriter()
         self.data_processor = DataProcessorFactory.create_processor(self.config, self.data_writer)
-        self.module_processor = DataProcessorFactory.get_module_processor(self.config.framework) \
-            if self.config.framework == Const.PT_FRAMEWORK else None
+        self.module_processor = DataProcessorFactory.get_module_processor(self.config.framework)
         self.module_count = {}
         if self.config.task == Const.FREE_BENCHMARK:
             self.scope = build_scope(ListScope, self.config.scope, self.config.list)
@@ -59,16 +73,16 @@ class DataCollector:
     def write_json(self):
         self.data_writer.write_json()
 
-    def update_data(self, data_info, msg=''):
+    def update_data(self, name, data_info):
+        msg = f"msprobe is collecting data on {name}."
         if self.config.task == Const.OVERFLOW_CHECK:
             if self.data_processor.has_overflow:
+                msg += " Overflow detected."
+                logger.warning(msg)
                 self.data_writer.update_data(data_info)
-                msg += "Overflow detected."
-            else:
-                msg += "No Overflow, OK."
-        else:
-            self.data_writer.update_data(data_info)
-        return msg
+            return
+        logger.debug(msg)
+        self.data_writer.update_data(data_info)
 
     def pre_forward_data_collect(self, name, module, pid, module_input_output):
         backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
@@ -78,7 +92,7 @@ class DataCollector:
             return
         logger.info(f"API {name} is inplace.")
         data_info = self.data_processor.analyze_pre_forward_inplace(name, module_input_output)
-        self.handle_data(name, data_info)
+        self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
     def forward_data_collect(self, name, module, pid, module_input_output):
         self.update_construct(name)
@@ -92,13 +106,7 @@ class DataCollector:
         if self.config.level == "L2":
             return
         self.data_writer.update_stack(self.data_processor.analyze_api_call_stack(name))
-        if self.config.framework == Const.MS_FRAMEWORK:
-            self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
-        else:
-            if self.data_processor.is_terminated:
-                self.handle_data(name, data_info, flush=True)
-                raise Exception(f"[{Const.TOOL_NAME}] exit")
-            self.handle_data(name, data_info)
+        self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
     def backward_data_collect(self, name, module, pid, module_input_output):
         self.update_construct(name)
@@ -106,13 +114,7 @@ class DataCollector:
             return
 
         data_info = self.data_processor.analyze_backward(name, module, module_input_output)
-        if self.config.framework == Const.MS_FRAMEWORK:
-            self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
-        else:
-            if self.data_processor.is_terminated:
-                self.handle_data(name, data_info, flush=True)
-                raise Exception(f"[{Const.TOOL_NAME}] exit")
-            self.handle_data(name, data_info)
+        self.handle_data(name, data_info, flush=self.data_processor.is_terminated)
 
     def backward_input_data_collect(self, name, module, pid, module_input_output):
         self.update_construct(name)
@@ -131,18 +133,15 @@ class DataCollector:
         self.handle_data(name, data_info)
 
     def update_construct(self, name):
-        if self.config.framework == Const.PT_FRAMEWORK and \
-           self.config.level not in DataCollector.level_without_construct:
+        if self.config.level not in DataCollector.level_without_construct:
             self.data_writer.update_construct({name: self.module_processor.api_parent_node})
             self.data_writer.update_construct(self.module_processor.module_node)
 
     def handle_data(self, name, data_info, flush=False):
         if data_info:
-            msg = f"msprobe is collecting data on {name}. "
-            msg = self.update_data(data_info, msg)
-            logger.info(MsgConst.CLEAR_SYMBOL + msg, end='\r')
+            self.update_data(name, data_info)
         if not flush:
-            self.data_writer.flush_data_when_buffer_is_full()
+            self.data_writer.flush_data_periodically()
         else:
             self.write_json()
 

msprobe/core/data_dump/data_processor/base.py

@@ -1,11 +1,27 @@
-import os
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0  (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import inspect
+import os
 from dataclasses import dataclass
 from typing import Tuple, Dict, Optional, Any
+
 import numpy as np
-from msprobe.core.common.log import logger
-from msprobe.core.common.utils import convert_tuple
 from msprobe.core.common.const import Const
+from msprobe.core.common.log import logger
+from msprobe.core.common.utils import convert_tuple, CompareException
 
 
 @dataclass
@@ -69,8 +85,11 @@ class TensorStatInfo:
 
 class BaseDataProcessor:
     _recursive_key_stack = []
-    special_type = (np.integer, np.floating, np.bool_, np.complexfloating, np.str_, np.byte, np.unicode_,
-                    bool, int, float, str, slice, type(Ellipsis))
+    special_type = (
+        np.integer, np.floating, np.bool_, np.complexfloating, np.str_, np.byte, np.unicode_,
+        bool, int, float, str, slice,
+        type(Ellipsis)
+    )
 
     def __init__(self, config, data_writer):
         self.data_writer = data_writer
@@ -86,26 +105,27 @@ class BaseDataProcessor:
     @property
     def data_path(self):
         return self.data_writer.dump_tensor_data_dir
-    
+
     @property
     def is_terminated(self):
         return False
 
     @staticmethod
     def analyze_api_call_stack(name):
+        try:
+            api_stack = inspect.stack()[5:]
+        except Exception as e:
+            logger.warning(f"The call stack of <{name}> failed to retrieve, {e}.")
+            api_stack = None
         stack_str = []
-        for (_, path, line, func, code, _) in inspect.stack()[5:]:
-            if not code:
-                continue
-            stack_line = " ".join([
-                "File", ", ".join([
-                    path,
-                    " ".join(["line", str(line)]),
-                    " ".join(["in", func]),
-                    " ".join(["\n", code[0].strip()])
-                ])
-            ])
-            stack_str.append(stack_line)
+        if api_stack:
+            for (_, path, line, func, code, _) in api_stack:
+                if not code:
+                    continue
+                stack_line = f"File {path}, line {str(line)}, in {func}, \n {code[0].strip()}"
+                stack_str.append(stack_line)
+        else:
+            stack_str.append(Const.WITHOUT_CALL_STACK)
         stack_info_struct = {name: stack_str}
         return stack_info_struct
 
@@ -167,7 +187,10 @@ class BaseDataProcessor:
         return cls.special_type
 
     @classmethod
-    def recursive_apply_transform(cls, args, transform):
+    def recursive_apply_transform(cls, args, transform, depth=0):
+        if depth > Const.MAX_DEPTH:
+            logger.error(f"The maximum depth of recursive transform, {Const.MAX_DEPTH} is reached.")
+            raise CompareException(CompareException.RECURSION_LIMIT_ERROR)
         if isinstance(args, cls.get_special_types()):
             arg_transform = transform(args, cls._recursive_key_stack)
             return arg_transform
@@ -175,14 +198,14 @@ class BaseDataProcessor:
             result_list = []
             for i, arg in enumerate(args):
                 cls._recursive_key_stack.append(str(i))
-                result_list.append(cls.recursive_apply_transform(arg, transform))
+                result_list.append(cls.recursive_apply_transform(arg, transform, depth=depth + 1))
                 cls._recursive_key_stack.pop()
             return type(args)(result_list)
         elif isinstance(args, dict):
             result_dict = {}
             for k, arg in args.items():
                 cls._recursive_key_stack.append(str(k))
-                result_dict[k] = cls.recursive_apply_transform(arg, transform)
+                result_dict[k] = cls.recursive_apply_transform(arg, transform, depth=depth + 1)
                 cls._recursive_key_stack.pop()
             return result_dict
         elif args is not None:
@@ -222,7 +245,7 @@ class BaseDataProcessor:
 
     def analyze_pre_forward(self, name, module, module_input_output: ModuleForwardInputsOutputs):
         pass
-    
+
     def analyze_element(self, element):
         return self.recursive_apply_transform(element, self.analyze_single_element)