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

msprobe/pytorch/pt_config.py

@@ -1,12 +1,33 @@
-import json
+# 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_config import CommonConfig, BaseConfig
-from msprobe.core.common.file_utils import FileOpen
 from msprobe.core.common.const import Const
-from msprobe.pytorch.hook_module.utils import get_ops
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.file_utils import FileOpen, load_json
+from msprobe.core.common.log import logger
+from msprobe.core.common_config import BaseConfig, CommonConfig
 from msprobe.core.grad_probe.constant import level_adp
-from msprobe.core.grad_probe.utils import check_numeral_list_ascend
+from msprobe.core.grad_probe.utils import check_bounds
+from msprobe.pytorch.free_benchmark.common.enums import (
+    DeviceType,
+    HandlerType,
+    PytorchFreeBenchmarkConst,
+)
+from msprobe.pytorch.hook_module.utils import get_ops
 
 
 class TensorConfig(BaseConfig):
@@ -16,7 +37,7 @@ class TensorConfig(BaseConfig):
         self.nfs_path = json_config.get("nfs_path", "")
         self.host = json_config.get("host", "")
         self.port = json_config.get("port", -1)
-        self.tls_path = json_config.get("tls_path", "")
+        self.tls_path = json_config.get("tls_path", "./")
         self.check_config()
         self._check_file_format()
         self._check_tls_path_config()
@@ -26,13 +47,8 @@ class TensorConfig(BaseConfig):
             raise Exception("file_format is invalid")
 
     def _check_tls_path_config(self):
-        if self.tls_path:
-            if not os.path.exists(self.tls_path):
-                raise Exception("tls_path: %s does not exist" % self.tls_path)
-            if not os.path.exists(os.path.join(self.tls_path, "client.key")):
-                raise Exception("tls_path does not contain client.key")
-            if not os.path.exists(os.path.join(self.tls_path, "client.crt")):
-                raise Exception("tls_path does not contain client.crt")
+        if self.tls_path and not os.path.exists(self.tls_path):
+            raise Exception("tls_path: %s does not exist" % self.tls_path)
 
 
 class StatisticsConfig(BaseConfig):
@@ -61,23 +77,142 @@ class OverflowCheckConfig(BaseConfig):
 
 
 class FreeBenchmarkCheckConfig(BaseConfig):
+
     def __init__(self, json_config):
         super().__init__(json_config)
-        self.fuzz_device = json_config.get("fuzz_device")
-        self.pert_mode = json_config.get("pert_mode")
-        self.handler_type = json_config.get("handler_type")
-        self.fuzz_level = json_config.get("fuzz_level")
-        self.fuzz_stage = json_config.get("fuzz_stage")
-        self.if_preheat = json_config.get("if_preheat")
-        self.preheat_step = json_config.get("preheat_step")
-        self.max_sample = json_config.get("max_sample")
+        self.fuzz_device = json_config.get("fuzz_device", PytorchFreeBenchmarkConst.DEFAULT_DEVICE)
+        self.pert_mode = json_config.get("pert_mode", PytorchFreeBenchmarkConst.DEFAULT_MODE)
+        self.handler_type = json_config.get("handler_type", PytorchFreeBenchmarkConst.DEFAULT_HANDLER)
+        self.fuzz_level = json_config.get("fuzz_level", PytorchFreeBenchmarkConst.DEFAULT_FUZZ_LEVEL)
+        self.fuzz_stage = json_config.get("fuzz_stage", PytorchFreeBenchmarkConst.DEFAULT_FUZZ_STAGE)
+        self.if_preheat = json_config.get("if_preheat", False)
+        self.preheat_step = json_config.get("preheat_step", PytorchFreeBenchmarkConst.DEFAULT_PREHEAT_STEP)
+        self.max_sample = json_config.get("max_sample", PytorchFreeBenchmarkConst.DEFAULT_PREHEAT_STEP)
         self.check_freebenchmark_config()
 
     def check_freebenchmark_config(self):
-        if self.if_preheat and self.handler_type == "fix":
-            raise Exception("Preheating is not supported in fix handler type")
-        if self.preheat_step and self.preheat_step == 0:
-            raise Exception("preheat_step cannot be 0")
+        self._check_pert_mode()
+        self._check_fuzz_device()
+        self._check_handler_type()
+        self._check_fuzz_stage()
+        self._check_fuzz_level()
+        self._check_if_preheat()
+        if self.handler_type == HandlerType.FIX:
+            self._check_fix_config()
+        if self.if_preheat:
+            self._check_preheat_config()
+
+    def _check_pert_mode(self):
+        if self.pert_mode not in PytorchFreeBenchmarkConst.PERTURBATION_MODE_LIST:
+            msg = (
+                f"pert_mode is invalid, it should be one of"
+                f" {PytorchFreeBenchmarkConst.PERTURBATION_MODE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_fuzz_device(self):
+        if self.fuzz_device not in PytorchFreeBenchmarkConst.DEVICE_LIST:
+            msg = (
+                f"fuzz_device is invalid, it should be one of"
+                f" {PytorchFreeBenchmarkConst.DEVICE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if (self.fuzz_device == DeviceType.CPU) ^ (
+                self.pert_mode in PytorchFreeBenchmarkConst.CPU_MODE_LIST
+        ):
+            msg = (
+                f"You neet to and can only set fuzz_device as {DeviceType.CPU} "
+                f"when pert_mode in {PytorchFreeBenchmarkConst.CPU_MODE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_handler_type(self):
+        if self.handler_type not in PytorchFreeBenchmarkConst.HANDLER_LIST:
+            msg = (
+                f"handler_type is invalid, it should be one of"
+                f" {PytorchFreeBenchmarkConst.HANDLER_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_fuzz_stage(self):
+        if self.fuzz_stage not in PytorchFreeBenchmarkConst.FUZZ_STAGE_LIST:
+            msg = (
+                f"fuzz_stage is invalid, it should be one of"
+                f" {PytorchFreeBenchmarkConst.FUZZ_STAGE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_fuzz_level(self):
+        if self.fuzz_level not in PytorchFreeBenchmarkConst.FUZZ_LEVEL_LIST:
+            msg = (
+                f"fuzz_level is invalid, it should be one of"
+                f" {PytorchFreeBenchmarkConst.FUZZ_LEVEL_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_if_preheat(self):
+        if not isinstance(self.if_preheat, bool):
+            msg = "if_preheat is invalid, it should be a boolean"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_preheat_config(self):
+        if not isinstance(self.preheat_step, int):
+            msg = "preheat_step is invalid, it should be an integer"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if self.preheat_step <= 0:
+            msg = "preheat_step must be greater than 0"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if not isinstance(self.max_sample, int):
+            msg = "max_sample is invalid, it should be an integer"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if self.max_sample <= 0:
+            msg = "max_sample must be greater than 0"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+
+    def _check_fix_config(self):
+        if self.if_preheat:
+            msg = f"Preheating is not supported for {HandlerType.FIX} handler type"
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if self.fuzz_stage not in PytorchFreeBenchmarkConst.FIX_STAGE_LIST:
+            msg = (
+                f"The fuzz_stage when opening {HandlerType.FIX} handler must be one of "
+                f"{PytorchFreeBenchmarkConst.FIX_STAGE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
+        if self.pert_mode not in PytorchFreeBenchmarkConst.FIX_MODE_LIST:
+            msg = (
+                f"The pert_mode when opening {HandlerType.FIX} handler must be one of "
+                f"{PytorchFreeBenchmarkConst.FIX_MODE_LIST}"
+            )
+            logger.error_log_with_exp(
+                msg, MsprobeException(MsprobeException.INVALID_PARAM_ERROR, msg)
+            )
 
 
 class RunUTConfig(BaseConfig):
@@ -93,7 +228,7 @@ class RunUTConfig(BaseConfig):
         self.host = json_config.get("host", "")
         self.port = json_config.get("port", -1)
         self.rank_list = json_config.get("rank_list", Const.DEFAULT_LIST)
-        self.tls_path = json_config.get("tls_path", "")
+        self.tls_path = json_config.get("tls_path", "./")
         self.check_run_ut_config()
 
     @classmethod
@@ -118,13 +253,8 @@ class RunUTConfig(BaseConfig):
 
     @classmethod
     def check_tls_path_config(cls, tls_path):
-        if tls_path:
-            if not os.path.exists(tls_path):
-                raise Exception("tls_path: %s does not exist" % tls_path)
-            if not os.path.exists(os.path.join(tls_path, "server.key")):
-                raise Exception("tls_path does not contain server.key")
-            if not os.path.exists(os.path.join(tls_path, "server.crt")):
-                raise Exception("tls_path does not contain server.crt")
+        if tls_path and not os.path.exists(tls_path):
+            raise Exception("tls_path: %s does not exist" % tls_path)
 
     def check_run_ut_config(self):
         RunUTConfig.check_filter_list_config(Const.WHITE_LIST, self.white_list)
@@ -141,13 +271,13 @@ class GradToolConfig(BaseConfig):
         self.param_list = json_config.get("param_list", [])
         self.bounds = json_config.get("bounds", [-1, 0, 1])
         self._check_config()
-    
+
     def _check_config(self):
         if self.grad_level not in level_adp.keys():
             raise Exception(f"grad_level must be one of {level_adp.keys()}")
         if not isinstance(self.param_list, list):
             raise Exception(f"param_list must be a list")
-        check_numeral_list_ascend(self.bounds)
+        check_bounds(self.bounds)
 
 
 def parse_task_config(task, json_config):
@@ -178,10 +308,9 @@ def parse_json_config(json_file_path, task):
     if not json_file_path:
         config_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
         json_file_path = os.path.join(config_dir, "config.json")
-    with FileOpen(json_file_path, 'r') as file:
-        json_config = json.load(file)
+    json_config = load_json(json_file_path)
     common_config = CommonConfig(json_config)
-    if task and task in Const.TASK_LIST:
+    if task:
         task_config = parse_task_config(task, json_config)
     else:
         task_config = parse_task_config(common_config.task, json_config)

msprobe/pytorch/service.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 functools
 import os
 
@@ -6,6 +21,7 @@ import torch
 from msprobe.core.common.const import Const
 from msprobe.core.common.exceptions import DistributedNotInitializedError, MsprobeException
 from msprobe.core.common.file_utils import create_directory
+from msprobe.core.common.utils import print_tools_ends_info
 from msprobe.core.data_dump.data_collector import build_data_collector
 from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
 from msprobe.core.data_dump.scope import BaseScope
@@ -16,7 +32,10 @@ from msprobe.pytorch.hook_module.api_registry import api_register
 from msprobe.pytorch.hook_module.hook_module import HOOKModule
 from msprobe.pytorch.module_processer import ModuleProcesser
 from msprobe.pytorch.api_accuracy_checker.common.utils import ApiData
+
 torch_version_above_or_equal_2 = torch.__version__.split('+')[0] >= '2.0'
+if torch_version_above_or_equal_2:
+    from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.dump_dispatch import run_ut_dispatch
 
 HookFn = namedtuple('hookFn', ['pre_hook', 'forward_hook', 'backward_hook', 'forward_hook_torch_version_below_2'])
 
@@ -32,6 +51,7 @@ class Service:
         self.first_start = True
         self.current_rank = None
         self.dump_iter_dir = None
+        self.should_stop_service = False
         self.attl = None
 
     @staticmethod
@@ -39,14 +59,29 @@ class Service:
         logger.info_on_rank_0("Data needed ends here.")
         api_register.api_originality()
 
+    @staticmethod
+    def is_registered_backward_hook(module):
+        if hasattr(module, '_backward_hooks') and \
+                len(module._backward_hooks) > 0 and \
+                module._is_full_backward_hook is False:
+            return True
+        return False
+
+    def check_register_full_backward_hook(self, module):
+        if self.is_registered_backward_hook(module):
+            module._backward_hooks.clear()
+            module._is_full_backward_hook = None
+            logger.warning("Found deprecated backward hooks. Removing them and switching to full backward hooks.")
+
     def build_hook(self, module_type, name):
         def pre_hook(api_or_module_name, module, args, kwargs):
+            if not self.should_execute_hook():
+                return args, kwargs
+
             if module_type == BaseScope.Module_Type_Module:
                 api_or_module_name = module.mindstudio_reserved_name
             self.data_collector.update_api_or_module_name(api_or_module_name)
 
-            if not self.switch:
-                return args, kwargs
             if self.config.online_run_ut:
                 return None, None
             if self.data_collector:
@@ -55,13 +90,13 @@ class Service:
             return args, kwargs
 
         def forward_hook(api_or_module_name, module, args, kwargs, output):
+            if not self.should_execute_hook():
+                return None
+
             if module_type == BaseScope.Module_Type_Module:
                 api_or_module_name = module.mindstudio_reserved_name
             self.data_collector.update_api_or_module_name(api_or_module_name)
 
-            if not self.switch:
-                return None
-
             if self.config.online_run_ut:
                 if self.data_collector.scope and not self.data_collector.scope.check(api_or_module_name):
                     return None
@@ -80,18 +115,14 @@ class Service:
             return forward_hook(api_or_module_name, module, args, {}, output)
 
         def backward_hook(api_or_module_name, module, grad_input, grad_output):
+            if not self.should_execute_hook():
+                return
+
             if module_type == BaseScope.Module_Type_Module:
                 api_or_module_name = module.mindstudio_reserved_name
             self.data_collector.update_api_or_module_name(api_or_module_name)
 
-            if not self.switch:
-                return
-
             if self.config.online_run_ut:
-                if self.data_collector.scope and not self.data_collector.scope.check(api_or_module_name):
-                    return
-                api_data = ApiData(name[:-1], grad_input, {}, grad_output, self.current_iter, self.current_rank)
-                self.attl_send(api_data)
                 return
 
             if self.data_collector:
@@ -105,26 +136,15 @@ class Service:
         pre_forward_hook_fn = functools.partial(pre_hook, forward_name_template)
         forward_hook_fn = functools.partial(forward_hook, forward_name_template)
         backward_hook_fn = functools.partial(backward_hook, backward_name_template)
-        forward_hook_torch_version_below_2_fn = functools.partial(forward_hook_torch_version_below_2, forward_name_template)
+        forward_hook_torch_version_below_2_fn = functools.partial(forward_hook_torch_version_below_2,
+                                                                  forward_name_template)
         return HookFn(pre_forward_hook_fn, forward_hook_fn, backward_hook_fn, forward_hook_torch_version_below_2_fn)
 
-    def step(self):
-        self.current_iter += 1
-        self.data_collector.update_iter(self.current_iter)
-
-        ModuleProcesser.reset_module_stats()
-        HOOKModule.reset_module_stats()
-
     def start(self, model, api_origin=False):
-        self.model = model
-        if self.config.step and self.current_iter > max(self.config.step):
-            if self.config.online_run_ut:
-                # send stop signal if online_run_ut
-                self.attl_stop()
-            self.stop()
-            raise Exception("msprobe: exit after iteration {}".format(max(self.config.step)))
-        if self.config.step and self.current_iter not in self.config.step:
+        if self.need_stop_service():
             return
+
+        self.model = model
         if self.first_start:
             try:
                 self.current_rank = get_rank_if_initialized()
@@ -138,6 +158,8 @@ class Service:
             self.first_start = False
         if api_origin:
             api_register.api_modularity()
+        if self.config.online_run_ut and torch_version_above_or_equal_2:
+            run_ut_dispatch(self.attl, True)
         self.switch = True
         logger.info_on_rank_0(f"Dump switch is turned on at step {self.current_iter}. ")
         if self.config.level != "L2" and not self.config.online_run_ut:
@@ -145,6 +167,8 @@ class Service:
             logger.info_on_rank_0(f"Dump data will be saved in {self.dump_iter_dir}.")
 
     def stop(self):
+        if self.should_stop_service:
+            return
         if self.config.level == "L2":
             return
         if self.config.step and self.current_iter not in self.config.step:
@@ -152,10 +176,47 @@ class Service:
         if self.config.rank and self.current_rank not in self.config.rank:
             return
         self.switch = False
-        if self.config.online_run_ut:
+        if self.config.online_run_ut and torch_version_above_or_equal_2:
+            run_ut_dispatch(self.attl, False)
             return
         self.data_collector.write_json()
 
+    def step(self):
+        if self.should_stop_service:
+            return
+        self.current_iter += 1
+        self.data_collector.update_iter(self.current_iter)
+
+        ModuleProcesser.reset_module_stats()
+        HOOKModule.reset_module_stats()
+        self.data_collector.data_writer.reset_cache()
+
+    def need_stop_service(self):
+        if self.should_stop_service:
+            return True
+        end_service = self.config.step and self.current_iter > max(self.config.step) or \
+                      self.data_collector and self.data_collector.data_processor.is_terminated
+        if end_service:
+            if self.config.online_run_ut:
+                # send stop signal if online_run_ut
+                self.attl_stop()
+            if self.config.level in [Const.LEVEL_L1, Const.LEVEL_L2, Const.LEVEL_MIX]:
+                api_register.api_originality()
+            self.switch = False
+            self.should_stop_service = True
+            print_tools_ends_info()
+            return True
+        if self.config.step and self.current_iter not in self.config.step:
+            return True
+        return False
+
+    def should_execute_hook(self):
+        if not self.switch:
+            return False
+        if self.data_collector and self.data_collector.data_processor.is_terminated:
+            return False
+        return True
+
     def create_dirs(self):
         create_directory(self.config.dump_path)
         self.dump_iter_dir = os.path.join(self.config.dump_path, f"step{self.current_iter}")
@@ -187,14 +248,16 @@ class Service:
                 prefix = BaseScope.Module_Type_Module + Const.SEP + name + Const.SEP + \
                          module.__class__.__name__ + Const.SEP
 
-                pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 \
-                    = self.build_hook(BaseScope.Module_Type_Module, prefix)
+                pre_forward_hook, forward_hook, backward_hook, forward_hook_torch_version_below_2 = self.build_hook(
+                    BaseScope.Module_Type_Module, prefix)
                 if torch_version_above_or_equal_2:
                     module.register_forward_hook(forward_hook, with_kwargs=True)
                 else:
+                    self.check_register_full_backward_hook(module)
                     module.register_full_backward_hook(
                         self.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
                     module.register_forward_hook(forward_hook_torch_version_below_2)
+                self.check_register_full_backward_hook(module)
                 module.register_full_backward_hook(backward_hook)
 
                 module.register_forward_pre_hook(
@@ -204,11 +267,13 @@ class Service:
                 if torch_version_above_or_equal_2:
                     module.register_full_backward_pre_hook(
                         self.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
+                    self.check_register_full_backward_hook(module)
                     module.register_full_backward_hook(
                         self.module_processor.node_hook(prefix + Const.BACKWARD, Const.STOP))
 
         if self.config.level in ["mix", "L1", "L2"]:
-            api_register.initialize_hook(functools.partial(self.build_hook, BaseScope.Module_Type_API))
+            api_register.initialize_hook(functools.partial(self.build_hook, BaseScope.Module_Type_API),
+                                         self.config.online_run_ut)
             api_register.api_modularity()
 
         if Const.STATISTICS == self.config.task or Const.TENSOR == self.config.task: