mindstudio-probe 8.1.2__py3-none-any.whl → 8.2.1__py3-none-any.whl

msprobe/README.md

@@ -35,17 +35,17 @@ export MSPROBE_LOG_LEVEL={x}
 
 ## 环境和依赖
 
-- 硬件环境请参见《[昇腾产品形态说明](https://gitee.com/link?target=https%3A%2F%2Fwww.hiascend.com%2Fdocument%2Fdetail%2Fzh%2Fcanncommercial%2F80RC22%2Fquickstart%2Fquickstart%2Fquickstart_18_0002.html)》。
-- 软件环境请参见《[CANN 软件安装指南](https://gitee.com/link?target=https%3A%2F%2Fwww.hiascend.com%2Fdocument%2Fdetail%2Fzh%2Fcanncommercial%2F80RC22%2Fsoftwareinst%2Finstg%2Finstg_0000.html%3FMode%3DPmIns%26OS%3DUbuntu%26Software%3DcannToolKit)》安装昇腾设备开发或运行环境，即toolkit软件包。
+- 硬件环境请参见《[昇腾产品形态说明](https://www.hiascend.com/document/detail/zh/canncommercial/80RC22/quickstart/quickstart/quickstart_18_0002.html)》。
+- 软件环境请参见《[CANN 软件安装指南](https://www.hiascend.com/document/detail/zh/canncommercial/80RC22/softwareinst/instg/instg_0000.html?Mode=PmIns&OS=Ubuntu&Software=cannToolKit)》安装昇腾设备开发或运行环境，即toolkit软件包。
 
 以上环境依赖请根据实际环境选择适配的版本。
 
 ## 版本配套说明
 
-- msprobe支持AscendPyTorch 1.11.0或更高版本，支持的PyTorch和CANN以及PyTorch和python软件版本配套关系请参见《[Ascend Extension for PyTorch插件](https://gitee.com/ascend/pytorch)》。
+- msprobe支持AscendPyTorch 1.11.0或更高版本，支持的PyTorch和CANN以及PyTorch和python软件版本配套关系请参见《[Ascend Extension for PyTorch插件](https://gitcode.com/Ascend/pytorch)》。
 - msprobe支持MindSpore 2.4.0或更高版本，支持的MindSpore和CANN以及MindSpore和python软件版本配套关系请参见《[MindSpore版本发布列表](https://www.mindspore.cn/versions)》。
 - msprobe支持MSAdapter 2.1.0。
-- msprobe支持的固件驱动版本与配套CANN软件支持的固件驱动版本相同，开发者可通过“[昇腾社区-固件与驱动](https://gitee.com/link?target=https%3A%2F%2Fwww.hiascend.com%2Fhardware%2Ffirmware-drivers%2Fcommunity%3Fproduct%3D2%26model%3D28%26cann%3D8.0.RC3.alpha003%26driver%3D1.0.25.alpha)”页面根据产品型号与CANN软件版本获取配套的固件与驱动。
+- msprobe支持的固件驱动版本与配套CANN软件支持的固件驱动版本相同，开发者可通过“[昇腾社区-固件与驱动](https://www.hiascend.com/hardware/firmware-drivers/community?product=2&model=28&cann=8.0.RC3.alpha003&driver=1.0.25.alpha)”页面根据产品型号与CANN软件版本获取配套的固件与驱动。
 
 
 ## 🚨 工具限制与注意事项
@@ -84,7 +84,7 @@ msprobe 通过在训练脚本中添加 PrecisionDebugger 接口的方式对 API
 
 精度预检旨在昇腾 NPU 上扫描训练模型中的所有 API 进行 API 复现，给出精度情况的诊断和分析。对应 config.json 中的 "run_ut" task。
 
-PyTorch 场景的[离线预检](./docs/07.accuracy_checker_PyTorch.md)和[在线预检](./docs/08.accuracy_checker_online_PyTorch.md)
+PyTorch 场景的[离线预检](./docs/07.accuracy_checker_PyTorch.md)
 
 MindSpore 动态图场景的[离线预检](./docs/09.accuracy_checker_MindSpore.md)
 
@@ -165,7 +165,7 @@ MindSpore 动态图场景的[离线预检](./docs/09.accuracy_checker_MindSpore.
 
 训练前或精度比对前，对比两个环境下可能影响训练精度的配置差异。
 
-[PyTorch 训练前配置检查](./docs/31.config_check.md)
+[训练前配置检查](./docs/31.config_check.md)
 
 训练过程中或结束后，比较两个不同的checkpoint，评估模型相似度。
 

msprobe/core/common/const.py

@@ -24,6 +24,8 @@ class Const:
     Class for const
     """
     TOOL_NAME = "msprobe"
+    MD5_INDEX = "md5_index"
+    MD5 = "md5"
 
     ipv4_pattern = "([1-9]?\d|1\d{2}|2[0-4]\d|25[0-5])(\.([1-9]?\d|1\d{2}|2[0-4]\d|25[0-5])){3}$"
     SEP = "."
@@ -52,9 +54,9 @@ class Const:
     SIX_SEGMENT = 6
     SEVEN_SEGMENT = 7
 
-    MAX_DEPTH = 10
+    MAX_DEPTH = 400
     CPU_QUARTER = 4
-    DUMP_MAX_DEPTH = 50
+    DUMP_MAX_DEPTH = 400
 
     EXTERN_INPUT_LIST_MAX_LEN = 100
     MAX_PROCESS_NUM = 128
@@ -72,6 +74,7 @@ class Const:
     ONLINE_DUMP_MODE = [ALL, LIST, AUTO, OFF]
     SUMMARY = "summary"
     MD5 = "md5"
+    HASH = "hash"
     VALUE = "value"
     SUMMARY_MODE = ["statistics", "md5"]
 
@@ -113,9 +116,13 @@ class Const:
     RUN_UT = "run_ut"
     GRAD_PROBE = "grad_probe"
     STRUCTURE = "structure"
-    TASK_LIST = [TENSOR, STATISTICS, OVERFLOW_CHECK, FREE_BENCHMARK, RUN_UT, GRAD_PROBE, STRUCTURE]
+    EXCEPTION_DUMP = "exception_dump"
+    DUMP_PRECISION_HIGH = "high"
+    DUMP_PRECISION_LOW = "low"
+    TASK_LIST = [TENSOR, STATISTICS, OVERFLOW_CHECK, FREE_BENCHMARK, RUN_UT, GRAD_PROBE, STRUCTURE, EXCEPTION_DUMP]
     DUMP_DATA_COLLECTION_LIST = [STATISTICS, TENSOR, STRUCTURE]
     DUMP_DATA_MODE_LIST = [ALL, INPUT, OUTPUT, FORWARD, BACKWARD]
+    DUMP_PRECISION_LIST = [DUMP_PRECISION_LOW, DUMP_PRECISION_HIGH]
     LEVEL_L0 = "L0"
     LEVEL_L1 = "L1"
     LEVEL_L2 = "L2"
@@ -237,7 +244,11 @@ class Const:
     MEAN = 'Mean'
     NORM = 'Norm'
     DATA_NAME = 'data_name'
+    STATE = 'state'
+    REQ_GRAD = 'requires_grad'
+    API_ORIGIN_NAME = 'api_origin_name'
     TENSOR_STAT_INDEX = 'tensor_stat_index'
+    SUMMARY_METRICS_LIST = [MAX, MIN, MEAN, NORM]
 
     CODE_STACK = 'Code Stack'
     OP_NAME = 'Op Name'
@@ -260,8 +271,15 @@ class Const:
 
     TENSOR_STAT_LEN = 2
 
+    TENSOR_TYPE = "torch.Tensor"
+    DTENSOR_TYPE = "torch.distributed.tensor.DTensor"
+    FAKE_TENSOR_TYPE = "torch._subclasses.fake_tensor.FakeTensor"
+    AC_TENSOR_TYPE = "torch.distributed._functional_collectives.AsyncCollectiveTensor"
+
     SUPPORT_API_FILE_NAME = "support_wrap_ops.yaml"
 
+    API_ATTR_LIST = ["__name__", "default"]
+
     PT_API_TYPE_FUNCTIONAL = "functional"
     PT_API_TYPE_TENSOR = "tensor"
     PT_API_TYPE_TORCH = "torch"
@@ -355,22 +373,22 @@ class Const:
     }
 
     def _fused_adamw_(
-        self,
-        grads,
-        exp_avgs,
-        exp_avg_sqs,
-        max_exp_avg_sqs,
-        state_steps,
-        *,
-        lr,
-        beta1,
-        beta2,
-        weight_decay,
-        eps,
-        amsgrad,
-        maximize,
-        grad_scale=None,
-        found_inf=None
+            self,
+            grads,
+            exp_avgs,
+            exp_avg_sqs,
+            max_exp_avg_sqs,
+            state_steps,
+            *,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+            eps,
+            amsgrad,
+            maximize,
+            grad_scale=None,
+            found_inf=None
     ):
         pass
 
@@ -382,6 +400,13 @@ class Const:
     MATCH_MODE_NAME = "pure name"
     MATCH_MODE_MAPPING = "mapping"
     MATCH_MODE_SIMILARITY = "similarity"
+    CONFIG_CHECK_PASS = "pass"
+    CONFIG_CHECK_WARNING = "warning"
+    CONFIG_CHECK_ERROR = "error"
+
+    MIX_DUMP_NAMES = {'graph', 'pynative'}
+
+    MEGATRON_MICRO_STEP_NUMBER = 'megatron_micro_step_number'
 
 
 class CompareConst:
@@ -397,10 +422,14 @@ class CompareConst:
     BENCH_DTYPE = "Bench Dtype"
     NPU_SHAPE = "NPU Tensor Shape"
     BENCH_SHAPE = "Bench Tensor Shape"
+    NPU_CSV_FILE = "NPU CSV File"
+    BENCH_CSV_FILE = "Bench CSV File"
     NPU_MAX = "NPU max"
     NPU_MIN = "NPU min"
     NPU_MEAN = "NPU mean"
     NPU_NORM = "NPU l2norm"
+    NPU_P2POP_PEER = "NPU P2POp peer"
+
     BENCH_MAX = "Bench max"
     BENCH_MIN = "Bench min"
     BENCH_MEAN = "Bench mean"
@@ -416,6 +445,9 @@ class CompareConst:
     MIN_RELATIVE_ERR = "MinRelativeErr"
     MEAN_RELATIVE_ERR = "MeanRelativeErr"
     NORM_RELATIVE_ERR = "NormRelativeErr"
+    REQ_GRAD_CONSIST = "Requires_grad Consistent"
+    NPU_REQ_GRAD = "NPU Requires_grad"
+    BENCH_REQ_GRAD = "Bench Requires_grad"
     ACCURACY = "Accuracy Reached or Not"
     STACK = "NPU_Stack_Info"
     DATA_NAME = "Data_name"
@@ -437,7 +469,7 @@ class CompareConst:
     SUMMARY = "summary"
     COMPARE_RESULT = "compare_result"
     COMPARE_MESSAGE = "compare_message"
-    MAX_EXCEL_LENGTH = 1048576
+    MAX_EXCEL_LENGTH = 1048500
     YES = "Yes"
     NO = "No"
     STATISTICS_INDICATOR_NUM = 4
@@ -485,21 +517,21 @@ class CompareConst:
 
     ULP_ERR_STATUS = "ulp_err_status"
 
-    COMPARE_RESULT_HEADER = [
-        NPU_NAME, BENCH_NAME, NPU_DTYPE, BENCH_DTYPE, NPU_SHAPE, BENCH_SHAPE, COSINE, EUC_DIST,
-        MAX_ABS_ERR, MAX_RELATIVE_ERR, ONE_THOUSANDTH_ERR_RATIO, FIVE_THOUSANDTHS_ERR_RATIO,
-        NPU_MAX, NPU_MIN, NPU_MEAN, NPU_NORM, BENCH_MAX, BENCH_MIN, BENCH_MEAN, BENCH_NORM, ACCURACY, ERROR_MESSAGE
-    ]
+    ALL_COMPARE_INDEX = [COSINE, EUC_DIST, MAX_ABS_ERR, MAX_RELATIVE_ERR,
+                         ONE_THOUSANDTH_ERR_RATIO, FIVE_THOUSANDTHS_ERR_RATIO]
+    SUMMARY_COMPARE_INDEX = [MAX_DIFF, MIN_DIFF, MEAN_DIFF, NORM_DIFF,
+                             MAX_RELATIVE_ERR, MIN_RELATIVE_ERR, MEAN_RELATIVE_ERR, NORM_RELATIVE_ERR]
+    MD5_COMPARE_INDEX = [RESULT]
 
-    SUMMARY_COMPARE_RESULT_HEADER = [
-        NPU_NAME, BENCH_NAME, NPU_DTYPE, BENCH_DTYPE, NPU_SHAPE, BENCH_SHAPE, MAX_DIFF, MIN_DIFF, MEAN_DIFF, NORM_DIFF,
-        MAX_RELATIVE_ERR, MIN_RELATIVE_ERR, MEAN_RELATIVE_ERR, NORM_RELATIVE_ERR,
-        NPU_MAX, NPU_MIN, NPU_MEAN, NPU_NORM, BENCH_MAX, BENCH_MIN, BENCH_MEAN, BENCH_NORM, RESULT, ERROR_MESSAGE
-    ]
+    BASIC_INFO = [NPU_NAME, BENCH_NAME, NPU_DTYPE, BENCH_DTYPE, NPU_SHAPE, BENCH_SHAPE, NPU_REQ_GRAD, BENCH_REQ_GRAD]
+    SUMMARY_INFO = [NPU_MAX, NPU_MIN, NPU_MEAN, NPU_NORM, BENCH_MAX, BENCH_MIN, BENCH_MEAN, BENCH_NORM]
 
-    MD5_COMPARE_RESULT_HEADER = [
-        NPU_NAME, BENCH_NAME, NPU_DTYPE, BENCH_DTYPE, NPU_SHAPE, BENCH_SHAPE, NPU_MD5, BENCH_MD5, RESULT
-    ]
+    COMPARE_RESULT_HEADER = BASIC_INFO + ALL_COMPARE_INDEX + SUMMARY_INFO + [REQ_GRAD_CONSIST, ACCURACY, ERROR_MESSAGE]
+
+    SUMMARY_COMPARE_RESULT_HEADER = BASIC_INFO + SUMMARY_COMPARE_INDEX + SUMMARY_INFO + [REQ_GRAD_CONSIST, RESULT,
+                                                                                         ERROR_MESSAGE]
+
+    MD5_COMPARE_RESULT_HEADER = BASIC_INFO + [NPU_MD5, BENCH_MD5, REQ_GRAD_CONSIST] + MD5_COMPARE_INDEX
 
     COMPARE_RESULT_HEADER_STACK = COMPARE_RESULT_HEADER + [STACK]
 
@@ -513,11 +545,6 @@ class CompareConst:
         Const.MD5: MD5_COMPARE_RESULT_HEADER
     }
 
-    ALL_COMPARE_INDEX = [COSINE, EUC_DIST, MAX_ABS_ERR, MAX_RELATIVE_ERR, ONE_THOUSANDTH_ERR_RATIO,
-                         FIVE_THOUSANDTHS_ERR_RATIO]
-    SUMMARY_COMPARE_INDEX = [MAX_DIFF, MIN_DIFF, MEAN_DIFF, NORM_DIFF,
-                             MAX_RELATIVE_ERR, MIN_RELATIVE_ERR, MEAN_RELATIVE_ERR, NORM_RELATIVE_ERR]
-
     # dtype match
 
     DTYPE_MATCH_GROUPS = [
@@ -554,6 +581,8 @@ class CompareConst:
     ULP_FLOAT16_THRESHOLD = 1
 
     # compare result data
+    NO_REAL_DATA = 'No real data'
+    API_UNMATCH = 'api unmatched'
     READ_NONE = 'No data'
     NONE = 'None'
     SHAPE_UNMATCH = 'shape unmatched'
@@ -622,6 +651,9 @@ class CompareConst:
         MAX_DIFF: None, MIN_DIFF: None, MEAN_DIFF: None, NORM_DIFF: None, MAX_RELATIVE_ERR: None,
         MIN_RELATIVE_ERR: None, MEAN_RELATIVE_ERR: None, NORM_RELATIVE_ERR: None
     }
+    MS_GRAPH_CSV = {
+        NPU_CSV_FILE: None, BENCH_CSV_FILE: None
+    }
 
     API_MAPPING_KEYS_TO_COMPARE = [
         ('ms_args', 'pt_args'),
@@ -641,9 +673,11 @@ class CompareConst:
 
     OP_NAME_X = 'op_name_x'
     MATCH_RESULT_COLUMNS = [
-        OP_NAME_X, 'dtype_x', 'shape_x', 'summary_x', 'stack_info_x', 'data_name_x',
+        OP_NAME_X, 'dtype_x', 'shape_x', 'summary_x', 'stack_info_x', 'state_x', 'api_origin_name_x',
+        'requires_grad_x', 'data_name_x',
         CMP_KEY, CMP_SHAPE,
-        'op_name_y', 'dtype_y', 'shape_y', 'summary_y', 'stack_info_y', 'data_name_y',
+        'op_name_y', 'dtype_y', 'shape_y', 'summary_y', 'stack_info_y', 'state_y', 'api_origin_name_y',
+        'requires_grad_y', 'data_name_y'
     ]
 
     INTERNAL_API_MAPPING_FILE = 'ms_to_pt_api.yaml'
@@ -674,6 +708,8 @@ class FileCheckConst:
     IR_SUFFIX = ".ir"
     ZIP_SUFFIX = ".zip"
     SHELL_SUFFIX = ".sh"
+    LOG_SUFFIX = ".log"
+    DB_SUFFIX = '.db'
     MAX_PKL_SIZE = 1073741824  # 1 * 1024 * 1024 * 1024
     MAX_NUMPY_SIZE = 10737418240  # 10 * 1024 * 1024 * 1024
     MAX_JSON_SIZE = 1073741824  # 1 * 1024 * 1024 * 1024
@@ -686,6 +722,8 @@ class FileCheckConst:
     MAX_FILE_IN_ZIP_SIZE = 1073741824  # 1 * 1024 * 1024 * 1024
     MAX_FILE_SIZE = 1073741824  # 1 * 1024 * 1024 * 1024
     COMMOM_FILE_SIZE = 1048576  # 1 * 1024 * 1024
+    MAX_LOG_SIZE = 10737418240  # 1 * 1024 * 1024 * 1024
+    MAX_DB_SIZE = 10737418240  # 10 * 1024 * 1024 * 1024
     DIR = "dir"
     FILE = "file"
     DATA_DIR_AUTHORITY = 0o750
@@ -699,7 +737,9 @@ class FileCheckConst:
         XLSX_SUFFIX: MAX_XLSX_SIZE,
         YAML_SUFFIX: MAX_YAML_SIZE,
         IR_SUFFIX: MAX_IR_SIZE,
-        ZIP_SUFFIX: MAX_ZIP_SIZE
+        ZIP_SUFFIX: MAX_ZIP_SIZE,
+        LOG_SUFFIX: MAX_LOG_SIZE,
+        DB_SUFFIX: MAX_DB_SIZE
     }
     CSV_BLACK_LIST = r'^[＋－＝％＠\+\-=%@]|;[＋－＝％＠\+\-=%@]'
 
@@ -758,6 +798,11 @@ class MonitorConst:
     DEFAULT_STEP_INTERVAL = 1
 
     OP_LIST = ["norm", "min", "max", "zeros", "nans", "id", "mean", "shape", "dtype"]
+    OP_MONVIS_SUPPORTED = [
+        "norm", "min", "max", "zeros", "nans", "mean",
+        "entropy", "softmax_max", "sr", "kernel_norm", "std_x", "jacobian",
+        "proxy", "token_similarity"
+    ]
     MONITOR_OUTPUT_DIR = "MONITOR_OUTPUT_DIR"
     DEFAULT_MONITOR_OUTPUT_DIR = "./monitor_output"
     DATABASE = "database"
@@ -770,6 +815,8 @@ class MonitorConst:
     )
     DEEPSPEED_ZERO_OPT_FILTER = "DeepSpeedZeroOptimizer"
     RULE_NAME = ['AnomalyTurbulence', 'AnomalyNan']
+    L2_HOOKS = ["linear_hook", "attention_hook"]
+    SA_ORDERS = ["s,b,h,d", "b,s,h,d"]
 
     SLICE_SIZE = 20480
     # used for name
@@ -781,6 +828,7 @@ class MonitorConst:
     ACTV_OUT = "output"
     ACTVGRAD_IN = "input_grad"
     ACTVGRAD_OUT = "output_grad"
+    FSDP_FLAT_SEP = "_fsdp_wrapped_module."
     # used for tasks
     ACTV = "actv"
     ACTVGRAD = "actv_grad"
@@ -820,3 +868,12 @@ class MonitorConst:
         TRAIN_STAGE[key] = BACKWARD_STAGE
     for key in OPTIMIZER_KEY:
         TRAIN_STAGE[key] = OPTIMIZER_STAGE
+
+    # csv2db
+    DEFAULT_INT_VALUE = 0
+    MAX_PROCESS_NUM = 128
+    CSV_FILE_PATTERN = r"_(\d+)-(\d+)\.csv"
+    BATCH_SIZE = 10000
+    MAX_PARTITION = 10_000_000
+    MIN_PARTITION = 10
+    

msprobe/core/common/db_manager.py

@@ -0,0 +1,256 @@
+# 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.
+import re
+import sqlite3
+from typing import List, Tuple, Dict, Any
+from functools import wraps
+
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.file_utils import check_path_before_create, change_mode
+from msprobe.core.common.const import FileCheckConst
+
+SAFE_SQL_PATTERN = re.compile(r'^[a-zA-Z0-9_]+$')
+
+
+def check_identifier_safety(name):
+    """验证标识符是否安全（防止SQL注入）"""
+    if not isinstance(name, str) or SAFE_SQL_PATTERN.match(name) is None:
+        raise ValueError(f"Invalid SQL identifier: {name}, potential SQL injection risk!")
+
+
+def _db_operation(func):
+    """数据库操作装饰器，自动管理连接"""
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        conn, curs = None, None
+        try:
+            conn, curs = self._get_connection()
+            result = func(self, conn, curs, *args, **kwargs)
+            return result  # 显式返回正常结果
+            
+        except sqlite3.Error as err:
+            logger.error(f"Database operation failed: {err}")
+            if conn:
+                conn.rollback()
+            return None  # 显式返回错误情况下的None
+            
+        finally:
+            self._release_connection(conn, curs)
+    return wrapper
+
+
+class DBManager:
+    """
+    数据库管理类，封装常用数据库操作
+    """
+
+    DEFAULT_FETCH_SIZE = 10000
+    DEFAULT_INSERT_SIZE = 10000
+    MAX_ROW_COUNT = 100000000
+
+    def __init__(self, db_path: str):
+        """
+        初始化DBManager
+        :param db_path: 数据库文件路径
+        :param table_config: 表配置对象
+        """
+        self.db_path = db_path
+
+    @staticmethod
+    def _get_where_sql(where_list):
+        if not where_list:
+            return "", tuple()
+
+        where_clauses = []
+        where_values = []
+        if where_list:
+            for col, val in where_list.items():
+                check_identifier_safety(col)
+                where_clauses.append(f"{col} = ?")
+                where_values.append(val)
+            if where_clauses:
+                where_sql = " WHERE " + " AND ".join(where_clauses)
+        return where_sql, tuple(where_values)
+
+    @_db_operation
+    def insert_data(self, conn: sqlite3.Connection, curs: sqlite3.Cursor,
+                    table_name: str, data: List[Tuple], key_list: List[str] = None) -> int:
+        """
+        批量插入数据
+        :param table_name: 表名
+        :param data: 要插入的数据列表
+        :param batch_size: 每批插入的大小
+        :return: 插入的行数
+        """
+        check_identifier_safety(table_name)
+
+        if not data:
+            return 0
+        columns = len(data[0])
+        if key_list:
+            if not isinstance(key_list, list):
+                raise TypeError(
+                    f"key_list must be a list, got {type(key_list)}"
+                )
+            if columns != len(key_list):
+                raise ValueError(
+                    f"When inserting into table {table_name}, the length of key list ({key_list})"
+                    f"does not match the data({columns}).")
+            for key in key_list:
+                check_identifier_safety(key)
+
+        batch_size = self.DEFAULT_INSERT_SIZE
+        placeholders = ", ".join(["?"] * columns)
+        if key_list:
+            keys = ", ".join(key_list)
+            sql = f"INSERT OR IGNORE INTO {table_name} ({keys}) VALUES ({placeholders})"
+        else:
+            sql = f"INSERT OR IGNORE INTO {table_name} VALUES ({placeholders})"
+
+        inserted_rows = 0
+        for i in range(0, len(data), batch_size):
+            batch = data[i:i + batch_size]
+            curs.executemany(sql, batch)
+            inserted_rows += curs.rowcount
+
+        conn.commit()
+        return inserted_rows
+
+    @_db_operation
+    def select_data(self, conn: sqlite3.Connection, curs: sqlite3.Cursor,
+                    table_name: str,
+                    columns: List[str] = None,
+                    where: dict = None) -> List[Dict]:
+        """
+        查询数据
+        :param table_name: 表名
+        :param columns: 要查询的列
+        :param where: WHERE条件
+        :return: 查询结果列表(字典形式)
+        """
+        check_identifier_safety(table_name)
+
+        if not columns:
+            raise ValueError("columns parameter cannot be empty, specify columns to select (e.g. ['id', 'name'])")
+        if not isinstance(columns, list) or not all(isinstance(col, str) for col in columns):
+            raise TypeError("columns must be a list of strings (e.g. ['id', 'name'])")
+        
+        for col in columns:
+            check_identifier_safety(col)
+        
+        cols = ", ".join(columns)
+        sql = f"SELECT {cols} FROM {table_name}"
+
+        where_sql, where_parems = self._get_where_sql(where)
+        curs.execute(sql + where_sql, where_parems)
+
+        return [dict(row) for row in curs.fetchall()]
+
+    @_db_operation
+    def update_data(self, conn: sqlite3.Connection, curs: sqlite3.Cursor,
+                    table_name: str, updates: Dict[str, Any],
+                    where: dict = None) -> int:
+        """
+        更新数据
+        :param table_name: 表名
+        :param updates: 要更新的字段和值
+        :param where: WHERE条件
+        :param where_params: WHERE条件参数
+        :return: 影响的行数
+        """
+        check_identifier_safety(table_name)
+        if not updates:
+            raise ValueError("columns parameter cannot be empty, specify it to update (e.g. {'name': 'xxx'}")
+        if not isinstance(updates, dict):
+            raise TypeError(f"updates must be a dictionary, got: {type(updates)}")
+        for key in updates.keys():
+            check_identifier_safety(key)
+
+        set_clause = ", ".join([f"{k} = ?" for k in updates.keys()])
+        sql = f"UPDATE {table_name} SET {set_clause}"
+
+        params = tuple(updates.values())
+
+        where_sql, where_parems = self._get_where_sql(where)
+
+        curs.execute(sql + where_sql, params + where_parems)
+        conn.commit()
+        return curs.rowcount
+
+    @_db_operation
+    def execute_sql(self, conn: sqlite3.Connection, curs: sqlite3.Cursor,
+                    sql: str, params: Tuple = None) -> List[Dict]:
+        """
+        执行自定义SQL查询
+        :param sql: SQL语句
+        :param params: 参数
+        :return: 查询结果
+        """
+        curs.execute(sql, params or ())
+        if sql.strip().upper().startswith("SELECT"):
+            return [dict(row) for row in curs.fetchall()]
+        conn.commit()
+        return []
+
+    def table_exists(self, table_name: str) -> bool:
+        """
+        :param table_name: 表名
+        :return: 查询结果
+        """
+        result = self.select_data(
+            table_name="sqlite_master",
+            columns=["name"],
+            where={"type": "table", "name": table_name}
+        )
+        return len(result) > 0
+
+    @_db_operation
+    def execute_multi_sql(self, conn: sqlite3.Connection, curs: sqlite3.Cursor,
+                          sql_commands: List[str]) -> List[List[Dict]]:
+        """
+        批量执行多个SQL语句
+        :param sql_commands: [sql1, sql2, ...]
+        :return: 每个SELECT语句的结果列表
+        """
+        results = []
+        for sql in sql_commands:
+            curs.execute(sql)
+            if sql.strip().upper().startswith("SELECT"):
+                results.append([dict(row) for row in curs.fetchall()])
+        conn.commit()
+        return results
+
+    def _get_connection(self) -> Tuple[sqlite3.Connection, sqlite3.Cursor]:
+        """获取数据库连接和游标"""
+        check_path_before_create(self.db_path)
+        try:
+            conn = sqlite3.connect(self.db_path)
+            conn.row_factory = sqlite3.Row  # 使用Row工厂获取字典形式的结果
+            curs = conn.cursor()
+            return conn, curs
+        except sqlite3.Error as err:
+            logger.error(f"Database connection failed: {err}")
+            raise
+
+    def _release_connection(self, conn: sqlite3.Connection, curs: sqlite3.Cursor) -> None:
+        """释放数据库连接"""
+        try:
+            if curs is not None:
+                curs.close()
+            if conn is not None:
+                conn.close()
+        except sqlite3.Error as err:
+            logger.error(f"Failed to release database connection: {err}")
+        change_mode(self.db_path, FileCheckConst.DATA_FILE_AUTHORITY)

msprobe/core/common/file_utils.py

@@ -33,7 +33,7 @@ import pandas as pd
 from msprobe.core.common.decorator import recursion_depth_decorator
 from msprobe.core.common.log import logger
 from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.common.const import FileCheckConst, CompareConst
+from msprobe.core.common.const import FileCheckConst, CompareConst, Const
 from msprobe.core.common.global_lock import global_lock, is_main_process
 
 proc_lock = multiprocessing.Lock()
@@ -172,7 +172,7 @@ def check_path_exists(path):
     if not os.path.exists(path):
         logger.error('The file path %s does not exist.' % path)
         raise FileCheckException(FileCheckException.ILLEGAL_PATH_ERROR)
-    
+
 
 def check_path_not_exists(path):
     if os.path.exists(path):
@@ -259,8 +259,8 @@ def check_path_type(file_path, file_type):
 def check_others_writable(directory):
     dir_stat = os.stat(directory)
     is_writable = (
-        bool(dir_stat.st_mode & stat.S_IWGRP) or  # 组可写
-        bool(dir_stat.st_mode & stat.S_IWOTH)     # 其他用户可写
+            bool(dir_stat.st_mode & stat.S_IWGRP) or  # 组可写
+            bool(dir_stat.st_mode & stat.S_IWOTH)  # 其他用户可写
     )
     return is_writable
 
@@ -319,7 +319,7 @@ def check_dirpath_before_read(path):
             check_path_owner_consistent(dirpath)
         except FileCheckException:
             logger.warning(f"The directory {dirpath} is not yours.")
-    
+
 
 def check_file_or_directory_path(path, isdir=False):
     """
@@ -422,6 +422,26 @@ def load_json(json_path):
     return data
 
 
+def load_construct_json(json_path):
+    construct_dict_o = load_json(json_path)
+    if Const.MEGATRON_MICRO_STEP_NUMBER in construct_dict_o:
+        construct_dict = {}
+        micro_step_dict = {Const.MEGATRON_MICRO_STEP_NUMBER: construct_dict_o.get(Const.MEGATRON_MICRO_STEP_NUMBER)}
+        del construct_dict_o[Const.MEGATRON_MICRO_STEP_NUMBER]
+        for key, value in construct_dict_o.items():
+            if isinstance(value, list):
+                if len(value) != 2:
+                    logger.error(f'Parse construct json file "{os.path.basename(json_path)}" failed.')
+                    raise RuntimeError()
+                construct_dict[key] = value[0]
+                micro_step_dict[key] = value[1]
+            else:
+                construct_dict[key] = value
+                micro_step_dict[key] = 0
+        return construct_dict, micro_step_dict
+    return construct_dict_o, {}
+
+
 def save_json(json_path, data, indent=None, mode="w"):
     check_path_before_create(json_path)
     json_path = os.path.realpath(json_path)
@@ -520,6 +540,9 @@ def move_directory(src_path, dst_path):
     check_file_or_directory_path(src_path, isdir=True)
     check_path_before_create(dst_path)
     try:
+        if os.path.exists(dst_path):
+            logger.warning(f"The destination directory {dst_path} already exists, it will be removed.")
+            shutil.rmtree(dst_path)
         shutil.move(src_path, dst_path)
     except Exception as e:
         logger.error(f"move directory {src_path} to {dst_path} failed")

msprobe/core/common/log.py

@@ -89,6 +89,13 @@ class BaseLogger:
         self.error(msg)
         raise exception
 
+    def warning_log_with_exp(self, msg, exception):
+        """
+        打印警告日志并抛出指定异常
+        """
+        self.warning(msg)
+        raise exception
+
     def _print_log(self, level, msg, end='\n'):
         current_rank = self.get_rank()
         current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())