mindstudio-probe 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

msprobe/pytorch/parse_tool/lib/compare.py

@@ -22,6 +22,8 @@ from collections import namedtuple
 from msprobe.pytorch.parse_tool.lib.utils import Util
 from msprobe.pytorch.parse_tool.lib.config import Const
 from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
+from msprobe.core.common.utils import create_directory, write_csv, save_npy_to_txt
+from msprobe.core.common.file_check import FileChecker
 
 
 class Compare:
@@ -36,7 +38,7 @@ class Compare:
         self.log.info("Compare finished!!")
 
     def compare_vector(self, my_dump_path, golden_dump_path, result_dir, msaccucmp_path):
-        self.util.create_dir(result_dir)
+        create_directory(result_dir)
         self.util.check_path_valid(result_dir)
         call_msaccucmp = self.util.check_msaccucmp(msaccucmp_path)
         cmd = '%s %s compare -m %s -g %s -out %s' % (
@@ -65,7 +67,7 @@ class Compare:
             self.util.print_panel("\n".join(summary_txt))
 
     def convert(self, dump_file, data_format, output, msaccucmp_path):
-        self.util.create_dir(output)
+        create_directory(output)
         self.util.check_path_valid(output)
         call_msaccucmp = self.util.check_msaccucmp(msaccucmp_path)
         if data_format:
@@ -83,21 +85,22 @@ class Compare:
         (left, right, save_txt, rl, al, diff_count) = args
         if left is None or right is None:
             raise ParseException("invalid input or output")
-        try:
-            left_data = np.load(left)
-            right_data = np.load(right)
-        except UnicodeError as e:
-            self.log.error("%s %s" % ("UnicodeError", str(e)))
-            self.log.warning("Please check the npy file")
-            raise ParseException(ParseException.PARSE_UNICODE_ERROR) from e
-        except IOError:
-            self.log.error("Failed to load npy %s or %s." % (left, right))
-            raise ParseException(ParseException.PARSE_LOAD_NPY_ERROR) from e
+        if self.util.check_path_valid(left) and self.util.check_path_valid(right):
+            try:
+                left_data = np.load(left)
+                right_data = np.load(right)
+            except UnicodeError as e:
+                self.log.error("%s %s" % ("UnicodeError", str(e)))
+                self.log.warning("Please check the npy file")
+                raise ParseException(ParseException.PARSE_UNICODE_ERROR) from e
+            except IOError:
+                self.log.error("Failed to load npy %s or %s." % (left, right))
+                raise ParseException(ParseException.PARSE_LOAD_NPY_ERROR) from e
 
         # save to txt
         if save_txt:
-            self.util.save_npy_to_txt(left_data, left + ".txt")
-            self.util.save_npy_to_txt(right_data, right + ".txt")
+            save_npy_to_txt(left_data, left + ".txt")
+            save_npy_to_txt(right_data, right + ".txt")
         # compare data
         (total_cnt, all_close, cos_sim, err_percent) = self.do_compare_data(left_data, right_data, rl, al, diff_count)
         content = ['Left:', ' ├─ NpyFile: %s' % left]
@@ -157,8 +160,10 @@ class Compare:
         return res
 
     def compare_npy(self, file, bench_file, output_path):
-        data = np.load(file)
-        bench_data = np.load(bench_file)
+        if self.util.check_path_valid(file):
+            data = np.load(file)
+        if self.util.check_path_valid(bench_file):
+            bench_data = np.load(bench_file)
         shape, dtype = data.shape, data.dtype
         bench_shape, bench_dtype = bench_data.shape, bench_data.dtype
         filename = os.path.basename(file)
@@ -181,7 +186,7 @@ class Compare:
         rel_diff_max = np.max(rel_error)
         compare_result = [[filename, bench_filename, data_mean, bench_data_mean, md5_consistency, abs_diff_max,
                            rel_diff_max]]
-        self.util.write_csv(compare_result, output_path)
+        write_csv(compare_result, output_path)
 
     def compare_all_file_in_directory(self, my_dump_dir, golden_dump_dir, output_path):
         if not (self.util.is_subdir_count_equal(my_dump_dir, golden_dump_dir)
@@ -228,7 +233,7 @@ class Compare:
             "Max Abs Error",
             "Max Relative Error"
         ]]
-        self.util.write_csv(title_rows, output_path)
+        write_csv(title_rows, output_path)
 
         my_ordered_subdirs = self.util.get_sorted_subdirectories_names(my_dump_dir)
         golden_ordered_subdirs = self.util.get_sorted_subdirectories_names(golden_dump_dir)
@@ -246,7 +251,9 @@ class Compare:
 
     def convert_api_dir_to_npy(self, dump_dir, param, output_dir, msaccucmp_path):
         dump_dir = self.util.path_strip(dump_dir)
-        for root, _, files in os.walk(dump_dir):
+        for root, _, files in os.walk(dump_dir, topdown=True):
+            path_checker = FileChecker(root)
+            path_checker.common_check()
             for file in files:
                 file_path = os.path.join(root, file)
                 file_name = os.path.basename(file_path)
@@ -257,3 +264,8 @@ class Compare:
                 timestamp = parts[-1]
                 output_path = os.path.join(output_dir, op_name, timestamp)
                 self.convert_dump_to_npy(file_path, param, output_path, msaccucmp_path)
+            path_depth = root.count(os.sep)
+            if path_depth <= Const.MAX_TRAVERSAL_DEPTH:
+                yield root, _, files
+            else:
+                _[:] = []

msprobe/pytorch/parse_tool/lib/config.py

@@ -33,11 +33,12 @@ class Const:
     OFFLINE_DUMP_CONVERT_PATTERN = \
         r"^([A-Za-z0-9_-]+)\.([A-Za-z0-9_-]+)\.([0-9]+)(\.[0-9]+)?\.([0-9]{1,255})" \
         r"\.([a-z]+)\.([0-9]{1,255})(\.[x0-9]+)?\.npy$"
-    NUMPY_PATTERN = r".*\.npy$"
+    NUMPY_PATTERN = r"^[\w\-_-]\.npy$"
     NPY_SUFFIX = ".npy"
     PKL_SUFFIX = ".pkl"
     DIRECTORY_LENGTH = 4096
     FILE_NAME_LENGTH = 255
+    MAX_TRAVERSAL_DEPTH = 5
     FILE_PATTERN = r'^[a-zA-Z0-9_./-]+$'
     ONE_GB = 1 * 1024 * 1024 * 1024
     TEN_GB = 10 * 1024 * 1024 * 1024

msprobe/pytorch/parse_tool/lib/parse_tool.py

@@ -23,7 +23,7 @@ from msprobe.pytorch.parse_tool.lib.utils import Util
 from msprobe.pytorch.parse_tool.lib.compare import Compare
 from msprobe.pytorch.parse_tool.lib.visualization import Visualization
 from msprobe.pytorch.parse_tool.lib.parse_exception import catch_exception, ParseException
-
+from msprobe.core.common.utils import create_directory
 
 class ParseTool:
     def __init__(self):
@@ -33,7 +33,7 @@ class ParseTool:
 
     @catch_exception
     def prepare(self):
-        self.util.create_dir(Const.DATA_ROOT_DIR)
+        create_directory(Const.DATA_ROOT_DIR)
 
     @catch_exception
     def do_vector_compare(self, args):
@@ -112,8 +112,8 @@ class ParseTool:
         args = parser.parse_args(argv)
         self.util.check_path_valid(args.my_dump_path)
         self.util.check_path_valid(args.golden_dump_path)
-        self.util.check_path_format(args.my_dump_path, Const.NPY_SUFFIX)
-        self.util.check_path_format(args.golden_dump_path, Const.NPY_SUFFIX)
+        self.util.check_file_path_format(args.my_dump_path, Const.NPY_SUFFIX)
+        self.util.check_file_path_format(args.golden_dump_path, Const.NPY_SUFFIX)
         compare_data_args = namedtuple('compare_data_args', ['my_dump_path', 'golden_dump_path', 'save', 'rtol', 'atol', 'count'])
         compare_data_args.__new__.__defaults__ = (False, 0.001, 0.001, 20)
         res = compare_data_args(args.my_dump_path, args.golden_dump_path, args.save, args.rtol, args.atol, args.count)

msprobe/pytorch/parse_tool/lib/utils.py

@@ -31,7 +31,7 @@ from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
 from msprobe.core.common.file_check import change_mode, check_other_user_writable,\
     check_path_executable, check_path_owner_consistent
 from msprobe.core.common.const import FileCheckConst
-from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.file_check import FileOpen, FileChecker
 from msprobe.core.common.utils import check_file_or_directory_path
 from msprobe.pytorch.common.log import logger
 
@@ -57,12 +57,7 @@ except ImportError as err:
 class Util:
     def __init__(self):
         self.ms_accu_cmp = None
-        logging.basicConfig(
-            level=Const.LOG_LEVEL,
-            format="%(asctime)s (%(process)d) -[%(levelname)s]%(message)s",
-            datefmt="%Y-%m-%d %H:%M:%S"
-        )
-        self.log = logging.getLogger()
+        self.log = logger
         self.python = sys.executable
 
     @staticmethod
@@ -82,6 +77,8 @@ class Util:
     @staticmethod
     def get_subdir_count(self, directory):
         subdir_count = 0
+        path_checker = FileChecker(directory)
+        path_checker.common_check()
         for _, dirs, _ in os.walk(directory):
             subdir_count += len(dirs)
             break
@@ -90,8 +87,15 @@ class Util:
     @staticmethod
     def get_subfiles_count(self, directory):
         file_count = 0
-        for _, _, files in os.walk(directory):
+        for root, _, files in os.walk(directory, topdown=True):
+            path_checker = FileChecker(root)
+            path_checker.common_check()
             file_count += len(files)
+            path_depth = root.count(os.sep)
+            if path_depth <= Const.MAX_TRAVERSAL_DEPTH:
+                yield root, _, files
+            else:
+                _[:] = []
         return file_count
 
     @staticmethod
@@ -128,17 +132,6 @@ class Util:
         md5_hash = hashlib.md5(np_bytes)
         return md5_hash.hexdigest()
 
-    @staticmethod
-    def write_csv(self, data, filepath):
-        need_change_mode = False
-        if not os.path.exists(filepath):
-            need_change_mode = True
-        with FileOpen(filepath, 'a') as f:
-            writer = csv.writer(f)
-            writer.writerows(data)
-        if need_change_mode:
-            change_mode(filepath, FileCheckConst.DATA_FILE_AUTHORITY)
-
     @staticmethod
     def deal_with_dir_or_file_inconsistency(self, output_path):
         if os.path.exists(output_path):
@@ -160,10 +153,17 @@ class Util:
     
     @staticmethod
     def dir_contains_only(self, path, endfix):
-        for _, _, files in os.walk(path):
+        for root, _, files in os.walk(path, topdown=True):
+            path_checker = FileChecker(root)
+            path_checker.common_check()
             for file in files:
                 if not file.endswith(endfix):
                     return False
+            path_depth = root.count(os.sep)
+            if path_depth <= Const.MAX_TRAVERSAL_DEPTH:
+                yield root, _, files
+            else:
+                _[:] = []
         return True
     
     @staticmethod
@@ -188,7 +188,7 @@ class Util:
         if not cmd:
             self.log.error("Commond is None")
             return -1
-        self.log.debug("[RUN CMD]: %s", cmd)
+        self.log.info("[RUN CMD]: %s", cmd)
         cmd = cmd.split(" ")
         complete_process = subprocess.run(cmd, shell=False)
         return complete_process.returncode
@@ -208,7 +208,7 @@ class Util:
                 "Check msaccucmp failed in dir %s. This is not a correct msaccucmp file" % target_file)
             raise ParseException(ParseException.PARSE_MSACCUCMP_ERROR)
         result = subprocess.run(
-            [self.python, target_file, "--help"], stdout=subprocess.PIPE)
+            [self.python, target_file, "--help"], stdout=subprocess.PIPE, shell=False)
         if result.returncode == 0:
             self.log.info("Check [%s] success.", target_file)
         else:
@@ -217,37 +217,12 @@ class Util:
             raise ParseException(ParseException.PARSE_MSACCUCMP_ERROR)
         return target_file
 
-    def create_dir(self, path):
-        path = self.path_strip(path)
-        if os.path.exists(path):
-            return
-        self.check_path_name(path)
-        try:
-            os.makedirs(path, mode=FileCheckConst.DATA_DIR_AUTHORITY)
-        except OSError as e:
-            self.log.error("Failed to create %s.", path)
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR) from e
-
     def gen_npy_info_txt(self, source_data):
         (shape, dtype, max_data, min_data, mean) = \
             self.npy_info(source_data)
         return \
                 '[Shape: %s] [Dtype: %s] [Max: %s] [Min: %s] [Mean: %s]' % (shape, dtype, max_data, min_data, mean)
 
-    def save_npy_to_txt(self, data, dst_file='', align=0):
-        if os.path.exists(dst_file):
-            self.log.info("Dst file %s exists, will not save new one.", dst_file)
-            return
-        shape = data.shape
-        data = data.flatten()
-        if align == 0:
-            align = 1 if len(shape) == 0 else shape[-1]
-        elif data.size % align != 0:
-            pad_array = np.zeros((align - data.size % align,))
-            data = np.append(data, pad_array)
-        np.savetxt(dst_file, data.reshape((-1, align)), delimiter=' ', fmt='%g')
-        change_mode(dst_file, FileCheckConst.DATA_FILE_AUTHORITY)
-
     def list_convert_files(self, path, external_pattern=""):
         return self.list_file_with_pattern(
             path, Const.OFFLINE_DUMP_CONVERT_PATTERN, external_pattern, self._gen_npu_dump_convert_file_info
@@ -274,27 +249,8 @@ class Util:
 
     def check_path_valid(self, path):
         path = self.path_strip(path)
-        if not path or not os.path.exists(path):
-            self.log.error("The path %s does not exist." % path)
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-        if os.path.islink(path):
-            self.log.error('The file path {} is a soft link.'.format(path))
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-        if len(os.path.realpath(path)) > Const.DIRECTORY_LENGTH or len(os.path.basename(path)) > \
-                Const.FILE_NAME_LENGTH:
-            self.log.error('The file path length exceeds limit.')
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-        if not re.match(Const.FILE_PATTERN, os.path.realpath(path)):
-            self.log.error('The file path {} contains special characters.'.format(path))
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-        if os.path.isfile(path):
-            file_size = os.path.getsize(path)
-            if path.endswith(Const.PKL_SUFFIX) and file_size > Const.ONE_GB:
-                self.log.error('The file {} size is greater than 1GB.'.format(path))
-                raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-            if path.endswith(Const.NPY_SUFFIX) and file_size > Const.TEN_GB:
-                self.log.error('The file {} size is greater than 10GB.'.format(path))
-                raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
+        path_checker = FileChecker(path)
+        path_checker.common_check()
         return True
 
     def check_files_in_path(self, path):
@@ -322,17 +278,24 @@ class Util:
         self.check_path_valid(path)
         file_list = {}
         re_pattern = re.compile(pattern)
-        for dir_path, _, file_names in os.walk(path, followlinks=True):
+        for dir_path, _, file_names in os.walk(path, topdown=True):
+            path_checker = FileChecker(dir)
+            path_checker.common_check()
             for name in file_names:
                 match = re_pattern.match(name)
                 if not match:
                     continue
-                if extern_pattern != '' and not re.match(extern_pattern, name):
+                if extern_pattern != '' and re_pattern.match(extern_pattern) and not re.match(extern_pattern, name):
                     continue
                 file_list[name] = gen_info_func(name, match, dir_path)
+            path_depth = dir_path.count(os.sep)
+            if path_depth <= Const.MAX_TRAVERSAL_DEPTH:
+                yield dir_path, _, file_names
+            else:
+                _[:] = []
         return file_list
 
-    def check_path_format(self, path, suffix):
+    def check_file_path_format(self, path, suffix):
         if os.path.isfile(path):
             if not path.endswith(suffix):
                 self.log.error("%s is not a %s file." % (path, suffix))
@@ -344,15 +307,6 @@ class Util:
             self.log.error("The file path %s is invalid" % path)
             raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
 
-    def check_path_name(self, path):
-        if len(os.path.realpath(path)) > Const.DIRECTORY_LENGTH or len(os.path.basename(path)) > \
-                Const.FILE_NAME_LENGTH:
-            self.log.error('The file path length exceeds limit.')
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-        if not re.match(Const.FILE_PATTERN, os.path.realpath(path)):
-            self.log.error('The file path {} contains special characters.'.format(path))
-            raise ParseException(ParseException.PARSE_INVALID_PATH_ERROR)
-
     def check_str_param(self, param):
         if len(param) > Const.FILE_NAME_LENGTH:
             self.log.error('The parameter length exceeds limit')

msprobe/pytorch/parse_tool/lib/visualization.py

@@ -21,6 +21,7 @@ from msprobe.pytorch.parse_tool.lib.config import Const
 from msprobe.pytorch.parse_tool.lib.utils import Util
 from msprobe.pytorch.parse_tool.lib.parse_exception import ParseException
 from msprobe.core.common.file_check import FileOpen
+from msprobe.core.common.utils import save_npy_to_txt
 
 
 class Visualization:
@@ -43,18 +44,18 @@ class Visualization:
         summary = ['[yellow]%s[/yellow]' % self.util.gen_npy_info_txt(np_data), 'Path: %s' % target_file,
                    "TextFile: %s.txt" % target_file]
         self.util.print_panel(self.util.create_columns([table, "\n".join(summary)]), target_file)
-        self.util.save_npy_to_txt(np_data, target_file + ".txt")
+        save_npy_to_txt(np_data, target_file + ".txt")
 
     def print_npy_data(self, file_name):
         file_name = self.util.path_strip(file_name)
         self.util.check_path_valid(file_name)
-        self.util.check_path_format(file_name, Const.NPY_SUFFIX)
+        self.util.check_file_path_format(file_name, Const.NPY_SUFFIX)
         return self.print_npy_summary(file_name)
 
     def parse_pkl(self, path, api_name):
         path = self.util.path_strip(path)
         self.util.check_path_valid(path)
-        self.util.check_path_format(path, Const.PKL_SUFFIX)
+        self.util.check_file_path_format(path, Const.PKL_SUFFIX)
         self.util.check_str_param(api_name)
         with FileOpen(path, "r") as pkl_handle:
             title_printed = False

msprobe/pytorch/pt_config.py

@@ -4,18 +4,36 @@ 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
+from msprobe.pytorch.hook_module.utils import get_ops
+from msprobe.core.grad_probe.constant import level_adp
+from msprobe.core.grad_probe.utils import check_numeral_list_ascend
 
 
 class TensorConfig(BaseConfig):
     def __init__(self, json_config):
         super().__init__(json_config)
+        self.online_run_ut = json_config.get("online_run_ut", False)
+        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.check_config()
         self._check_file_format()
+        self._check_tls_path_config()
 
     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")
 
+    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")
+
 
 class StatisticsConfig(BaseConfig):
     def __init__(self, json_config):
@@ -31,12 +49,12 @@ class StatisticsConfig(BaseConfig):
 class OverflowCheckConfig(BaseConfig):
     def __init__(self, json_config):
         super().__init__(json_config)
-        self.overflow_num = json_config.get("overflow_nums")
+        self.overflow_nums = 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):
+        if self.overflow_nums is not None and not isinstance(self.overflow_nums, 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")
@@ -61,20 +79,96 @@ class FreeBenchmarkCheckConfig(BaseConfig):
         if self.preheat_step and self.preheat_step == 0:
             raise Exception("preheat_step cannot be 0")
 
+
+class RunUTConfig(BaseConfig):
+    WrapApi = get_ops()
+
+    def __init__(self, json_config):
+        super().__init__(json_config)
+        self.white_list = json_config.get("white_list", Const.DEFAULT_LIST)
+        self.black_list = json_config.get("black_list", Const.DEFAULT_LIST)
+        self.error_data_path = json_config.get("error_data_path", Const.DEFAULT_PATH)
+        self.is_online = json_config.get("is_online", False)
+        self.nfs_path = json_config.get("nfs_path", "")
+        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.check_run_ut_config()
+
+    @classmethod
+    def check_filter_list_config(cls, key, filter_list):
+        if not isinstance(filter_list, list):
+            raise Exception("%s must be a list type" % key)
+        if not all(isinstance(item, str) for item in filter_list):
+            raise Exception("All elements in %s must be string type" % key)
+        invalid_api = [item for item in filter_list if item not in cls.WrapApi]
+        if invalid_api:
+            raise Exception("Invalid api in %s: %s" % (key, invalid_api))
+
+    @classmethod
+    def check_error_data_path_config(cls, error_data_path):
+        if not os.path.exists(error_data_path):
+            raise Exception("error_data_path: %s does not exist" % error_data_path)
+
+    @classmethod
+    def check_nfs_path_config(cls, nfs_path):
+        if nfs_path and not os.path.exists(nfs_path):
+            raise Exception("nfs_path: %s does not exist" % nfs_path)
+
+    @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")
+
+    def check_run_ut_config(self):
+        RunUTConfig.check_filter_list_config(Const.WHITE_LIST, self.white_list)
+        RunUTConfig.check_filter_list_config(Const.BLACK_LIST, self.black_list)
+        RunUTConfig.check_error_data_path_config(self.error_data_path)
+        RunUTConfig.check_nfs_path_config(self.nfs_path)
+        RunUTConfig.check_tls_path_config(self.tls_path)
+
+
+class GradToolConfig(BaseConfig):
+    def __init__(self, json_config):
+        super().__init__(json_config)
+        self.grad_level = json_config.get("grad_level", "L1")
+        self.param_list = json_config.get("param_list", [])
+        self.bounds = json_config.get("bounds", [])
+    
+    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)
+
+
 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
+        config_dic = json_config.get(Const.TENSOR, 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
+        config_dic = json_config.get(Const.STATISTICS, 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
+        config_dic = json_config.get(Const.OVERFLOW_CHECK, 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
+        config_dic = json_config.get(Const.FREE_BENCHMARK, default_dic)
         return FreeBenchmarkCheckConfig(config_dic)
+    elif task == Const.RUN_UT:
+        config_dic = json_config.get(Const.RUN_UT, default_dic)
+        return RunUTConfig(config_dic)
+    elif task == Const.GRAD_PROBE:
+        config_dic = json_config.get(Const.GRAD_PROBE, default_dic)
+        return GradToolConfig(config_dic)
     else:
         return StatisticsConfig(default_dic)