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

msprobe/core/data_dump/json_writer.py

@@ -1,116 +1,117 @@
-import os
-import csv
-import fcntl
-import json
-from pathlib import Path
-
-from msprobe.core.common.file_check import change_mode, FileOpen
-from msprobe.core.common.log import logger
-from msprobe.core.common.const import Const, FileCheckConst
-
-
-class DataWriter:
-
-    def __init__(self, init_json=None) -> None:
-        self.dump_count = 0
-        self.init_json = init_json
-        self.dump_file_path = None  # os.path.join(dump_dir, DataWriter.dump_json_name)
-        self.stack_file_path = None  # os.path.join(dump_dir, DataWriter.stack_json_name)
-        self.construct_file_path = None  # os.path.join(dump_dir, DataWriter.construct_json_name)
-        self.free_benchmark_file_path = None
-        self.dump_tensor_data_dir = None
-        self.buffer_size = 1000
-        self.cache_data = {Const.DATA: {}}
-        self.cache_stack = {}
-        self.cache_construct = {}
-
-    @staticmethod
-    def write_data_to_csv(result: list, result_header: tuple, file_path: str):
-        if not result:
-            return
-        is_exists = os.path.exists(file_path)
-        append = "a+" if is_exists else "w+"
-        with FileOpen(file_path, append) as csv_file:
-            spawn_writer = csv.writer(csv_file)
-            if not is_exists:
-                spawn_writer.writerow(result_header)
-            spawn_writer.writerows([result,])
-        is_new_file = not is_exists
-        if is_new_file:
-            change_mode(file_path, FileCheckConst.DATA_FILE_AUTHORITY)
-
-    def initialize_json_file(self, **kwargs):
-        kwargs.update({"dump_data_dir": self.dump_tensor_data_dir, Const.DATA: {}})
-        with FileOpen(self.dump_file_path, 'w') as f:
-            json.dump(kwargs, f)
-        change_mode(self.dump_file_path, FileCheckConst.DATA_FILE_AUTHORITY)
-
-        if os.path.exists(self.stack_file_path):
-            os.remove(self.stack_file_path)
-        Path(self.stack_file_path).touch()
-        change_mode(self.stack_file_path, FileCheckConst.DATA_FILE_AUTHORITY)
-
-        if os.path.exists(self.construct_file_path):
-            os.remove(self.construct_file_path)
-        Path(self.construct_file_path).touch()
-        change_mode(self.construct_file_path, FileCheckConst.DATA_FILE_AUTHORITY)
-
-    def update_dump_paths(self, dump_file_path, stack_file_path, construct_file_path, dump_data_dir, 
-                          free_benchmark_file_path):
-        self.dump_file_path = dump_file_path
-        self.stack_file_path = stack_file_path
-        self.construct_file_path = construct_file_path
-        self.dump_tensor_data_dir = dump_data_dir
-        self.free_benchmark_file_path = free_benchmark_file_path
-
-    def update_data(self, new_data):
-        key = next(iter(new_data.keys()))  # assert len(new_data.keys()) == 1
-        if key in self.cache_data[Const.DATA]:
-            self.cache_data[Const.DATA][key].update(new_data[key])
-        else:
-            self.cache_data[Const.DATA].update(new_data)
-
-    def flush_data_when_buffer_is_full(self):
-        if len(self.cache_data[Const.DATA]) >= self.buffer_size:
-            self.write_data_json(self.dump_file_path)
-
-    def update_stack(self, new_data):
-        self.cache_stack.update(new_data)
-
-    def update_construct(self, new_data):
-        self.cache_construct.update(new_data)
-
-    def write_data_json(self, file_path):
-        logger.info(f"dump.json is at {os.path.dirname(os.path.dirname(file_path))}. ")
-        if Path(file_path).exists() and os.path.getsize(file_path) > 0:
-            with FileOpen(file_path, "r+") as f:
-                fcntl.flock(f, fcntl.LOCK_EX)
-                data_to_write = json.load(f)
-                fcntl.flock(f, fcntl.LOCK_UN)
-        else:
-            self.init_json['data_path'] = self.dump_tensor_data_dir
-            data_to_write = self.init_json
-        data_to_write[Const.DATA].update(self.cache_data[Const.DATA])
-        with FileOpen(file_path, 'w+') as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
-            json.dump(data_to_write, f, indent=1)
-            fcntl.flock(f, fcntl.LOCK_UN)
-
-        self.cache_data[Const.DATA].clear()
-
-    def write_stack_info_json(self, file_path):
-        with FileOpen(file_path, 'w+') as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
-            json.dump(self.cache_stack, f, indent=1)
-            fcntl.flock(f, fcntl.LOCK_UN)
-
-    def write_construct_info_json(self, file_path):
-        with FileOpen(file_path, 'w+') as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
-            json.dump(self.cache_construct, f, indent=1)
-            fcntl.flock(f, fcntl.LOCK_UN)
-
-    def write_json(self):
-        self.write_data_json(self.dump_file_path)
-        self.write_stack_info_json(self.stack_file_path)
-        self.write_construct_info_json(self.construct_file_path)
+# 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 csv
+import os
+
+from msprobe.core.common.const import Const, FileCheckConst
+from msprobe.core.common.file_utils import change_mode, FileOpen, save_json
+from msprobe.core.common.log import logger
+
+
+class DataWriter:
+
+    def __init__(self) -> None:
+        self.dump_file_path = None
+        self.stack_file_path = None
+        self.construct_file_path = None
+        self.free_benchmark_file_path = None
+        self.dump_tensor_data_dir = None
+        self.flush_size = 1000
+        self.cache_data = {}
+        self.cache_stack = {}
+        self.cache_construct = {}
+
+    @staticmethod
+    def write_data_to_csv(result: list, result_header: tuple, file_path: str):
+        if not result:
+            return
+        is_exists = os.path.exists(file_path)
+        append = "a+" if is_exists else "w+"
+        with FileOpen(file_path, append) as csv_file:
+            spawn_writer = csv.writer(csv_file)
+            if not is_exists:
+                spawn_writer.writerow(result_header)
+            spawn_writer.writerows([result,])
+        is_new_file = not is_exists
+        if is_new_file:
+            change_mode(file_path, FileCheckConst.DATA_FILE_AUTHORITY)
+
+    def reset_cache(self):
+        self.cache_data = {}
+        self.cache_stack = {}
+        self.cache_construct = {}
+
+    def initialize_json_file(self, **kwargs):
+        if not self.cache_data:
+            kwargs.update({"dump_data_dir": self.dump_tensor_data_dir, Const.DATA: {}})
+            self.cache_data = kwargs
+            save_json(self.dump_file_path, self.cache_data, indent=1)
+        if not self.cache_stack:
+            save_json(self.stack_file_path, self.cache_stack, indent=1)
+        if not self.cache_construct:
+            save_json(self.construct_file_path, self.cache_construct, indent=1)
+
+    def update_dump_paths(self, dump_file_path, stack_file_path, construct_file_path, dump_data_dir,
+                          free_benchmark_file_path):
+        self.dump_file_path = dump_file_path
+        self.stack_file_path = stack_file_path
+        self.construct_file_path = construct_file_path
+        self.dump_tensor_data_dir = dump_data_dir
+        self.free_benchmark_file_path = free_benchmark_file_path
+
+    def flush_data_periodically(self):
+        dump_data = self.cache_data.get(Const.DATA)
+        if dump_data and isinstance(dump_data, dict) and len(dump_data) % self.flush_size == 0:
+            self.write_json()
+
+    def update_data(self, new_data):
+        if not isinstance(new_data, dict) or len(new_data.keys()) != 1:
+            logger.warning(f"The data info({new_data}) should be a dict with only one outer key.")
+            return
+        dump_data = self.cache_data.get(Const.DATA)
+        if not isinstance(dump_data, dict):
+            logger.warning(f"The dump data({dump_data}) should be a dict.")
+            return
+
+        key = next(iter(new_data.keys()))
+        if key in dump_data:
+            dump_data.get(key).update(new_data.get(key))
+        else:
+            dump_data.update(new_data)
+
+    def update_stack(self, new_data):
+        self.cache_stack.update(new_data)
+
+    def update_construct(self, new_data):
+        self.cache_construct.update(new_data)
+
+    def write_data_json(self, file_path):
+        logger.info(f"dump.json is at {os.path.dirname(os.path.dirname(file_path))}. ")
+        save_json(file_path, self.cache_data, indent=1)
+
+    def write_stack_info_json(self, file_path):
+        save_json(file_path, self.cache_stack, indent=1)
+
+    def write_construct_info_json(self, file_path):
+        save_json(file_path, self.cache_construct, indent=1)
+
+    def write_json(self):
+        if self.cache_data:
+            self.write_data_json(self.dump_file_path)
+        if self.cache_stack:
+            self.write_stack_info_json(self.stack_file_path)
+        if self.cache_construct:
+            self.write_construct_info_json(self.construct_file_path)

msprobe/core/data_dump/scope.py

@@ -1,178 +1,194 @@
-from abc import ABC, abstractmethod
-from msprobe.core.common.exceptions import ScopeException
-from msprobe.core.common.const import Const
-
-
-def build_scope(scope_class, scope=None, api_list=None):
-    if not scope and not api_list:
-        return None
-    if scope is None:
-        scope = []
-    if api_list is None:
-        api_list = []
-    if scope_class:
-        return scope_class(scope, api_list)
-    return build_range_scope_according_to_scope_name(scope, api_list)
-
-
-def build_range_scope_according_to_scope_name(scope, api_list):
-    api_range_scope = APIRangeScope(scope, api_list)
-    module_range_scope = ModuleRangeScope(scope, api_list)
-    if not scope:  # 如果没有scope参数则用哪类scope都一样
-        return api_range_scope
-    if api_range_scope.is_valid and module_range_scope.is_valid:
-        raise ScopeException(ScopeException.InvalidScope, f"scope={scope}.")
-    elif api_range_scope.is_valid:
-        return api_range_scope
-    elif module_range_scope.is_valid:
-        return module_range_scope
-    else:
-        raise ScopeException(ScopeException.InvalidScope, f"scope={scope}")
-
-
-class BaseScope(ABC):
-    Module_Type_Module = "Module"
-    Module_Type_API = "api"
-
-    def __init__(self, scope, api_list):
-        scope, api_list = self.rectify_args(scope, api_list)
-        self.scope = scope
-        self.api_list = api_list
-
-    @staticmethod
-    def rectify_args(scope, api_list):
-        if not isinstance(api_list, list):
-            raise ScopeException(ScopeException.InvalidApiStr,
-                f"api_list参数须配置为列表，实际类型为{type(api_list)}.")
-        for api in api_list:
-            if not isinstance(api, str):
-                raise ScopeException(ScopeException.InvalidApiStr,
-                    f"api_list中的元素须配置为字符串，实际类型为{type(api)}.")
-        if isinstance(scope, str):
-            scope = [scope]
-            return scope, api_list
-        if not isinstance(scope, list):
-            raise ScopeException(ScopeException.InvalidScope,
-                f"scope参数须配置为字符串或列表，实际类型为{type(scope)}.")
-        for s in scope:
-            if not isinstance(s, str):
-                raise ScopeException(ScopeException.InvalidScope,
-                f"scope列表元素要求类型为字符串，实际类型为{type(s)}.")
-        return scope, api_list
-
-    @abstractmethod
-    def check(self, name):
-        pass
-
-    def check_api_list(self, api_name):
-        if not self.api_list:
-            return True
-        for api_str in self.api_list:
-            if api_str in api_name:
-                return True
-        return False
-
-
-class ListScope(BaseScope):
-    @staticmethod
-    def rectify_args(scope, api_list):
-        if scope and api_list:
-            raise ScopeException(ScopeException.ArgConflict,
-                f"scope和api_list不可以同时配置，实际配置为scope={scope}, api_list={api_list}.")
-        return super(ListScope, ListScope).rectify_args(scope, api_list)
-
-    def check(self, module_name):
-        if not self.scope or module_name in self.scope:
-            return self.check_api_list(module_name)
-        return False
-
-
-class RangeScope(BaseScope, ABC):
-
-    def __init__(self, *args):
-        super().__init__(*args)
-        self.in_scope = False
-        self.is_valid = self.check_scope_is_valid()
-
-
-    @staticmethod
-    def rectify_args(scope, api_list):
-        scope, api_list = super(RangeScope, RangeScope).rectify_args(scope, api_list)
-        if isinstance(scope, list):
-            if len(scope) == 1:
-                scope.append(scope[0])
-            elif len(scope) > 2:
-                raise ScopeException(ScopeException.InvalidScope,
-                    f"scope参数指定区间断点，须传入长度为1或2的列表，实际长度为{len(scope)}.")
-
-        return scope, api_list
-
-    @abstractmethod
-    def check_scope_is_valid(self):
-        pass
-
-    def begin_module(self, module_name):
-        pass
-
-    def end_module(self, module_name):
-        pass
-
-
-class APIRangeScope(RangeScope):
-    def check_scope_is_valid(self):
-        if not self.scope:
-            return True
-        scope_start_type = self.scope[0].split(Const.SEP)[0]
-        if scope_start_type == BaseScope.Module_Type_Module:
-            return False
-        scope_stop_type = self.scope[1].split(Const.SEP)[0]
-        if scope_stop_type == BaseScope.Module_Type_Module:
-            return False
-        return True
-
-    def check(self, api_name):
-        if self.scope and api_name == self.scope[0]:
-            self.in_scope = True
-
-        if not self.scope or self.in_scope:
-            result = self.check_api_list(api_name)
-        else:
-            result = False
-
-        if self.scope and api_name == self.scope[1]:
-            self.in_scope = False
-        return result
-
-
-class ModuleRangeScope(RangeScope):
-    """
-        模块与api不同的是，模块内部还有子结构需要dump，
-        需要用pre_hook和full_backward_hook来精确控制module的开始和结束，
-        在这些hook触发时调用begin_module和end_module做区间控制
-    """
-    def check_scope_is_valid(self):
-        if not self.scope:
-            return True
-        scope_start_type = self.scope[0].split(Const.SEP)[0]
-        scope_stop_type = self.scope[1].split(Const.SEP)[0]
-        if scope_start_type == BaseScope.Module_Type_Module and \
-                scope_stop_type == BaseScope.Module_Type_Module:
-            return True
-        return False
-
-    def begin_module(self, module_name):
-        if not self.scope:
-            return
-        if module_name == self.scope[0]:
-            self.in_scope = True
-
-    def end_module(self, module_name):
-        if not self.scope:
-            return
-        if module_name == self.scope[1]:
-            self.in_scope = False
-
-    def check(self, module_name):
-        if not self.scope or self.in_scope:
-            return self.check_api_list(module_name)
-        return False
+# 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 abc import ABC, abstractmethod
+
+from msprobe.core.common.const import Const
+from msprobe.core.common.exceptions import ScopeException
+
+
+def build_scope(scope_class, scope=None, api_list=None):
+    if not scope and not api_list:
+        return None
+    if scope is None:
+        scope = []
+    if api_list is None:
+        api_list = []
+    if scope_class:
+        return scope_class(scope, api_list)
+    return build_range_scope_according_to_scope_name(scope, api_list)
+
+
+def build_range_scope_according_to_scope_name(scope, api_list):
+    api_range_scope = APIRangeScope(scope, api_list)
+    module_range_scope = ModuleRangeScope(scope, api_list)
+    if not scope:  # 如果没有scope参数则用哪类scope都一样
+        return api_range_scope
+    if api_range_scope.is_valid and module_range_scope.is_valid:
+        raise ScopeException(ScopeException.InvalidScope, f"scope={scope}.")
+    elif api_range_scope.is_valid:
+        return api_range_scope
+    elif module_range_scope.is_valid:
+        return module_range_scope
+    else:
+        raise ScopeException(ScopeException.InvalidScope, f"scope={scope}")
+
+
+class BaseScope(ABC):
+    Module_Type_Module = "Module"
+    Module_Type_API = "api"
+    module_type = ["Module", "Cell"]
+
+    def __init__(self, scope, api_list):
+        scope, api_list = self.rectify_args(scope, api_list)
+        self.scope = scope
+        self.api_list = api_list
+
+    @staticmethod
+    def rectify_args(scope, api_list):
+        if not isinstance(api_list, list):
+            raise ScopeException(ScopeException.InvalidApiStr,
+                f"api_list参数须配置为列表，实际类型为{type(api_list)}.")
+        for api in api_list:
+            if not isinstance(api, str):
+                raise ScopeException(ScopeException.InvalidApiStr,
+                    f"api_list中的元素须配置为字符串，实际类型为{type(api)}.")
+        if isinstance(scope, str):
+            scope = [scope]
+            return scope, api_list
+        if not isinstance(scope, list):
+            raise ScopeException(ScopeException.InvalidScope,
+                f"scope参数须配置为字符串或列表，实际类型为{type(scope)}.")
+        for s in scope:
+            if not isinstance(s, str):
+                raise ScopeException(ScopeException.InvalidScope,
+                f"scope列表元素要求类型为字符串，实际类型为{type(s)}.")
+        return scope, api_list
+
+    @abstractmethod
+    def check(self, name):
+        pass
+
+    def check_api_list(self, api_name):
+        if not self.api_list:
+            return True
+        for api_str in self.api_list:
+            if api_str in api_name:
+                return True
+        return False
+
+
+class ListScope(BaseScope):
+    @staticmethod
+    def rectify_args(scope, api_list):
+        if scope and api_list:
+            raise ScopeException(ScopeException.ArgConflict,
+                f"scope和api_list不可以同时配置，实际配置为scope={scope}, api_list={api_list}.")
+        return super(ListScope, ListScope).rectify_args(scope, api_list)
+
+    def check(self, name):
+        if not self.scope or name in self.scope:
+            return self.check_api_list(name)
+        return False
+
+
+class RangeScope(BaseScope, ABC):
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        self.in_scope = False
+        self.is_valid = self.check_scope_is_valid()
+
+    @staticmethod
+    def rectify_args(scope, api_list):
+        scope, api_list = super(RangeScope, RangeScope).rectify_args(scope, api_list)
+        if isinstance(scope, list):
+            if len(scope) == 1:
+                scope.append(scope[0])
+            elif len(scope) > 2:
+                raise ScopeException(ScopeException.InvalidScope,
+                    f"scope参数指定区间断点，须传入长度为1或2的列表，实际长度为{len(scope)}.")
+        return scope, api_list
+
+    @abstractmethod
+    def check_scope_is_valid(self):
+        pass
+
+    def begin_module(self, module_name):
+        pass
+
+    def end_module(self, module_name):
+        pass
+
+
+class APIRangeScope(RangeScope):
+    def check_scope_is_valid(self):
+        if not self.scope:
+            return True
+        scope_start_type = self.scope[0].split(Const.SEP)[0]
+        if scope_start_type in BaseScope.module_type:
+            return False
+        scope_stop_type = self.scope[1].split(Const.SEP)[0]
+        if scope_stop_type in BaseScope.module_type:
+            return False
+        return True
+
+    def check(self, name):
+        if self.scope and name == self.scope[0]:
+            self.in_scope = True
+
+        if not self.scope or self.in_scope:
+            result = self.check_api_list(name)
+        else:
+            result = False
+
+        if self.scope and name == self.scope[1]:
+            self.in_scope = False
+        return result
+
+
+class ModuleRangeScope(RangeScope):
+    """
+        模块与api不同的是，模块内部还有子结构需要dump，
+        需要用pre_hook和full_backward_hook来精确控制module的开始和结束，
+        在这些hook触发时调用begin_module和end_module做区间控制
+    """
+
+    def check_scope_is_valid(self):
+        if not self.scope:
+            return True
+        scope_start_type = self.scope[0].split(Const.SEP)[0]
+        scope_stop_type = self.scope[1].split(Const.SEP)[0]
+        if scope_start_type in BaseScope.module_type and \
+                scope_stop_type in BaseScope.module_type:
+            return True
+        return False
+
+    def begin_module(self, module_name):
+        if not self.scope:
+            return
+        if module_name == self.scope[0]:
+            self.in_scope = True
+
+    def end_module(self, module_name):
+        if not self.scope:
+            return
+        if module_name == self.scope[1]:
+            self.in_scope = False
+
+    def check(self, name):
+        if not self.scope or self.in_scope:
+            return self.check_api_list(name)
+        return False