mindspore 2.7.0__cp310-cp310-win_amd64.whl → 2.7.1__cp310-cp310-win_amd64.whl

mindspore/profiler/common/path_manager.py

@@ -135,7 +135,8 @@ class PathManager:
             msg = f"The path does not exist: {path}"
             raise ProfilerPathErrorException(msg)
         if os.name != 'nt' and os.stat(path).st_uid != os.getuid():
-            msg = f"Path {path} owner[{os.stat(path).st_uid}] does not match the current user[{os.getuid()}]."
+            msg = (f"Path {path} owner[{os.stat(path).st_uid}] does not match the current user[{os.getuid()}]."
+                   f"Please execute chown -R $(id -un) {path}")
             raise ProfilerPathErrorException(msg)
 
     @classmethod
@@ -153,7 +154,7 @@ class PathManager:
             msg = f"Invalid path is a soft link: {path}"
             raise ProfilerPathErrorException(msg)
         if not os.access(path, os.W_OK):
-            msg = f"The path writeable permission check failed: {path}"
+            msg = f"The path writeable permission check failed: {path}. Please execute chmod -R 755 {path}"
             raise ProfilerPathErrorException(msg)
 
     @classmethod
@@ -171,7 +172,7 @@ class PathManager:
             msg = f"Invalid path is a soft link: {path}"
             raise ProfilerPathErrorException(msg)
         if not os.access(path, os.R_OK):
-            msg = f"The path readable permission check failed: {path}"
+            msg = f"The path readable permission check failed: {path}. Please execute chmod -R 755 {path}"
             raise ProfilerPathErrorException(msg)
 
     @classmethod
@@ -246,26 +247,6 @@ class PathManager:
         except Exception as err:
             raise ProfilerPathErrorException(f"Failed to make directory: {path}, err: {err}") from err
 
-    @classmethod
-    def create_file_safety(cls, path: str):
-        """
-        Function Description:
-            create file safety
-        Parameter:
-            path: the file to remove
-        Exception Description:
-            when invalid data throw exception
-        """
-        if os.path.islink(path):
-            raise RuntimeError(f"Failed to create file: {path}, is a soft link")
-        if os.path.exists(path):
-            logger.warning("File already exists: %s", path)
-            return
-        try:
-            os.close(os.open(path, os.O_WRONLY | os.O_CREAT, cls.DATA_FILE_AUTHORITY))
-        except Exception as err:
-            raise RuntimeError(f"Failed to create file: {path}, err: {err}") from err
-
     @classmethod
     def _input_path_common_check(cls, path: str):
         """
@@ -399,6 +380,57 @@ class PathManager:
             return False
         if os.name == 'nt':
             return False
-        if os.stat(path).st_uid == 0 or os.stat(path).st_uid == os.getuid():
+        if os.stat(lib_path).st_uid == 0 or os.stat(lib_path).st_uid == os.getuid():
             return True
         return False
+
+    @classmethod
+    def check_path_is_other_writable(cls, path):
+        """Check whether the file or directory in the specified path has writable permissions for others."""
+        file_stat = os.stat(path)
+        if file_stat.st_mode & (stat.S_IWGRP | stat.S_IWOTH):
+            msg = (f"File path {path} has group or others writable permissions, which is not allowed."
+                   f"Please execute chmod -R 755 {path}")
+            raise ProfilerPathErrorException(msg)
+
+    @classmethod
+    def check_path_is_owner_or_root(cls, path):
+        """Check path is owner or root."""
+        if not os.path.exists(path):
+            msg = f"The path does not exist: {path}"
+            raise ProfilerPathErrorException(msg)
+        file_stat = os.stat(path)
+        current_uid = os.getuid()
+        file_uid = file_stat.st_uid
+        if file_uid not in (0, current_uid):
+            return False
+        return True
+
+    @classmethod
+    def check_path_is_executable(cls, path):
+        """Check path is executable"""
+        return os.access(path, os.X_OK)
+
+    @classmethod
+    def check_path_is_readable(cls, path):
+        """Check path is readable"""
+        if os.path.islink(path):
+            msg = f"Invalid path is a soft link: {path}"
+            raise ProfilerPathErrorException(msg)
+        if not os.access(path, os.R_OK):
+            msg = f"The path readable permission check failed: {path}."
+            raise ProfilerPathErrorException(msg)
+
+    @classmethod
+    def walk_with_depth(cls, path, *args, max_depth=10, **kwargs):
+        """walk path depth"""
+        if not isinstance(path, str):
+            return
+        base_depth = path.count(os.sep)
+        if path.endswith(os.sep):
+            base_depth -= 1
+        for root, dirs, files in os.walk(path, *args, **kwargs):
+            if root.count(os.sep) - base_depth > max_depth:
+                dirs.clear()
+                continue
+            yield root, dirs, files

mindspore/profiler/common/profiler_context.py

@@ -24,7 +24,6 @@ from typing import (
 )
 
 from mindspore.communication.management import GlobalComm
-from mindspore.communication.management import get_local_rank
 from mindspore.communication.management import get_rank
 from mindspore.profiler.common.constant import (
     DeviceTarget,
@@ -44,6 +43,7 @@ from mindspore import context
 from mindspore import log as logger
 from mindspore.profiler.common.profiler_info import ProfilerInfo
 from mindspore.profiler.experimental_config import _ExperimentalConfig
+from mindspore.profiler.common.util import get_device_id
 
 
 @Singleton
@@ -488,17 +488,7 @@ class ProfilerContext:
         """
         Initialize the device ID.
         """
-        self._device_id = str(context.get_context("device_id"))
-
-        if not self._device_id or not self._device_id.isdigit():
-            if GlobalComm.INITED and self._device_target == DeviceTarget.NPU.value:
-                self._device_id = str(get_local_rank())
-            else:
-                self._device_id = os.getenv("DEVICE_ID")
-
-        if not self._device_id or not self._device_id.isdigit():
-            self._device_id = "0"
-            logger.warning("Fail to get DEVICE_ID, use 0 instead.")
+        self._device_id = get_device_id()
 
     def _init_rank_id(self) -> None:
         """

mindspore/profiler/common/profiler_info.py

@@ -97,13 +97,14 @@ class ProfilerInfo:
         Load time parameters from msprof profile and host start log.
         This method should be called before TimeConverter.init_parameters.
         """
+        msprof_info = MsprofCmdTool(msprof_profile_path).get_msprof_info()
         if not msprof_profile_path or not msprof_profile_host_path:
             raise ValueError(
                 "msprof_profile_path and msprof_profile_host_path must be provided"
             )
         self._read_host_start_log(msprof_profile_host_path)
         self._read_start_info(msprof_profile_host_path)
-        self._get_freq_from_msprof(msprof_profile_path)
+        self._get_freq_from_msprof(msprof_info)
 
     @property
     def time_parameters(self) -> Dict[str, Any]:
@@ -237,7 +238,7 @@ class ProfilerInfo:
             self._collection_time_begin * self.US_TO_NS - self._clock_monotonic_raw_info
         )
 
-    def _get_freq_from_msprof(self, msprof_profile_path: str) -> None:
+    def _get_freq_from_msprof(self, msprof_info: str) -> None:
         """
         Get frequency from get_msprof_info.py script
 
@@ -250,7 +251,6 @@ class ProfilerInfo:
             }
         }
         """
-        msprof_info = MsprofCmdTool(msprof_profile_path).get_msprof_info()
 
         if not isinstance(msprof_info, dict):
             raise RuntimeError("msprof_info must be a dictionary")

mindspore/profiler/common/profiler_path_manager.py

@@ -143,6 +143,19 @@ class ProfilerPathManager:
                     new_file_path = os.path.join(self._prof_ctx.ascend_profiler_output_path, new_file_name)
                     shutil.move(db_file, new_file_path)
 
+    def remove_db_file(self):
+        """
+        Remove the db file in the output path.
+        """
+        if not self._prof_ctx.msprof_profile_output_path:
+            return
+        db_files = glob.glob(os.path.join(
+            os.path.dirname(self._prof_ctx.msprof_profile_output_path),
+            'msprof*.db'
+        ))
+        for db_file in db_files:
+            if os.path.isfile(db_file):
+                os.remove(db_file)
 
     def create_output_path(self):
         """

mindspore/profiler/common/util.py

@@ -25,10 +25,15 @@ import re
 import shutil
 import stat
 
+from mindspore.communication.management import GlobalComm
+from mindspore.communication.management import get_local_rank
+from mindspore import context
 from mindspore import log as logger
 from mindspore.profiler.common.path_manager import PathManager
 from mindspore.profiler.common.exceptions.exceptions import ProfilerPathErrorException
 
+from mindspore.profiler.common.constant import DeviceTarget
+
 
 def no_exception_func(
         default_ret: Any = None,
@@ -75,12 +80,16 @@ def get_cann_version():
     ascend_home_path = os.environ.get("ASCEND_HOME_PATH", "")
     cann_version = "not known"
     try:
-        PathManager.check_directory_path_readable(os.path.realpath(ascend_home_path))
-        for dirpath, _, filenames in os.walk(os.path.realpath(ascend_home_path)):
+        if not PathManager.check_path_is_owner_or_root(ascend_home_path):
+            raise PermissionError(f"PermissionError, CANN package user id: {os.stat(ascend_home_path).st_uid}, "
+                                  f"current user id: {os.getuid()}. "
+                                  f"Ensure CANN package user id and current user id consistency")
+        PathManager.check_path_is_readable(os.path.realpath(ascend_home_path))
+        for dirpath, _, filenames in PathManager.walk_with_depth(os.path.realpath(ascend_home_path)):
             install_files = [file for file in filenames if re.match(r"ascend_.{1,20}_install\.info", file)]
             if install_files:
                 filepath = os.path.realpath(os.path.join(dirpath, install_files[0]))
-                PathManager.check_directory_path_readable(filepath)
+                PathManager.check_path_is_readable(filepath)
                 with open(filepath, "r") as f:
                     for line in f:
                         if line.find("version") != -1:
@@ -441,6 +450,24 @@ def get_newest_file(file_list):
     return newest_file_list
 
 
+def get_device_id():
+    """
+    Get device ID.
+    """
+    device_id = str(context.get_context("device_id"))
+
+    if not device_id or not device_id.isdigit():
+        if GlobalComm.INITED and context.get_context("device_target") == DeviceTarget.NPU.value:
+            device_id = str(get_local_rank())
+        else:
+            device_id = os.getenv("DEVICE_ID")
+
+    if not device_id or not device_id.isdigit():
+        logger.warning("Fail to get DEVICE_ID, use 0 instead.")
+        device_id = "0"
+    return device_id
+
+
 class ProfilerPathManager:
     """A path manager to manage profiler path"""
 

mindspore/profiler/experimental_config.py

@@ -84,7 +84,8 @@ class _ExperimentalConfig:
             HCCS data, PCIe data, and Stars Chip Trans. Default: ``False``.
         host_sys (list, optional): Collect the data of system call classes on the host side.
             Default: ``[]``, indicating that system class data on the host side is not collected.
-            You need to set `start_profile` of :class:`mindspore.profiler.profile` to ``False``.When collecting DISK or
+            You need to set `start_profile` of :class:`mindspore.profiler.profile` to ``False``.Currently, only
+            the **root user** supports collecting DISK or OSRT data, when collecting DISK or
             OSRT data, it is necessary to install the iotop, perf, and ltrace third-party tools in advance.
             For detailed steps, please refer to `Installing Third-party Tools
             <https://www.hiascend.com/document/detail/zh/mindstudio/80RC1/T&ITools/Profiling/atlasprofiling_16_0136.

mindspore/profiler/platform/npu_profiler.py

@@ -36,6 +36,8 @@ from mindspore.profiler.platform.base_profiler import BaseProfiler
 from mindspore.profiler.common.profiler_path_manager import ProfilerPathManager
 from mindspore.profiler.common.profiler_info import ProfilerInfo
 from mindspore.profiler.common.process_pool import MultiProcessPool
+from mindspore.profiler.common.constant import MsprofModeName
+from mindspore.profiler.common.util import no_exception_func
 from mindspore.profiler.analysis.task_manager import TaskManager
 from mindspore.profiler.analysis.time_converter import TimeConverter
 from mindspore.profiler.analysis.parser.ascend_cann_parser import AscendMsprofParser
@@ -58,6 +60,7 @@ from mindspore.profiler.analysis.viewer.ms_operator_details_viewer import MsOper
 from mindspore.profiler.common.util import print_msg_with_pid
 from mindspore.profiler.common.log import ProfilerLogger
 from mindspore.profiler.mstx import Mstx
+from mindspore.profiler.common.util import get_device_id
 
 
 @PROFILERS.register_module(DeviceTarget.NPU.value)
@@ -68,6 +71,7 @@ class NpuProfiler(BaseProfiler):
 
     def __init__(self) -> None:
         super().__init__()
+        self._is_env_not_valid = self._is_environment_not_valid()
         self._prof_ctx = ProfilerContext()
         self._prof_info = ProfilerInfo()
         self._prof_path_mgr = ProfilerPathManager()
@@ -78,7 +82,7 @@ class NpuProfiler(BaseProfiler):
         # initialize profiler backend
         self._profiler.init(
             self._prof_ctx.ascend_ms_dir,
-            int(self._prof_ctx.device_id),
+            int(get_device_id()),
             json.dumps(self._prof_ctx.npu_profiler_params),
         )
         self._logger.info("NpuProfiler init profiler backend params %s",
@@ -97,6 +101,8 @@ class NpuProfiler(BaseProfiler):
 
     def start(self) -> None:
         """Start profiling."""
+        if self._is_env_not_valid:
+            return
         self._logger.info("NpuProfiler start.")
 
         Mstx.enable = self._prof_ctx.npu_profiler_params.get("mstx", False)
@@ -115,6 +121,8 @@ class NpuProfiler(BaseProfiler):
 
     def stop(self) -> None:
         """Stop profiling."""
+        if self._is_env_not_valid:
+            return
         self._logger.info("NpuProfiler stop.")
 
         Mstx.enable = False
@@ -149,16 +157,29 @@ class NpuProfiler(BaseProfiler):
 
     def analyse(self, **kwargs) -> None:
         """Analyse the profiling data."""
+        if self._is_env_not_valid:
+            return
         self._logger.info("NpuProfiler analyse.")
 
         NPUProfilerAnalysis.online_analyse(async_mode=kwargs.get('async_mode'))
 
     def finalize(self) -> None:
         """Finalize profiling data."""
+        if self._is_env_not_valid:
+            return
         self._logger.info("NpuProfiler finalize.")
         if self._profiler:
             self._profiler.finalize()
 
+    @staticmethod
+    def _is_environment_not_valid() -> bool:
+        # check msprof dynamic environment variable
+        if os.getenv(MsprofModeName.MSPROF_DYNAMIC_ENV) is not None:
+            logger.error(f"The environment variable '{MsprofModeName.MSPROF_DYNAMIC_ENV}' has been set."
+                         f"Please execute 'unset {MsprofModeName.MSPROF_DYNAMIC_ENV}'.")
+            return True
+        return False
+
 
 class NPUProfilerAnalysis:
     """
@@ -166,6 +187,7 @@ class NPUProfilerAnalysis:
     """
 
     @classmethod
+    @no_exception_func()
     def online_analyse(cls, async_mode: bool = False):
         """
         Online analysis for NPU
@@ -179,6 +201,7 @@ class NPUProfilerAnalysis:
             cls._run_tasks(**ProfilerContext().to_dict())
 
     @classmethod
+    @no_exception_func()
     def offline_analyse(
             cls,
             path: str,
@@ -262,11 +285,15 @@ class NPUProfilerAnalysis:
         task_mgr = cls._construct_task_mgr(**kwargs)
         task_mgr.run()
         ProfilerLogger.get_instance().info(json.dumps(task_mgr.cost_time, indent=4))
-        activities = kwargs.get("activities")
-        if activities and ProfilerActivity.NPU.value in activities:
-            ProfilerPathManager().move_db_file()
-        if kwargs.get("data_simplification") and ProfilerActivity.NPU.value in kwargs.get("activities"):
-            ProfilerPathManager().simplify_data()
+        activities = kwargs.get("activities", [])
+        export_type = kwargs.get("export_type", [])
+        if ProfilerActivity.NPU.value in activities:
+            if ExportType.Db.value in export_type:
+                ProfilerPathManager().move_db_file()
+            else:
+                ProfilerPathManager().remove_db_file()
+            if kwargs.get("data_simplification"):
+                ProfilerPathManager().simplify_data()
 
     @classmethod
     def _construct_task_mgr(cls, **kwargs) -> TaskManager:

mindspore/run_check/_check_version.py

@@ -44,6 +44,23 @@ class EnvChecker(metaclass=ABCMeta):
     def check_version(self):
         pass
 
+    @staticmethod
+    def _concat_variable(env_name, env_value):
+        """concat value to the beginning of env specified by env_name"""
+        if not os.getenv(env_name, ""):
+            os.environ[env_name] = env_value
+        else:
+            paths = os.environ[env_name].split(':')
+            if paths and paths[0] == env_value:
+                return
+            if env_value not in paths:
+                os.environ[env_name] = env_value + ':' + os.environ[env_name]
+            else:
+                # move env_value to beginning
+                new_paths = [p for p in paths if p != env_value]
+                new_paths.insert(0, env_value)
+                os.environ[env_name] = ':'.join(new_paths)
+
 
 class CPUEnvChecker(EnvChecker):
     """CPU environment check."""
@@ -61,10 +78,7 @@ class CPUEnvChecker(EnvChecker):
         """set env for cpu"""
         plugin_dir = os.path.dirname(self.library_path)
         akg_dir = os.path.join(plugin_dir, "plugin/cpu")
-        if os.getenv('LD_LIBRARY_PATH'):
-            os.environ['LD_LIBRARY_PATH'] = akg_dir + ":" + os.environ['LD_LIBRARY_PATH']
-        else:
-            os.environ['LD_LIBRARY_PATH'] = akg_dir
+        EnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)
 
 
 class GPUEnvChecker(EnvChecker):
@@ -142,10 +156,7 @@ class GPUEnvChecker(EnvChecker):
         v_str = str(v.major) + "." + str(v.minor)
         plugin_dir = os.path.dirname(self.library_path)
         akg_dir = os.path.join(plugin_dir, "gpu" + v_str)
-        if os.getenv('LD_LIBRARY_PATH'):
-            os.environ['LD_LIBRARY_PATH'] = akg_dir + ":" + os.environ['LD_LIBRARY_PATH']
-        else:
-            os.environ['LD_LIBRARY_PATH'] = akg_dir
+        EnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)
         os.environ['CUDA_CACHE_MAXSIZE'] = "4000000000"
 
     def _get_bin_path(self, bin_name):
@@ -258,7 +269,7 @@ class AscendEnvChecker(EnvChecker):
 
     def __init__(self, library_path):
         self.library_path = library_path
-        self.version = ["7.7", "7.8", "8.2"]
+        self.version = ["7.8", "8.2", "8.3"]
 
         # env
         self.path = os.getenv("PATH")
@@ -278,13 +289,6 @@ class AscendEnvChecker(EnvChecker):
         self.ascend_opp_kernel_path_check = "/opp_kernel"
         self.v = ""
 
-    @staticmethod
-    def _concat_variable(env_name, env_value):
-        if os.getenv(env_name) is None:
-            os.environ[env_name] = env_value
-        else:
-            os.environ[env_name] = env_value + ":" + os.environ[env_name]
-
     def check_custom_version(self):
         """custom op version check"""
 
@@ -359,11 +363,6 @@ class AscendEnvChecker(EnvChecker):
         """
             opp kernel install check
         """
-        from mindspore._c_expression import MSContext
-        soc_version = MSContext.get_instance().get_ascend_soc_version()
-        if soc_version == "ascend310":
-            return
-
         opp_kernel_path = self.ascend_opp_path.replace("opp", "opp_kernel")
         if not os.path.exists(opp_kernel_path):
             logger.critical("MindSpore relies on \"Ascend opp_kernel\" folder of the Ascend AI software package ("
@@ -389,7 +388,7 @@ class AscendEnvChecker(EnvChecker):
         os.environ['IGNORE_INFER_ERROR'] = "1"
         plugin_dir = os.path.dirname(self.library_path)
         akg_dir = os.path.join(plugin_dir, "ascend")
-        AscendEnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)
+        EnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)
 
         self._check_env()
 
@@ -493,6 +492,85 @@ def _set_pb_env():
         os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 
+def _check_dir_path_safety(dir_path):
+    """Check safety of env directory path."""
+    if not os.path.exists(dir_path):
+        logger.warning(f"Path {dir_path} not exists.")
+        return False
+
+    if not os.path.isdir(dir_path):
+        logger.warning(f"Path {dir_path} is not a directory.")
+        return False
+
+    # check if path is suspicious
+    suspicious_patterns = [
+        "/tmp/", "/var/tmp/", "/dev/", "/proc/",
+        "\\temp\\", "\\windows\\temp\\",
+        "appdata", "local\\temp"
+    ]
+    lower_path = dir_path.lower()
+    for pattern in suspicious_patterns:
+        if pattern in lower_path:
+            logger.warning(f"Path {dir_path} is suspicious.")
+            return False
+
+    # check whether the path points to a system-critical directory
+    critical_dirs = [
+        "/bin", "/sbin", "/usr/bin", "/usr/sbin",
+        "/windows", "/system32", "c:\\windows"
+    ]
+    for critical_dir in critical_dirs:
+        if critical_dir in lower_path:
+            logger.warning(f"Path {dir_path} points to a system-critical directory.")
+            return False
+
+    return True
+
+
+def check_cuda_path_safety(cuda_path):
+    """Check safety of cuda path."""
+    if not _check_dir_path_safety(cuda_path):
+        return False
+
+    expected_files = ["nvcc", "cudart.dll", "cudart.so"]
+    has_expected_content = False
+    for expected_file in expected_files:
+        if os.path.exists(os.path.join(cuda_path, "bin", expected_file)):
+            has_expected_content = True
+            break
+
+    if not has_expected_content:
+        logger.warning(f"The directory {cuda_path} does not contain the typical file structure of CUDA")
+        return False
+
+    return True
+
+
+def check_cudnn_path_safety(cudnn_path):
+    """Check safety of cudnn path."""
+    if not _check_dir_path_safety(cudnn_path):
+        return False
+
+    expected_files = [
+        "include/cudnn.h",
+        "lib64/libcudnn.so",  # Linux
+        "lib/libcudnn.dylib",  # macOS
+        "lib/x64/cudnn.lib",  # Windows
+        "bin/cudnn64_7.dll"   # Windows
+    ]
+    found_files = []
+    for expected_file in expected_files:
+        full_path = os.path.join(cudnn_path, expected_file)
+        if os.path.exists(full_path):
+            found_files.append(expected_file)
+
+    if not found_files:
+        logger.warning(f"The directory {cudnn_path} does not contain the typical file structure of CUDNN")
+        return False
+
+    return True
+
+
 def _add_cuda_path():
     """add cuda path on windows."""
     if platform.system().lower() == 'windows':
@@ -500,15 +578,21 @@ def _add_cuda_path():
         if cuda_home is None:
             pass
         else:
+            if not check_cuda_path_safety(cuda_home):
+                logger.error(f"CUDA_PATH {cuda_home} is not safe, skip add cuda path.")
+                return
             cuda_bin_path = os.path.join(os.environ['CUDA_PATH'], 'bin')
             if sys.version_info >= (3, 8):
                 os.add_dll_directory(cuda_bin_path)
             else:
                 os.environ['PATH'] += os.pathsep + cuda_bin_path
-        cudann_home = os.environ.get('CUDNN_HOME')
-        if cudann_home is None:
+        cudnn_home = os.environ.get('CUDNN_HOME')
+        if cudnn_home is None:
             pass
         else:
+            if not check_cudnn_path_safety(cudnn_home):
+                logger.error(f"CUDNN_HOME {cuda_home} is not safe, skip add cudnn home.")
+                return
             cuda_home_bin_path = os.path.join(os.environ['CUDNN_HOME'], 'bin')
             if sys.version_info >= (3, 8):
                 os.add_dll_directory(cuda_home_bin_path)

mindspore/runtime/__init__.py

@@ -21,7 +21,8 @@ from mindspore.runtime.executor import launch_blocking, dispatch_threads_num, se
                                     set_kernel_launch_group, set_kernel_launch_capture
 from mindspore.runtime.memory import set_memory, memory_stats, memory_reserved, max_memory_reserved, empty_cache,\
                                  memory_replay, reset_peak_memory_stats, memory_summary, memory_allocated,\
-                                 max_memory_allocated, reset_max_memory_reserved, reset_max_memory_allocated
+                                 max_memory_allocated, reset_max_memory_reserved, reset_max_memory_allocated,\
+                                 PluggableAllocator, MemPool, use_mem_pool
 from mindspore.runtime.stream import Stream, synchronize, set_cur_stream, current_stream, \
     default_stream, communication_stream, StreamCtx
 from mindspore.runtime.event import Event
@@ -33,7 +34,7 @@ __all__ = [
     "Stream", "communication_stream", "synchronize", "set_cur_stream", "current_stream", "default_stream", "StreamCtx",
     "set_memory", "memory_stats", "memory_reserved", "max_memory_reserved", "empty_cache", "memory_replay",
     "reset_peak_memory_stats", "memory_summary", "memory_allocated", "max_memory_allocated",
-    "reset_max_memory_reserved", "reset_max_memory_allocated", "Event"
+    "reset_max_memory_reserved", "reset_max_memory_allocated", "Event", "PluggableAllocator", "MemPool", "use_mem_pool"
 ]
 
 __all__.sort()

mindspore/runtime/executor.py

@@ -196,7 +196,7 @@ def set_kernel_launch_group(thread_num=2, kernel_group_num=8):
 
 
 @args_type_check(enable_capture_graph=bool)
-def set_kernel_launch_capture(enable_capture_graph):
+def set_kernel_launch_capture(enable_capture_graph, op_capture_skip=None):
     """
     In O0/O1 mode, the incremental inference scenario supports graph capture.
     By capturing the CPU-side operator dispatch behavior into a graph,
@@ -208,12 +208,20 @@ def set_kernel_launch_capture(enable_capture_graph):
     Args:
         enable_capture_graph (bool): Whether to enable graph capture.
             It can be turned on or off at any position in the script.
+        op_capture_skip (list): Custom non-captured operator names. Default: ``None``.
 
     Examples:
         >>> import mindspore as ms
-        >>> ms.runtime.set_kernel_launch_capture(enable_capture_graph=True)
+        >>> op_capture_skip = ['matmul', 'addn']
+        >>> ms.runtime.set_kernel_launch_capture(True, op_capture_skip)
     """
     if RuntimeConf.get_instance().is_kernel_launch_group_configured():
         raise RuntimeError("The kernel launch group and kernel launch capture can not be set together")
 
-    return RuntimeConf.get_instance().set_kernel_launch_capture(enable_capture_graph)
+    if op_capture_skip is None:
+        op_capture_skip = []
+
+    if not isinstance(op_capture_skip, list):
+        raise TypeError("op_capture_skip must be a list")
+
+    return RuntimeConf.get_instance().set_kernel_launch_capture(enable_capture_graph, op_capture_skip)