mindstudio-probe 1.3.0__py3-none-any.whl → 8.1.1__py3-none-any.whl

msprobe/mindspore/dump/jit_dump.py

@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 from collections import defaultdict
+import os
+import types
 
 import mindspore
+from mindspore import nn
 from mindspore._c_expression import PyNativeExecutor_
 try:
     from mindspore.common.api import _MindsporeFunctionExecutor
@@ -25,7 +27,9 @@ except ImportError:
 
 from msprobe.core.common.log import logger
 from msprobe.core.common.const import Const
+from msprobe.core.common.runtime import Runtime
 from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
+from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.mindspore.dump.hook_cell.api_register import get_api_register
 
 
@@ -34,24 +38,20 @@ _api_register = get_api_register()
 
 def dump_jit(name, in_feat, out_feat, is_forward):
     pid = os.getpid()
-    ori_args = str(name)
-    index = ori_args.find("<")
-    if index != 0 and index != -1:
-        result = ori_args[0:index]
-    elif name is not None and "<" not in str(name):
-        result = str(name)
-    else:
-        result = "JitFunction"
+    name = name if name else "JitFunction"
     if JitDump.need_dump():
         if is_forward:
-            JitDump.jit_count[result] += 1
-            name_template = (Const.JIT + Const.SEP + result + Const.SEP +
-                             str(JitDump.jit_count[result]) + Const.SEP + Const.FORWARD)
+            if name in JitDump.jit_count:
+                JitDump.jit_count[name] += 1
+            else:
+                JitDump.jit_count[name] = 0
+            name_template = (Const.JIT + Const.SEP + name + Const.SEP +
+                             str(JitDump.jit_count[name]) + Const.SEP + Const.FORWARD)
             JitDump.data_collector.update_api_or_module_name(name_template)
             module_input_output = ModuleForwardInputsOutputs(args=in_feat, kwargs={}, output=out_feat)
             JitDump.data_collector.forward_data_collect(name_template, None, pid, module_input_output)
         else:
-            name_template = Const.JIT + Const.SEP + result + Const.SEP + str(JitDump.jit_count[result]) + Const.SEP + \
+            name_template = Const.JIT + Const.SEP + name + Const.SEP + str(JitDump.jit_count[name]) + Const.SEP + \
                             Const.BACKWARD
             JitDump.data_collector.update_api_or_module_name(name_template)
             module_input_output = ModuleBackwardInputsOutputs(grad_input=in_feat, grad_output=out_feat)
@@ -74,11 +74,11 @@ class JitDump(_MindsporeFunctionExecutor):
     def __call__(self, *args, **kwargs):
         _api_register.restore_all_api()
         out = super().__call__(*args, **kwargs)
-        if JitDump.jit_dump_switch and len(args) > 0:
-            if self.name and self.name != "construct":
+        if JitDump.jit_dump_switch and len(args) > 0 and self.name:
+            if self.name != "construct":
                 dump_jit(self.name, args, out, True)
-            else:
-                dump_jit(args[0], args, out, True)
+            elif Runtime.run_mode != MsConst.PYNATIVE_GRAPH_MODE and isinstance(args[0], nn.Cell):
+                dump_jit(args[0].__class__.__name__, args, out, True)
             JitDump.jit_enable = True
         elif len(args) == 0:
             logger.warning(f"The jit function {self.name} has no input arguments, nothing will be dumped.")
@@ -109,6 +109,9 @@ class JitDump(_MindsporeFunctionExecutor):
         else:
             output = self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
         if JitDump.jit_dump_switch and JitDump.jit_enable:
-            dump_jit(obj, args, None, False)
+            if isinstance(obj, types.FunctionType):
+                dump_jit(obj.__name__, args, None, False)
+            elif Runtime.run_mode != MsConst.PYNATIVE_GRAPH_MODE and isinstance(obj, nn.Cell):
+                dump_jit(obj.__class__.__name__, args, None, False)
             _api_register.register_all_api()
         return output

msprobe/mindspore/dump/kernel_kbyk_dump.py

@@ -39,9 +39,12 @@ class KernelKbykDump:
         common_set["input_output"] = 0
         common_set["kernels"] = []
         common_set["support_device"] = [0, 1, 2, 3, 4, 5, 6, 7]
-        e2e_set = dict()
-        e2e_set["enable"] = True
-        e2e_set["trans_flag"] = True
+        e2e_set = {
+                    "enable": not config.async_dump,
+                    "trans_flag": True,
+                    "stat_calc_mode": config.stat_cal_mode,
+                    "device_stat_precision_mode": config.device_stat_precision_mode,
+                  }
 
         if config.list:
             common_set["dump_mode"] = 1

msprobe/mindspore/dym_loader/hook_dynamic_loader.cpp

@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2024-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.
+ */
+
+#include "hook_dynamic_loader.h"
+#include <sys/stat.h>
+#include <cstdlib>
+#include <cstring>
+#include <pybind11/embed.h>
+#include "utils/log_adapter.h"
+
+namespace py = pybind11;
+
+HookDynamicLoader &HookDynamicLoader::GetInstance() 
+{
+    static HookDynamicLoader instance;
+    return instance;
+}
+
+bool HookDynamicLoader::LoadFunction(void *handle, const std::string &functionName) {
+    void *func = dlsym(handle, functionName.c_str());
+    if (!func) {
+        MS_LOG(WARNING) << "Could not load function: " << functionName << ", error: " << dlerror();
+        return false;
+    }
+    funcMap_[functionName] = func;
+    return true;
+}
+
+bool HookDynamicLoader::LoadLibrary()
+{
+    std::string msprobePath = "";
+    // 获取gil锁
+    py::gil_scoped_acquire acquire;
+    try {
+        py::module msprobeMod = py::module::import("msprobe.lib._msprobe_c");
+        if (!py::hasattr(msprobeMod, "__file__")) {
+            MS_LOG(WARNING) << "Adump mod not found";
+            return false;
+    }
+    msprobePath = msprobeMod.attr("__file__").cast<std::string>();
+    } catch (const std::exception& e) {
+    MS_LOG(WARNING) << "Adump mod path unable to get: " << e.what();
+    return false;
+    }
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (handle_) {
+        MS_LOG(WARNING) << "Hook library already loaded!";
+        return false;
+    }
+    if (msprobePath == "") {
+        MS_LOG(WARNING) << "Adump path not loaded";
+        return false;
+    }
+    handle_ = dlopen(msprobePath.c_str(), RTLD_LAZY | RTLD_LOCAL);
+    if (!handle_) {
+        MS_LOG(WARNING) << "Failed to load Hook library: " << dlerror();
+        return false;
+    }
+
+    for (const auto &functionName : functionList_) {
+        if (!LoadFunction(handle_, functionName)) {
+            MS_LOG(WARNING) << "Failed to load adump function";
+            dlclose(handle_);
+            handle_ = nullptr;
+            return false;
+        }
+    }
+
+    MS_LOG(INFO) << "Hook library loaded successfully.";
+    return true;
+}
+
+bool HookDynamicLoader::UnloadLibrary() 
+{
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (!handle_) {
+        MS_LOG(WARNING) << "Hook library hasn't been loaded.";
+        return false;
+    }
+
+    dlclose(handle_);
+    handle_ = nullptr;
+    funcMap_.clear();
+    MS_LOG(INFO) << "Library unloaded successfully.";
+    return true;
+}
+
+void *HookDynamicLoader::GetHooker(const std::string &funcName)
+{
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto iter = funcMap_.find(funcName);
+    if (iter == funcMap_.end()) {
+      MS_LOG(WARNING) << "Function not found: " << funcName;
+      return nullptr;
+    }
+    return iter->second;
+}

msprobe/mindspore/dym_loader/hook_dynamic_loader.h

@@ -27,26 +27,26 @@ constexpr auto kHookBegin = "MS_DbgOnStepBegin";
 constexpr auto kHookEnd = "MS_DbgOnStepEnd";
 
 class HookDynamicLoader {
- public:
-  static HookDynamicLoader &GetInstance();
+public:
+    static HookDynamicLoader &GetInstance();
 
-  HookDynamicLoader(const HookDynamicLoader &) = delete;
-  HookDynamicLoader &operator=(const HookDynamicLoader &) = delete;
+    HookDynamicLoader(const HookDynamicLoader &) = delete;
+    HookDynamicLoader &operator=(const HookDynamicLoader &) = delete;
 
-  bool LoadLibrary();
-  bool UnloadLibrary();
-  void *GetHooker(const std::string &funcName);
+    bool LoadLibrary();
+    bool UnloadLibrary();
+    void *GetHooker(const std::string &funcName);
 
- private:
-  // Helper functions
-  bool loadFunction(void *handle, const std::string &functionName);
+private:
+    // Helper functions
+    bool LoadFunction(void *handle, const std::string &functionName);
 
-  HookDynamicLoader() = default;
+    HookDynamicLoader() = default;
 
-  void *handle_ = nullptr;
-  std::vector<std::string> functionList_ = {kHookBegin, kHookEnd};
-  std::map<std::string, void *> funcMap_;
-  std::mutex mutex_;
+    void *handle_ = nullptr;
+    std::vector<std::string> functionList_ = {kHookBegin, kHookEnd};
+    std::map<std::string, void *> funcMap_;
+    std::mutex mutex_;
 };
 
 #endif  // HOOK_DYNAMIC_LOADER_H

msprobe/mindspore/free_benchmark/api_pynative_self_check.py

@@ -23,6 +23,8 @@ import mindspore as ms
 from msprobe.core.common.const import Const
 from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import check_path_length, load_yaml
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.hook_manager import HookSet
 from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.common.log import logger
@@ -35,7 +37,6 @@ from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
 from msprobe.mindspore.free_benchmark.common.utils import Tools
 from msprobe.mindspore.free_benchmark.handler.handler_factory import HandlerFactory
 from msprobe.mindspore.free_benchmark.perturbation.perturbation_factory import PerturbationFactory
-from msprobe.mindspore.runtime import Runtime
 
 
 _api_register = get_api_register()
@@ -75,7 +76,7 @@ class ApiPyNativeSelfCheck:
             ret = None
 
             if not need_wrapper_func():
-                del cell.input_kwargs
+                del cell.msprobe_input_kwargs
                 return ret
 
             api_name_with_id = api_name_with_id[:-1]
@@ -84,9 +85,9 @@ class ApiPyNativeSelfCheck:
                         api_name_with_id[api_name_with_id.find(Const.SEP) + 1:api_name_with_id.rfind(Const.SEP)])
             if api_name in self.api_list:
                 ret = check_self(api_name_with_id, output_data, self.ori_func.get(api_name),
-                                 *input_data, **cell.input_kwargs)
+                                 *input_data, **cell.msprobe_input_kwargs)
 
-            del cell.input_kwargs
+            del cell.msprobe_input_kwargs
             return ret
 
         def backward_hook(cell, grad_input, grad_output):
@@ -105,8 +106,13 @@ class ApiPyNativeSelfCheck:
 
         def pre_backward_hook(cell, grad_input):
             return None
-
-        return pre_hook, wrap_forward_hook, wrap_backward_hook, pre_backward_hook
+        
+        return HookSet(
+            forward_hook=wrap_forward_hook,
+            forward_pre_hook=pre_hook,
+            backward_hook=wrap_backward_hook,
+            backward_pre_hook=pre_backward_hook
+        )
 
     def store_original_func(self):
         for api_name in self.api_list:

msprobe/mindspore/free_benchmark/common/utils.py

@@ -19,10 +19,10 @@ from typing import Any, Optional
 import mindspore as ms
 from mindspore import Tensor, ops
 
+from msprobe.core.common.runtime import Runtime
 from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
-from msprobe.mindspore.runtime import Runtime
 
 
 class Tools:

msprobe/mindspore/grad_probe/global_context.py

@@ -41,8 +41,12 @@ class GlobalContext:
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
             cls._instance_lock.acquire()
-            cls._instance = object.__new__(cls)
-            cls._instance_lock.release()
+            try:
+                cls._instance = object.__new__(cls)
+            except Exception as e:
+                raise RuntimeError("grad_probe global context init failed") from e
+            finally:
+                cls._instance_lock.release()
         return cls._instance
 
     def init_context(self, config_dict: Dict):
@@ -69,6 +73,7 @@ class GlobalContext:
             create_directory(self._setting.get(GradConst.OUTPUT_PATH))
         else:
             logger.warning("The output_path exists, the data will be covered.")
+
         self._setting[GradConst.TIME_STAMP] = str(int(time.time()))
 
     def get_context(self, key: str):

msprobe/mindspore/grad_probe/grad_stat_csv.py

@@ -15,6 +15,7 @@
 
 import hashlib
 from abc import ABC, abstractmethod
+import zlib
 
 import mindspore
 from mindspore import ops
@@ -76,8 +77,8 @@ class CsvMd5(CsvItem):
     def generate_csv_content(csv_input):
         grad = csv_input.grad
         tensor_bytes = grad.float().numpy().tobytes()
-        md5_hash = hashlib.md5(tensor_bytes)
-        return [md5_hash.hexdigest()]
+        md5_hash = f"{zlib.crc32(tensor_bytes):08x}"
+        return [md5_hash]
 
 
 @register_csv_item(GradConst.DISTRIBUTION)

msprobe/mindspore/mindspore_service.py

@@ -0,0 +1,114 @@
+# 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.
+
+from collections import defaultdict
+import mindspore as ms
+from mindspore.ops.primitive import Primitive
+
+from msprobe.core.common.utils import Const
+from msprobe.core.service import BaseService
+from msprobe.mindspore.cell_processor import CellProcessor
+from msprobe.mindspore.common.log import logger
+from msprobe.mindspore.common.utils import (
+    get_rank_if_initialized,
+    is_mindtorch,
+    get_cells_and_names_with_index
+)
+from msprobe.mindspore.dump.hook_cell.api_register import get_api_register, ApiTemplate
+from msprobe.mindspore.dump.hook_cell.ms_hook_manager import MindsproeHookManager
+from msprobe.mindspore.dump.hook_cell.primitive_hooks import PrimitiveHookService
+from msprobe.mindspore.dump.jit_dump import JitDump
+
+try:
+    from mindspore.common._pijit_context import PIJitCaptureContext
+except ImportError:
+    pijit_label = False
+else:
+    pijit_label = True
+
+
+class MindsporeService(BaseService):    
+    @property
+    def _get_framework_type(self):
+        return Const.MT_FRAMEWORK if is_mindtorch() else Const.MS_FRAMEWORK
+    
+    @staticmethod
+    def _get_current_rank():
+        return get_rank_if_initialized()
+
+    def empty(self, *args, **kwargs):
+        pass
+
+    def reset_status(self):
+        self._reset_status()
+
+    def _init_specific_components(self):
+        self.logger = logger
+        self.api_register = get_api_register()
+        self.primitive_hook_service = PrimitiveHookService(self)
+        self.cell_processor = CellProcessor(self.data_collector.scope)
+        self.hook_manager = MindsproeHookManager(self.data_collector, self.config)
+        self._setup_jit_context()
+        self.api_template = ApiTemplate
+
+    def _setup_jit_context(self):
+        if self.config.level in [Const.LEVEL_MIX, Const.LEVEL_L1]:
+            JitDump.set_config(self.config)
+            JitDump.set_data_collector(self.data_collector)
+            if hasattr(ms.common.api, "_MindsporeFunctionExecutor"):
+                ms.common.api._MindsporeFunctionExecutor = JitDump
+            else:
+                ms.common.api._JitExecutor = JitDump
+            ms.common.api._PyNativeExecutor.grad = JitDump.grad
+            if pijit_label:
+                PIJitCaptureContext.__enter__ = self.empty
+                PIJitCaptureContext.__exit__ = self.empty
+
+    def _register_module_hook(self):
+        self.cell_processor.register_cell_hook(self.model, self.build_hook, self.config)
+        self.logger.info_on_rank_0(f"The module {self.config.task} hook function is successfully mounted to the model.")
+
+    def _register_hook(self):
+        self._register_primitive_hook()
+
+    def _register_primitive_hook(self):
+        if self.config.level not in [Const.LEVEL_MIX, Const.LEVEL_L1]:
+            return
+        if not self.model or self.config.task not in Const.DUMP_DATA_COLLECTION_LIST:
+            return
+
+        primitive_set = set()
+        cells_and_names_with_index, _ = get_cells_and_names_with_index(self.model)
+        for cells_and_names in cells_and_names_with_index.values():
+            for _, cell in cells_and_names:
+                for attribute, value in vars(cell).items():
+                    if isinstance(value, Primitive):
+                        primitive_set.add((attribute, value))
+
+        for pname, primitive in primitive_set:
+            primitive_class_name = primitive.__class__.__name__
+            primitive_combined_name = pname + Const.SEP + primitive_class_name
+            new_primitive = type('NewPrimitive', (primitive.__class__,),
+                                    {'__call__': self.primitive_hook_service.wrap_primitive(primitive.__call__,
+                                                                                            primitive_combined_name)})
+            primitive.__class__ = new_primitive
+
+    def _reset_status(self):
+        super()._reset_status()
+        self.primitive_hook_service.primitive_counters.clear()
+        JitDump.jit_count = defaultdict(int)
+
+    def _change_jit_switch(self, status):
+        JitDump.jit_dump_switch = status

msprobe/mindspore/monitor/common_func.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2024-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.
+
+
+from mindspore import nn
+from mindspore import communication
+from msprobe.mindspore.monitor.utils import logger
+from msprobe.mindspore.common.utils import is_mindtorch
+if is_mindtorch():
+    import torch
+
+
+def is_valid_instance(model):
+    return isinstance(model, torch.nn.Module) if is_mindtorch() else isinstance(model, nn.Cell)
+
+
+def get_submodules(model):
+    if not is_valid_instance(model):
+        logger.info("Counter invalid model, nothing to hook")
+        return {}
+    return model.named_modules() if is_mindtorch() else model.cells_and_names()
+
+
+def get_parameters(model):
+    if not is_valid_instance(model):
+        return {}
+    if is_mindtorch():
+        return model.named_parameters()
+    else:
+        return model.parameters_and_names()
+
+
+def get_rank():
+    if comm_is_initialized():
+        return communication.get_rank()
+    return 0
+
+
+def comm_is_initialized():
+    return communication.GlobalComm.INITED