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

msprobe/visualization/db_utils.py

@@ -1,27 +1,29 @@
-# 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 os
 import sqlite3
 import json
 import re
 import time
+from typing import Dict
 from msprobe.core.common.log import logger
 from msprobe.core.common.file_utils import change_mode, check_path_before_create, FileChecker
-from msprobe.core.common.const import FileCheckConst
-from msprobe.visualization.utils import GraphConst
+from msprobe.core.common.const import FileCheckConst, Const
+from msprobe.visualization.utils import GraphConst, update_shared_dict, update_pbar_info, post_process_db_pbar
 from msprobe.visualization.builder.msprobe_adapter import format_node_data
 
 TEXT_PRIMARY_KEY = 'TEXT PRIMARY KEY'
@@ -29,6 +31,7 @@ TEXT_NOT_NULL = 'TEXT NOT NULL'
 INTEGER_NOT_NULL = 'INTEGER NOT NULL'
 TEXT = 'TEXT'
 INTEGER = 'INTEGER'
+DB_INSERT_SIZE = 1000
 
 node_columns = {
     'id': TEXT_PRIMARY_KEY,
@@ -51,7 +54,8 @@ node_columns = {
     'data_source': TEXT,
     'dump_data_dir': TEXT,
     'step': INTEGER_NOT_NULL,
-    'rank': INTEGER_NOT_NULL
+    'rank': INTEGER_NOT_NULL,
+    'is_distributed': INTEGER
 }
 
 config_columns = {
@@ -77,7 +81,8 @@ indexes = {
     "index3": ["step", "rank", "data_source", "node_order"],
     "index4": ["step", "rank", "node_order"],
     "index5": ["step", "rank", "micro_step_id", "node_order"],
-    "index6": ["step", "rank", "modified", "matched_node_link"]
+    "index6": ["step", "rank", "modified", "matched_node_link"],
+    "index7": ["is_distributed"]
 }
 
 SAFE_NAME_PATTERN = re.compile(r'^[a-zA-Z0-9_]+$')
@@ -133,7 +138,7 @@ def create_insert_sql_from_dict(table_name, columns_dict, ignore_insert=False):
     return sql
 
 
-def to_db(db_path, create_table_sql, insert_sql, data, db_insert_size=1000):
+def to_db(db_path, create_table_sql, insert_sql, data, pbar_info=None):
     max_retries = 10
     initial_delay = 0.1
     if not os.path.exists(db_path):
@@ -144,6 +149,11 @@ def to_db(db_path, create_table_sql, insert_sql, data, db_insert_size=1000):
 
     retry_count = 0
     current_delay = initial_delay
+    total_data = len(data)
+    commited_data = 0
+
+    if pbar_info:
+        update_shared_dict(pbar_info.current_stage_dict, pbar_info.task_id, 1)
 
     while retry_count <= max_retries:
         conn = None
@@ -154,9 +164,13 @@ def to_db(db_path, create_table_sql, insert_sql, data, db_insert_size=1000):
             cursor.execute("PRAGMA journal_mode=WAL")
             cursor.execute("PRAGMA synchronous=NORMAL")
             cursor.execute(create_table_sql)
-            for i in range(0, len(data), db_insert_size):
-                batch = data[i:i + db_insert_size]
+            for i in range(0, total_data, DB_INSERT_SIZE):
+                batch = data[i:i + DB_INSERT_SIZE]
                 cursor.executemany(insert_sql, batch)
+
+                if pbar_info:
+                    commited_data += len(batch)
+                    update_pbar_info(pbar_info, commited_data, total_data)
             conn.commit()
             return
         except sqlite3.OperationalError as e:
@@ -217,12 +231,16 @@ def add_table_index(db_path):
         conn.close()
 
 
-def post_process_db(db_path):
+def post_process_db(db_path, pbar_info=None, is_parallel_merge=False):
+    logger.info('Start adding index to db file, please wait...')
     add_table_index(db_path)
     change_mode(db_path, FileCheckConst.DATA_FILE_AUTHORITY)
+    if pbar_info:
+        post_process_db_pbar(pbar_info, is_parallel_merge)
+    logger.info('Adding index to db file completed.')
 
 
-def node_to_db(graph, db_name):
+def node_to_db(graph, db_name, pbar_info=None):
     create_table_sql = create_table_sql_from_dict('tb_nodes', node_columns)
     insert_sql = create_insert_sql_from_dict('tb_nodes', node_columns)
     data = []
@@ -241,8 +259,9 @@ def node_to_db(graph, db_name):
                      json.dumps(node.matched_distributed), 0,
                      json.dumps(format_node_data(node.input_data, node.id, graph.compare_mode)),
                      json.dumps(format_node_data(node.output_data, node.id, graph.compare_mode)),
-                     graph.data_source, graph.data_path, graph.step, graph.rank))
-    to_db(db_name, create_table_sql, insert_sql, data)
+                     graph.data_source, graph.data_path, graph.step, graph.rank,
+                     1 if node.id.startswith(Const.DISTRIBUTED) else 0))
+    to_db(db_name, create_table_sql, insert_sql, data, pbar_info=pbar_info)
     stack_to_db(stack_dict, db_name)
 
 
@@ -274,3 +293,95 @@ def get_node_unique_id(graph, node):
 
 def get_stack_unique_id(graph, stack_dict):
     return f'{get_graph_unique_id(graph)}_{len(stack_dict)}'
+
+
+class DBActuator:
+    def __init__(self, db_path):
+        self.db_path = db_path
+        self._init_connection()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.conn:
+            try:
+                self.conn.close()
+                logger.info(f"Database connection to {self.db_path} closed successfully")
+            except sqlite3.Error as e:
+                logger.error(f"Failed to close database connection: {e}")
+        return False
+
+    def query_distributed_nodes_info(self):
+        query = """
+ 	         SELECT
+ 	             node_name,
+ 	             input_data,
+ 	             rank,
+ 	             step,
+ 	             data_source,
+ 	             precision_index,
+ 	             overflow_level
+ 	         FROM
+ 	             tb_nodes
+ 	         WHERE
+ 	             is_distributed = 1
+ 	         """
+        if not self.conn:
+            logger.warning("Database connection is not initialized.")
+            return []
+        try:
+            with self.conn as conn:
+                cursor = conn.execute(query)
+                rows = cursor.fetchall()
+                cursor.close()
+
+            return rows
+        except Exception as e:
+            logger.error(f'Failed to query distributed nodes: {e}')
+            return []
+
+    def update_matched_distributed(self, distributed_info: Dict[int, dict], step=0,
+                                   data_source=GraphConst.JSON_NPU_KEY, batch_size=1000):
+        update_data_list = []
+        update_sql = """
+     	         UPDATE tb_nodes 
+     	         SET matched_distributed = ? 
+     	         WHERE step = ? AND rank = ? AND data_source = ? AND node_name = ?;
+
+     	         """
+        if not self.conn:
+            logger.warning("Database connection is not initialized.")
+            return
+        for rank, nodes_dict in distributed_info.items():
+            for node in nodes_dict.values():
+                if not node.matched_distributed:
+                    continue
+                try:
+
+                    matched_distributed = json.dumps(node.matched_distributed)
+                except Exception as e:
+                    logger.warning(f"Failed to serialize matched_distributed for node {node.id}: {e}")
+                    continue
+                update_data_list.append((matched_distributed, step, rank, data_source, node.id))
+
+        try:
+            with self.conn as conn:
+                cursor = conn.cursor()
+                for i in range(0, len(update_data_list), batch_size):
+                    batch_data = update_data_list[i:i + batch_size]
+                    cursor.executemany(update_sql, batch_data)
+                cursor.close()
+
+        except Exception as e:
+            logger.error(f"Failed to update matched_distributed: {e}")
+            raise
+
+    def _init_connection(self):
+        FileChecker(self.db_path, FileCheckConst.FILE, FileCheckConst.READ_WRITE_ABLE,
+                    FileCheckConst.DB_SUFFIX).common_check()
+        try:
+            self.conn = sqlite3.connect(self.db_path)
+        except sqlite3.Error as e:
+            logger.error(f"Failed to connect to database: {e}")
+            raise e

msprobe/visualization/graph/__init__.py

@@ -1,14 +1,15 @@
-# Copyright (c) 2024-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.
+# -------------------------------------------------------------------------

msprobe/visualization/graph/base_node.py

@@ -1,21 +1,21 @@
-# Copyright (c) 2024-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.
+# -------------------------------------------------------------------------
 
 from msprobe.core.overflow_check.level import OverflowLevel
 from msprobe.visualization.utils import GraphConst
-from msprobe.visualization.builder.msprobe_adapter import format_node_data, compare_data, compare_data_fuzzy
 from msprobe.core.common.log import logger
 
 
@@ -44,21 +44,9 @@ class BaseNode:
         return info
 
     def __eq__(self, other):
-        """
-        用来判断两个节点是否可以被匹配上，认为结构上是否一致
-        """
-        if not compare_data(self.input_data, other.input_data):
-            return False
-        if not compare_data(self.output_data, other.output_data):
-            return False
-        return True
-
-    def fuzzy_eq(self, other):
-        if not compare_data_fuzzy(self.input_data, other.input_data):
-            return False
-        if not compare_data_fuzzy(self.output_data, other.output_data):
-            return False
-        return True
+        if self.id == other.id:
+            return True
+        return False
 
     def set_input_output(self, input_data, output_data):
         self.input_data = input_data

msprobe/visualization/graph/distributed_analyzer.py

@@ -1,21 +1,27 @@
-# Copyright (c) 2024-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 json
 from enum import Enum
-from msprobe.visualization.utils import GraphConst
+from msprobe.visualization.utils import GraphConst, post_process_db_pbar
 from msprobe.core.common.const import Const, CompareConst
 from msprobe.core.common.log import logger
+from msprobe.visualization.db_utils import DBActuator
+from msprobe.visualization.graph.base_node import BaseNode
+from msprobe.visualization.graph.node_op import NodeOp
 
 
 class CommunicationType(Enum):
@@ -40,8 +46,7 @@ CANNOT_MATCH = 'cannot match distributed node in rank'
 
 class DistributedAnalyzer:
 
-    def __init__(self, graphs: dict, overflow_check: bool):
-        self.graphs = graphs
+    def __init__(self, distributed_info: dict, overflow_check: bool):
         self.overflow_check = overflow_check
         self.config = {
             # 当前通信api名称: 匹配目标通信api名称, 获取rank信息的位置参数或关键字参数, 通信类型, 分布式类型
@@ -55,6 +60,7 @@ class DistributedAnalyzer:
             'reduce': ['reduce', '1', CommunicationType.RECEIVE.value, DistributedType.COLLECTIVE]
         }
         self.group_node_mapping = {}
+        self.distributed_info = distributed_info
         self._make_group_node_mapping()
 
     @staticmethod
@@ -101,18 +107,30 @@ class DistributedAnalyzer:
         if not group_ranks:
             logger.debug(f'The group_ranks of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
             return None, None
+        group_ranks = DistributedAnalyzer._check_and_convert_group_ranks(group_ranks)
+        if not group_ranks:
+            return None, None
         group_id = group.get('group_id')
         if not group_id:
             logger.debug(f'The group_id of node {node.id} does not exist, {CANNOT_MATCH}{rank}')
             return None, None
         return group_ranks, group_id
 
+    @staticmethod
+    def _check_and_convert_group_ranks(group_ranks):
+        if not isinstance(group_ranks, list):
+            try:
+                group_ranks = json.loads(group_ranks)
+            except json.decoder.JSONDecodeError as e:
+                logger.debug(f'Error occurred, group_rank cannot be converted to list: {e}')
+                return []
+        return group_ranks
+
     def distributed_match(self):
-        for rank, graph in self.graphs.items():
-            nodes = graph.node_map
-            for node_id, node in nodes.items():
-                # 不是通信节点或者已经匹配过了
-                if not node_id.startswith(Const.DISTRIBUTED) or node.matched_distributed:
+        for rank, nodes_dict in self.distributed_info.items():
+            for node_id, node in nodes_dict.items():
+                # 已经匹配过了
+                if node.matched_distributed:
                     continue
                 api_name, distributed_type = self._get_distributed_name_and_type(node_id)
                 if api_name == GraphConst.BATCH_P2P:
@@ -137,14 +155,11 @@ class DistributedAnalyzer:
             "2": {}
         }
         """
-        for rank, graph in self.graphs.items():
+        for rank, nodes_dict in self.distributed_info.items():
             group_count = {}
             group_info = {}
             batch_p2p_count = {}
-            nodes = graph.node_map
-            for node_id, node in nodes.items():
-                if not node_id.startswith(Const.DISTRIBUTED):
-                    continue
+            for node_id, node in nodes_dict.items():
                 api_name, distributed_type = self._get_distributed_name_and_type(node_id)
                 if api_name == GraphConst.BATCH_P2P:
                     self._make_batch_p2p_mapping(node, rank, batch_p2p_count)
@@ -213,18 +228,18 @@ class DistributedAnalyzer:
         :param target_api_name: 与当前节点产生通信的节点api名称, 仅p2p通信需要配置
         :return: 目标节点
         """
-        target_graph = self.graphs.get(target_rank)
-        if not target_graph:
-            logger.debug(f'Graph data does not exist, {CANNOT_MATCH}{target_rank}')
+        target_nodes_dict = self.distributed_info.get(str(target_rank))
+        if not target_nodes_dict:
+            logger.debug(f'Node data does not exist, {CANNOT_MATCH}{target_rank}')
             return None
-        target_group_mapping = self.group_node_mapping.get(target_rank)
+        target_group_mapping = self.group_node_mapping.get(str(target_rank))
         # p2p通信，想要获取目标节点，需要替换unique_group_id中的rank和api name,
         # 例如isend发送到rank1，对应的irecv接收自rank0, isend_rank1与irecv_rank0对应
         target_unique_group_id = (unique_group_id
                                   .replace(Const.RANK + str(target_rank), Const.RANK + str(rank))
                                   .replace(api_name, target_api_name)) if target_api_name else unique_group_id
         target_node_id = target_group_mapping.get(target_unique_group_id, '')
-        target_node = target_graph.node_map.get(target_node_id)
+        target_node = target_nodes_dict.get(target_node_id)
         if not target_node:
             logger.debug(f'Node {target_node_id} does not exist, {CANNOT_MATCH}{target_rank}')
             return None
@@ -244,6 +259,7 @@ class DistributedAnalyzer:
             else communications_type
         index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
             else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+        index = CompareConst.NAN if index is None else index
         matched_distributed = {
             'communications_type': communications_type,
             'nodes_info': {target_rank: [str(index), target_node.id]}
@@ -264,11 +280,10 @@ class DistributedAnalyzer:
         """
         config_info = self.config.get(api_name)
         target_api_name = config_info[0]
-        #
         target_rank = self._get_target_rank(node, rank, config_info[1])
         if target_rank is None:
             return
-        unique_group_id = self.group_node_mapping.get(rank, {}).get(node.id, '')
+        unique_group_id = self.group_node_mapping.get(str(rank), {}).get(node.id, '')
         target_node = self._get_target_node(rank, unique_group_id, api_name, target_rank, target_api_name)
         if not target_node:
             return
@@ -288,13 +303,6 @@ class DistributedAnalyzer:
                 f'The rank is inconsistent, cannot match distributed node')
             return
 
-        # 点对点通信，两个匹配节点的输出数据要一致
-        if not DistributedAnalyzer._node_output_all_equal(node.output_data.get(node.id + '.output.0'),
-                                                          target_node.output_data.get(target_node.id + '.output.0')):
-            logger.debug(f'{node.id} output of rank{rank} is different from the {target_node.id} '
-                           f'output of rank{target_rank}, cannot match distributed node')
-            return
-
         self._add_node_matched_distributed(node, target_node, api_name, target_rank)
         self._add_node_matched_distributed(target_node, node, target_api_name, rank)
 
@@ -319,7 +327,10 @@ class DistributedAnalyzer:
         group_ranks, group_id = self._get_group_info(node, rank)
         if not group_ranks or not group_id:
             return
-        unique_group_id = self.group_node_mapping.get(rank, {}).get(node.id, '')
+        group_ranks = DistributedAnalyzer._check_and_convert_group_ranks(group_ranks)
+        if not group_ranks:
+            return
+        unique_group_id = self.group_node_mapping.get(str(rank), {}).get(node.id, '')
         matched_distributed = {'communications_type': communications_type}
         nodes_info = {}
         for target_rank in group_ranks:
@@ -340,6 +351,7 @@ class DistributedAnalyzer:
             # 给当前通信节点添加matched_distributed字段信息
             index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
                 else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+            index = CompareConst.NAN if index is None else index
             nodes_info[target_rank] = [str(index), target_node.id]
             if config_info:
                 # 给匹配上的目标节点也添加matched_distributed字段信息
@@ -357,7 +369,7 @@ class DistributedAnalyzer:
         :param rank: 当前节点所属rank
         :return:
         """
-        unique_group_ids = self.group_node_mapping.get(rank, {}).get(node.id)
+        unique_group_ids = self.group_node_mapping.get(str(rank), {}).get(node.id)
         if not unique_group_ids:
             return
         matched_distributed = [] if len(unique_group_ids) > 1 else {}
@@ -376,6 +388,7 @@ class DistributedAnalyzer:
             communications_type = self.config.get(api_name)[2]
             index = target_node.data.get(GraphConst.OVERFLOW_LEVEL, CompareConst.NAN) if self.overflow_check \
                 else target_node.data.get(GraphConst.JSON_INDEX_KEY, CompareConst.NAN)
+            index = CompareConst.NAN if index is None else index
             matched_info = {
                 'communications_type': communications_type,
                 'nodes_info': {target_rank: [str(index), target_node.id]}
@@ -384,3 +397,47 @@ class DistributedAnalyzer:
                 else matched_distributed.update(matched_info)
         if matched_distributed:
             node.matched_distributed = matched_distributed
+
+
+def distributed_analyse(db_path: str, overflow_check: bool, pbar_info=None):
+    distributed_info_data_source = {}
+    db_actuator = DBActuator(db_path)
+    rows = db_actuator.query_distributed_nodes_info()
+    if not rows:
+        if pbar_info:
+            post_process_db_pbar(pbar_info, True)
+        return
+    logger.info('Start analyzing distributed nodes...')
+    for row in rows:
+        if len(row) < 7:
+            continue
+        node_id = row[0]
+        input_data = row[1]
+        rank = row[2]
+        step = row[3]
+        data_source = row[4]
+        precision_index = row[5]
+        overflow_level = row[6]
+        node = BaseNode(NodeOp.function_api, node_id)
+        node.data[GraphConst.JSON_INDEX_KEY] = precision_index
+        node.data[GraphConst.OVERFLOW_LEVEL] = overflow_level
+        try:
+            node.input_data = json.loads(input_data)
+        except json.JSONDecodeError as e:
+            logger.debug(f"JSON parsing of input data for node {node_id} failed: {e}, data: {input_data}")
+            continue
+
+        # data_source -> step -> rank -> {node_id: node} 分类
+        distributed_info_data_source.setdefault(data_source, {}) \
+            .setdefault(step, {}) \
+            .setdefault(str(rank), {})[node_id] = node
+
+    for data_source, distributed_info_steps in distributed_info_data_source.items():
+        for step, distributed_info in distributed_info_steps.items():
+            analyzer = DistributedAnalyzer(distributed_info, overflow_check)
+            analyzer.distributed_match()
+            db_actuator.update_matched_distributed(analyzer.distributed_info, step, data_source)
+
+    if pbar_info:
+        post_process_db_pbar(pbar_info, True)
+    logger.info('Successfully analyzed distributed nodes.')

msprobe/visualization/graph/graph.py

@@ -1,17 +1,19 @@
-# Copyright (c) 2024-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.
+# -------------------------------------------------------------------------
+
 from msprobe.core.overflow_check.checker import AnomalyDetector
 from msprobe.visualization.graph.base_node import BaseNode
 from msprobe.visualization.graph.node_op import NodeOp
@@ -50,8 +52,6 @@ class Graph:
         if not node_n or node_n.id not in graph_b.node_map:
             return None, []
         node_b = graph_b.node_map.get(node_n.id)
-        if node_n != node_b:
-            return None, []
         ancestors_n = node_n.get_ancestors()
         ancestors_b = node_b.get_ancestors()
         if ancestors_n != ancestors_b:
@@ -71,11 +71,9 @@ class Graph:
         return node_b, ancestors_n, ancestors_b
 
     @staticmethod
-    def fuzzy_match(node_n, node_b, check_shape=True):
+    def fuzzy_match(node_n, node_b):
         if not node_n or not node_b:
             return None, [], []
-        if check_shape and not node_n.fuzzy_eq(node_b):
-            return None, [], []
         ancestors_n = node_n.get_ancestors()
         ancestors_b = node_b.get_ancestors()
         return node_b, ancestors_n, ancestors_b