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

msprobe/core/data_dump/data_collector.py

@@ -1,140 +1,154 @@
-
-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
-from msprobe.core.data_dump.data_processor.factory import DataProcessorFactory
-
-
-def build_data_collector(config):
-    return DataCollector(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"]
-
-    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_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)
-
-    @property
-    def dump_data_dir(self):
-        return self.data_writer.dump_tensor_data_dir
-
-    @property
-    def dump_file_path(self):
-        return self.data_writer.dump_file_path
-    
-    @staticmethod
-    def check_scope_and_pid(scope, name, pid):
-        return (not scope or scope.check(name)) and pid == os.getpid()
-
-    @staticmethod
-    def is_inplace(module):
-        return getattr(module, "op_is_inplace", False)
-    
-    def if_return_forward_new_output(self):
-        return self.data_processor.if_return_forward_new_output()
-    
-    def get_forward_new_output(self):
-        return self.data_processor.get_forward_new_output()
-
-    def visit_and_clear_overflow_status(self, api_or_module_name):
-        self.data_processor.visit_and_clear_overflow_status(api_or_module_name)
-
-    def write_json(self):
-        self.data_writer.write_json()
-
-    def update_data(self, data_info, msg=''):
-        if self.config.task == Const.OVERFLOW_CHECK:
-            if self.data_processor.has_overflow:
-                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
-
-    def pre_forward_data_collect(self, name, module, pid, module_input_output):
-        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)
-        if not self.is_inplace(module):
-            return
-        logger.info(f"API {name} is inplace.")
-        if self.check_scope_and_pid(self.scope, name, pid):
-            data_info = self.data_processor.analyze_pre_forward_inplace(name, module_input_output)
-            self.update_data(data_info)
-
-    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 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))
-        self.handle_data(name, data_info)
-
-    def backward_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
-
-        data_info = self.data_processor.analyze_backward(name, module, module_input_output)
-        self.handle_data(name, data_info)
-
-    def update_construct(self, name):
-        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):
-        msg = f"msProbe is collecting data on {name}. "
-        if data_info:
-            msg = self.update_data(data_info, msg)
-            logger.info(msg)
-        self.data_writer.flush_data_when_buffer_is_full()
-
-    def module_count_func(self, name, name_template):
-        module_name = name.split(Const.SEP)[-3]
-        if "forward" in name_template:
-            if module_name not in self.module_count:
-                self.module_count[module_name] = [0, [0]]
-            else:
-                if self.module_count[module_name][-1] and \
-                        self.module_count[module_name][0] != self.module_count[module_name][-1][-1]:
-                    self.module_count[module_name][-1].pop()
-                self.module_count[module_name][0] += 1
-                self.module_count[module_name][-1].append(self.module_count[module_name][0])
-            index = self.module_count[module_name][0]
-        else:
-            backward_stack = self.module_count[module_name][-1] if module_name in self.module_count else []
-            if not backward_stack:
-                index = "abnormal"
-            else:
-                index = backward_stack.pop()
-        return index
-
-    def update_dump_paths(self, *args):
-        self.data_writer.update_dump_paths(*args)
-        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level)
-    
-    def update_iter(self, current_iter):
-        self.data_processor.update_iter(current_iter)
+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.data_dump.data_processor.factory import DataProcessorFactory
+
+
+def build_data_collector(config):
+    return DataCollector(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"]
+
+    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_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()
+
+    @property
+    def dump_data_dir(self):
+        return self.data_writer.dump_tensor_data_dir
+
+    @property
+    def dump_file_path(self):
+        return self.data_writer.dump_file_path
+
+    @staticmethod
+    def check_scope_and_pid(scope, name, pid):
+        return (not scope or scope.check(name)) and pid == os.getpid()
+
+    @staticmethod
+    def is_inplace(module):
+        return getattr(module, "op_is_inplace", False)
+
+    def if_return_forward_new_output(self):
+        return self.data_processor.if_return_forward_new_output()
+
+    def get_forward_new_output(self):
+        return self.data_processor.get_forward_new_output()
+
+    def update_api_or_module_name(self, api_or_module_name):
+        self.data_processor.update_api_or_module_name(api_or_module_name)
+
+    def write_json(self):
+        self.data_writer.write_json()
+
+    def update_data(self, data_info, msg=''):
+        if self.config.task == Const.OVERFLOW_CHECK:
+            if self.data_processor.has_overflow:
+                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
+
+    def pre_forward_data_collect(self, name, module, pid, module_input_output):
+        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)
+        if not self.is_inplace(module) or not self.check_scope_and_pid(self.scope, name, pid):
+            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)
+
+    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 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)
+
+    def backward_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
+
+        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)
+
+    def backward_input_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
+
+        data_info = self.data_processor.analyze_backward_input(name, module, module_input_output)
+        self.handle_data(name, data_info)
+
+    def backward_output_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
+
+        data_info = self.data_processor.analyze_backward_output(name, module, module_input_output)
+        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:
+            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')
+        if not flush:
+            self.data_writer.flush_data_when_buffer_is_full()
+        else:
+            self.write_json()
+
+    def update_dump_paths(self, *args):
+        self.data_writer.update_dump_paths(*args)
+        self.data_writer.initialize_json_file(task=self.config.task, level=self.config.level)
+
+    def update_iter(self, current_iter):
+        self.data_processor.update_iter(current_iter)