mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/core/common/global_lock.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2025-2025, 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 multiprocessing.shared_memory import SharedMemory
+import random
+import time
+import atexit
+import os
+
+from msprobe.core.common.log import logger
+
+
+def is_main_process():
+    return multiprocessing.current_process().name == 'MainProcess'
+
+
+class GlobalLock:
+    def __init__(self):
+        self.name = self.get_lock_name()
+        try:
+            self._shm = SharedMemory(create=False, name=self.name)
+            time.sleep(random.randint(0, 500) / 10000) # 等待随机时长以避免同时获得锁
+        except FileNotFoundError:
+            try:
+                self._shm = SharedMemory(create=True, name=self.name, size=1)
+                self._shm.buf[0] = 0
+                logger.debug(f'{self.name} is created.')
+            except FileExistsError:
+                self.__init__()
+
+    @classmethod
+    def get_lock_name(cls):
+        if is_main_process():
+            return f'global_lock_{os.getpid()}'
+        return f'global_lock_{os.getppid()}'
+
+    @classmethod
+    def is_lock_exist(cls):
+        try:
+            SharedMemory(create=False, name=cls.get_lock_name()).close()
+            return True
+        except FileNotFoundError:
+            return False
+
+    def cleanup(self):
+        self._shm.close()
+        if is_main_process():
+            try:
+                self._shm.unlink()
+                logger.debug(f'{self.name} is unlinked.')
+            except FileNotFoundError:
+                logger.warning(f'{self.name} has already been unlinked.')
+
+    def acquire(self, timeout=180):
+        """
+        acquire global lock, default timeout is 3 minutes.
+
+        :param float timeout: timeout(seconds), default value is 180.
+        """
+        start = time.time()
+        while time.time() - start < timeout:
+            if self._shm.buf[0] == 0:
+                self._shm.buf[0] = 1
+                return
+            time.sleep(random.randint(10, 500) / 10000)  # 自旋，等待1-50ms
+        self._shm.buf[0] = 1
+
+    def release(self):
+        self._shm.buf[0] = 0
+
+
+global_lock = GlobalLock()
+atexit.register(global_lock.cleanup)

msprobe/core/common/runtime.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2025, 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
+
+
+class Runtime:
+    step_count: int = 0
+    rank_id: int = -1
+    is_running: bool = False
+    run_mode: str = Const.PYNATIVE_MODE
+    current_iter: int = 0
+    current_rank: None

msprobe/core/common/utils.py

@@ -18,9 +18,8 @@ import os
 import re
 import subprocess
 import time
-from collections import defaultdict
+import inspect
 from datetime import datetime, timezone
-from functools import wraps
 
 import numpy as np
 
@@ -28,10 +27,15 @@ from msprobe.core.common.file_utils import (FileOpen, check_file_or_directory_pa
 from msprobe.core.common.const import Const, CompareConst
 from msprobe.core.common.log import logger
 from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.decorator import recursion_depth_decorator
 
 
 device = collections.namedtuple('device', ['type', 'index'])
 prefixes = ['api_stack', 'list', 'range', 'acl']
+file_suffix_to_file_type = {
+    "dump.json": Const.DUMP_JSON_FILE,
+    "debug.json": Const.DEBUG_JSON_FILE,
+}
 
 
 class MsprobeBaseException(Exception):
@@ -75,6 +79,8 @@ class MsprobeBaseException(Exception):
     MERGE_COMPARE_RESULT_ERROR = 33
     NAMES_STRUCTS_MATCH_ERROR = 34
     INVALID_STATE_ERROR = 35
+    INVALID_API_NAME_ERROR = 36
+    CROSS_FRAME_ERROR = 37
 
     def __init__(self, code, error_info: str = ""):
         super(MsprobeBaseException, self).__init__()
@@ -191,27 +197,6 @@ def check_regex_prefix_format_valid(prefix):
         raise ValueError(f"prefix contains invalid characters, prefix pattern {Const.REGEX_PREFIX_PATTERN}")
 
 
-def execute_command(cmd):
-    """
-    Function Description:
-        run the following command
-    Parameter:
-        cmd: command
-    Exception Description:
-        when invalid command throw exception
-    """
-    logger.info('Execute command:%s' % cmd)
-    process = subprocess.Popen(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    while process.poll() is None:
-        line = process.stdout.readline()
-        line = line.strip()
-        if line:
-            logger.info(line)
-    if process.returncode != 0:
-        logger.error('Failed to execute command:%s' % " ".join(cmd))
-        raise CompareException(CompareException.INVALID_DATA_ERROR)
-
-
 def add_time_as_suffix(name):
     return '{}_{}.csv'.format(name, time.strftime("%Y%m%d%H%M%S", time.localtime(time.time())))
 
@@ -232,21 +217,41 @@ def format_value(value):
     return float('{:.12f}'.format(value))
 
 
-def md5_find(data):
-    for key_op in data:
-        for api_info in data[key_op]:
-            if isinstance(data[key_op][api_info], list):
-                for data_detail in data[key_op][api_info]:
-                    if data_detail and 'md5' in data_detail:
-                        return True
-            if isinstance(data[key_op][api_info], bool):
-                continue
-            elif data[key_op][api_info] and 'md5' in data[key_op][api_info]:
+@recursion_depth_decorator('msprobe.core.common.utils.md5_find', max_depth=Const.DUMP_MAX_DEPTH)
+def md5_find(data, json_type=Const.DUMP_JSON_FILE):
+    if json_type == Const.DUMP_JSON_FILE:
+        for key_op in data:
+            for api_info in data[key_op]:
+                if isinstance(data[key_op][api_info], list):
+                    for data_detail in data[key_op][api_info]:
+                        if data_detail and Const.MD5 in data_detail:
+                            return True
+                if isinstance(data[key_op][api_info], bool):
+                    continue
+                elif data[key_op][api_info] and Const.MD5 in data[key_op][api_info]:
+                    return True
+    elif json_type == Const.DEBUG_JSON_FILE:
+        if isinstance(data, dict):
+            if Const.MD5 in data:
                 return True
+            else:
+                for _, data_info in data.items():
+                    if md5_find(data_info, Const.DEBUG_JSON_FILE):
+                        return True
+        elif isinstance(data, list):
+            for data_info in data:
+                if md5_find(data_info, Const.DEBUG_JSON_FILE):
+                    return True
+        else:
+            return False
     return False
 
 
 def detect_framework_by_dump_json(file_path):
+    json_data = load_json(file_path)
+    framework = json_data.get("framework", None)
+    if framework in [Const.PT_FRAMEWORK, Const.MS_FRAMEWORK]:
+        return framework
     pattern_ms = r'"type":\s*"mindspore'
     pattern_pt = r'"type":\s*"torch'
     with FileOpen(file_path, 'r') as file:
@@ -276,13 +281,26 @@ def get_stack_construct_by_dump_json_path(dump_json_path):
 def set_dump_path(input_param):
     npu_path = input_param.get("npu_json_path", None)
     bench_path = input_param.get("bench_json_path", None)
-    npu_path_valid = npu_path is not None and npu_path.endswith("dump.json")
-    bench_path_valid = bench_path is not None and bench_path.endswith("dump.json")
-    if not npu_path_valid or not bench_path_valid:
-        logger.error(f"Please check the json path is valid. npu_path: {npu_path}, bench_path: {bench_path}")
+    dump_json_path_valid = npu_path is not None and npu_path.endswith("dump.json") and \
+        bench_path is not None and bench_path.endswith("dump.json")
+    debug_json_path_valid = npu_path is not None and npu_path.endswith("debug.json") and \
+        bench_path is not None and bench_path.endswith("debug.json")
+    if not dump_json_path_valid and not debug_json_path_valid:
+        logger.error(f"Please check the json path is valid and ensure that neither npu_path nor bench_path is None.")
+        raise CompareException(CompareException.INVALID_PATH_ERROR)
+    input_param[CompareConst.NPU_DUMP_DATA_DIR] = os.path.join(os.path.dirname(npu_path), Const.DUMP_TENSOR_DATA)
+    input_param[CompareConst.BENCH_DUMP_DATA_DIR] = os.path.join(os.path.dirname(bench_path), Const.DUMP_TENSOR_DATA)
+
+
+def get_file_type(file_path):
+    if not isinstance(file_path, str):
+        logger.error("get_file_type failed, check the type of file_path.")
+        raise CompareException(CompareException.INVALID_PATH_ERROR)
+    file_type = file_suffix_to_file_type.get(file_path.split(Const.SCOPE_SEPARATOR)[-1])
+    if file_type is None:
+        logger.error("get_file_type failed, file_path is neither dump.json nor debug.json.")
         raise CompareException(CompareException.INVALID_PATH_ERROR)
-    input_param['npu_dump_data_dir'] = os.path.join(os.path.dirname(npu_path), Const.DUMP_TENSOR_DATA)
-    input_param['bench_dump_data_dir'] = os.path.join(os.path.dirname(bench_path), Const.DUMP_TENSOR_DATA)
+    return file_type
 
 
 def get_dump_mode(input_param):
@@ -290,6 +308,7 @@ def get_dump_mode(input_param):
     bench_path = input_param.get("bench_json_path", None)
     npu_json_data = load_json(npu_path)
     bench_json_data = load_json(bench_path)
+    json_type = get_file_type(file_path=npu_path)
 
     npu_task = npu_json_data.get('task', None)
     bench_task = bench_json_data.get('task', None)
@@ -309,8 +328,8 @@ def get_dump_mode(input_param):
         return Const.STRUCTURE
 
     if npu_task == Const.STATISTICS:
-        npu_md5_compare = md5_find(npu_json_data['data'])
-        bench_md5_compare = md5_find(bench_json_data['data'])
+        npu_md5_compare = md5_find(npu_json_data['data'], json_type)
+        bench_md5_compare = md5_find(bench_json_data['data'], json_type)
         if npu_md5_compare == bench_md5_compare:
             return Const.MD5 if npu_md5_compare else Const.SUMMARY
         else:
@@ -424,6 +443,37 @@ def get_real_step_or_rank(step_or_rank_input, obj):
     return real_step_or_rank
 
 
+def check_init_step(step):
+    if not is_int(step):
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                        f"{step} must be an integer")
+    if not step >= 0:
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                f"{step} must be greater than or equal to 0")
+
+
+def check_token_range(token_range):
+    if token_range is None:
+        return
+    if not isinstance(token_range, (list, tuple)):
+        logger.error("Token_range must be a list or tuple.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if len(token_range) != 2:
+        logger.error("Token_range must contains exactly 2 elements.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+    start, end = token_range
+    if not isinstance(start, int) or not isinstance(end, int):
+        logger.error("Start and end in token_range must be integer.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if start > end:
+        logger.error("Start in token_range must less than the end.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if start < 0:
+        logger.error("Start in token_range must >= 0.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+
 def check_seed_all(seed, mode, rm_dropout):
     if is_int(seed):
         if seed < 0 or seed > Const.MAX_SEED_VALUE:
@@ -467,36 +517,6 @@ def safe_get_value(container, index, container_name, key=None):
         raise MsprobeBaseException(MsprobeBaseException.INVALID_OBJECT_TYPE_ERROR) from e
 
 
-# 记录工具函数递归的深度
-recursion_depth = defaultdict(int)
-
-
-# 装饰一个函数，当函数递归调用超过限制时，抛出异常并打印函数信息。
-def recursion_depth_decorator(func_info):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            func_id = id(func)
-            recursion_depth[func_id] += 1
-            if recursion_depth[func_id] > Const.MAX_DEPTH:
-                msg = f"call {func_info} exceeds the recursion limit."
-                logger.error_log_with_exp(
-                    msg,
-                    MsprobeException(
-                        MsprobeException.RECURSION_LIMIT_ERROR, msg
-                    ),
-                )
-            try:
-                result = func(*args, **kwargs)
-            finally:
-                recursion_depth[func_id] -= 1
-            return result
-
-        return wrapper
-
-    return decorator
-
-
 def check_str_param(param):
     if not re.match(Const.REGEX_PREFIX_PATTERN, param):
         logger.error('The parameter {} contains special characters.'.format(param))
@@ -509,4 +529,60 @@ class DumpPathAggregation:
     construct_file_path = None
     dump_tensor_data_dir = None
     free_benchmark_file_path = None
-    debug_file_path = None
+    debug_file_path = None
+
+
+def is_save_variable_valid(variable, valid_special_types, depth=0):
+    if depth > Const.DUMP_MAX_DEPTH:
+        return False
+    if isinstance(variable, valid_special_types):
+        return True
+    elif isinstance(variable, (list, tuple)):
+        return all(is_save_variable_valid(item, valid_special_types, depth + 1) for item in variable)
+    elif isinstance(variable, dict):
+        return all(isinstance(key, str) and is_save_variable_valid(value, valid_special_types, depth + 1)
+                   for key, value in variable.items())
+    else:
+        return False
+
+
+def replace_last_occurrence(text, old, new):
+    if text is None:
+        return text
+    index = text.rfind(old)
+    if index != -1:
+        return text[:index] + text[index:].replace(old, new, 1)
+    return text
+
+
+def load_stack_json(stack_path):
+    stack_dict = load_json(stack_path)
+    if not stack_dict.get(Const.NEW_STACK_FLAG):
+        return stack_dict
+
+    new_stack_dict = {}
+    for stack_info in stack_dict.values():
+        if len(stack_info) != 2:
+            continue
+        api_list, stack_str = stack_info
+        for api_name in api_list:
+            new_stack_dict.update({api_name: stack_str})
+    return new_stack_dict
+
+
+def analyze_api_call_stack(name):
+    try:
+        api_stack = inspect.stack()[2:]
+    except Exception as e:
+        logger.warning(f"The call stack of {name} failed to retrieve, {e}.")
+        api_stack = None
+    stack_str = []
+    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()} \n"
+            stack_str.append(stack_line)
+    else:
+        stack_str.append(Const.WITHOUT_CALL_STACK)
+    return "".join(stack_str)

msprobe/core/common_config.py

@@ -111,3 +111,10 @@ class BaseConfig:
                         f"The element '{mode}' of data_mode {self.data_mode} is not in {Const.DUMP_DATA_MODE_LIST}.",
                         MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
                     )
+
+    def _check_summary_mode(self):
+        if self.summary_mode and self.summary_mode not in Const.SUMMARY_MODE:
+            logger.error_log_with_exp(
+                        f"summary_mode is invalid, summary_mode is not in {Const.SUMMARY_MODE}.",
+                        MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+                    )