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

msprobe/core/common/utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-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
 #
@@ -13,14 +12,17 @@
 # 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 collections
 import os
 import re
 import subprocess
 import time
-import json
+from collections import defaultdict
 from datetime import datetime, timezone
+from functools import wraps
+
+import numpy as np
 
 from msprobe.core.common.file_utils import (FileOpen, check_file_or_directory_path, load_json)
 from msprobe.core.common.const import Const, CompareConst
@@ -68,6 +70,11 @@ class MsprobeBaseException(Exception):
     FUNCTION_CALL_ERROR = 28
     FORWARD_DATA_COLLECTION_ERROR = 29
     BACKWARD_DATA_COLLECTION_ERROR = 30
+    INVALID_KEY_ERROR = 31
+    MISSING_HEADER_ERROR = 32
+    MERGE_COMPARE_RESULT_ERROR = 33
+    NAMES_STRUCTS_MATCH_ERROR = 34
+    INVALID_STATE_ERROR = 35
 
     def __init__(self, code, error_info: str = ""):
         super(MsprobeBaseException, self).__init__()
@@ -99,7 +106,14 @@ class DumpException(MsprobeBaseException):
         return f"Dump Error Code {self.code}: {self.error_info}"
 
 
-def check_compare_param(input_param, output_path, summary_compare=False, md5_compare=False):
+def is_json_file(file_path):
+    if isinstance(file_path, str) and file_path.lower().endswith('.json'):
+        return True
+    else:
+        return False
+
+
+def check_compare_param(input_param, output_path, dump_mode, stack_mode):
     if not isinstance(input_param, dict):
         logger.error(f"Invalid input parameter 'input_param', the expected type dict but got {type(input_param)}.")
         raise CompareException(CompareException.INVALID_PARAM_ERROR)
@@ -107,18 +121,31 @@ def check_compare_param(input_param, output_path, summary_compare=False, md5_com
         logger.error(f"Invalid input parameter 'output_path', the expected type str but got {type(output_path)}.")
         raise CompareException(CompareException.INVALID_PARAM_ERROR)
 
-    check_file_or_directory_path(input_param.get("npu_json_path"), False)
-    check_file_or_directory_path(input_param.get("bench_json_path"), False)
-    check_file_or_directory_path(input_param.get("stack_json_path"), False)
-    if not summary_compare and not md5_compare:
+    def check_json_path(json_path_str):
+        json_path = input_param.get(json_path_str)
+        check_file_or_directory_path(json_path, False)
+        json_type_check = is_json_file(json_path)
+        if not json_type_check:
+            logger.error(f"Invalid {json_path_str}: {json_path}, please check!")
+            raise CompareException(CompareException.INVALID_PATH_ERROR)
+
+    check_json_path("npu_json_path")
+    check_json_path("bench_json_path")
+    if stack_mode:
+        check_json_path("stack_json_path")
+
+    if dump_mode == Const.ALL:
         check_file_or_directory_path(input_param.get("npu_dump_data_dir"), True)
         check_file_or_directory_path(input_param.get("bench_dump_data_dir"), True)
     check_file_or_directory_path(output_path, True)
 
     with FileOpen(input_param.get("npu_json_path"), "r") as npu_json, \
-            FileOpen(input_param.get("bench_json_path"), "r") as bench_json, \
-            FileOpen(input_param.get("stack_json_path"), "r") as stack_json:
-        check_json_file(input_param, npu_json, bench_json, stack_json)
+            FileOpen(input_param.get("bench_json_path"), "r") as bench_json:
+        _check_json(npu_json, input_param.get("npu_json_path"))
+        _check_json(bench_json, input_param.get("bench_json_path"))
+    if stack_mode:
+        with FileOpen(input_param.get("stack_json_path"), "r") as stack_json:
+            _check_json(stack_json, input_param.get("stack_json_path"))
 
 
 def check_configuration_param(stack_mode=False, auto_analyze=True, fuzzy_match=False, is_print_compare_log=True):
@@ -179,7 +206,7 @@ def execute_command(cmd):
         line = process.stdout.readline()
         line = line.strip()
         if line:
-            print(line)
+            logger.info(line)
     if process.returncode != 0:
         logger.error('Failed to execute command:%s' % " ".join(cmd))
         raise CompareException(CompareException.INVALID_DATA_ERROR)
@@ -212,25 +239,29 @@ def md5_find(data):
                 for data_detail in data[key_op][api_info]:
                     if data_detail and 'md5' in data_detail:
                         return True
-            elif 'md5' in data[key_op][api_info]:
+            elif data[key_op][api_info] and 'md5' in data[key_op][api_info]:
                 return True
     return False
 
 
-def struct_json_get(input_param, framework):
-    if framework == Const.PT_FRAMEWORK:
-        prefix = "bench"
-    elif framework == Const.MS_FRAMEWORK:
-        prefix = "npu"
-    else:
-        logger.error("Error framework found.")
-        raise CompareException(CompareException.INVALID_PARAM_ERROR)
+def detect_framework_by_dump_json(file_path):
+    pattern_ms = r'"type":\s*"mindspore'
+    pattern_pt = r'"type":\s*"torch'
+    with FileOpen(file_path, 'r') as file:
+        for line in file:
+            if re.search(pattern_ms, line):
+                return Const.MS_FRAMEWORK
+            if re.search(pattern_pt, line):
+                return Const.PT_FRAMEWORK
+    logger.error(f"{file_path} must be based on the MindSpore or PyTorch framework.")
+    raise CompareException(CompareException.INVALID_PARAM_ERROR)
 
-    frame_json_path = input_param.get(f"{prefix}_json_path", None)
-    if not frame_json_path:
-        logger.error(f"Please check the json path is valid.")
+
+def get_stack_construct_by_dump_json_path(dump_json_path):
+    if not dump_json_path:
+        logger.error("The path is empty. Please enter a valid path.")
         raise CompareException(CompareException.INVALID_PATH_ERROR)
-    directory = os.path.dirname(frame_json_path)
+    directory = os.path.dirname(dump_json_path)
     check_file_or_directory_path(directory, True)
     stack_json = os.path.join(directory, "stack.json")
     construct_json = os.path.join(directory, "construct.json")
@@ -240,41 +271,57 @@ def struct_json_get(input_param, framework):
     return stack, construct
 
 
-def task_dumppath_get(input_param):
+def set_dump_path(input_param):
     npu_path = input_param.get("npu_json_path", None)
     bench_path = input_param.get("bench_json_path", None)
-    if not npu_path or not bench_path:
-        logger.error(f"Please check the json path is valid.")
+    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}")
         raise CompareException(CompareException.INVALID_PATH_ERROR)
-    with FileOpen(npu_path, 'r') as npu_f:
-        npu_json_data = json.load(npu_f)
-    with FileOpen(bench_path, 'r') as bench_f:
-        bench_json_data = json.load(bench_f)
-    if npu_json_data['task'] != bench_json_data['task']:
+    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)
+
+
+def get_dump_mode(input_param):
+    npu_path = input_param.get("npu_json_path", None)
+    bench_path = input_param.get("bench_json_path", None)
+    npu_json_data = load_json(npu_path)
+    bench_json_data = load_json(bench_path)
+
+    npu_task = npu_json_data.get('task', None)
+    bench_task = bench_json_data.get('task', None)
+
+    if not npu_task or not bench_task:
+        logger.error(f"Please check the dump task is correct, npu's task is {npu_task}, bench's task is {bench_task}.")
+        raise CompareException(CompareException.INVALID_TASK_ERROR)
+
+    if npu_task != bench_task:
         logger.error(f"Please check the dump task is consistent.")
         raise CompareException(CompareException.INVALID_TASK_ERROR)
-    if npu_json_data['task'] == Const.TENSOR:
-        summary_compare = False
-        md5_compare = False
-    elif npu_json_data['task'] == Const.STATISTICS:
-        md5_compare = md5_find(npu_json_data['data'])
-        if md5_compare:
-            summary_compare = False
+
+    if npu_task == Const.TENSOR:
+        return Const.ALL
+
+    if npu_task == Const.STATISTICS:
+        npu_md5_compare = md5_find(npu_json_data['data'])
+        bench_md5_compare = md5_find(bench_json_data['data'])
+        if npu_md5_compare == bench_md5_compare:
+            return Const.MD5 if npu_md5_compare else Const.SUMMARY
         else:
-            summary_compare = True
-    else:
-        logger.error(f"Compare is not required for overflow_check or free_benchmark.")
-        raise CompareException(CompareException.INVALID_TASK_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 summary_compare, md5_compare
+            logger.error(f"Please check the dump task is consistent, "
+                         f"dump mode of npu and bench should both be statistics or md5.")
+            raise CompareException(CompareException.INVALID_TASK_ERROR)
 
+    logger.error(f"Compare applies only to task is tensor or statistics")
+    raise CompareException(CompareException.INVALID_TASK_ERROR)
 
-def get_header_index(header_name, summary_compare=False):
-    if summary_compare:
-        header = CompareConst.SUMMARY_COMPARE_RESULT_HEADER[:]
-    else:
-        header = CompareConst.COMPARE_RESULT_HEADER[:]
+
+def get_header_index(header_name, dump_mode):
+    header = CompareConst.HEAD_OF_COMPARE_MODE.get(dump_mode)
+    if not header:
+        logger.error(f"{dump_mode} not in {CompareConst.HEAD_OF_COMPARE_MODE}")
+        raise CompareException(CompareException.INVALID_PARAM_ERROR)
     if header_name not in header:
         logger.error(f"{header_name} not in data name")
         raise CompareException(CompareException.INVALID_PARAM_ERROR)
@@ -282,7 +329,7 @@ def get_header_index(header_name, summary_compare=False):
 
 
 def convert_tuple(data):
-    return data if isinstance(data, tuple) else (data, )
+    return data if isinstance(data, tuple) else (data,)
 
 
 def check_op_str_pattern_valid(string, op_name=None, stack=False):
@@ -302,6 +349,10 @@ def is_invalid_pattern(string):
     return re.search(pattern, string)
 
 
+def is_int(x):
+    return isinstance(x, int) and not isinstance(x, bool)
+
+
 def print_tools_ends_info():
     total_len = len(Const.TOOL_ENDS_SUCCESSFULLY) + Const.FILL_CHAR_NUMS
     logger.info('*' * total_len)
@@ -315,51 +366,61 @@ def get_step_or_rank_from_string(step_or_rank, obj):
         try:
             borderlines = int(splited[0]), int(splited[1])
         except (ValueError, IndexError) as e:
-            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    "The hyphen(-) must start and end with decimal numbers.") from e
     else:
-        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
-                               f'The string parameter for {obj} only supports formats like "3-5". Now string parameter for {obj} is "{step_or_rank}".')
-    if all(Const.STEP_RANK_MAXIMUM_RANGE[0] <= b <= Const.STEP_RANK_MAXIMUM_RANGE[1] for b in borderlines):
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                               f'The string parameter for {obj} only supports formats like "3-5". '
+                               f'Now string parameter for {obj} is "{step_or_rank}".')
+    if all(Const.STEP_RANK_MINIMUM_VALUE <= b <= Const.STEP_RANK_MAXIMUM_VALUE for b in borderlines):
         if borderlines[0] <= borderlines[1]:
             continual_step_or_rank = list(range(borderlines[0], borderlines[1] + 1))
         else:
-            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
-                               f'For the hyphen(-) in {obj}, the left boundary ({borderlines[0]}) cannot be greater than the right boundary ({borderlines[1]}).')
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                                   f'For the hyphen(-) in {obj}, the left boundary ({borderlines[0]}) cannot be '
+                                   f'greater than the right boundary ({borderlines[1]}).')
     else:
-        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
-                               f"The boundaries must fall within the range of [{Const.STEP_RANK_MAXIMUM_RANGE[0]}, {Const.STEP_RANK_MAXIMUM_RANGE[1]}].")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                               f"The boundaries must fall within the range of "
+                               f"[{Const.STEP_RANK_MINIMUM_VALUE}, {Const.STEP_RANK_MAXIMUM_VALUE}].")
     return continual_step_or_rank
 
 
 def get_real_step_or_rank(step_or_rank_input, obj):
     if obj not in [Const.STEP, Const.RANK]:
-        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                f"Only support parsing {[Const.STEP, Const.RANK]}, the current parsing object is {obj}.")
     if step_or_rank_input is None:
         return []
     if not isinstance(step_or_rank_input, list):
         raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, f"{obj} is invalid, it should be a list")
+    if len(step_or_rank_input) > Const.STEP_RANK_MAXIMUM_VALUE:
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
+                               f"{obj} is invalid, its length cannot exceed {Const.STEP_RANK_MAXIMUM_VALUE}")
+
     real_step_or_rank = []
     for element in step_or_rank_input:
-        if not isinstance(element, (int, str)):
-            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
+        if not is_int(element) and not isinstance(element, str):
+            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR,
                                    f"{obj} element {element} must be an integer or string.")
-        if isinstance(element, int) and element < 0:
-            raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR, 
-                                   f"Each element of {obj} must be non-negative, currently it is {element}.")
-        if isinstance(element, int) and Const.STEP_RANK_MAXIMUM_RANGE[0] <= element <= Const.STEP_RANK_MAXIMUM_RANGE[1]:
+        if is_int(element):
+            if not Const.STEP_RANK_MINIMUM_VALUE <= element <= Const.STEP_RANK_MAXIMUM_VALUE:
+                raise MsprobeException(
+                    MsprobeException.INVALID_PARAM_ERROR,
+                    f"Each element of {obj} must be between {Const.STEP_RANK_MINIMUM_VALUE} and "
+                    f"{Const.STEP_RANK_MAXIMUM_VALUE}, currently it is {element}."
+                )
             real_step_or_rank.append(element)
-        elif isinstance(element, str) and Const.HYPHEN in element:
-            continual_step_or_rank = get_step_or_rank_from_string(element, obj)
-            real_step_or_rank.extend(continual_step_or_rank)
+            continue
+        continual_step_or_rank = get_step_or_rank_from_string(element, obj)
+        real_step_or_rank.extend(continual_step_or_rank)
     real_step_or_rank = list(set(real_step_or_rank))
     real_step_or_rank.sort()
     return real_step_or_rank
 
 
-def check_seed_all(seed, mode):
-    if isinstance(seed, int):
+def check_seed_all(seed, mode, rm_dropout):
+    if is_int(seed):
         if seed < 0 or seed > Const.MAX_SEED_VALUE:
             logger.error(f"Seed must be between 0 and {Const.MAX_SEED_VALUE}.")
             raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
@@ -369,3 +430,69 @@ def check_seed_all(seed, mode):
     if not isinstance(mode, bool):
         logger.error("seed_all mode must be bool.")
         raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+    if not isinstance(rm_dropout, bool):
+        logger.error("The rm_dropout parameter must be bool.")
+        raise MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+
+
+def safe_get_value(container, index, container_name, key=None):
+    try:
+        # 处理字典情况
+        if isinstance(container, dict):
+            return container.get(key)[index]
+        # 处理列表、元组、numpy情况
+        elif isinstance(container, (list, tuple, np.ndarray)):
+            return container[index]
+        else:
+            err_msg = f"Unsupported container type for '{container_name}': {type(container)}"
+            logger.error(err_msg)
+            raise MsprobeBaseException(MsprobeBaseException.INVALID_OBJECT_TYPE_ERROR)
+    except IndexError as e:
+        err_msg = "index out of bounds error occurs, please check!\n" \
+                  f"{container_name} is {container}\n" \
+                  f"index is {index}"
+        logger.error(err_msg)
+        raise MsprobeBaseException(MsprobeBaseException.INDEX_OUT_OF_BOUNDS_ERROR) from e
+    except TypeError as e:
+        err_msg = "wrong type, please check!\n" \
+                  f"{container_name} is {container}\n" \
+                  f"index is {index}\n" \
+                  f"key is {key}"
+        logger.error(err_msg)
+        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))
+        raise MsprobeBaseException(MsprobeBaseException.INVALID_CHAR_ERROR)

msprobe/core/common_config.py

@@ -1,7 +1,21 @@
+# 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, FileCheckConst
 from msprobe.core.common.log import logger
 from msprobe.core.common.exceptions import MsprobeException
-from msprobe.core.common.file_utils import FileChecker
 from msprobe.core.common.utils import get_real_step_or_rank
 
 
@@ -12,8 +26,8 @@ class CommonConfig:
         self.rank = get_real_step_or_rank(json_config.get('rank'), Const.RANK)
         self.step = get_real_step_or_rank(json_config.get('step'), Const.STEP)
         self.level = json_config.get('level')
-        self.acl_config = json_config.get('acl_config')
         self.enable_dataloader = json_config.get('enable_dataloader', False)
+        self.async_dump = json_config.get("async_dump", False)
         self._check_config()
 
     def _check_config(self):
@@ -29,16 +43,11 @@ class CommonConfig:
         if not isinstance(self.enable_dataloader, bool):
             logger.error_log_with_exp("enable_dataloader is invalid, it should be a boolean",
                                       MsprobeException(MsprobeException.INVALID_PARAM_ERROR))
-        if self.acl_config:
-            self._check_acl_config()
-
-    def _check_acl_config(self):
-        if not isinstance(self.acl_config, str):
-            logger.error_log_with_exp("acl_config is invalid, it should be a string",
+        if not isinstance(self.async_dump, bool):
+            logger.error_log_with_exp("async_dump is invalid, it should be a boolean",
                                       MsprobeException(MsprobeException.INVALID_PARAM_ERROR))
-        file_checker = FileChecker(
-            file_path=self.acl_config, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
-        file_checker.common_check()
+        elif self.async_dump:
+            logger.warning("async_dump is True, it may cause OOM when dumping large tensor.")
 
 
 class BaseConfig:
@@ -46,7 +55,6 @@ class BaseConfig:
         self.scope = json_config.get('scope')
         self.list = json_config.get('list')
         self.data_mode = json_config.get('data_mode')
-        self.backward_input = json_config.get("backward_input")
         self.file_format = json_config.get("file_format")
         self.summary_mode = json_config.get("summary_mode")
         self.overflow_nums = json_config.get("overflow_nums")
@@ -74,5 +82,32 @@ class BaseConfig:
     def check_config(self):
         self._check_str_list_config(self.scope, "scope")
         self._check_str_list_config(self.list, "list")
-        self._check_str_list_config(self.data_mode, "data_mode")
-        self._check_str_list_config(self.backward_input, "backward_input")
+        self._check_data_mode()
+
+    def _check_data_mode(self):
+        if self.data_mode is not None:
+            if not isinstance(self.data_mode, list):
+                logger.error_log_with_exp("data_mode is invalid, it should be a list[str]",
+                                          MsprobeException(MsprobeException.INVALID_PARAM_ERROR))
+
+            if Const.ALL in self.data_mode and len(self.data_mode) != 1:
+                logger.error_log_with_exp(
+                    "'all' cannot be combined with other options in data_mode.",
+                    MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+                )
+
+            if len(self.data_mode) >= len(Const.DUMP_DATA_MODE_LIST):
+                logger.error_log_with_exp(
+                    f"The number of elements in the data_made cannot exceed {len(Const.DUMP_DATA_MODE_LIST) - 1}.",
+                    MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+                )
+
+            for mode in self.data_mode:
+                if not isinstance(mode, str):
+                    logger.error_log_with_exp("data_mode is invalid, it should be a list[str]",
+                                              MsprobeException(MsprobeException.INVALID_PARAM_ERROR))
+                if mode not in Const.DUMP_DATA_MODE_LIST:
+                    logger.error_log_with_exp(
+                        f"The element '{mode}' of data_mode {self.data_mode} is not in {Const.DUMP_DATA_MODE_LIST}.",
+                        MsprobeException(MsprobeException.INVALID_PARAM_ERROR)
+                    )