From a263757757f01fb7d57c634db1507ddd33f181e7 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Fri, 18 Jul 2025 05:07:55 -0400 Subject: [PATCH 1/6] use SafeUnpickler Signed-off-by: yiliu30 --- neural_compressor/utils/utility.py | 68 +++++++++++++++++++++++++----- 1 file changed, 58 insertions(+), 10 deletions(-) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index e934d580474..fdc597e4f53 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -23,6 +23,7 @@ """ import _thread import ast +import builtins import importlib import logging import os @@ -33,6 +34,7 @@ import sys import threading import time +from collections import OrderedDict from contextlib import contextmanager from enum import Enum from functools import wraps @@ -390,16 +392,62 @@ def get_all_fp32_data(data): return [float(i) for i in data.replace("[", " ").replace("]", " ").split(" ") if i.strip() and len(i) < 32] -def get_tuning_history(tuning_history_path): - """Get tuning history. - - Args: - tuning_history_path: The tuning history path, which need users to assign - """ - with open(tuning_history_path, "rb") as f: - strategy_object = pickle.load(f) - tuning_history = strategy_object.tuning_history - return tuning_history +def get_tuning_history(history_path): + class SafeUnpickler(pickle.Unpickler): + def find_class(self, module, name): + # Allowed built-in types + allowed_builtins = { + "dict", + "list", + "tuple", + "set", + "frozenset", + "str", + "bytes", + "int", + "float", + "complex", + "bool", + "NoneType", + "slice", + "type", + "object", + "bytearray", + "bytes", + "ellipsis", + "filter", + "map", + "range", + "reversed", + "zip", + } + if module == "builtins" and name in allowed_builtins: + return getattr(builtins, name) + + # Allow collections.OrderedDict + if module == "collections" and name == "OrderedDict": + return OrderedDict + + # Allow specific neural_compressor classes + if module.startswith("neural_compressor."): + # Validate class name exists in module + mod_path = module.replace(".__", " ") # Handle submodules + for part in mod_path.split(): + try: + __import__(part) + except ImportError: + continue + mod = sys.modules.get(module) + if mod and hasattr(mod, name): + return getattr(mod, name) + + # Block all other classes + raise pickle.UnpicklingError(f"Unsafe class: {module}.{name}") + + with open(history_path, "rb") as f: + strategy_object = SafeUnpickler(f).load() + tuning_history = strategy_object.tuning_history + return tuning_history def recover(fp32_model, tuning_history_path, num, **kwargs): From c42d13533dcf4c9e75391e22fda9545daf18c586 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Sat, 19 Jul 2025 15:01:04 +0800 Subject: [PATCH 2/6] Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- neural_compressor/utils/utility.py | 1 - 1 file changed, 1 deletion(-) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index fdc597e4f53..74303e1fca7 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -413,7 +413,6 @@ def find_class(self, module, name): "type", "object", "bytearray", - "bytes", "ellipsis", "filter", "map", From a35e39e26e7048ad4511cf238f1e7cdf75ab3883 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Sun, 20 Jul 2025 23:25:44 -0400 Subject: [PATCH 3/6] add numpy Signed-off-by: yiliu30 --- neural_compressor/utils/utility.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index fdc597e4f53..6a0405c2b20 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -393,6 +393,12 @@ def get_all_fp32_data(data): def get_tuning_history(history_path): + """Get tuning history. + + Args: + tuning_history_path: The tuning history path, which need users to assign + """ + class SafeUnpickler(pickle.Unpickler): def find_class(self, module, name): # Allowed built-in types @@ -441,6 +447,14 @@ def find_class(self, module, name): if mod and hasattr(mod, name): return getattr(mod, name) + # Allow all numpy classes + if module.startswith("numpy."): + import numpy + + mod = sys.modules.get(module) + if mod and hasattr(mod, name): + return getattr(mod, name) + # Block all other classes raise pickle.UnpicklingError(f"Unsafe class: {module}.{name}") From 66cd19492c3740f97d207d8a64c27c3f8ead4a82 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Mon, 21 Jul 2025 01:06:28 -0400 Subject: [PATCH 4/6] debug Signed-off-by: yiliu30 --- neural_compressor/utils/utility.py | 1 + 1 file changed, 1 insertion(+) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index b44b57f83d8..f693c2a20bd 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -401,6 +401,7 @@ def get_tuning_history(history_path): class SafeUnpickler(pickle.Unpickler): def find_class(self, module, name): + print(f"Unpickling module:{module} name: {name}...") # Allowed built-in types allowed_builtins = { "dict", From d9a741bfc11b716d0c0e8fff4e767fe3e967caa2 Mon Sep 17 00:00:00 2001 From: yiliu30 Date: Mon, 21 Jul 2025 02:54:09 -0400 Subject: [PATCH 5/6] fix numpy check Signed-off-by: yiliu30 --- neural_compressor/utils/utility.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index f693c2a20bd..a1bfaff1405 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -396,7 +396,7 @@ def get_tuning_history(history_path): """Get tuning history. Args: - tuning_history_path: The tuning history path, which need users to assign + history_path: The tuning history path, which need users to assign """ class SafeUnpickler(pickle.Unpickler): @@ -435,7 +435,7 @@ def find_class(self, module, name): return OrderedDict # Allow specific neural_compressor classes - if module.startswith("neural_compressor."): + if module.startswith("neural_compressor"): # Validate class name exists in module mod_path = module.replace(".__", " ") # Handle submodules for part in mod_path.split(): @@ -448,8 +448,7 @@ def find_class(self, module, name): return getattr(mod, name) # Allow all numpy classes - if module.startswith("numpy."): - import numpy + if module.startswith("numpy"): mod = sys.modules.get(module) if mod and hasattr(mod, name): From 7f46d0f601e332145416c86728ef1c5db64fc754 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Mon, 21 Jul 2025 15:33:43 +0800 Subject: [PATCH 6/6] Update utility.py --- neural_compressor/utils/utility.py | 1 - 1 file changed, 1 deletion(-) diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py index a1bfaff1405..6b3245b6659 100644 --- a/neural_compressor/utils/utility.py +++ b/neural_compressor/utils/utility.py @@ -401,7 +401,6 @@ def get_tuning_history(history_path): class SafeUnpickler(pickle.Unpickler): def find_class(self, module, name): - print(f"Unpickling module:{module} name: {name}...") # Allowed built-in types allowed_builtins = { "dict",