mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/pytorch/api_accuracy_checker/run_ut/run_ut.py

@@ -45,11 +45,11 @@ from msprobe.pytorch.api_accuracy_checker.compare.compare_column import CompareC
 from msprobe.pytorch.api_accuracy_checker.common.config import CheckerConfig
 from msprobe.pytorch.common.parse_json import parse_json_info_forward_backward
 from msprobe.core.common.file_utils import FileChecker, change_mode, \
-    create_directory, get_json_contents, read_csv, check_file_or_directory_path, check_crt_valid
+    create_directory, get_json_contents, read_csv, check_file_or_directory_path
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.pt_config import parse_json_config
 from msprobe.core.common.const import Const, FileCheckConst, CompareConst
-from msprobe.core.common.utils import safe_get_value, CompareException
+from msprobe.core.common.utils import safe_get_value, CompareException, is_int, check_op_str_pattern_valid
 from msprobe.pytorch.common.utils import seed_all
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import ATTL, ATTLConfig, move2device_exec
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.device_dispatch import ConsumerDispatcher
@@ -65,6 +65,8 @@ DETAILS_FILE_NAME = "accuracy_checking_details_" + current_time + ".csv"
 
 not_backward_list = ['repeat_interleave']
 unsupported_backward_list = ['masked_select']
+unsupported_api_list = ["to", "empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided", 
+                        "empty_with_format"]
 
 
 tqdm_params = {
@@ -83,6 +85,9 @@ tqdm_params = {
 }
 
 
+seed_all()
+
+
 def run_ut(config):
     logger.info("start UT test")
     if config.online_config.is_online:
@@ -93,7 +98,7 @@ def run_ut(config):
         logger.info(f"UT task details will be saved in {config.details_csv_path}")
 
     if config.save_error_data:
-        logger.info(f"UT task error_datas will be saved in {config.error_data_path}")
+        logger.info(f"UT task error_data will be saved in {config.error_data_path}")
     compare = Comparator(config.result_csv_path, config.details_csv_path, config.is_continue_run_ut, config=config)
 
     if config.online_config.is_online:
@@ -117,6 +122,7 @@ def run_ut(config):
 def run_api_offline(config, compare, api_name_set):
     err_column = CompareColumn()
     for _, (api_full_name, api_info_dict) in enumerate(tqdm(config.forward_content.items(), **tqdm_params)):
+        check_op_str_pattern_valid(api_full_name)
         if api_full_name in api_name_set:
             continue
         if is_unsupported_api(api_full_name):
@@ -218,6 +224,7 @@ def blacklist_and_whitelist_filter(api_name, black_list, white_list):
     If api is both in black_list and black_list, black_list first.
     return: False for exec api, True for not exec
     """
+    black_list.extend(unsupported_api_list)
     if black_list and api_name in black_list:
         return True
     if white_list and api_name not in white_list:
@@ -317,7 +324,8 @@ def run_torch_api_online(api_full_name, api_data, backward_content):
     if kwargs.get("device"):
         del kwargs["device"]
 
-    device_out = exec_api(api_type, api_name, Const.CUDA_LOWERCASE, args, kwargs)
+    device_exec_params = ExecParams(api_type, api_name, current_device, args, kwargs, False, None)
+    device_out = exec_api(device_exec_params)
     device_out = move2device_exec(device_out, "cpu")
     return UtDataInfo(None, None, out, device_out, None, in_fwd_data_list, None, rank=api_data.rank)
 
@@ -344,6 +352,9 @@ def need_to_backward(grad_index, out):
 
 def run_backward(args, grad, grad_index, out):
     if grad_index is not None:
+        if not is_int(grad_index):
+            logger.error(f"{grad_index} dtype is not int")
+            raise TypeError(f"{grad_index} dtype is not int")
         if grad_index >= len(out):
             logger.error(f"Run backward error when grad_index is {grad_index}")
             raise IndexError(f"Run backward error when grad_index is {grad_index}")
@@ -430,6 +441,7 @@ def preprocess_forward_content(forward_content):
     arg_cache = {}
 
     for key, value in forward_content.items():
+        check_op_str_pattern_valid(key)
         base_key = key.rsplit(Const.SEP, 1)[0]
 
         if key not in arg_cache:
@@ -469,7 +481,7 @@ def _run_ut(parser=None):
     _run_ut_parser(parser)
     args = parser.parse_args(sys.argv[1:])
     run_ut_command(args)
-
+    
 
 def checked_online_config(online_config):
     if not online_config.is_online:
@@ -491,7 +503,10 @@ def checked_online_config(online_config):
         check_file_or_directory_path(online_config.tls_path, isdir=True)
         check_file_or_directory_path(os.path.join(online_config.tls_path, "server.key"))
         check_file_or_directory_path(os.path.join(online_config.tls_path, "server.crt"))
-        check_crt_valid(os.path.join(online_config.tls_path, "server.crt"))
+        check_file_or_directory_path(os.path.join(online_config.tls_path, "ca.crt"))
+        crl_path = os.path.join(online_config.tls_path, "crl.pem")
+        if os.path.exists(crl_path):
+            check_file_or_directory_path(crl_path)
 
     # host and port
     if not isinstance(online_config.host, str) or not re.match(Const.ipv4_pattern, online_config.host):
@@ -561,7 +576,15 @@ def run_ut_command(args):
     error_data_path = checker_config.error_data_path
     if save_error_data:
         if args.result_csv_path:
-            time_info = result_csv_path.split('.')[0].split('_')[-1]
+            parts_by_dot = result_csv_path.split(Const.SEP)
+            if len(parts_by_dot) < 2 or not parts_by_dot[0]:
+                raise ValueError("result_csv_path does not contain a valid file name with an extension.")
+            file_name_part = parts_by_dot[0]
+            parts_by_underscore = file_name_part.split(Const.REPLACEMENT_CHARACTER)
+            if len(parts_by_underscore) < 2:
+                raise ValueError("File name part does not contain enough '_' separated segments.")
+            time_info = parts_by_underscore[-1]
+
             global UT_ERROR_DATA_DIR
             UT_ERROR_DATA_DIR = 'ut_error_data' + time_info
         error_data_path = initialize_save_error_data(error_data_path)
@@ -579,9 +602,8 @@ def run_ut_command(args):
     }
     run_ut_config = checker_config.get_run_ut_config(**config_params)
     run_ut(run_ut_config)
+    logger.info("UT task completed.")
 
 
 if __name__ == '__main__':
-    seed_all()
     _run_ut()
-    logger.info("UT task completed.")

msprobe/pytorch/api_accuracy_checker/run_ut/run_ut_utils.py

@@ -1,9 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
-# Licensed under the Apache License, Version 2.0  (the "License");
+# 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
 #
@@ -18,8 +16,8 @@
 import os
 from collections import namedtuple
 import re
-import torch
 
+import torch
 try:
     import torch_npu
 except ImportError:
@@ -33,11 +31,9 @@ from msprobe.core.common.const import FileCheckConst, Const, CompareConst
 from msprobe.core.common.file_utils import FileChecker
 from msprobe.core.common.log import logger
 from msprobe.core.common.utils import CompareException
+from msprobe.pytorch.hook_module.api_register import ApiTemplate, get_api_register
 from msprobe.pytorch.hook_module.wrap_aten import AtenOPTemplate
-from msprobe.pytorch.hook_module.wrap_functional import FunctionalOPTemplate
-from msprobe.pytorch.hook_module.wrap_npu_custom import NpuOPTemplate
-from msprobe.pytorch.hook_module.wrap_tensor import TensorOPTemplate
-from msprobe.pytorch.hook_module.wrap_torch import TorchOPTemplate
+
 
 hf_32_standard_api = ["conv1d", "conv2d"]
 not_detach_set = {'resize_', 'resize_as_', 'set_', 'transpose_', 't_', 'squeeze_', 'unsqueeze_'}
@@ -108,17 +104,28 @@ def exec_api(exec_params):
     kwargs = exec_params.kwargs
     is_autocast = exec_params.is_autocast
     autocast_dtype = exec_params.autocast_dtype
-
-    if api_type == "Functional":
-        torch_api = FunctionalOPTemplate(api_name, str, False)
-    if api_type == "Tensor":
-        torch_api = TensorOPTemplate(api_name, str, False)
-    if api_type == "Torch":
-        torch_api = TorchOPTemplate(api_name, str, False)
-    if api_type == "Aten":
+    out = None
+
+    prefix_map = Const.API_DATA_PREFIX.get(Const.PT_FRAMEWORK, {})
+    if not prefix_map or api_type not in prefix_map.values() or \
+        api_type not in (
+            Const.FUNCTIONAL_API_TYPE_PREFIX,
+            Const.TENSOR_API_TYPE_PREFIX,
+            Const.TORCH_API_TYPE_PREFIX,
+            Const.ATEN_API_TYPE_PREFIX,
+            Const.NPU_API_TYPE_PREFIX
+    ):
+        return out
+
+    if api_type == Const.ATEN_API_TYPE_PREFIX:
         torch_api = AtenOPTemplate(api_name, None, False)
-    if api_type == "NPU":
-        torch_api = NpuOPTemplate(api_name, None, False, device)
+    else:
+        api_register = get_api_register()
+        api_register.initialize_hook(None)
+        api_func_type = list(prefix_map.keys())[list(prefix_map.values()).index(api_type)]
+        api_func = api_register.ori_api_attr.get(Const.PT_FRAMEWORK + Const.SEP + api_func_type, {}).get(api_name)
+
+        torch_api = ApiTemplate(api_name, api_func, api_type, None, need_hook=False, device=device)
     if is_autocast:
         with autocast(dtype=autocast_dtype):
             out = torch_api.forward(*args, **kwargs)
@@ -248,7 +255,8 @@ def record_skip_info(api_full_name, compare, compare_alg_results):
 
 def is_unsupported_api(api_name, is_overflow_check=False):
     split_name = api_name.split(Const.SEP)[0]
-    flag = (split_name == Const.DISTRIBUTED) or (is_overflow_check and split_name == Const.NPU)
+    unsupport_type_list = [Const.DISTRIBUTED, Const.MINDSPEED_API_TYPE_PREFIX]
+    flag = (split_name in unsupport_type_list) or (is_overflow_check and split_name == Const.NPU)
     if flag:
         logger.info(f"{split_name} api is not supported for run ut. SKIP.")
     return flag

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/attl.py

@@ -27,6 +27,7 @@ from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.client import T
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.server import TCPServer
 from msprobe.core.common.file_utils import remove_path
 from msprobe.pytorch.common.utils import logger, save_api_data, load_api_data, save_pkl, load_pkl
+from msprobe.core.common.decorator import recursion_depth_decorator
 
 BufferType = Union[ApiData, Dict[str, Any], str]  # Union[Tensor, Tuple[Optional[Tensor]]]
 
@@ -168,11 +169,12 @@ class ATTL:
         return buffer
 
 
+@recursion_depth_decorator("move2device_exec")
 def move2device_exec(obj, device):
     if isinstance(obj, (tuple, list)):
         data_list = [move2device_exec(val, device) for val in obj]
         return data_list if isinstance(obj, list) else tuple(data_list)
-    if isinstance(obj, dict):
+    if isinstance(obj, dict): 
         return {key: move2device_exec(val, device) for key, val in obj.items()}
     elif isinstance(obj, torch.Tensor):
         obj = obj.detach()

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/client.py

@@ -12,23 +12,22 @@
 # 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 hashlib
+from functools import partial
+import zlib
 import io
 import struct
 import time
 import os
-import signal
 from queue import Queue
 from threading import Thread
 from typing import Union
 
-from twisted.internet import reactor, protocol, endpoints
+from twisted.internet import reactor, protocol, endpoints, ssl
 from twisted.protocols.basic import FileSender
 
 from msprobe.pytorch.common.utils import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.utils import STRUCT_UNPACK_MODE as unpack_mode, \
-    STR_TO_BYTES_ORDER as bytes_order
+    STR_TO_BYTES_ORDER as bytes_order, cipher_list, verify_callback, load_ssl_pem
 
 MAX_SENDING_QUEUE_SIZE = 20
 
@@ -104,11 +103,28 @@ class TCPClient:
         self.factory = MessageClientFactory()
         self.factory.protocol = cur_protocol
         if self.tls_path:
-            from twisted.internet import ssl
-            client_key = os.path.join(self.tls_path, "client.key")
-            client_crt = os.path.join(self.tls_path, "client.crt")
-            client_context_factory = ssl.DefaultOpenSSLContextFactory(client_key, client_crt)
-            endpoint = endpoints.SSL4ClientEndpoint(reactor, self.host, self.port, client_context_factory)
+            client_key, client_crt, ca_crt, crl_pem = load_ssl_pem(
+                key_file=os.path.join(self.tls_path, "client.key"),
+                cert_file=os.path.join(self.tls_path, "client.crt"),
+                ca_file=os.path.join(self.tls_path, "ca.crt"),
+                crl_file=os.path.join(self.tls_path, "crl.pem")
+            )
+
+            ssl_options = ssl.CertificateOptions(
+                privateKey=client_key,
+                certificate=client_crt,
+                method=ssl.SSL.TLSv1_2_METHOD,
+                verify=True,
+                requireCertificate=True,
+                caCerts=[ca_crt],  # 信任的CA证书列表
+            )
+            ssl_context = ssl_options.getContext()
+            ssl_context.set_cipher_list(cipher_list)
+            ssl_context.set_options(ssl.SSL.OP_NO_RENEGOTIATION)
+            ssl_context.set_verify(ssl.SSL.VERIFY_PEER | ssl.SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
+                                   partial(verify_callback, crl=crl_pem))
+
+            endpoint = endpoints.SSL4ClientEndpoint(reactor, self.host, self.port, ssl_options)
         else:
             endpoint = endpoints.TCP4ClientEndpoint(reactor, self.host, self.port)
         d = endpoint.connect(self.factory)
@@ -299,12 +315,12 @@ class ClientProtocol(protocol.Protocol):
 
     def send_wrapped_data(self, data, sequence_number: int = 0, rank: int = 0, step: int = 0):
         length = len(data)
-        md5_hash = hashlib.md5(data).hexdigest() if self.check_sum else ""
+        data_crc = f"{zlib.crc32(data):08x}" if self.check_sum else ""
         data_meaasge = length.to_bytes(8, byteorder=bytes_order) + \
                        sequence_number.to_bytes(8, byteorder=bytes_order) + \
                        rank.to_bytes(8, byteorder=bytes_order) + \
                        step.to_bytes(8, byteorder=bytes_order) + \
-                       md5_hash.encode() + \
+                       data_crc.encode() + \
                        data
         logger.debug(f"send 流水号: {sequence_number}; RANK: {rank}; STEP: {step}; LENGTH: {length}")
 
@@ -346,7 +362,7 @@ class ClientProtocol(protocol.Protocol):
     def connectionLost(self, reason):
         self.signal_exit = True
         self.factory.num_connections -= 1
-        logger.info(f"Lost connection with server, reason is : {reason}")
+        logger.info(f"Lost connection with server, reason is : {reason.value}")
 
 
 class MessageClientFactory(protocol.ClientFactory):

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/device_dispatch.py

@@ -29,7 +29,6 @@ from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import move2target_device
 from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import generate_cpu_params
 
-
 # NPU vs GPU api list
 CompareApi = set(absolute_standard_api) | set(binary_standard_api) | set(thousandth_standard_api)
 
@@ -43,6 +42,15 @@ OnlineApiPrecisionCompareConfig = namedtuple('OnlineApiPrecisionCompareConfig',
 CommonCompareConfig = namedtuple('CommonCompareConfig', ['compare', 'handle_func', 'config'])
 
 
+def get_gpu_device():
+    try:
+        import torch_npu
+        is_gpu = False
+    except ImportError:
+        is_gpu = True
+    return is_gpu
+
+
 def run_ut_process(xpu_id, consumer_queue, common_config, api_precision_csv_file):
     """ When consumer_queue(shared with ConsumerDispatcher) is not empty, consume api data from consumer_queue.
     :param xpu_id: int
@@ -51,7 +59,9 @@ def run_ut_process(xpu_id, consumer_queue, common_config, api_precision_csv_file
     :param api_precision_csv_file: list, length is 2, result file name and details file name
     :return:
     """
-    gpu_device = torch.device(f'cuda:{xpu_id}')
+    device_info = "cuda" if get_gpu_device() else "npu"
+    logger.info(f"Start run_ut_process for {device_info} device, rank: {xpu_id}.")
+    gpu_device = torch.device(f'{device_info}:{xpu_id}')
 
     while True:
         if consumer_queue.empty():

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/server.py

@@ -12,19 +12,19 @@
 # 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 os.path
+from functools import partial
+import os
 import struct
-import hashlib
+import zlib
 import time
 import io
 from threading import Thread
 
-from twisted.internet import reactor, protocol, endpoints
+from twisted.internet import reactor, protocol, endpoints, ssl
 
 from msprobe.pytorch.common.utils import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.utils import cipher_list, \
-    STRUCT_UNPACK_MODE as unpack_mode, STR_TO_BYTES_ORDER as bytes_order
+    STRUCT_UNPACK_MODE as unpack_mode, STR_TO_BYTES_ORDER as bytes_order, verify_callback, load_ssl_pem
 
 
 class TCPServer:
@@ -44,15 +44,28 @@ class TCPServer:
         self.factory.protocol = self.build_protocol
 
         if self.tls_path:
-            from OpenSSL import SSL
-            from twisted.internet import ssl
-            server_key = os.path.join(self.tls_path, "server.key")
-            server_crt = os.path.join(self.tls_path, "server.crt")
-            server_context_factory = ssl.DefaultOpenSSLContextFactory(server_key, server_crt, SSL.TLSv1_2_METHOD)
-            server_context_ = server_context_factory.getContext()
-            server_context_.set_cipher_list(cipher_list)
-            server_context_.set_options(SSL.OP_NO_RENEGOTIATION)
-            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, server_context_factory)
+            server_key, server_crt, ca_crt, crl_pem = load_ssl_pem(
+                key_file=os.path.join(self.tls_path, "server.key"),
+                cert_file=os.path.join(self.tls_path, "server.crt"),
+                ca_file=os.path.join(self.tls_path, "ca.crt"),
+                crl_file=os.path.join(self.tls_path, "crl.pem")
+            )
+
+            ssl_options = ssl.CertificateOptions(
+                privateKey=server_key,
+                certificate=server_crt,
+                method=ssl.SSL.TLSv1_2_METHOD,
+                verify=True,
+                requireCertificate=True,
+                caCerts=[ca_crt],  # 信任的CA证书列表
+            )
+            ssl_context = ssl_options.getContext()
+            ssl_context.set_cipher_list(cipher_list)
+            ssl_context.set_options(ssl.SSL.OP_NO_RENEGOTIATION)
+            ssl_context.set_verify(ssl.SSL.VERIFY_PEER | ssl.SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
+                                   partial(verify_callback, crl=crl_pem))
+
+            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, ssl_options)
         else:
             endpoint = endpoints.TCP4ServerEndpoint(reactor, self.port)
         endpoint.listen(self.factory)
@@ -85,10 +98,10 @@ class ServerProtocol(protocol.Protocol):
         self.consumer_queue = shared_queue
         self.check_sum = check_sum
         self.length_width = 8
-        self.md5_width = 32
+        self.crc_width = 8
         self.obj_length = None
         self.tell = 0
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
         self.sequence_number = -1
         self.rank = -1
@@ -99,7 +112,7 @@ class ServerProtocol(protocol.Protocol):
         self.buffer = io.BytesIO()
         self.obj_length = None
         self.tell = 0
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
         self.factory.transport_dict[self.transport] = 1
         self.factory.transport_list.append(self.transport)
@@ -132,11 +145,12 @@ class ServerProtocol(protocol.Protocol):
             time.sleep(0.1)
 
         obj_key = str(self.sequence_number) + "_" + str(self.rank) + "_" + str(self.step)
+        # get the crc value of a 16-bit string with a length of 8
+        recv_crc = f"{zlib.crc32(self.obj_body):08x}"
 
-        recv_md5 = hashlib.md5(self.obj_body).hexdigest()
-        if self.check_sum and recv_md5 != self.obj_md5:
-            # when needs check md5 and check no pass, indicates received data error, send b"ERROR" to client.
-            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_md5}, but get {recv_md5}")
+        if self.check_sum and recv_crc != self.obj_crc:
+            # when needs check hash value and check no pass, indicates received data error, send b"ERROR" to client.
+            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_crc}, but get {recv_crc}")
             self.send_ack(self.ACK_ERROR)
         else:
             if self.obj_body == self.ACK_STOP:
@@ -146,7 +160,7 @@ class ServerProtocol(protocol.Protocol):
             if obj_key in self.sequence_number_dict:
                 logger.debug(f"这是一次异常的重传，可以忽略。 {obj_key}, {self.sequence_number_dict}")
             else:
-                self.sequence_number_dict[obj_key] = self.obj_md5
+                self.sequence_number_dict[obj_key] = self.obj_crc
                 self.consumer_queue.put(self.obj_body, block=True)
 
         self.reset_env()
@@ -173,7 +187,7 @@ class ServerProtocol(protocol.Protocol):
         self.sequence_number = -1
         self.rank = -1
         self.step = -1
-        self.obj_md5 = None
+        self.obj_crc = None
         self.obj_body = None
 
     def dataReceived(self, data):
@@ -192,15 +206,15 @@ class ServerProtocol(protocol.Protocol):
             logger.debug(
                 f"流水号: {self.sequence_number}; RANK: {self.rank}; STEP: {self.step}; Length: {self.obj_length}")
 
-        # If needs check md5 but not parse md5 yet, read 32b md5 values
-        check_sum_and_md5 = (self.check_sum
+        # If needs check hash but not parse crc yet, read 8b crc values
+        check_sum_and_crc = (self.check_sum
                              and self.obj_length is not None
-                             and self.obj_md5 is None
-                             and len(self.buffer.getvalue()) - self.tell >= self.md5_width)
-        if check_sum_and_md5:
-            self.obj_md5 = self.buffer.read(self.md5_width).decode()
-            self.tell += self.md5_width
-            logger.debug(f"MD5: {self.obj_md5}")
+                             and self.obj_crc is None
+                             and len(self.buffer.getvalue()) - self.tell >= self.crc_width)
+        if check_sum_and_crc:
+            self.obj_crc = self.buffer.read(self.crc_width).decode()
+            self.tell += self.crc_width
+            logger.debug(f"Hash value: {self.obj_crc}")
 
         current_length = len(self.buffer.getvalue()) - self.tell
         if self.obj_length is not None and 0 < self.obj_length <= current_length:

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/utils.py

@@ -12,6 +12,16 @@
 # 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 os
+from datetime import datetime, timezone
+
+from OpenSSL import crypto
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+from dateutil import parser
+
+from msprobe.core.common.file_utils import FileOpen
+from msprobe.core.common.log import logger
 
 cipher_list = ":".join(
     ["TLS_DHE_RSA_WITH_AES_128_GCM_SHA256",
@@ -42,3 +52,147 @@ cipher_list = ":".join(
 
 STRUCT_UNPACK_MODE = "!Q"
 STR_TO_BYTES_ORDER = "big"
+
+
+def is_certificate_revoked(cert, crl):
+    # 获取证书的序列号
+    cert_serial_number = cert.get_serial_number()
+
+    # 检查证书是否在CRL中
+    revoked_serials = [revoked_cert.serial_number for revoked_cert in crl]
+    if cert_serial_number in revoked_serials:
+        logger.error(f"证书已吊销:{cert_serial_number:020x}")
+        return True
+
+    return False
+
+
+def verify_callback(conn, cert, errno, depth, preverify_ok, crl=None):
+    """
+    验证对端证书的有效性
+    :param conn: OpenSSL.SSL.Connection, SSL 连接对象
+    :param cert: OpenSSL.crypto.X509, 当前证书
+    :param errno: int, OpenSSL错误代码, 0:无错误 | 9:证书过期 | 18: 自签名证书
+    :param depth: int, 当前证书在证书链中的深度 (0=叶子节点), 1:中间CA证书 -1:根CA证书 2+:更高级别CA证书
+    :param preverify_ok: int, 验证结果 (1=通过， 0=失败)
+    :param crl: _CRLInternal, CRL证书对象
+    :return: bool, True表示接受证书, False表示拒绝
+    """
+
+    if not preverify_ok:
+        from OpenSSL import SSL
+        error_str = SSL._ffi.string(SSL._lib.X509_verify_cert_error_string(errno)).decode()
+        logger.error(f"证书验证失败 (depth={depth}, err={errno}): {error_str}")
+        return False
+
+    if crl and is_certificate_revoked(cert, crl):
+        return False
+
+    return preverify_ok
+
+
+def load_ssl_pem(key_file, cert_file, ca_file, crl_file):
+    """
+    Load SSL PEM files.
+
+    Args:
+        key_file (str): The path to the private key file.
+        cert_file (str): The path to the certificate file.
+        ca_file (str): The path to the CA certificate file.
+        crl_file (str): The path to the CRL file.
+
+    Returns:
+        tuple: (key, crt, ca_crt, crl)
+
+    Raises:
+        Exception: If the file paths are invalid or the file contents are incorrect, exceptions may be thrown.
+    """
+
+    try:
+        # your_private_key_password
+        passphrase = ""
+        if not passphrase:
+            import pwinput
+            passphrase = pwinput.pwinput("Enter your password: ")
+        with FileOpen(key_file, "rb") as f:
+            key = crypto.load_privatekey(crypto.FILETYPE_PEM, f.read(), passphrase.encode())
+            del passphrase
+        with FileOpen(cert_file, "rb") as f:
+            crt = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
+            check_crt_valid(crt)
+
+            crt_serial_number = hex(crt.get_serial_number())[2:]
+            logger.info(f"crt_serial_number: {crt_serial_number}")
+
+        check_certificate_match(crt, key)
+
+        with FileOpen(ca_file, "rb") as f:
+            ca_crt = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
+            check_crt_valid(ca_crt)
+
+            ca_serial_number = hex(ca_crt.get_serial_number())[2:]
+            logger.info(f"ca_serial_number: {ca_serial_number}")
+        crl = None
+        if os.path.exists(crl_file):
+            with FileOpen(crl_file, "rb") as f:
+                crl = x509.load_pem_x509_crl(f.read(), default_backend())
+                check_crl_valid(crl, ca_crt)
+            for revoked_cert in crl:
+                logger.info(f"Serial Number: {revoked_cert.serial_number}, "
+                            f"Revocation Date: {revoked_cert.revocation_date_utc}")
+
+    except Exception as e:
+        raise RuntimeError(f"The SSL certificate is invalid") from e
+
+    return key, crt, ca_crt, crl
+
+
+def check_crt_valid(pem):
+    """
+    Check the validity of the SSL certificate.
+
+    Raises:
+    RuntimeError: If the SSL certificate is invalid or expired.
+    """
+    try:
+        pem_start = parser.parse(pem.get_notBefore().decode("UTF-8"))
+        pem_end = parser.parse(pem.get_notAfter().decode("UTF-8"))
+        logger.info(f"The SSL certificate passes the verification and the validity period "
+                    f"starts from {pem_start} ends at {pem_end}.")
+    except Exception as e:
+        raise RuntimeError(f"The SSL certificate is invalid") from e
+
+    now_utc = datetime.now(tz=timezone.utc)
+    if pem.has_expired() or not (pem_start <= now_utc <= pem_end):
+        raise RuntimeError(f"The SSL certificate has expired.")
+
+
+def check_certificate_match(certificate, private_key):
+    """
+    Check certificate and private_key is match or not. if mismatched, an exception is thrown.
+    :param certificate:
+    :param private_key:
+    :return:
+    """
+    test_data = os.urandom(256)
+    try:
+        signature = crypto.sign(private_key, test_data, "sha256")
+        crypto.verify(
+            certificate,  # 包含公钥的证书
+            signature,  # 生成的签名
+            test_data,  # 原始数据
+            "sha256",  # 哈希算法
+        )
+        logger.info("公钥和私钥匹配")
+    except Exception as e:
+        raise RuntimeError("公钥和私钥不匹配") from e
+
+
+def check_crl_valid(crl, ca_crt):
+    # 验证CRL签名（确保CRL未被篡改）
+    if not crl.is_signature_valid(ca_crt.get_pubkey().to_cryptography_key()):
+        raise RuntimeError("CRL签名无效！")
+
+    # 检查CRL有效期
+    if not (crl.last_update <= datetime.utcnow() <= crl.next_update):
+        raise RuntimeError("CRL已过期或尚未生效!")