mindstudio-probe 8.3.2__py3-none-any.whl → 26.0.0a1__py3-none-any.whl

msprobe/msprobe.py

@@ -1,163 +1,134 @@
-# Copyright (c) 2024, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import argparse
-import importlib.util
 import sys
 
-from msprobe.core.common.const import Const
-from msprobe.core.common.file_utils import root_privilege_warning
 from msprobe.core.common.log import logger
+from msprobe.core.compare.utils import _compare_parser
 from msprobe.core.compare.compare_cli import compare_cli
 from msprobe.core.compare.merge_result.merge_result_cli import _merge_result_parser, merge_result_cli
-from msprobe.core.compare.utils import _compare_parser
 from msprobe.core.config_check.config_check_cli import _config_checking_parser, _run_config_checking_command
-
-
-def is_module_available(module_name):
-    spec = importlib.util.find_spec(module_name)
-    return spec is not None
+from msprobe.overflow_check.analyzer import _overflow_check_parser, _run_overflow_check
+from msprobe.core.acc_check.acc_check_cli import acc_check_cli, multi_acc_check_cli
+from msprobe.visualization.graph_service import _graph_service_parser, _graph_service_command
+from msprobe.core.dump.dump2db.dump2db import _data2db_service_parser, _data2db_command
+from msprobe.infer.offline.compare.msquickcmp.main import _offline_dump_parser, offline_dump_cli
+from msprobe.core.install_deps.install_deps import _install_deps_parser, install_deps_cli
 
 
 def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
         description="msprobe(mindstudio probe), [Powered by MindStudio].\n"
-                    "Providing one-site accuracy difference debugging toolkit for training on Ascend Devices.\n"
+                    "A full-process, all-scenario precision tool base on Ascend products.\n"
                     f"For any issue, refer README.md first",
     )
 
     parser.set_defaults(print_help=parser.print_help)
-    parser.add_argument('-f', '--framework', required=True, choices=[Const.PT_FRAMEWORK, Const.MS_FRAMEWORK],
-                        help='Deep learning framework.')
     subparsers = parser.add_subparsers()
-    subparsers.add_parser('parse')
-    compare_cmd_parser = subparsers.add_parser('compare')
-    run_ut_cmd_parser = subparsers.add_parser('run_ut')
-    multi_run_ut_cmd_parser = subparsers.add_parser('multi_run_ut')
-    api_precision_compare_cmd_parser = subparsers.add_parser('api_precision_compare')
-    run_overflow_check_cmd_parser = subparsers.add_parser('run_overflow_check')
-    code_mapping_cmd_parser = subparsers.add_parser('code_mapping')
-    graph_service_cmd_parser = subparsers.add_parser('graph')
-    op_generate_cmd_parser = subparsers.add_parser('op_generate')
+
+    compare_parser = subparsers.add_parser('compare')
+    _compare_parser(compare_parser)
+    acc_check_cmd_parser = subparsers.add_parser('acc_check')
+    multi_acc_check_cmd_parser = subparsers.add_parser('multi_acc_check')
     merge_result_parser = subparsers.add_parser('merge_result')
+    _merge_result_parser(merge_result_parser)
+    overflow_check_parse = subparsers.add_parser('overflow_check')
+    _overflow_check_parser(overflow_check_parse)
     config_checking_parser = subparsers.add_parser('config_check')
-    nan_analyze_parser = subparsers.add_parser('nan_analyze')
     _config_checking_parser(config_checking_parser)
-    _compare_parser(compare_cmd_parser)
-    _merge_result_parser(merge_result_parser)
-
-    is_torch_available = is_module_available("torch")
-
-    if len(sys.argv) < 4:
+    # api_precision_compare 是 PyTorch 相关能力，延迟加载
+    api_precision_compare_cmd_parser = subparsers.add_parser('api_precision_compare')
+    try:
+        from msprobe.pytorch.api_accuracy_checker.compare.api_precision_compare import (
+            _api_precision_compare_parser
+        )
+        _api_precision_compare_parser(api_precision_compare_cmd_parser)
+    except ImportError as e:
+        # torch 不存在时，parser 仍然存在，但给提示
+        def _no_torch_parser(_):
+            api_precision_compare_cmd_parser.set_defaults(
+                func=lambda *_: (
+                    logger.error(
+                        "api_precision_compare requires PyTorch environment. "
+                        "Please install torch / torch_npu first."
+                    ),
+                    sys.exit(1)
+                )
+            )
+        _no_torch_parser(api_precision_compare_cmd_parser)
+    graph_service_cmd_parser = subparsers.add_parser('graph_visualize')
+    _graph_service_parser(graph_service_cmd_parser)
+    graph_service_cmd_parser_deprecated = subparsers.add_parser('graph')
+    _graph_service_parser(graph_service_cmd_parser_deprecated)
+    data2db_parser = subparsers.add_parser('data2db')
+    _data2db_service_parser(data2db_parser)
+
+    offline_dump_parser = subparsers.add_parser('offline_dump')
+    _offline_dump_parser(offline_dump_parser)
+
+    install_deps_cmd_parser = subparsers.add_parser('install_deps')
+    _install_deps_parser(install_deps_cmd_parser)
+
+    if len(sys.argv) >= 2 and sys.argv[1] == "acc_check":
+        acc_check_cli(sys.argv[2:])
+        return
+    elif len(sys.argv) >= 2 and sys.argv[1] == "multi_acc_check":
+        multi_acc_check_cli(sys.argv[2:])
+        return
+    elif len(sys.argv) < 2:
         parser.print_help()
         sys.exit(0)
 
-    root_privilege_warning()
-    framework_args = parser.parse_args(sys.argv[1:3])
-    if framework_args.framework == Const.PT_FRAMEWORK:
-        from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut import _run_ut_parser, run_ut_command
-        from msprobe.pytorch.parse_tool.cli import parse as cli_parse
-        from msprobe.pytorch.api_accuracy_checker.run_ut.multi_run_ut import prepare_config, run_parallel_ut
-        from msprobe.pytorch.api_accuracy_checker.compare.api_precision_compare import _api_precision_compare_parser, \
-            _api_precision_compare_command
-        from msprobe.pytorch.api_accuracy_checker.run_ut.run_overflow_check import _run_overflow_check_parser, \
-            _run_overflow_check_command
-        from msprobe.visualization.graph_service import _pt_graph_service_parser, _pt_graph_service_command
-        from msprobe.pytorch.api_accuracy_checker.generate_op_script.op_generator import _op_generator_parser, \
-            _run_operator_generate_commond
-        from msprobe.nan_analyze.analyzer import _nan_analyze_parser, _run_nan_analyze
-
-        _run_ut_parser(run_ut_cmd_parser)
-        _run_ut_parser(multi_run_ut_cmd_parser)
-        multi_run_ut_cmd_parser.add_argument('-n', '--num_splits', type=int, choices=range(1, 65), default=8,
-                                             help='Number of splits for parallel processing. Range: 1-64')
-        _api_precision_compare_parser(api_precision_compare_cmd_parser)
-        _run_overflow_check_parser(run_overflow_check_cmd_parser)
-        _pt_graph_service_parser(graph_service_cmd_parser)
-        _op_generator_parser(op_generate_cmd_parser)
-        _nan_analyze_parser(nan_analyze_parser)
-    elif framework_args.framework == Const.MS_FRAMEWORK:
-        from msprobe.mindspore.api_accuracy_checker.cmd_parser import add_api_accuracy_checker_argument
-        from msprobe.visualization.graph_service import _ms_graph_service_parser, _ms_graph_service_command
-        add_api_accuracy_checker_argument(run_ut_cmd_parser)
-        from msprobe.mindspore.api_accuracy_checker.cmd_parser import multi_add_api_accuracy_checker_argument
-        multi_add_api_accuracy_checker_argument(multi_run_ut_cmd_parser)
-        from msprobe.mindspore.code_mapping.cmd_parser import add_ir_parser_arguments
-        add_ir_parser_arguments(code_mapping_cmd_parser)
-
-        _ms_graph_service_parser(graph_service_cmd_parser)
-
-        from msprobe.mindspore.api_accuracy_checker.generate_op_script.op_generator import _op_generator_parser, \
-            _run_operator_generate_commond
-        _op_generator_parser(op_generate_cmd_parser)
-
     args = parser.parse_args(sys.argv[1:])
-    if sys.argv[2] == Const.PT_FRAMEWORK:
-        if not is_torch_available:
-            logger.error("PyTorch does not exist, please install PyTorch library")
-            raise Exception("PyTorch does not exist, please install PyTorch library")
-        if sys.argv[3] == "run_ut":
-            run_ut_command(args)
-        elif sys.argv[3] == "parse":
-            cli_parse()
-        elif sys.argv[3] == "multi_run_ut":
-            config = prepare_config(args)
-            run_parallel_ut(config)
-        elif sys.argv[3] == "api_precision_compare":
+    if sys.argv[1] == "compare":
+        compare_cli(args, sys.argv[1:])
+    elif sys.argv[1] == "merge_result":
+        merge_result_cli(args)
+    elif sys.argv[1] == "overflow_check":
+        _run_overflow_check(args)
+    elif sys.argv[1] == "graph_visualize":
+        _graph_service_command(args)
+    elif sys.argv[1] == "api_precision_compare":
+        try:
+            from msprobe.pytorch.api_accuracy_checker.compare.api_precision_compare import (
+                _api_precision_compare_command
+            )
             _api_precision_compare_command(args)
-        elif sys.argv[3] == "run_overflow_check":
-            _run_overflow_check_command(args)
-        elif sys.argv[3] == "graph":
-            _pt_graph_service_command(args)
-        elif sys.argv[3] == 'op_generate':
-            _run_operator_generate_commond(args)
-        elif sys.argv[3] == "compare":
-            if args.cell_mapping is not None or args.api_mapping is not None:
-                logger.error("Argument -cm or -am is not supported in PyTorch framework")
-                raise Exception("Argument -cm or -am is not supported in PyTorch framework")
-            compare_cli(args)
-        elif sys.argv[3] == "merge_result":
-            merge_result_cli(args)
-        elif sys.argv[3] == "config_check":
-            _run_config_checking_command(args)
-        elif sys.argv[3] == "nan_analyze":
-            _run_nan_analyze(args)
-    else:
-        if not is_module_available(Const.MS_FRAMEWORK):
-            logger.error("MindSpore does not exist, please install MindSpore library")
-            raise Exception("MindSpore does not exist, please install MindSpore library")
-        if sys.argv[3] == "compare":
-            compare_cli(args)
-        elif sys.argv[3] == "merge_result":
-            merge_result_cli(args)
-        elif sys.argv[3] == "run_ut":
-            from msprobe.mindspore.api_accuracy_checker.main import api_checker_main
-            api_checker_main(args)
-        elif sys.argv[3] == "multi_run_ut":
-            from msprobe.mindspore.api_accuracy_checker.main import mul_api_checker_main
-            mul_api_checker_main(args)
-        elif sys.argv[3] == "graph":
-            _ms_graph_service_command(args)
-        elif sys.argv[3] == 'op_generate':
-            _run_operator_generate_commond(args)
-        elif sys.argv[3] == "code_mapping":
-            from msprobe.mindspore.code_mapping.main import code_mapping_main
-            code_mapping_main(args)
-        elif sys.argv[3] == "config_check":
-            _run_config_checking_command(args)
+        except ImportError:
+            logger.error(
+                "api_precision_compare requires PyTorch environment. "
+                "Please install torch / torch_npu first."
+            )
+            sys.exit(1)
+
+    elif sys.argv[1] == "graph":
+        logger.warning('The "graph" parameter has been deprecated and will be removed in future versions. '
+                       'Please use the "graph_visualize" parameter instead.')
+        _graph_service_command(args)
+    elif sys.argv[1] == "config_check":
+        _run_config_checking_command(args)
+    elif sys.argv[1] == "data2db":
+        _data2db_command(args)
+
+    elif sys.argv[1] == "offline_dump":
+        offline_dump_cli(args)
+    elif sys.argv[1] == "install_deps":
+        install_deps_cli(args)
 
 
 if __name__ == "__main__":

msprobe/overflow_check/__init__.py

@@ -0,0 +1,15 @@
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
+#
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
+#
+#          http://license.coscl.org.cn/MulanPSL2
+#
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------

msprobe/{nan_analyze → overflow_check}/analyzer.py

@@ -1,37 +1,39 @@
-# Copyright (c) 2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import time
 from collections import defaultdict
 import os
 from itertools import dropwhile, chain
 
-from msprobe.core.common.file_utils import check_file_or_directory_path, save_json, make_dir
+from msprobe.core.common.file_utils import check_file_or_directory_path, save_json, make_dir, load_json
 from msprobe.core.common.log import logger
 from msprobe.core.common.const import Const
-from msprobe.nan_analyze.utils import (RankPath, FileCache, is_communication_op, is_ignore_op, NanAnalyseConst,
-                                       analyze_anomaly_in_group)
-from msprobe.nan_analyze.graph import DataNode, CommunicationNode
+from msprobe.overflow_check.utils import (RankPath, FileCache, is_communication_op, is_ignore_op, OverFlowCheckConst,
+                                          analyze_anomaly_in_group)
+from msprobe.overflow_check.graph import DataNode, CommunicationNode
 
 
-class NanAnalyzer:
+class OverFlowCheck:
     def __init__(self, input_path, output_path):
         self._input_path = input_path
         self._output_path = output_path
         self._paths = {}
         self._resolve_input_path()
+        self._check_framework()
         self._anomaly_nodes = []  # 记录所有异常节点
         self._cache = FileCache()
         self._first_comm_nodes = {}  # 记录各rank下首个通信节点的node_id
@@ -51,18 +53,27 @@ class NanAnalyzer:
         for path in contents:
             if not path.startswith('rank'):
                 continue
-            rank_str = path.strip('rank')
+            rank_str = path[len('rank'):]
             if not rank_str:
                 rank = 0
             elif not rank_str.isdigit():
                 continue
             else:
                 rank = int(rank_str)
-            dump_path = os.path.join(self._input_path, path, NanAnalyseConst.DUMP_FILE)
-            construct_path = os.path.join(self._input_path, path, NanAnalyseConst.CONSTRUCT_FILE)
-            stack_path = os.path.join(self._input_path, path, NanAnalyseConst.STACK_FILE)
+            dump_path = os.path.join(self._input_path, path, OverFlowCheckConst.DUMP_FILE)
+            construct_path = os.path.join(self._input_path, path, OverFlowCheckConst.CONSTRUCT_FILE)
+            stack_path = os.path.join(self._input_path, path, OverFlowCheckConst.STACK_FILE)
             self._paths[rank] = RankPath(rank, dump_path, construct_path, stack_path)
 
+    def _check_framework(self):
+        for _, rankpath in self._paths.items():
+            data = load_json(rankpath.dump_path)
+            if data.get('framework', '') != "pytorch":
+                logger.error('overflow check is only available for PyTorch, though the data is stored in '
+                             'a different framework.')
+                raise RuntimeError(f'invalid dump_path {rankpath.dump_path}')
+            break
+
     def _pre_analyze(self):
         logger.info('Start searching anomaly node before communication.')
         for path in self._paths.values():
@@ -148,9 +159,9 @@ class NanAnalyzer:
     def _find_connection(self, conn_info, cur_node, searched_ranks, seen_nodes):
         def connect():
             seen_nodes.add(search_node.node_id)
-            if search_node.type == NanAnalyseConst.DST:
+            if search_node.type == OverFlowCheckConst.DST:
                 cur_node.add_dst(search_node)
-            elif search_node.type == NanAnalyseConst.SRC:
+            elif search_node.type == OverFlowCheckConst.SRC:
                 search_node.layer = cur_node.layer
                 search_node.add_dst(cur_node)
             else:
@@ -167,7 +178,7 @@ class NanAnalyzer:
                 if not (search_id.startswith(tar_id_prefix) and search_node.type == conn_info.get('type')):
                     continue
                 search_conn_ranks = search_node.find_connected_nodes().get('ranks')
-                if ((not search_conn_ranks and search_node.api not in NanAnalyseConst.DIRECTED_API) or
+                if ((not search_conn_ranks and search_node.api not in OverFlowCheckConst.DIRECTED_API) or
                     cur_node.rank in search_conn_ranks):  # 有些无向通信算子没有填ProcessGroup，默认连接所有rank
                     connect()
                     found = True
@@ -241,15 +252,15 @@ class NanAnalyzer:
         return self._rank_comm_nodes_dict.get(rank, {}).get(node_id)
 
 
-def _nan_analyze_parser(parser):
+def _overflow_check_parser(parser):
     parser.add_argument("-i", "--input_path", dest="input_path", default="", type=str,
                         help="<Required> The dump file path, over step level. eg: \"xxx/step_0/\".",
                         required=True)
     parser.add_argument("-o", "--output_path", dest="output_path", default="./output", type=str,
-                        help="<optional> The nan inf analyze result output file path.",
+                        help="<optional> The overflow check result output file path.",
                         required=False)
 
 
-def _run_nan_analyze(args):
+def _run_overflow_check(args):
     check_file_or_directory_path(args.input_path, True)
-    NanAnalyzer(args.input_path, args.output_path).analyze()
+    OverFlowCheck(args.input_path, args.output_path).analyze()

msprobe/{nan_analyze → overflow_check}/graph.py

@@ -1,22 +1,23 @@
-# Copyright (c) 2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 from dataclasses import dataclass
 from msprobe.core.common.const import Const
 from msprobe.core.common.log import logger
-from msprobe.nan_analyze.utils import FileCache, RankPath, is_ignore_op, check_item_anomaly, NanAnalyseConst
+from msprobe.overflow_check.utils import FileCache, RankPath, is_ignore_op, check_item_anomaly, OverFlowCheckConst
 from msprobe.core.common.exceptions import MsprobeException
 
 
@@ -46,6 +47,8 @@ class DataNode:
         seen = set(op_name)
         while True:
             op_name = construct_info.get(op_name)
+            if isinstance(op_name, list):
+                op_name = op_name[0]
             if not op_name or op_name in seen:
                 return construct
             construct.insert(0, op_name)
@@ -140,7 +143,7 @@ class CommunicationNode:
 
     def has_nan_inf(self):
         return self.input_has_nan_inf() or check_item_anomaly(self.data.outputs)
-    
+
     def input_has_nan_inf(self):
         return check_item_anomaly(self.data.input_args) or check_item_anomaly(self.data.input_kwargs)
 
@@ -148,15 +151,15 @@ class CommunicationNode:
         """
         根据 api/类型/入参/调用次数 确定相连接的node的op_name
         """
-        tar_api = NanAnalyseConst.P2P_API_MAPPING.get(self.api, self.api)
+        tar_api = OverFlowCheckConst.P2P_API_MAPPING.get(self.api, self.api)
         ranks = set()
-        for dst in [NanAnalyseConst.DST, NanAnalyseConst.DST_GROUP]:
+        for dst in [OverFlowCheckConst.DST, OverFlowCheckConst.DST_GROUP]:
             if dst in self.data.input_kwargs:
                 dst_value = self.data.input_kwargs.get(dst)
                 if dst_value:
                     ranks.add(dst_value.get('value'))
                 break
-        for src in [NanAnalyseConst.SRC, NanAnalyseConst.SRC_GROUP]:
+        for src in [OverFlowCheckConst.SRC, OverFlowCheckConst.SRC_GROUP]:
             if src in self.data.input_kwargs:
                 src_value = self.data.input_kwargs.get(src)
                 if src_value:
@@ -170,24 +173,24 @@ class CommunicationNode:
         if group:
             ranks.update(group.get('group_ranks'))
         return {'ranks': ranks, 'api': f'Distributed.{tar_api}',
-                'type': NanAnalyseConst.OPPOSITE_DIR.get(self.type, NanAnalyseConst.LINK)}
+                'type': OverFlowCheckConst.OPPOSITE_DIR.get(self.type, OverFlowCheckConst.LINK)}
 
     def _resolve_type(self):
-        for src in [NanAnalyseConst.SRC, NanAnalyseConst.SRC_GROUP]:
+        for src in [OverFlowCheckConst.SRC, OverFlowCheckConst.SRC_GROUP]:
             if src in self.data.input_kwargs and self.data.input_kwargs[src]:
                 if self.data.input_kwargs[src].get('value') == self.rank:
-                    return NanAnalyseConst.SRC
+                    return OverFlowCheckConst.SRC
                 else:
-                    return NanAnalyseConst.DST
-        for dst in [NanAnalyseConst.DST, NanAnalyseConst.DST_GROUP]:
+                    return OverFlowCheckConst.DST
+        for dst in [OverFlowCheckConst.DST, OverFlowCheckConst.DST_GROUP]:
             if dst in self.data.input_kwargs and self.data.input_kwargs[dst]:
                 if self.data.input_kwargs[dst].get('value') == self.rank:
-                    return NanAnalyseConst.DST
+                    return OverFlowCheckConst.DST
                 else:
-                    return NanAnalyseConst.SRC
-        if self.api in NanAnalyseConst.DIRECTED_API:
+                    return OverFlowCheckConst.SRC
+        if self.api in OverFlowCheckConst.DIRECTED_API:
             for item in self.data.input_args:
                 if item.get(Const.TYPE) == 'int':
-                    node_type = NanAnalyseConst.DIRECTED_API[self.api]
-                    return node_type if item.get('value') == self.rank else NanAnalyseConst.OPPOSITE_DIR[node_type]
-        return NanAnalyseConst.LINK
+                    node_type = OverFlowCheckConst.DIRECTED_API[self.api]
+                    return node_type if item.get('value') == self.rank else OverFlowCheckConst.OPPOSITE_DIR[node_type]
+        return OverFlowCheckConst.LINK

msprobe/{nan_analyze → overflow_check}/utils.py

@@ -1,17 +1,18 @@
-# Copyright (c) 2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 from collections import OrderedDict
 from dataclasses import dataclass
@@ -120,7 +121,7 @@ def is_communication_op(op_name):
     # 定义通信算子的关键字，覆盖各种通信操作，如all_reduce, send, broadcast等
     # 从wrap文件中读取，先硬编码在文件中
     return (op_name.startswith('Distributed.') and
-            any(keyword in op_name for keyword in NanAnalyseConst.COMMUNICATION_KEYWORDS))
+            any(keyword in op_name for keyword in OverFlowCheckConst.COMMUNICATION_KEYWORDS))
 
 
 def is_ignore_op(op_name):
@@ -163,7 +164,7 @@ def analyze_anomaly_in_group(nodes_group):
             anomaly_nodes.append(node.data)
 
     # 先看src或link中input是否有异常
-    src_list = list(filter(lambda node: node.type in [NanAnalyseConst.SRC, NanAnalyseConst.LINK], nodes_group))
+    src_list = list(filter(lambda node: node.type in [OverFlowCheckConst.SRC, OverFlowCheckConst.LINK], nodes_group))
     input_anomaly_nodes = list(filter(lambda node: node.input_has_nan_inf(), src_list))
     # 如果有异常回溯计算节点找到异常来源
     # 使用cpu模拟节点进行计算，查看结果是否有问题。需要对所有计算节点录入/映射，暂不实现。
@@ -174,7 +175,7 @@ def analyze_anomaly_in_group(nodes_group):
     return anomaly_nodes
 
 
-class NanAnalyseConst:
+class OverFlowCheckConst:
     COMMUNICATION_KEYWORDS = {
         'send',  # send 算子
         'recv',  # recv 算子

msprobe/pytorch/__init__.py

@@ -1,24 +1,30 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# All rights reserved.
+# -------------------------------------------------------------------------
+#  This file is part of the MindStudio project.
+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
 #
-# 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
+# MindStudio is licensed under Mulan PSL v2.
+# You can use this software according to the terms and conditions of the Mulan PSL v2.
+# You may obtain a copy of Mulan PSL v2 at:
 #
-# http://www.apache.org/licenses/LICENSE-2.0
+#          http://license.coscl.org.cn/MulanPSL2
 #
-# 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.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
+# See the Mulan PSL v2 for more details.
+# -------------------------------------------------------------------------
 
 import torch
-from .compare.distributed_compare import compare_distributed
-from .compare.pt_compare import compare
+
+from msprobe.pytorch.dump.debugger.precision_debugger import PrecisionDebugger, module_dump, module_dump_end
 from .common.utils import seed_all
-from .debugger.precision_debugger import PrecisionDebugger, module_dump, module_dump_end
+from .torchair_dump import set_fx_dump_config, set_ge_dump_config
 
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
 if torch_version_above_or_equal_2:
     from msprobe.pytorch.monitor.module_hook import TrainerMon
+
+try:
+    from msprobe.pytorch.aclgraph_dump import acl_save  # noqa: F401
+except Exception:
+    acl_save = None