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

msprobe/mindspore/debugger/precision_debugger.py

@@ -19,22 +19,33 @@ from collections import defaultdict, namedtuple
 import mindspore as ms
 from mindspore._c_expression import MSContext
 
-from msprobe.core.common.const import Const, FileCheckConst, MsgConst
-from msprobe.core.common.exceptions import MsprobeException
-from msprobe.core.common.file_utils import FileChecker
-from msprobe.core.common.utils import get_real_step_or_rank
+from msprobe.core.common.const import Const, MsgConst
+from msprobe.core.common.utils import check_token_range
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.debugger.precision_debugger import BasePrecisionDebugger
 from msprobe.mindspore.cell_processor import CellProcessor
 from msprobe.mindspore.common.const import Const as MsConst
-from msprobe.mindspore.common.utils import set_register_backward_hook_functions, check_save_param
+from msprobe.mindspore.common.utils import (
+    set_register_backward_hook_functions,
+    check_save_param,
+    is_graph_mode_cell_dump_allowed
+)
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.dump.hook_cell.api_registry import api_register
+from msprobe.mindspore.dump.graph_mode_cell_dump import GraphModeCellDump
+from msprobe.mindspore.dump.hook_cell.api_register import get_api_register
 from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.mindspore.grad_probe.grad_monitor import GradientMonitor
-from msprobe.mindspore.ms_config import parse_json_config
-from msprobe.mindspore.runtime import Runtime
-from msprobe.mindspore.service import Service
+from msprobe.mindspore.ms_config import parse_task_config
+from msprobe.mindspore.mindspore_service import MindsporeService
 from msprobe.mindspore.task_handler_factory import TaskHandlerFactory
 
+try:
+    from mindspore._c_expression import _dump_start, _dump_stop, _dump_step, _set_init_iter, _dump_set_dynamic
+except ImportError:
+    enable_dynamic_kbyk_dump = False
+else:
+    enable_dynamic_kbyk_dump = True
+
 try:
     from msprobe.lib import _msprobe_c
 except ImportError:
@@ -44,9 +55,7 @@ except ImportError:
 ConfigParameters = namedtuple("ConfigParameters", ["config_path", "task", "dump_path", "level"])
 
 
-class PrecisionDebugger:
-    _instance = None
-    task_not_need_service = [Const.GRAD_PROBE]
+class PrecisionDebugger(BasePrecisionDebugger):
 
     def __new__(cls, config_path=None, task=None, dump_path=None,
                 level=None, step=None, opt=None):
@@ -62,61 +71,33 @@ class PrecisionDebugger:
                  level=None, step=None):
         if self.initialized:
             return
-        self.initialized = True
-
         set_register_backward_hook_functions()
+        super().__init__(config_path, task, dump_path, level, step)
 
-        if not config_path:
-            config_path = os.path.join(os.path.dirname(__file__), "../../config.json")
-
-        config_params = ConfigParameters(config_path, task, dump_path, level)
-        self.check_input_params(config_params)
-
-        common_config, task_config = parse_json_config(config_path)
-        common_config.task = task if task else common_config.task
-        self.task = common_config.task
         if self.task == Const.GRAD_PROBE:
-            self.gm = GradientMonitor(common_config, task_config)
+            self.gm = GradientMonitor(self.common_config, self.task_config)
             return
-        common_config.step = get_real_step_or_rank(
-            step, Const.STEP) if step is not None else common_config.step
-        common_config.level = level if level else common_config.level
-        common_config.dump_path = dump_path if dump_path else common_config.dump_path
-        self.config = DebuggerConfig(common_config, task_config)
+        self.common_config.level = level if level else self.common_config.level
+        self.common_config.dump_path = dump_path if dump_path else self.common_config.dump_path
+        self.config = DebuggerConfig(self.common_config, self.task_config)
 
-        if _msprobe_c:
+        if self._need_msprobe_c() and _msprobe_c:
+            os.environ["MS_HOOK_ENABLE"] = "on"
             _msprobe_c._PrecisionDebugger(framework="MindSpore", config_path=config_path)
 
         self.config.execution_mode = self._get_execution_mode()
         if self._need_service():
             self.config.check_config_with_l2()
-            self.service = Service(self.config)
+            self.service = MindsporeService(self.config)
 
         Runtime.step_count = 0
         Runtime.is_running = False
+        if enable_dynamic_kbyk_dump:
+            _dump_set_dynamic()
 
     @staticmethod
-    def check_input_params(args):
-        if args.config_path is not None:
-            if not isinstance(args.config_path, str):
-                raise MsprobeException(
-                    MsprobeException.INVALID_PARAM_ERROR, f"config_path must be a string")
-            file_checker = FileChecker(
-                file_path=args.config_path, path_type=FileCheckConst.FILE, file_type=FileCheckConst.JSON_SUFFIX)
-            file_checker.common_check()
-
-        if args.task is not None and args.task not in Const.TASK_LIST:
-            raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"task must be one of {Const.TASK_LIST}")
-
-        if args.dump_path is not None:
-            if not isinstance(args.dump_path, str):
-                raise MsprobeException(
-                    MsprobeException.INVALID_PARAM_ERROR, f"dump_path must be a string")
-
-        if args.level is not None and args.level not in Const.LEVEL_LIST:
-            raise MsprobeException(
-                MsprobeException.INVALID_PARAM_ERROR, f"level must be one of {Const.LEVEL_LIST}")
+    def get_task_config(task, json_config):
+        return parse_task_config(task, json_config)
 
     @staticmethod
     def _get_execution_mode():
@@ -137,7 +118,7 @@ class PrecisionDebugger:
             return MsConst.PYNATIVE_MODE
 
     @staticmethod
-    def _is_graph_dump(config):
+    def _is_graph_dump(config: DebuggerConfig):
         if config.level != MsConst.KERNEL:
             return False
         if not config.list:
@@ -147,63 +128,68 @@ class PrecisionDebugger:
         return is_graph
 
     @classmethod
-    def start(cls, model=None):
-        instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if _msprobe_c:
-            _msprobe_c._PrecisionDebugger().start()
-        if instance.task in PrecisionDebugger.task_not_need_service:
+    def start(cls, model=None, token_range=None):
+        instance = cls.get_instance()
+        if instance is None:
             return
-
+        if cls._need_msprobe_c() and _msprobe_c:
+            _msprobe_c._PrecisionDebugger().start()
+        check_token_range(token_range)
         instance.config.execution_mode = cls._get_execution_mode()
         if cls._need_service():
             if not instance.service:
-                instance.service = Service(instance.config)
-            instance.service.start(model)
+                instance.service = MindsporeService(instance.config)
+            instance.config.check_model(model, token_range)
+            instance.service.start(model, token_range)
         else:
             if not instance.first_start:
-                api_register.api_set_ori_func()
-                handler = TaskHandlerFactory.create(instance.config)
+                get_api_register().restore_all_api()
+                handler = TaskHandlerFactory.create(instance.config, model)
                 handler.handle()
-
+                if enable_dynamic_kbyk_dump:
+                    _set_init_iter(0)
+            if enable_dynamic_kbyk_dump:
+                is_valid_rank = (not instance.config.rank or Runtime.rank_id in instance.config.rank)
+                is_valid_step = (not instance.config.step or Runtime.step_count in instance.config.step)
+                if is_valid_rank and is_valid_step:
+                    _dump_start()
+            Runtime.is_running = True
         instance.first_start = True
-        Runtime.is_running = True
-
-    @classmethod
-    def forward_backward_dump_end(cls):
-        instance = cls._instance
-        instance.stop()
 
     @classmethod
     def stop(cls):
-        instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if _msprobe_c:
-            _msprobe_c._PrecisionDebugger().stop()
+        instance = cls.get_instance()
+        if instance is None:
+            return
+
         if instance.task == Const.GRAD_PROBE:
             instance.gm.stop()
-        if instance.task in PrecisionDebugger.task_not_need_service:
-            return
         if instance.service:
             instance.service.stop()
-        Runtime.is_running = False
-
+        else:
+            Runtime.is_running = False
+        if enable_dynamic_kbyk_dump:
+            _dump_stop()
+        if cls._need_msprobe_c() and _msprobe_c:
+            _msprobe_c._PrecisionDebugger().stop()
+    
     @classmethod
     def step(cls):
-        instance = cls._instance
-        if not instance:
-            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
-        if _msprobe_c:
-            _msprobe_c._PrecisionDebugger().step()
-        if instance.task in PrecisionDebugger.task_not_need_service:
+        instance = cls.get_instance()
+        if instance is None:
             return
+
         if instance.service:
             instance.service.step()
+        if is_graph_mode_cell_dump_allowed(instance.config):
+            GraphModeCellDump.step()
+        if enable_dynamic_kbyk_dump:
+            _dump_step(1)
+        if cls._need_msprobe_c() and _msprobe_c:
+            _msprobe_c._PrecisionDebugger().step()
+
         HOOKCell.cell_count = defaultdict(int)
         CellProcessor.reset_cell_stats()
-
         Runtime.step_count += 1
 
     @classmethod
@@ -230,7 +216,7 @@ class PrecisionDebugger:
         instance.config.execution_mode = cls._get_execution_mode()
         if cls._need_service():
             if not instance.service:
-                instance.service = Service(instance.config)
+                instance.service = MindsporeService(instance.config)
             instance.service.save(variable, name, save_backward)
 
     @classmethod
@@ -241,4 +227,11 @@ class PrecisionDebugger:
         if instance.config.execution_mode != MsConst.PYNATIVE_MODE:
             return False
         else:
-            return instance.config.task != Const.FREE_BENCHMARK and not instance._is_graph_dump(instance.config)
+            return instance.config.task != Const.FREE_BENCHMARK and not instance._is_graph_dump(instance.config)
+
+    @classmethod
+    def _need_msprobe_c(cls):
+        instance = cls._instance
+        if not instance:
+            raise Exception(MsgConst.NOT_CREATED_INSTANCE)
+        return instance.config.level_ori == Const.LEVEL_L2