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

msprobe/core/data_dump/data_collector.py

@@ -1,9 +1,25 @@
+# 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 atexit
 import os
 
-from msprobe.core.data_dump.scope import build_scope, ListScope
+from msprobe.core.data_dump.scope import ScopeFactory
 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
 
 
@@ -12,24 +28,17 @@ 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)
-        else:
-            self.scope = build_scope(None, self.config.scope, self.config.list)
-
-    def __del__(self):
-        self.write_json()
+        self.scope = ScopeFactory(self.config).build_scope()
+        atexit.register(self.write_json)
 
     @property
     def dump_data_dir(self):
@@ -59,18 +68,22 @@ 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):
+        if self.config.level == Const.LEVEL_L2 and self.check_scope_and_pid(self.scope, name, pid):
+            self.data_processor.analyze_pre_forward(name, module, module_input_output)
+            return
+
         backward_name = name.replace(Const.FORWARD, Const.BACKWARD)
         if self.check_scope_and_pid(self.scope, backward_name, pid):
             self.data_processor.analyze_pre_forward(backward_name, module, module_input_output)
@@ -78,27 +91,22 @@ 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)
         if not self.check_scope_and_pid(self.scope, name, pid):
             return
+        if self.config.level == Const.LEVEL_L2:
+            self.data_processor.analyze_forward(name, module, module_input_output)
+            return
 
         if not self.is_inplace(module):
             data_info = self.data_processor.analyze_forward(name, module, module_input_output)
         else:
             data_info = self.data_processor.analyze_forward_inplace(name, module_input_output)
-        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,9 @@ 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)
+        if self.config.level == Const.LEVEL_L2:
+            return
+        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 +135,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,28 @@
-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
-from dataclasses import dataclass
+import os
+from dataclasses import dataclass, is_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 +86,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
@@ -82,30 +102,33 @@ class BaseDataProcessor:
         self.current_iter = 0
         self._return_forward_new_output = False
         self._forward_new_output = None
+        if hasattr(config, "data_mode"):
+            self.allowed_data_mode = self._get_allowed_data_mode(config.data_mode)
 
     @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
 
@@ -162,34 +185,66 @@ class BaseDataProcessor:
     def _analyze_numpy(value, numpy_type):
         return {"type": numpy_type, "value": value}
 
+    @staticmethod
+    def _get_allowed_data_mode(data_mode):
+        if Const.ALL in data_mode:
+            allowed_data_mode = [Const.FORWARD, Const.BACKWARD, Const.INPUT, Const.OUTPUT]
+        else:
+            allowed_data_mode = list(set(data_mode))
+            if Const.FORWARD not in allowed_data_mode and Const.BACKWARD not in allowed_data_mode:
+                allowed_data_mode += [Const.FORWARD, Const.BACKWARD]
+            if Const.INPUT not in allowed_data_mode and Const.OUTPUT not in allowed_data_mode:
+                allowed_data_mode += [Const.INPUT, Const.OUTPUT]
+        return allowed_data_mode
+
     @classmethod
     def get_special_types(cls):
         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
+        elif isinstance(args, tuple) and hasattr(args, '_fields'):
+            # namedtuple to dict
+            args_dict = {field: getattr(args, field) for field in args._fields}
+            return cls.apply_transform_dict(args_dict, transform, depth)
+        elif is_dataclass(args):
+            # dataclass to dict
+            args_dict = {field: getattr(args, field) for field in args.__dataclass_fields__}
+            return cls.apply_transform_dict(args_dict, transform, depth)
         elif isinstance(args, (list, tuple)):
-            result_list = []
-            for i, arg in enumerate(args):
-                cls._recursive_key_stack.append(str(i))
-                result_list.append(cls.recursive_apply_transform(arg, transform))
-                cls._recursive_key_stack.pop()
+            result_list = cls.apply_transform_list(args, transform, depth)
             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)
-                cls._recursive_key_stack.pop()
-            return result_dict
+            return cls.apply_transform_dict(args, transform, depth)
         elif args is not None:
             logger.warning(f"Data type {type(args)} is not supported.")
             return None
         else:
             return None
+    
+    @classmethod
+    def apply_transform_dict(cls, args, transform, depth):
+        result_dict = {}
+        for k, arg in args.items():
+            cls._recursive_key_stack.append(str(k))
+            result_dict[k] = cls.recursive_apply_transform(arg, transform, depth=depth + 1)
+            cls._recursive_key_stack.pop()
+        return result_dict
+
+    @classmethod
+    def apply_transform_list(cls, args, transform, depth):
+        result_list = []
+        for i, arg in enumerate(args):
+            cls._recursive_key_stack.append(str(i))
+            result_list.append(cls.recursive_apply_transform(arg, transform, depth=depth + 1))
+            cls._recursive_key_stack.pop()
+        return result_list
 
     def if_return_forward_new_output(self):
         return self._return_forward_new_output
@@ -216,13 +271,11 @@ class BaseDataProcessor:
         Return:
             bool: True if the parameters are in data_mode or data_mode is all, False otherwise.
         """
-        return (Const.ALL in self.config.data_mode or
-                forward_backward in self.config.data_mode or
-                input_output in self.config.data_mode)
+        return forward_backward in self.allowed_data_mode and input_output in self.allowed_data_mode
 
     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)
 

msprobe/core/data_dump/data_processor/factory.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.
+
 from msprobe.core.common.const import Const
 
 
@@ -34,14 +49,14 @@ class DataProcessorFactory:
     @classmethod
     def register_processors(cls, framework):
         if framework == Const.PT_FRAMEWORK:
-            from .pytorch_processor import (
+            from msprobe.core.data_dump.data_processor.pytorch_processor import (
                 StatisticsDataProcessor as PytorchStatisticsDataProcessor,
                 TensorDataProcessor as PytorchTensorDataProcessor,
                 OverflowCheckDataProcessor as PytorchOverflowCheckDataProcessor,
                 FreeBenchmarkDataProcessor as PytorchFreeBenchmarkDataProcessor,
                 KernelDumpDataProcessor as PytorchKernelDumpDataProcessor
             )
-            from ....pytorch.module_processer import ModuleProcesser
+            from msprobe.pytorch.module_processer import ModuleProcesser
             cls.register_processor(Const.PT_FRAMEWORK, Const.STATISTICS, PytorchStatisticsDataProcessor)
             cls.register_processor(Const.PT_FRAMEWORK, Const.TENSOR, PytorchTensorDataProcessor)
             cls.register_processor(Const.PT_FRAMEWORK, Const.OVERFLOW_CHECK, PytorchOverflowCheckDataProcessor)
@@ -49,11 +64,13 @@ class DataProcessorFactory:
             cls.register_processor(Const.PT_FRAMEWORK, Const.KERNEL_DUMP, PytorchKernelDumpDataProcessor)
             cls.register_module_processor(Const.PT_FRAMEWORK, ModuleProcesser)
         elif framework == Const.MS_FRAMEWORK:
-            from .mindspore_processor import (
+            from msprobe.core.data_dump.data_processor.mindspore_processor import (
                 StatisticsDataProcessor as MindsporeStatisticsDataProcessor,
                 TensorDataProcessor as MindsporeTensorDataProcessor,
                 OverflowCheckDataProcessor as MindsporeOverflowCheckDataProcessor
             )
+            from msprobe.mindspore.cell_processor import CellProcessor
             cls.register_processor(Const.MS_FRAMEWORK, Const.STATISTICS, MindsporeStatisticsDataProcessor)
             cls.register_processor(Const.MS_FRAMEWORK, Const.TENSOR, MindsporeTensorDataProcessor)
             cls.register_processor(Const.MS_FRAMEWORK, Const.OVERFLOW_CHECK, MindsporeOverflowCheckDataProcessor)
+            cls.register_module_processor(Const.MS_FRAMEWORK, CellProcessor)

msprobe/core/data_dump/data_processor/mindspore_processor.py

@@ -17,6 +17,7 @@ import zlib
 
 import mindspore as ms
 from mindspore import mint, ops
+from mindspore._c_expression.typing import Number
 import numpy as np
 
 from msprobe.core.common.const import Const
@@ -29,7 +30,7 @@ from msprobe.mindspore.dump.hook_cell.api_registry import api_register
 
 
 class MindsporeDataProcessor(BaseDataProcessor):
-    mindspore_special_type = tuple([ms.Tensor])
+    mindspore_special_type = tuple([ms.Tensor, Number])
 
     def __init__(self, config, data_writer):
         super().__init__(config, data_writer)
@@ -40,7 +41,7 @@ class MindsporeDataProcessor(BaseDataProcessor):
     @staticmethod
     def get_md5_for_tensor(x):
         x = convert_bf16_to_fp32(x)
-        tensor_bytes = x.asnumpy().tobytes()
+        tensor_bytes = x.contiguous().asnumpy().tobytes()
         crc32_hash = zlib.crc32(tensor_bytes)
         return f"{crc32_hash:08x}"
 
@@ -57,25 +58,28 @@ class MindsporeDataProcessor(BaseDataProcessor):
         if data.numel() == 0:
             return tensor_stat
         elif data.dtype == ms.bool_:
-            data_np = data.asnumpy()
+            data_np = data.contiguous().asnumpy()
             tensor_stat.max = np.max(data_np).item()
             tensor_stat.min = np.min(data_np).item()
         elif not data.shape:
             tensor_stat.max = tensor_stat.min = tensor_stat.mean = tensor_stat.norm = data.item()
         elif data.dtype == ms.complex64 or data.dtype == ms.complex128:
-            data_abs = np.abs(data.asnumpy())
+            data_abs = np.abs(data.contiguous().asnumpy())
             tensor_stat.max = np.max(data_abs).item()
             tensor_stat.min = np.min(data_abs).item()
             tensor_stat.mean = np.mean(data_abs).item()
             tensor_stat.norm = np.linalg.norm(data_abs).item()
         else:
-            if data.dtype == ms.bfloat16 or not ops.is_floating_point(data):
+            if not ops.is_floating_point(data) or data.dtype == ms.float64:
                 data = data.to(ms.float32)
             api_register.norm_inner_op_set_ori_func()
             get_max_value = api_register.mint_ops_ori_attr.get("max", mint.max)
             get_min_value = api_register.mint_ops_ori_attr.get("min", mint.min)
             get_mean_value = api_register.mint_ops_ori_attr.get("mean", mint.mean)
-            get_norm_value = api_register.functional_ori_attr.get("norm", ops.norm)
+            if hasattr(mint, "norm"):
+                get_norm_value = api_register.mint_ops_ori_attr.get("norm", mint.norm)
+            else:
+                get_norm_value = api_register.functional_ori_attr.get("norm", ops.norm)
             tensor_stat.max = get_max_value(data).item()
             tensor_stat.min = get_min_value(data).item()
             tensor_stat.mean = get_mean_value(data).item()
@@ -90,9 +94,10 @@ class MindsporeDataProcessor(BaseDataProcessor):
         converted_numpy, numpy_type = self._convert_numpy_to_builtin(element)
         if converted_numpy is not element:
             return self._analyze_numpy(converted_numpy, numpy_type)
+        if isinstance(element, Number):
+            return self.analyze_dtype_in_kwargs(element)
         if isinstance(element, ms.Tensor):
             return self._analyze_tensor(element, Const.SEP.join(suffix_stack))
-
         if isinstance(element, (bool, int, float, str, slice, type(Ellipsis))):
             return self._analyze_builtin(element)
         return {}
@@ -163,7 +168,8 @@ class OverflowCheckDataProcessor(MindsporeDataProcessor):
                 save_tensor_as_npy(tensor, file_path)
             self.real_overflow_nums += 1
             if self.overflow_nums != -1 and self.real_overflow_nums >= self.overflow_nums:
-                logger.info(f"[{Const.TOOL_NAME}] 超过预设溢出次数 当前溢出次数: {self.real_overflow_nums}")
+                logger.info(f"[{Const.TOOL_NAME}] Reached the preset overflow times, "
+                            f"current overflow times: {self.real_overflow_nums}.")
         self.cached_tensors_and_file_paths = {}
 
     def _analyze_maybe_overflow_tensor(self, tensor_json):