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

msprobe/msaccucmp/dump_parse/big_dump_data.py

@@ -0,0 +1,317 @@
+# coding=utf-8
+# -------------------------------------------------------------------------
+#  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.
+# -------------------------------------------------------------------------
+
+"""
+Function:
+This file mainly involves the dump function.
+"""
+import os
+import time
+import struct
+import warnings
+import ctypes
+import tempfile
+import json
+from typing.io import BinaryIO
+
+import numpy as np
+from google.protobuf.message import DecodeError
+
+from msprobe.msaccucmp.dump_parse.proto_dump_data import DumpData
+from msprobe.msaccucmp.cmp_utils import path_check
+from msprobe.msaccucmp.cmp_utils import log
+from msprobe.msaccucmp.cmp_utils import common
+from msprobe.msaccucmp.cmp_utils.constant.const_manager import ConstManager, DD
+from msprobe.msaccucmp.cmp_utils.constant.compare_error import CompareError
+
+
+class BigDumpDataParser:
+    """
+    The class for big dump data parser
+    """
+    warnings.filterwarnings("ignore")
+
+    def __init__(self: any, dump_file_path: str) -> None:
+        self.dump_file_path = dump_file_path
+        self.file_size = 0
+        self.header_length = 0
+        self.dump_data = DumpData()
+        self.dump_json_data = {}
+        self.data_types = ['input', 'output', 'buffer', 'space']
+        self.parse_dump_so = "libascend_dump_parser.so"
+
+    @staticmethod
+    def find_shared_library(lib_name: str, relative_path: str) -> str:
+        """
+        优先从LD_LIBRARY_PATH查找 so 文件，如果没有则从相对路径中查找。
+        """
+        # 从环境变量 LD_LIBRARY_PATH 中查找
+        ld_paths = os.environ.get("LD_LIBRARY_PATH", "").split(":")
+        for path in ld_paths:
+            candidate = os.path.join(path, lib_name)
+            if os.path.exists(candidate):
+                return os.path.realpath(candidate)
+
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+        candidate = os.path.realpath(os.path.join(current_dir, relative_path, lib_name))
+
+        if os.path.exists(candidate):
+            return candidate
+
+        raise OSError(
+            f"Shared library {lib_name} not found."
+        )
+
+    def parse(self: any):
+        """
+        Parse the dump file path by big dump data format
+        :return: DumpData
+        :exception when read or parse file error
+        """
+        self.check_argument_valid()
+        try:
+            with open(self.dump_file_path, 'rb') as dump_file:
+                # read header length
+                self._read_header_length(dump_file)
+                self._parse_dump_to_json()
+                self._parse_binary_to_json_data(dump_file)
+                return self.dump_data.from_dict(self.dump_json_data)
+        except (OSError, IOError) as io_error:
+            log.print_error_log('Failed to read the dump file %r. %s'
+                                % (self.dump_file_path, str(io_error)))
+            raise CompareError(CompareError.MSACCUCMP_OPEN_FILE_ERROR) from io_error
+        finally:
+            pass
+
+    def check_argument_valid(self: any) -> None:
+        """
+        check argument valid
+        :exception when invalid
+        """
+        ret = path_check.check_path_valid(self.dump_file_path, True)
+        if ret != CompareError.MSACCUCMP_NONE_ERROR:
+            raise CompareError(ret)
+        # get file size
+        try:
+            self.file_size = os.path.getsize(self.dump_file_path)
+        except (OSError, IOError) as error:
+            log.print_error_log('get the size of dump file %r failed. %s'
+                                % (self.dump_file_path, str(error)))
+            raise CompareError(CompareError.MSACCUCMP_DUMP_FILE_ERROR) from error
+        finally:
+            pass
+
+        if self.file_size <= ConstManager.UINT64_SIZE:
+            log.print_error_log(
+                'The size of %r must be greater than %d, but the file size'
+                ' is %d. Please check the dump file.'
+                % (self.dump_file_path, ConstManager.UINT64_SIZE, self.file_size))
+            raise CompareError(CompareError.MSACCUCMP_UNMATCH_STANDARD_DUMP_SIZE)
+        if self.file_size > ConstManager.ONE_GB:
+            log.print_warn_log(
+                'The size (%d) of %r exceeds 1GB, it may task more time to run, please wait.'
+                % (self.file_size, self.dump_file_path))
+
+    def _read_header_length(self: any, dump_file: BinaryIO) -> None:
+        # read header length
+        header_length = dump_file.read(ConstManager.UINT64_SIZE)
+        self.header_length = struct.unpack(ConstManager.UINT64_FMT, header_length)[0]
+        # check header_length <= file_size - 8
+        if self.header_length > self.file_size - ConstManager.UINT64_SIZE:
+            log.print_warn_log(
+                'The header content size (%d) of %r must be less than or'
+                ' equal to %d (file size) - %d (header length).'
+                ' Please check the dump file.'
+                % (self.header_length, self.dump_file_path, self.file_size, ConstManager.UINT64_SIZE))
+            raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR)
+        dump_file.read(self.header_length)
+
+    def _parse_dump_to_json(self):
+        # 读取 dump 二进制数据
+        with open(self.dump_file_path, 'rb') as dump_file:
+            binary_data = dump_file.read()
+        # 加载 C 解析库
+        try:
+            self.parse_dump_so = self.find_shared_library(self.parse_dump_so, "../../../adump/lib64")
+            ret = path_check.check_path_valid(self.parse_dump_so, True, False)
+            if ret != CompareError.MSACCUCMP_NONE_ERROR:
+                raise CompareError(ret)
+            if path_check.is_group_and_others_writable(self.parse_dump_so):
+                log.print_error_log(f"Failed to load {self.parse_dump_so}, this file is not safe. Please check.")
+                raise CompareError(CompareError.MSACCUCMP_INVALID_FILE_ERROR)
+            dump_parse_cdll = ctypes.CDLL(self.parse_dump_so)
+        except (OSError, IOError) as e:
+            log.print_error_log(f"Failed to load {self.parse_dump_so}:{e}")
+            raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR) from e
+        data_ptr = ctypes.c_char_p(binary_data)
+
+        # 使用临时文件替代落盘 JSON
+        fd, tmp_json_path = tempfile.mkstemp(suffix=".json")
+        # 关闭 fd，让 C 库写文件
+        os.close(fd)
+        try:
+            res = dump_parse_cdll.ParseDumpProtoToJson(
+                data_ptr, ctypes.c_size_t(len(binary_data)), tmp_json_path.encode('utf-8'))
+
+            if res != 0 or not os.path.isfile(tmp_json_path):
+                log.print_error_log(f"Parse dump file to json failed.")
+                raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR)
+            # 从临时文件加载 JSON
+            with open(tmp_json_path, 'r') as load_f:
+                self.dump_json_data = json.load(load_f)
+        finally:
+            # 删除临时文件，避免落盘
+            if os.path.exists(tmp_json_path):
+                os.remove(tmp_json_path)
+
+    def _parse_binary_to_json_data(self, dump_file: BinaryIO):
+        used_size = self.header_length + ConstManager.UINT64_SIZE
+        for data_type in self.data_types:
+            for item in self.dump_json_data.get(data_type, []):
+                size = int(item.get('size', 0))
+                used_size += size
+                if used_size > self.file_size:
+                    log.print_error_log(f'The size of {self.dump_file_path} is invalid, please check the dump file.')
+                    raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR)
+                item['data'] = dump_file.read(size)
+
+
+class DumpDataHandler:
+    """
+    Handle dump data
+    """
+
+    def __init__(self: any, dump_file_path: str) -> None:
+        self.dump_file_path = dump_file_path
+        self.file_size = 0
+
+    def check_argument_valid(self: any) -> None:
+        """
+        check argument valid
+        :exception when invalid
+        """
+        ret = path_check.check_path_valid(self.dump_file_path, True)
+        if ret != CompareError.MSACCUCMP_NONE_ERROR:
+            raise CompareError(ret)
+        # get file size
+        try:
+            self.file_size = os.path.getsize(self.dump_file_path)
+        except (OSError, IOError) as error:
+            log.print_error_log('get the size of dump file %r failed. %s'
+                                % (self.dump_file_path, str(error)))
+            raise CompareError(CompareError.MSACCUCMP_DUMP_FILE_ERROR) from error
+        finally:
+            pass
+        if self.file_size == 0:
+            message = 'Failed to parse dump file %r. The file size is zero. Please check the dump file.' \
+                      % self.dump_file_path
+            log.print_error_log(message)
+            raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR, message)
+        if self.file_size > ConstManager.ONE_GB:
+            log.print_warn_log(
+                'The size (%d) of %r exceeds 1GB, it may task more time to run, please wait.'
+                % (self.file_size, self.dump_file_path))
+
+    def read_numpy_file(self: any) -> np.ndarray:
+        """
+        Read numpy file
+        :return: numpy data
+        """
+        self.check_argument_valid()
+        try:
+            if self.dump_file_path.endswith(".txt"):
+                numpy_data = np.loadtxt(self.dump_file_path)
+            else:
+                numpy_data = np.load(self.dump_file_path)
+        except (OSError, SystemError, ValueError, TypeError, RuntimeError, MemoryError, DecodeError) as error:
+            log.print_error_log('Failed to parse dump file "%r". Only data of the numpy format is supported. %s'
+                                % (self.dump_file_path, str(error)))
+            raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR) from error
+        finally:
+            pass
+        return numpy_data
+
+    def parse_dump_data(self: any):
+        """
+        Parse dump file
+        :param dump_version: the dump version
+        :return: DumpData
+        """
+        self.check_argument_valid()
+        if self.dump_file_path.endswith('.npy'):
+            numpy_data = self.read_numpy_file()
+            dump_data, _ = _convert_numpy_to_dump(numpy_data, only_header=False)
+            return dump_data
+        try:
+            with open(self.dump_file_path, 'rb') as dump_file:
+                file_content = dump_file.read()
+        except (OSError, SystemError, ValueError, TypeError, RuntimeError, MemoryError) as error:
+            message = 'Failed to open dump file %r. Please check the dump file. %s' \
+                      % (self.dump_file_path, str(error))
+            log.print_error_log(message)
+            raise CompareError(CompareError.MSACCUCMP_INVALID_DUMP_DATA_ERROR, message) from error
+        finally:
+            pass
+        return BigDumpDataParser(self.dump_file_path).parse()
+
+
+def _convert_numpy_to_dump(numpy_data: np.ndarray, only_header: bool):
+    dump_data = DumpData()
+    dump_data.version = '2.0'
+    dump_data.dump_time = int(round(time.time() * ConstManager.TIME_LENGTH))
+    output = dump_data.output.add()
+    output.data_type = common.get_data_type_by_dtype(numpy_data.dtype)
+    output.format = DD.FORMAT_RESERVED
+    total_size = 1
+    for dim in numpy_data.shape:
+        output.shape.dim.append(dim)
+        total_size *= dim
+    # make output data
+    struct_format = common.get_struct_format_by_data_type(output.data_type)
+    data = struct.pack('%d%s' % (total_size, struct_format), *numpy_data.flatten())
+    output.size = len(output.data)
+    if not only_header:
+        output.data = data
+    return dump_data, data
+
+
+def write_dump_data(numpy_data: np.ndarray, output_dump_path: str) -> None:
+    """
+    write numpy data to dump data file
+    :param numpy_data: the numpy data
+    :param output_dump_path: the output dump file path
+    :exception when write file error
+    """
+    # make content of dump data header
+    dump_data, data = _convert_numpy_to_dump(numpy_data, only_header=True)
+    dump_data_ser = dump_data.SerializeToString()
+    try:
+        path_check.check_write_path_secure(output_dump_path)
+        with os.fdopen(os.open(output_dump_path, ConstManager.WRITE_FLAGS,
+                               ConstManager.WRITE_MODES), 'wb') as dump_file:
+            # write the header length
+            dump_file.write(struct.pack(ConstManager.UINT64_FMT, len(dump_data_ser)))
+            # write the header content
+            dump_file.write(dump_data_ser)
+            # write output data
+            dump_file.write(data)
+    except IOError as io_error:
+        log.print_error_log('Failed to write dump file %r. %s'
+                            % (output_dump_path, str(io_error)))
+        raise CompareError(CompareError.MSACCUCMP_WRITE_FILE_ERROR) from io_error
+    finally:
+        pass