mindstudio-probe 1.0.1__py3-none-any.whl

msprobe/pytorch/pt_config.py

@@ -0,0 +1,93 @@
+import json
+import os
+
+from msprobe.core.common_config import CommonConfig, BaseConfig
+from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.const import Const
+
+
+class TensorConfig(BaseConfig):
+    def __init__(self, json_config):
+        super().__init__(json_config)
+        self.check_config()
+        self._check_file_format()
+
+    def _check_file_format(self):
+        if self.file_format is not None and self.file_format not in ["npy", "bin"]:
+            raise Exception("file_format is invalid")
+
+
+class StatisticsConfig(BaseConfig):
+    def __init__(self, json_config):
+        super().__init__(json_config)
+        self.check_config()
+        self._check_summary_mode()
+
+    def _check_summary_mode(self):
+        if self.summary_mode and self.summary_mode not in ["statistics", "md5"]:
+            raise Exception("summary_mode is invalid")
+
+
+class OverflowCheckConfig(BaseConfig):
+    def __init__(self, json_config):
+        super().__init__(json_config)
+        self.overflow_num = json_config.get("overflow_nums")
+        self.check_mode = json_config.get("check_mode")
+        self.check_overflow_config()
+
+    def check_overflow_config(self):
+        if self.overflow_num is not None and not isinstance(self.overflow_num, int):
+            raise Exception("overflow_num is invalid")
+        if self.check_mode is not None and self.check_mode not in ["all", "aicore", "atomic"]:
+            raise Exception("check_mode is invalid")
+
+
+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.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")
+
+def parse_task_config(task, json_config):
+    default_dic = {}
+    if task == Const.TENSOR:
+        config_dic = json_config.get(Const.TENSOR) if json_config.get(Const.TENSOR) else default_dic
+        return TensorConfig(config_dic)
+    elif task == Const.STATISTICS:
+        config_dic = json_config.get(Const.STATISTICS) if json_config.get(Const.STATISTICS) else default_dic
+        return StatisticsConfig(config_dic)
+    elif task == Const.OVERFLOW_CHECK:
+        config_dic = json_config.get(Const.OVERFLOW_CHECK) if json_config.get(Const.OVERFLOW_CHECK) else default_dic
+        return OverflowCheckConfig(config_dic)
+    elif task == Const.FREE_BENCHMARK:
+        config_dic = json_config.get(Const.FREE_BENCHMARK) if json_config.get(Const.FREE_BENCHMARK) else default_dic
+        return FreeBenchmarkCheckConfig(config_dic)
+    else:
+        return StatisticsConfig(default_dic)
+
+
+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(os.path.join(config_dir, "config"), "config.json")
+    with FileOpen(json_file_path, 'r') as file:
+        json_config = json.load(file)
+    common_config = CommonConfig(json_config)
+    if task and task in Const.TASK_LIST:
+        task_config = parse_task_config(task, json_config)
+    else:
+        task_config = parse_task_config(common_config.task, json_config)
+    return common_config, task_config

msprobe/pytorch/service.py

@@ -0,0 +1,167 @@
+import functools
+import os
+from pathlib import Path
+
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.file_check import FileChecker, check_path_before_create
+from msprobe.core.common.const import Const, FileCheckConst
+from msprobe.core.common.exceptions import DistributedNotInitializedError, MsaccException
+from msprobe.core.data_dump.data_collector import build_data_collector
+from msprobe.core.data_dump.scope import BaseScope
+from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
+from msprobe.pytorch.common.utils import get_rank_if_initialized
+from msprobe.pytorch.module_processer import ModuleProcesser
+from msprobe.pytorch.hook_module import remove_dropout
+from msprobe.pytorch.hook_module.api_registry import api_register
+
+
+class Service:
+    def __init__(self, config):
+        self.model = None
+        self.config = config
+        self.data_collector = build_data_collector(config)
+        self.module_processor = ModuleProcesser(self.data_collector.scope)
+        self.switch = False
+        self.current_iter = 0
+        self.first_start = True
+        self.current_rank = None
+        self.dump_iter_dir = None
+
+    def build_hook(self, module_type, name):
+        def pre_hook(api_or_module_name, module, args, kwargs):
+            if module_type == BaseScope.Module_Type_Module:
+                api_or_module_name = module.mindstudio_reserved_name
+            self.data_collector.visit_and_clear_overflow_status(api_or_module_name)
+
+            if not self.switch:
+                return args, kwargs
+            if self.data_collector:
+                module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+                self.data_collector.pre_forward_data_collect(api_or_module_name, module, pid, module_input_output)
+            return args, kwargs
+
+        def forward_hook(api_or_module_name, module, args, kwargs, output):
+            if module_type == BaseScope.Module_Type_Module:
+                api_or_module_name = module.mindstudio_reserved_name
+            self.data_collector.visit_and_clear_overflow_status(api_or_module_name)
+
+            if not self.switch:
+                return None
+            if self.data_collector:
+                module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+                self.data_collector.forward_data_collect(api_or_module_name, module, pid, module_input_output)
+                if self.data_collector.if_return_forward_new_output():
+                    return self.data_collector.get_forward_new_output()
+            return output
+
+        def backward_hook(api_or_module_name, module, grad_input, grad_output):
+            if module_type == BaseScope.Module_Type_Module:
+                api_or_module_name = module.mindstudio_reserved_name
+            self.data_collector.visit_and_clear_overflow_status(api_or_module_name)
+
+            if not self.switch:
+                return
+            if self.data_collector:
+                module_input_output = ModuleBackwardInputsOutputs(grad_input=grad_input, grad_output=grad_output)
+                self.data_collector.backward_data_collect(api_or_module_name, module, pid, module_input_output)
+
+        pid = os.getpid()
+        forward_name_template = name + Const.FORWARD
+        backward_name_template = name + Const.BACKWARD
+        pre_forward_hook = functools.partial(pre_hook, forward_name_template)
+        forward_hook = functools.partial(forward_hook, forward_name_template)
+        backward_hook = functools.partial(backward_hook, backward_name_template)
+        return pre_forward_hook, forward_hook, backward_hook
+
+    def step(self):
+        self.current_iter += 1
+        self.data_collector.update_iter(self.current_iter)
+
+    def start(self, model):
+        self.model = model
+        if self.config.step and self.current_iter > max(self.config.step):
+            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:
+            return
+        if self.first_start:
+            try:
+                self.current_rank = get_rank_if_initialized()
+            except DistributedNotInitializedError:
+                self.current_rank = None
+
+            if self.config.rank and self.current_rank not in self.config.rank:
+                return
+            self.register_hook_new()
+            self.first_start = False
+        self.switch = True
+        logger.info_on_rank_0(f"Dump switch is turned on at step {self.current_iter}. ")
+        if self.config.level != "L2":
+            self.create_dirs()
+            logger.info_on_rank_0(f"Dump data will be saved in {self.dump_iter_dir}.")
+
+    def stop(self):
+        if self.config.level == "L2":
+            return
+        if self.config.step and self.current_iter not in self.config.step:
+            return
+        if self.config.rank and self.current_rank not in self.config.rank:
+            return
+        self.switch = False
+        self.data_collector.write_json()
+
+    def create_dirs(self):
+        check_path_before_create(self.config.dump_path)
+        if not os.path.exists(self.config.dump_path):
+            Path(self.config.dump_path).mkdir(mode=0o750, exist_ok=True)
+        file_check = FileChecker(self.config.dump_path, FileCheckConst.DIR)
+        file_check.common_check()
+        self.dump_iter_dir = os.path.join(self.config.dump_path, f"step{self.current_iter}")
+        cur_rank = self.current_rank if self.current_rank is not None else ''
+        dump_dir = os.path.join(self.dump_iter_dir, f"rank{cur_rank}")
+        if not os.path.exists(dump_dir):
+            Path(dump_dir).mkdir(mode=0o750, parents=True, exist_ok=True)
+        if self.config.task in self.data_collector.tasks_need_tensor_data:
+            dump_data_dir = os.path.join(dump_dir, "dump_tensor_data")
+            Path(dump_data_dir).mkdir(mode=0o750, exist_ok=True)
+        else:
+            dump_data_dir = None
+
+        dump_file_path = os.path.join(dump_dir, "dump.json")
+        stack_file_path = os.path.join(dump_dir, "stack.json")
+        construct_file_path = os.path.join(dump_dir, "construct.json")
+        free_benchmark_file_path = os.path.join(self.config.dump_path, "free_benchmark.csv")
+        self.data_collector.update_dump_paths(
+            dump_file_path, stack_file_path, construct_file_path, dump_data_dir, free_benchmark_file_path)
+
+    def register_hook_new(self):
+        logger.info_on_rank_0("The {} hook function is successfully mounted to the model.".format(self.config.task))
+        if self.config.level in ["L0", "mix"]:
+            if self.model is None:
+                logger.error_log_with_exp("The model is None.", MsaccException.INVALID_PARAM_ERROR)
+            logger.info_on_rank_0("The init dump mode is enabled, and the module dump function will not be available")
+            for name, module in self.model.named_modules():
+                if module == self.model:
+                    continue
+                prefix = BaseScope.Module_Type_Module + Const.SEP + name + Const.SEP + \
+                         module.__class__.__name__ + Const.SEP
+
+                pre_forward_hook, forward_hook, backward_hook = self.build_hook(BaseScope.Module_Type_Module, prefix)
+                module.register_forward_hook(forward_hook, with_kwargs=True)
+                module.register_full_backward_hook(backward_hook)
+
+                module.register_forward_pre_hook(
+                    self.module_processor.node_hook(prefix + Const.FORWARD, Const.START))
+                module.register_forward_hook(
+                    self.module_processor.node_hook(prefix + Const.FORWARD, Const.STOP))
+                module.register_full_backward_pre_hook(
+                    self.module_processor.node_hook(prefix + Const.BACKWARD, Const.START))
+                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.api_modularity()
+
+        if Const.STATISTICS == self.config.task or Const.TENSOR == self.config.task:
+            remove_dropout()

msprobe/test/core_ut/common/test_utils.py

@@ -0,0 +1,345 @@
+#!/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");
+# 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 uuid
+
+from unittest import TestCase
+from unittest.mock import patch, MagicMock, mock_open
+
+from msprobe.core.common.log import logger
+from msprobe.core.common.const import Const
+from msprobe.core.common.utils import (CompareException,
+                                    check_seed_all,
+                                    check_inplace_op,
+                                    make_dump_path_if_not_exists,
+                                    check_mode_valid,
+                                    check_switch_valid,
+                                    check_dump_mode_valid,
+                                    check_summary_mode_valid,
+                                    check_summary_only_valid,
+                                    check_file_or_directory_path,
+                                    check_compare_param,
+                                    check_configuration_param,
+                                    is_starts_with,
+                                    _check_json,
+                                    check_json_file,
+                                    check_file_size,
+                                    check_regex_prefix_format_valid,
+                                    get_dump_data_path,
+                                    task_dumppath_get)
+from msprobe.core.common.file_check import FileCheckConst
+
+
+class TestUtils(TestCase):
+    @patch.object(logger, "error")
+    def test_check_seed_all(self, mock_error):
+        self.assertIsNone(check_seed_all(1234, True))
+        self.assertIsNone(check_seed_all(0, True))
+        self.assertIsNone(check_seed_all(Const.MAX_SEED_VALUE, True))
+
+        with self.assertRaises(CompareException) as context:
+            check_seed_all(-1, True)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with(f"Seed must be between 0 and {Const.MAX_SEED_VALUE}.")
+
+        with self.assertRaises(CompareException) as context:
+            check_seed_all(Const.MAX_SEED_VALUE + 1, True)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with(f"Seed must be between 0 and {Const.MAX_SEED_VALUE}.")
+
+        with self.assertRaises(CompareException) as context:
+            check_seed_all("1234", True)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("Seed must be integer.")
+
+        with self.assertRaises(CompareException) as context:
+            check_seed_all(1234, 1)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("seed_all mode must be bool.")
+
+    def test_check_inplace_op(self):
+        test_prefix_1 = "Distributed.broadcast.0.forward.input.0"
+        self.assertTrue(check_inplace_op(test_prefix_1))
+        test_prefix_2 = "Distributed_broadcast_0_forward_input_0"
+        self.assertFalse(check_inplace_op(test_prefix_2))
+        test_prefix_3 = "Torch.sum.0.backward.output.0"
+        self.assertFalse(check_inplace_op(test_prefix_3))
+
+    @patch.object(logger, "error")
+    def test_make_dump_path_if_not_exists(self, mock_error):
+        file_path = os.path.realpath(__file__)
+        dirname = os.path.dirname(file_path) + str(uuid.uuid4())
+
+        def test_mkdir(self, **kwargs):
+            raise OSError
+
+        if not os.path.exists(dirname):
+            with patch("msprobe.core.common.utils.Path.mkdir", new=test_mkdir):
+                with self.assertRaises(CompareException) as context:
+                    make_dump_path_if_not_exists(dirname)
+                self.assertEqual(context.exception.code, CompareException.INVALID_PATH_ERROR)
+
+        make_dump_path_if_not_exists(file_path)
+        mock_error.assert_called_with(f"{file_path} already exists and is not a directory.")
+
+    def test_check_mode_valid(self):
+        with self.assertRaises(ValueError) as context:
+            check_mode_valid("all", scope="scope")
+        self.assertEqual(str(context.exception), "scope param set invalid, it's must be a list.")
+
+        with self.assertRaises(ValueError) as context:
+            check_mode_valid("all", api_list="api_list")
+        self.assertEqual(str(context.exception), "api_list param set invalid, it's must be a list.")
+
+        mode = "all_list"
+        with self.assertRaises(CompareException) as context:
+            check_mode_valid(mode)
+        self.assertEqual(context.exception.code, CompareException.INVALID_DUMP_MODE)
+        self.assertEqual(str(context.exception),
+                         f"Current mode '{mode}' is not supported. Please use the field in {Const.DUMP_MODE}")
+
+        mode = "list"
+        with self.assertRaises(ValueError) as context:
+            check_mode_valid(mode)
+        self.assertEqual(str(context.exception),
+                         "set_dump_switch, scope param set invalid, it's should not be an empty list.")
+
+    @patch.object(logger, "error")
+    def test_check_switch_valid(self, mock_error):
+        with self.assertRaises(CompareException) as context:
+            check_switch_valid("Close")
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("Please set switch with 'ON' or 'OFF'.")
+
+    @patch.object(logger, "warning")
+    def test_check_dump_mode_valid(self, mock_warning):
+        dump_mode = check_dump_mode_valid("all")
+        mock_warning.assert_called_with("Please set dump_mode as a list.")
+        self.assertEqual(dump_mode, ["forward", "backward", "input", "output"])
+
+        with self.assertRaises(ValueError) as context:
+            check_dump_mode_valid("all_forward")
+        self.assertEqual(str(context.exception),
+                         "Please set dump_mode as a list containing one or more of the following: " +
+                         "'all', 'forward', 'backward', 'input', 'output'.")
+
+    def test_check_summary_mode_valid(self):
+        with self.assertRaises(CompareException) as context:
+            check_summary_mode_valid("MD5")
+        self.assertEqual(context.exception.code, CompareException.INVALID_SUMMARY_MODE)
+        self.assertEqual(str(context.exception), "The summary_mode is not valid")
+
+    @patch.object(logger, "error")
+    def test_check_summary_only_valid(self, mock_error):
+        summary_only = check_summary_only_valid(True)
+        self.assertTrue(summary_only)
+
+        with self.assertRaises(CompareException) as context:
+            check_summary_only_valid("True")
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("Params summary_only only support True or False.")
+
+    def test_check_file_or_directory_path(self):
+        class TestFileChecker:
+            file_path = ""
+            path_type = ""
+            ability = ""
+            checked = False
+
+            def __init__(self, file_path, path_type, ability=None):
+                TestFileChecker.file_path = file_path
+                TestFileChecker.path_type = path_type
+                TestFileChecker.ability = ability
+
+            def common_check(self):
+                TestFileChecker.checked = True
+
+        file_path = os.path.realpath(__file__)
+        dirname = os.path.dirname(file_path)
+
+        with patch("msprobe.core.common.utils.FileChecker", new=TestFileChecker):
+            check_file_or_directory_path(file_path, isdir=False)
+        self.assertTrue(TestFileChecker.checked)
+        self.assertEqual(TestFileChecker.file_path, file_path)
+        self.assertEqual(TestFileChecker.path_type, FileCheckConst.FILE)
+        self.assertEqual(TestFileChecker.ability, FileCheckConst.READ_ABLE)
+
+        TestFileChecker.checked = False
+        with patch("msprobe.core.common.utils.FileChecker", new=TestFileChecker):
+            check_file_or_directory_path(dirname, isdir=True)
+        self.assertTrue(TestFileChecker.checked)
+        self.assertEqual(TestFileChecker.file_path, dirname)
+        self.assertEqual(TestFileChecker.path_type, FileCheckConst.DIR)
+        self.assertEqual(TestFileChecker.ability, FileCheckConst.WRITE_ABLE)
+
+    @patch.object(logger, "error")
+    def test_check_compare_param(self, mock_error):
+        params = {
+            "npu_json_path": "npu_json_path",
+            "bench_json_path": "bench_json_path",
+            "stack_json_path": "stack_json_path",
+            "npu_dump_data_dir": "npu_dump_data_dir",
+            "bench_dump_data_dir": "bench_dump_data_dir"
+        }
+
+        call_args = [
+            ("npu_json_path", False),
+            ("bench_json_path", False),
+            ("stack_json_path", False),
+            ("npu_dump_data_dir", True),
+            ("bench_dump_data_dir", True),
+            ("output_path", True),
+            ("npu_json_path", False),
+            ("bench_json_path", False),
+            ("stack_json_path", False),
+            ("output_path", True)
+        ]
+
+        with self.assertRaises(CompareException) as context:
+            check_compare_param("npu_json_path", "output_path")
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("Invalid input parameters")
+
+        mock_check_file_or_directory_path = MagicMock()
+        mock_check_json_file = MagicMock()
+        with patch("msprobe.core.common.utils.FileOpen", mock_open(read_data="")), \
+             patch("msprobe.core.common.utils.check_json_file", new=mock_check_json_file), \
+             patch("msprobe.core.common.utils.check_file_or_directory_path", new=mock_check_file_or_directory_path):
+            check_compare_param(params, "output_path")
+            check_compare_param(params, "output_path", summary_compare=False, md5_compare=True)
+        for i in range(len(call_args)):
+            self.assertEqual(mock_check_file_or_directory_path.call_args_list[i][0], call_args[i])
+        self.assertEqual(len(mock_check_json_file.call_args[0]), 4)
+        self.assertEqual(mock_check_json_file.call_args[0][0], params)
+
+    @patch.object(logger, "error")
+    def test_check_configuration_param(self, mock_error):
+        with self.assertRaises(CompareException) as context:
+            check_configuration_param(stack_mode="False", auto_analyze=True, fuzzy_match=False)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PARAM_ERROR)
+        mock_error.assert_called_with("Invalid input parameters which should be only bool type.")
+
+    def test_is_starts_with(self):
+        string = "input_slot0"
+        self.assertFalse(is_starts_with(string, []))
+        self.assertFalse(is_starts_with("", ["input"]))
+        self.assertFalse(is_starts_with(string, ["output"]))
+        self.assertTrue(is_starts_with(string, ["input", "output"]))
+
+    @patch.object(logger, "error")
+    def test__check_json(self, mock_error):
+        class TestOpen:
+            def __init__(self, string):
+                self.string = string
+
+            def readline(self):
+                return self.string
+
+            def seek(self, begin, end):
+                self.string = str(begin) + "_" + str(end)
+
+        with self.assertRaises(CompareException) as context:
+            _check_json(TestOpen(""), "test.json")
+        self.assertEqual(context.exception.code, CompareException.INVALID_DUMP_FILE)
+        mock_error.assert_called_with("dump file test.json have empty line!")
+
+        handler = TestOpen("jons file\n")
+        _check_json(handler, "test.json")
+        self.assertEqual(handler.string, "0_0")
+
+    @patch("msprobe.core.common.utils._check_json")
+    def test_check_json_file(self, _mock_check_json):
+        input_param = {
+            "npu_json_path": "npu_json_path",
+            "bench_json_path": "bench_json_path",
+            "stack_json_path": "stack_json_path"
+        }
+        check_json_file(input_param, "npu_json", "bench_json", "stack_json")
+        self.assertEqual(_mock_check_json.call_args_list[0][0], ("npu_json", "npu_json_path"))
+        self.assertEqual(_mock_check_json.call_args_list[1][0], ("bench_json", "bench_json_path"))
+        self.assertEqual(_mock_check_json.call_args_list[2][0], ("stack_json", "stack_json_path"))
+
+    @patch.object(logger, "error")
+    def test_check_file_size(self, mock_error):
+        with patch("msprobe.core.common.utils.os.path.getsize", return_value=120):
+            with self.assertRaises(CompareException) as context:
+                check_file_size("input_file", 100)
+        self.assertEqual(context.exception.code, CompareException.INVALID_FILE_ERROR)
+        mock_error.assert_called_with("The size (120) of input_file exceeds (100) bytes, tools not support.")
+
+    def test_check_regex_prefix_format_valid(self):
+        prefix = "A" * 21
+        with self.assertRaises(ValueError) as context:
+            check_regex_prefix_format_valid(prefix)
+        self.assertEqual(str(context.exception), f"Maximum length of prefix is {Const.REGEX_PREFIX_MAX_LENGTH}, "
+                         f"while current length is {len(prefix)}")
+
+        prefix = "(prefix)"
+        with self.assertRaises(ValueError) as context:
+            check_regex_prefix_format_valid(prefix)
+        self.assertEqual(str(context.exception), f"prefix contains invalid characters, "
+                         f"prefix pattern {Const.REGEX_PREFIX_PATTERN}")
+
+    @patch("msprobe.core.common.utils.check_file_or_directory_path")
+    def test_get_dump_data_path(self, mock_check_file_or_directory_path):
+        file_path = os.path.realpath(__file__)
+        dirname = os.path.dirname(file_path)
+
+        dump_data_path, file_is_exist = get_dump_data_path(dirname)
+        self.assertEqual(mock_check_file_or_directory_path.call_args[0], (dirname, True))
+        self.assertEqual(dump_data_path, dirname)
+        self.assertTrue(file_is_exist)
+
+    @patch.object(logger, "error")
+    def test_task_dumppath_get(self, mock_error):
+        input_param = {
+            "npu_json_path": None,
+            "bench_json_path": "bench_json_path"
+        }
+        npu_json = {
+            "task": Const.TENSOR,
+            "dump_data_dir": "dump_data_dir",
+            "data": "data"
+        }
+
+        with self.assertRaises(CompareException) as context:
+            task_dumppath_get(input_param)
+        self.assertEqual(context.exception.code, CompareException.INVALID_PATH_ERROR)
+        mock_error.assert_called_with("Please check the json path is valid.")
+
+        input_param["npu_json_path"] = "npu_json_path"
+        with patch("msprobe.core.common.utils.FileOpen", mock_open(read_data="")), \
+             patch("msprobe.core.common.utils.json.load", return_value=npu_json):
+            summary_compare, md5_compare = task_dumppath_get(input_param)
+        self.assertFalse(summary_compare)
+        self.assertFalse(md5_compare)
+
+        npu_json["task"] = Const.STATISTICS
+        with patch("msprobe.core.common.utils.FileOpen", mock_open(read_data="")), \
+             patch("msprobe.core.common.utils.json.load", return_value=npu_json), \
+             patch("msprobe.core.common.utils.md5_find", return_value=True):
+            summary_compare, md5_compare = task_dumppath_get(input_param)
+        self.assertFalse(summary_compare)
+        self.assertTrue(md5_compare)
+
+        npu_json["task"] = Const.OVERFLOW_CHECK
+        with patch("msprobe.core.common.utils.FileOpen", mock_open(read_data="")), \
+             patch("msprobe.core.common.utils.json.load", return_value=npu_json):
+            with self.assertRaises(CompareException) as context:
+                task_dumppath_get(input_param)
+            self.assertEqual(context.exception.code, CompareException.INVALID_TASK_ERROR)
+            mock_error.assert_called_with("Compare is not required for overflow_check or free_benchmark.")

msprobe/test/core_ut/data_dump/test_data_collector.py

@@ -0,0 +1,47 @@
+import unittest
+from unittest.mock import patch, mock_open, MagicMock
+
+from msprobe.core.common.utils import Const
+from msprobe.core.data_dump.data_collector import DataCollector
+from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
+from msprobe.pytorch.pt_config import parse_json_config
+
+
+class TestDataCollector(unittest.TestCase):
+    def setUp(self):
+        mock_json_data = {
+            "dump_path": "./ut_dump",
+        }
+        with patch("msprobe.pytorch.pt_config.FileOpen", mock_open(read_data='')), \
+                patch("msprobe.pytorch.pt_config.json.load", return_value=mock_json_data):
+            common_config, task_config = parse_json_config("./config.json", Const.STATISTICS)
+        config = DebuggerConfig(common_config, task_config, Const.STATISTICS, "./ut_dump", "L1")
+        self.data_collector = DataCollector(config)
+
+    def test_update_data(self):
+        self.data_collector.config.task = Const.OVERFLOW_CHECK
+        self.data_collector.data_processor.has_overflow = True
+        with patch("msprobe.core.data_dump.json_writer.DataWriter.update_data", return_value=None):
+            result1 = self.data_collector.update_data("test message", "test1:")
+        self.assertEqual(result1, "test1:Overflow detected.")
+
+        self.data_collector.data_processor.has_overflow = False
+        result2 = self.data_collector.update_data("test message", "test2:")
+        self.assertEqual(result2, "test2:No Overflow, OK.")
+
+        self.data_collector.config.task = Const.STATISTICS
+        self.data_collector.data_processor.has_overflow = True
+        with patch("msprobe.core.data_dump.json_writer.DataWriter.update_data", return_value=None):
+            result3 = self.data_collector.update_data("test message", "test3")
+        self.assertEqual(result3, "test3")
+
+    def test_pre_forward_data_collect(self):
+        self.data_collector.check_scope_and_pid = MagicMock(return_value=False)
+        self.data_collector.is_inplace = MagicMock(return_value=False)
+        self.data_collector.data_processor.analyze_pre_forward = MagicMock()
+        name = "TestModule.forward"
+        pid = 123
+
+        self.data_collector.pre_forward_data_collect(name, None, pid, None)
+        self.data_collector.check_scope_and_pid.assert_called_once_with(
+            self.data_collector.scope, "TestModule.backward", 123)