mindstudio-probe 1.1.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

msprobe/core/common/file_utils.py

@@ -1,8 +1,7 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-# Copyright (C) 2022-2023. Huawei Technologies Co., Ltd. All rights reserved.
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Copyright (c) 2024-2024, 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
 #
@@ -13,13 +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 csv
 import fcntl
 import os
+import stat
 import json
 import re
 import shutil
+from datetime import datetime, timezone
+from dateutil import parser
 import yaml
 import numpy as np
 import pandas as pd
@@ -67,9 +69,11 @@ class FileChecker:
         self.check_path_ability()
         if self.is_script:
             check_path_owner_consistent(self.file_path)
-        check_path_pattern_vaild(self.file_path)
+        check_path_pattern_valid(self.file_path)
         check_common_file_size(self.file_path)
         check_file_suffix(self.file_path, self.file_type)
+        if self.path_type == FileCheckConst.FILE:
+            check_dirpath_before_read(self.file_path)
         return self.file_path
 
     def check_path_ability(self):
@@ -122,9 +126,10 @@ class FileOpen:
         self.file_path = os.path.realpath(self.file_path)
         check_path_length(self.file_path)
         self.check_ability_and_owner()
-        check_path_pattern_vaild(self.file_path)
+        check_path_pattern_valid(self.file_path)
         if os.path.exists(self.file_path):
             check_common_file_size(self.file_path)
+            check_dirpath_before_read(self.file_path)
 
     def check_ability_and_owner(self):
         if self.mode in self.SUPPORT_READ_MODE:
@@ -193,7 +198,7 @@ def check_path_owner_consistent(path):
         raise FileCheckException(FileCheckException.FILE_PERMISSION_ERROR)
 
 
-def check_path_pattern_vaild(path):
+def check_path_pattern_valid(path):
     if not re.match(FileCheckConst.FILE_VALID_PATTERN, path):
         logger.error('The file path %s contains special characters.' % (path))
         raise FileCheckException(FileCheckException.ILLEGAL_PATH_ERROR)
@@ -217,7 +222,6 @@ def check_common_file_size(file_path):
                 check_file_size(file_path, max_size)
                 return
         check_file_size(file_path, FileCheckConst.COMMOM_FILE_SIZE)
-        
 
 
 def check_file_suffix(file_path, file_suffix):
@@ -238,9 +242,18 @@ def check_path_type(file_path, file_type):
             raise FileCheckException(FileCheckException.INVALID_FILE_ERROR)
 
 
+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)     # 其他用户可写
+    )
+    return is_writable
+
+
 def make_dir(dir_path):
-    dir_path = os.path.realpath(dir_path)
     check_path_before_create(dir_path)
+    dir_path = os.path.realpath(dir_path)
     if os.path.isdir(dir_path):
         return
     try:
@@ -262,8 +275,9 @@ def create_directory(dir_path):
     Exception Description:
         when invalid data throw exception
     """
-    dir_path = os.path.realpath(dir_path)
+    check_link(dir_path)
     check_path_before_create(dir_path)
+    dir_path = os.path.realpath(dir_path)
     parent_dir = os.path.dirname(dir_path)
     if not os.path.isdir(parent_dir):
         create_directory(parent_dir)
@@ -271,6 +285,7 @@ def create_directory(dir_path):
 
 
 def check_path_before_create(path):
+    check_link(path)
     if path_len_exceeds_limit(path):
         raise FileCheckException(FileCheckException.ILLEGAL_PATH_ERROR, 'The file path length exceeds limit.')
 
@@ -279,6 +294,17 @@ def check_path_before_create(path):
                                  'The file path {} contains special characters.'.format(path))
 
 
+def check_dirpath_before_read(path):
+    path = os.path.realpath(path)
+    dirpath = os.path.dirname(path)
+    if check_others_writable(dirpath):
+        logger.warning(f"The directory is writable by others: {dirpath}.")
+    try:
+        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):
     """
     Function Description:
@@ -344,7 +370,7 @@ def load_yaml(yaml_path):
 def load_npy(filepath):
     check_file_or_directory_path(filepath)
     try:
-        npy = np.load(filepath)
+        npy = np.load(filepath, allow_pickle=False)
     except Exception as e:
         logger.error(f"The numpy file failed to load. Please check the path: {filepath}.")
         raise RuntimeError(f"Load numpy file {filepath} failed.") from e
@@ -354,7 +380,7 @@ def load_npy(filepath):
 def load_json(json_path):
     try:
         with FileOpen(json_path, "r") as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
+            fcntl.flock(f, fcntl.LOCK_SH)
             data = json.load(f)
             fcntl.flock(f, fcntl.LOCK_UN)
     except Exception as e:
@@ -363,11 +389,11 @@ def load_json(json_path):
     return data
 
 
-def save_json(json_path, data, indent=None):
-    json_path = os.path.realpath(json_path)
+def save_json(json_path, data, indent=None, mode="w"):
     check_path_before_create(json_path)
+    json_path = os.path.realpath(json_path)
     try:
-        with FileOpen(json_path, 'w') as f:
+        with FileOpen(json_path, mode) as f:
             fcntl.flock(f, fcntl.LOCK_EX)
             json.dump(data, f, indent=indent)
             fcntl.flock(f, fcntl.LOCK_UN)
@@ -378,8 +404,8 @@ def save_json(json_path, data, indent=None):
 
 
 def save_yaml(yaml_path, data):
-    yaml_path = os.path.realpath(yaml_path)
     check_path_before_create(yaml_path)
+    yaml_path = os.path.realpath(yaml_path)
     try:
         with FileOpen(yaml_path, 'w') as f:
             fcntl.flock(f, fcntl.LOCK_EX)
@@ -391,6 +417,37 @@ def save_yaml(yaml_path, data):
     change_mode(yaml_path, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
+def save_excel(path, data):
+    def validate_data(data):
+        """Validate that the data is a DataFrame or a list of (DataFrame, sheet_name) pairs."""
+        if isinstance(data, pd.DataFrame):
+            return "single"
+        elif isinstance(data, list):
+            if all(isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], pd.DataFrame) for item in data):
+                return "list"
+        raise ValueError("Data must be a DataFrame or a list of (DataFrame, sheet_name) pairs.")
+
+    check_path_before_create(path)
+    path = os.path.realpath(path)
+
+    # 验证数据类型
+    data_type = validate_data(data)
+
+    try:
+        if data_type == "single":
+            data.to_excel(path, index=False)
+        elif data_type == "list":
+            with pd.ExcelWriter(path) as writer:
+                for data_df, sheet_name in data:
+                    data_df.to_excel(writer, sheet_name=sheet_name, index=False)
+    except Exception as e:
+        logger.error(f'Save excel file "{os.path.basename(path)}" failed.')
+        raise RuntimeError(f"Save excel file {path} failed.") from e
+    change_mode(path, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
+
+
 def move_file(src_path, dst_path):
     check_file_or_directory_path(src_path)
     check_path_before_create(dst_path)
@@ -403,8 +460,8 @@ def move_file(src_path, dst_path):
 
 
 def save_npy(data, filepath):
-    filepath = os.path.realpath(filepath)
     check_path_before_create(filepath)
+    filepath = os.path.realpath(filepath)
     try:
         np.save(filepath, data)
     except Exception as e:
@@ -425,6 +482,7 @@ def save_npy_to_txt(data, dst_file='', align=0):
         pad_array = np.zeros((align - data.size % align,))
         data = np.append(data, pad_array)
     check_path_before_create(dst_file)
+    dst_file = os.path.realpath(dst_file)
     try:
         np.savetxt(dst_file, data.reshape((-1, align)), delimiter=' ', fmt='%g')
     except Exception as e:
@@ -438,8 +496,8 @@ def save_workbook(workbook, file_path):
     workbook: 要保存的工作簿对象
     file_path: 文件保存路径
     """
-    file_path = os.path.realpath(file_path)
     check_path_before_create(file_path)
+    file_path = os.path.realpath(file_path)
     try:
         workbook.save(file_path)
     except Exception as e:
@@ -451,7 +509,7 @@ def save_workbook(workbook, file_path):
 def write_csv(data, filepath, mode="a+", malicious_check=False):
     def csv_value_is_valid(value: str) -> bool:
         if not isinstance(value, str):
-            return True  
+            return True
         try:
             # -1.00 or +1.00 should be consdiered as digit numbers
             float(value)
@@ -459,16 +517,16 @@ def write_csv(data, filepath, mode="a+", malicious_check=False):
             # otherwise, they will be considered as formular injections
             return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
         return True
-    
+
     if malicious_check:
         for row in data:
             for cell in row:
                 if not csv_value_is_valid(cell):
-                    raise RuntimeError(f"Malicious value [{cell}] is not allowed " \
+                    raise RuntimeError(f"Malicious value [{cell}] is not allowed "
                                        f"to be written into the csv: {filepath}.")
 
-    file_path = os.path.realpath(filepath)
     check_path_before_create(filepath)
+    file_path = os.path.realpath(filepath)
     try:
         with FileOpen(filepath, mode, encoding='utf-8-sig') as f:
             writer = csv.writer(f)
@@ -479,16 +537,54 @@ def write_csv(data, filepath, mode="a+", malicious_check=False):
     change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
 
 
-def read_csv(filepath):
+def read_csv(filepath, as_pd=True, header='infer'):
     check_file_or_directory_path(filepath)
     try:
-        csv_data = pd.read_csv(filepath)
+        if as_pd:
+            csv_data = pd.read_csv(filepath, header=header)
+        else:
+            with FileOpen(filepath, 'r', encoding='utf-8-sig') as f:
+                csv_reader = csv.reader(f, delimiter=',')
+                csv_data = list(csv_reader)
     except Exception as e:
         logger.error(f"The csv file failed to load. Please check the path: {filepath}.")
         raise RuntimeError(f"Read csv file {filepath} failed.") from e
     return csv_data
 
 
+def write_df_to_csv(data, filepath, mode="w", header=True, malicious_check=False):
+    def csv_value_is_valid(value: str) -> bool:
+        if not isinstance(value, str):
+            return True
+        try:
+            # -1.00 or +1.00 should be consdiered as digit numbers
+            float(value)
+        except ValueError:
+            # otherwise, they will be considered as formular injections
+            return not bool(re.compile(FileCheckConst.CSV_BLACK_LIST).search(value))
+        return True
+
+    if not isinstance(data, pd.DataFrame):
+        raise ValueError("The data type of data is not supported. Only support pd.DataFrame.")
+
+    if malicious_check:
+        for i in range(len(data)):
+            for j in range(len(data.columns)):
+                cell = data.iloc[i, j]
+                if not csv_value_is_valid(cell):
+                    raise RuntimeError(f"Malicious value [{cell}] is not allowed "
+                                       f"to be written into the csv: {filepath}.")
+
+    check_path_before_create(filepath)
+    file_path = os.path.realpath(filepath)
+    try:
+        data.to_csv(filepath, mode=mode, header=header, index=False)
+    except Exception as e:
+        logger.error(f'Save csv file "{os.path.basename(file_path)}" failed')
+        raise RuntimeError(f"Save csv file {file_path} failed.") from e
+    change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
+
+
 def remove_path(path):
     if not os.path.exists(path):
         return
@@ -521,3 +617,57 @@ def get_json_contents(file_path):
 def get_file_content_bytes(file):
     with FileOpen(file, 'rb') as file_handle:
         return file_handle.read()
+
+
+# 对os.walk设置遍历深度
+def os_walk_for_files(path, depth):
+    res = []
+    for root, _, files in os.walk(path, topdown=True):
+        check_file_or_directory_path(root, isdir=True)
+        if root.count(os.sep) - path.count(os.sep) >= depth:
+            _[:] = []
+        else:
+            for file in files:
+                res.append({"file": file, "root": root})
+    return res
+
+
+def check_crt_valid(pem_path):
+    """
+    Check the validity of the SSL certificate.
+
+    Load the SSL certificate from the specified path, parse and check its validity period.
+    If the certificate is expired or invalid, raise a RuntimeError.
+
+    Parameters:
+    pem_path (str): The file path of the SSL certificate.
+
+    Raises:
+    RuntimeError: If the SSL certificate is invalid or expired.
+    """
+    import OpenSSL
+    try:
+        with FileOpen(pem_path, "r") as f:
+            pem_data = f.read()
+        cert = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, pem_data)
+        pem_start = parser.parse(cert.get_notBefore().decode("UTF-8"))
+        pem_end = parser.parse(cert.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:
+        logger.error("Failed to parse the SSL certificate. Check the certificate.")
+        raise RuntimeError(f"The SSL certificate is invalid, {pem_path}") from e
+
+    now_utc = datetime.now(tz=timezone.utc)
+    if cert.has_expired() or not (pem_start <= now_utc <= pem_end):
+        raise RuntimeError(f"The SSL certificate has expired and needs to be replaced, {pem_path}")
+
+
+def read_xlsx(file_path):
+    check_file_or_directory_path(file_path)
+    try:
+        result_df = pd.read_excel(file_path, keep_default_na=False)
+    except Exception as e:
+        logger.error(f"The xlsx file failed to load. Please check the path: {file_path}.")
+        raise RuntimeError(f"Read xlsx file {file_path} failed.") from e
+    return result_df

msprobe/core/common/inplace_op_checker.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 os
 from msprobe.core.common.file_utils import load_yaml
 

msprobe/core/common/inplace_ops.yaml

@@ -157,6 +157,9 @@ inplace_tensor_op:
   - trunc_
   - unsqueeze_
   - xlogy_
+  - bitwise_left_shift_
+  - bitwise_right_shift_
+  - arctan2_
 
 inplace_torch_op:
   - _add_relu_

msprobe/core/common/log.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2024-2024, 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 os
 import time
 import sys
@@ -5,6 +20,16 @@ from functools import wraps
 from msprobe.core.common.const import MsgConst
 
 
+def filter_special_chars(func):
+    @wraps(func)
+    def func_level(self, msg, **kwargs):
+        for char in MsgConst.SPECIAL_CHAR:
+            msg = msg.replace(char, '_')
+        return func(self, msg, **kwargs)
+
+    return func_level
+
+
 class BaseLogger:
     def __init__(self):
         self.rank = None
@@ -21,14 +46,6 @@ class BaseLogger:
     def get_rank(self):
         return self.rank
 
-    def filter_special_chars(func):
-        @wraps(func)
-        def func_level(self, msg, **kwargs):
-            for char in MsgConst.SPECIAL_CHAR:
-                msg = msg.replace(char, '_')
-            return func(self, msg, **kwargs)
-        return func_level
-    
     @filter_special_chars
     def error(self, msg):
         if self.level <= MsgConst.LogLevel.ERROR.value:
@@ -56,6 +73,7 @@ class BaseLogger:
                 return func(*args, **kwargs)
             else:
                 return None
+
         return func_rank_0
 
     def info_on_rank_0(self, msg):
@@ -66,7 +84,7 @@ class BaseLogger:
 
     def warning_on_rank_0(self, msg):
         return self.on_rank_0(self.warning)(msg)
-    
+
     def error_log_with_exp(self, msg, exception):
         self.error(msg)
         raise exception