fixing PR comments

Author: root

superbench/benchmarks/micro_benchmarks/huggingface_model_loader.py

@@ -4,6 +4,7 @@
 """Hugging Face model loader for benchmarking."""
 
 import os
+import re
 from pathlib import Path
 from typing import Optional, Tuple
 
@@ -20,6 +21,39 @@
 from superbench.common.utils import logger
 from superbench.benchmarks.micro_benchmarks.model_source_config import ModelSourceConfig
 
+# Strict allow-list for HuggingFace model identifiers. Accepts either a bare
+# repo name ('bert-base-uncased') or 'namespace/name' form, restricted to the
+# character set HF itself uses and bounded in length. Rejects '..', backslash,
+# colon, control chars, absolute paths, and anything that could be interpreted
+# as a local filesystem path by AutoConfig.from_pretrained (which silently
+# loads from disk when given a path that exists).
+_SAFE_MODEL_ID_RE = re.compile(r'^[A-Za-z0-9][A-Za-z0-9._-]{0,127}(/[A-Za-z0-9._-]{1,128})?$')
+
+
+def validate_model_identifier(model_identifier: Optional[str]) -> str:
+    """Validate a HuggingFace model identifier against a strict allow-list.
+
+    Args:
+        model_identifier: The identifier to validate (typically from CLI input).
+
+    Returns:
+        The validated identifier (unchanged) for convenient inline use.
+
+    Raises:
+        ValueError: If the identifier is missing or does not match the
+            permitted ``[namespace/]name`` shape. The check intentionally
+            rejects path-traversal sequences and characters that could let
+            ``from_pretrained`` load attacker-staged files from disk.
+    """
+    if not model_identifier or not _SAFE_MODEL_ID_RE.match(model_identifier):
+        raise ValueError(
+            f'Invalid model_identifier {model_identifier!r}. '
+            'Must be a HuggingFace repo id matching '
+            "'^[A-Za-z0-9][A-Za-z0-9._-]{0,127}(/[A-Za-z0-9._-]{1,128})?$' "
+            '(e.g. "bert-base-uncased" or "meta-llama/Llama-2-7b-hf").'
+        )
+    return model_identifier
+
 
 class ModelLoadError(Exception):
     """Exception raised when model loading fails."""
@@ -46,23 +80,41 @@ class HuggingFaceModelLoader:
     Attributes:
         cache_dir: Directory to cache downloaded models.
         token: HuggingFace authentication token for private/gated models.
+        allow_remote_code: Whether to allow HuggingFace to download and execute
+            repository-provided Python (``trust_remote_code=True``). Default
+            ``False``; enabling this turns ``--model_identifier`` into an RCE
+            sink, so it is opt-in only.
     """
-    def __init__(self, cache_dir: Optional[str] = None, token: Optional[str] = None):
+    def __init__(
+        self,
+        cache_dir: Optional[str] = None,
+        token: Optional[str] = None,
+        allow_remote_code: bool = False,
+    ):
         """Initialize the HuggingFace model loader.
 
         Args:
             cache_dir: Directory to cache downloaded models. If None, uses HF default.
             token: HuggingFace authentication token for private/gated models.
+            allow_remote_code: If True, allow execution of model-repo Python via
+                ``trust_remote_code=True``. Default False. Only enable for
+                trusted ``--model_identifier`` values; pin ``--revision <sha>``.
         """
         self.cache_dir = cache_dir or os.getenv('HF_HOME') or os.path.expanduser('~/.cache/huggingface')
         self.token = token or os.getenv('HF_TOKEN') or os.getenv('HUGGING_FACE_HUB_TOKEN')
+        self.allow_remote_code = bool(allow_remote_code)
 
         # Ensure cache directory exists
         Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
 
         logger.info(f'HuggingFaceModelLoader initialized with cache_dir: {self.cache_dir}')
         if self.token:
             logger.info('Authentication token provided for private/gated models (token not logged)')
+        if self.allow_remote_code:
+            logger.warning(
+                'allow_remote_code=True: HuggingFace may download and execute arbitrary Python '
+                'from model repositories. Only enable for trusted model identifiers; pin --revision.'
+            )
 
     def load_model(
         self,
@@ -94,6 +146,9 @@ def load_model(
         """
         logger.info(f'Loading model: {model_identifier}')
 
+        # Reject malformed / path-like identifiers before any network or disk activity.
+        validate_model_identifier(model_identifier)
+
         try:
             # Convert torch_dtype string to torch dtype
             dtype = self._get_torch_dtype(torch_dtype) if torch_dtype else None
@@ -112,22 +167,26 @@ def load_model(
             # Load config (use pre-downloaded config if provided)
             if config is None:
                 logger.info('Loading model configuration...')
-                config = AutoConfig.from_pretrained(model_identifier, trust_remote_code=True, **load_kwargs)
+                config = AutoConfig.from_pretrained(
+                    model_identifier, trust_remote_code=self.allow_remote_code, **load_kwargs
+                )
             else:
                 logger.info('Using pre-downloaded model configuration.')
 
             # Load tokenizer (may fail for some models, that's ok)
             tokenizer = None
             try:
                 logger.info('Loading tokenizer...')
-                tokenizer = AutoTokenizer.from_pretrained(model_identifier, trust_remote_code=True, **load_kwargs)
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_identifier, trust_remote_code=self.allow_remote_code, **load_kwargs
+                )
             except Exception as e:
                 logger.warning(f'Could not load tokenizer: {e}. Continuing without tokenizer.')
 
             # Load model
             logger.info(f'Loading model weights (dtype={torch_dtype}, device={device})...')
             model_kwargs = load_kwargs.copy()
-            model_kwargs['trust_remote_code'] = True
+            model_kwargs['trust_remote_code'] = self.allow_remote_code
 
             # Handle device mapping for large models
             effective_device_map = device_map
@@ -202,16 +261,31 @@ def load_model_from_config(
         if device is None:
             device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-        # Extract loading parameters
-        return self.load_model(
-            model_identifier=config.identifier,
-            torch_dtype=config.torch_dtype,
-            device=device,
-            revision=config.revision,
-            device_map=config.device_map,
-            config=config_pretrained,
-            **config.additional_kwargs
-        )
+        # Honor explicit per-call hf_token / cache_dir from the config without permanently
+        # mutating the loader instance. This makes ModelSourceConfig the single source of
+        # truth for callers that don't rely on HF_TOKEN / HF_HOME env vars.
+        original_token = self.token
+        original_cache_dir = self.cache_dir
+        try:
+            if config.hf_token:
+                self.token = config.hf_token
+            if config.cache_dir:
+                self.cache_dir = config.cache_dir
+                Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+
+            # Extract loading parameters
+            return self.load_model(
+                model_identifier=config.identifier,
+                torch_dtype=config.torch_dtype,
+                device=device,
+                revision=config.revision,
+                device_map=config.device_map,
+                config=config_pretrained,
+                **config.additional_kwargs
+            )
+        finally:
+            self.token = original_token
+            self.cache_dir = original_cache_dir
 
     def _get_torch_dtype(self, dtype_str: str) -> torch.dtype:
         """Convert dtype string to torch.dtype.
@@ -242,7 +316,7 @@ def _get_torch_dtype(self, dtype_str: str) -> torch.dtype:
         }
 
         if normalized_dtype not in dtype_map:
-            raise ValueError(f"Invalid dtype '{dtype_str}'.Must be one of {list(dtype_map.keys())}")
+            raise ValueError(f"Invalid dtype '{dtype_str}'. Must be one of {list(dtype_map.keys())}")
 
         return dtype_map[normalized_dtype]
 

superbench/benchmarks/micro_benchmarks/model_source_config.py

@@ -47,8 +47,10 @@ def __post_init__(self):
         if self.source not in ['in-house', 'huggingface']:
             raise ValueError(f"Invalid model source '{self.source}'. Must be 'in-house' or 'huggingface'.")
 
-        # Validate torch_dtype
-        valid_dtypes = ['float32', 'float16', 'bfloat16', 'int8']
+        # Validate torch_dtype. NOTE: 'int8' is intentionally excluded here — it is handled
+        # post-export via quantize_dynamic (see ort_inference_performance.py) rather than via
+        # the HF torch_dtype loading path, which does not accept torch.int8.
+        valid_dtypes = ['float32', 'float16', 'bfloat16']
         if self.torch_dtype not in valid_dtypes:
             raise ValueError(f"Invalid torch_dtype '{self.torch_dtype}'. Must be one of {valid_dtypes}.")
 

superbench/benchmarks/micro_benchmarks/ort_inference_performance.py

@@ -12,10 +12,13 @@
 import numpy as np
 
 from superbench.common.utils import logger
-from superbench.benchmarks import BenchmarkRegistry, Platform, Precision
+from superbench.benchmarks import BenchmarkRegistry, Platform, Precision, ReturnCode
 from superbench.benchmarks.micro_benchmarks import MicroBenchmark
 from superbench.benchmarks.micro_benchmarks.model_source_config import ModelSourceConfig
-from superbench.benchmarks.micro_benchmarks.huggingface_model_loader import HuggingFaceModelLoader
+from superbench.benchmarks.micro_benchmarks.huggingface_model_loader import (
+    HuggingFaceModelLoader,
+    validate_model_identifier,
+)
 
 
 class ORTInferenceBenchmark(MicroBenchmark):
@@ -42,6 +45,9 @@ def __init__(self, name, parameters=''):
         ]
         self.__graph_opt_level = None
         self.__model_cache_path = Path(torch.hub.get_dir()) / 'checkpoints'
+        # Stashed HF config (populated in _preprocess_huggingface_models) so that
+        # __inference() can derive vocab_size / dynamic input shapes from it.
+        self._hf_config = None
 
     def add_parser_arguments(self):
         """Add the specified arguments."""
@@ -124,6 +130,24 @@ def add_parser_arguments(self):
             help='Sequence length for transformer models.',
         )
 
+        self._parser.add_argument(
+            '--require_cuda',
+            action='store_true',
+            default=False,
+            required=False,
+            help='Fail if CUDAExecutionProvider is not available. '
+            'Default: warn and fall back to other registered ORT providers (CPU/ROCm/etc.).',
+        )
+
+        self._parser.add_argument(
+            '--allow_remote_code',
+            action='store_true',
+            default=False,
+            required=False,
+            help='Allow HuggingFace to execute model-repo Python (trust_remote_code=True). '
+            'SECURITY: enables RCE from --model_identifier. Pin --revision <sha> when used.',
+        )
+
     def _preprocess(self):
         """Preprocess/preparation operations before the benchmarking.
 
@@ -179,8 +203,19 @@ def _preprocess_huggingface_models(self):
 
         if not self._args.model_identifier:
             logger.error('--model_identifier is required when using --model_source huggingface')
+            self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
             return False
 
+        # Reject malformed / path-like identifiers up front, before any network or disk activity.
+        try:
+            validate_model_identifier(self._args.model_identifier)
+        except ValueError as e:
+            logger.error(str(e))
+            self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
+            return False
+
+        allow_remote_code = bool(getattr(self._args, 'allow_remote_code', False))
+
         try:
             logger.info(f'Loading HuggingFace model: {self._args.model_identifier}')
 
@@ -190,13 +225,18 @@ def _preprocess_huggingface_models(self):
             load_kwargs = {}
             if hf_token:
                 load_kwargs['token'] = hf_token
-            hf_config = AutoConfig.from_pretrained(self._args.model_identifier, trust_remote_code=True, **load_kwargs)
+            hf_config = AutoConfig.from_pretrained(
+                self._args.model_identifier, trust_remote_code=allow_remote_code, **load_kwargs
+            )
+            # Stash for __inference() to read vocab_size / other model metadata later.
+            self._hf_config = hf_config
 
             precision_str = self._args.precision.value if self._args.precision != Precision.INT8 else 'float32'
             fits, param_m, est_gb, avail_gb = HuggingFaceModelLoader.check_memory_fits(
                 self._args.model_identifier, hf_config, precision_str, mode='inference', token=hf_token
             )
             if not fits:
+                self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
                 return False
 
             # Step 2: Proceed with model download and ONNX export
@@ -217,7 +257,7 @@ def _preprocess_huggingface_models(self):
             )
 
             # Load model from HuggingFace on CPU
-            loader = HuggingFaceModelLoader()
+            loader = HuggingFaceModelLoader(allow_remote_code=allow_remote_code)
             hf_model, _, _ = loader.load_model_from_config(model_config, device='cpu')
             from superbench.benchmarks.micro_benchmarks._export_torch_to_onnx import torch2onnxExporter
             exporter = torch2onnxExporter()
@@ -237,6 +277,15 @@ def _preprocess_huggingface_models(self):
                 export_precision = self._args.precision.value
             model_name_with_precision = f'{model_name}.{export_precision}'
 
+            # Defense-in-depth: confirm the resolved output path stays inside the rank
+            # directory even though validate_model_identifier already rejected '..' / '\\'.
+            proc_root = proc_output_path.resolve()
+            resolved_out = (proc_output_path / f'{model_name_with_precision}.onnx').resolve()
+            if proc_root not in resolved_out.parents:
+                logger.error(f'Refusing to write ONNX outside rank dir: {resolved_out} not under {proc_root}')
+                self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
+                return False
+
             # Export directly to final destination to avoid path issues with external data
             onnx_path = exporter.export_huggingface_model(
                 model=hf_model,
@@ -248,6 +297,7 @@ def _preprocess_huggingface_models(self):
 
             if not onnx_path:
                 logger.error(f'Failed to export {self._args.model_identifier} to ONNX')
+                self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
                 return False
 
             # Apply INT8 quantization if requested (matching in-house model behavior)
@@ -268,25 +318,32 @@ def _preprocess_huggingface_models(self):
             logger.error(f'Failed to prepare HuggingFace model: {str(e)}')
             import traceback
             logger.error(traceback.format_exc())
+            self._result.set_return_code(ReturnCode.MICROBENCHMARK_EXECUTION_FAILURE)
             return False
 
     def _benchmark(self):
         """Implementation for benchmarking."""
         import onnxruntime as ort
         precision_metric = {'float16': 'fp16', 'float32': 'fp32', 'int8': 'int8'}
 
-        # Require CUDAExecutionProvider — this benchmark targets GPU inference
         available = ort.get_available_providers()
-        if 'CUDAExecutionProvider' not in available:
-            logger.error(f'CUDAExecutionProvider is not available (available: {available}).')
-            return False
+        cuda_available = 'CUDAExecutionProvider' in available
+        if not cuda_available:
+            msg = f'CUDAExecutionProvider is not available (available providers: {available}).'
+            if getattr(self._args, 'require_cuda', False):
+                logger.error(msg + ' --require_cuda was set, aborting.')
+                return False
+            logger.warning(
+                msg + ' Falling back to registered providers; pass --require_cuda to fail instead.'
+            )
+        providers = ['CUDAExecutionProvider'] if cuda_available else available
 
         for model in self._args.pytorch_models:
             sess_options = ort.SessionOptions()
             sess_options.graph_optimization_level = self.__graph_opt_level[self._args.graph_opt_level]
             file_name = '{model}.{precision}.onnx'.format(model=model, precision=self._args.precision)
             ort_sess = ort.InferenceSession(
-                f'{self.__model_cache_path / file_name}', sess_options, providers=['CUDAExecutionProvider']
+                f'{self.__model_cache_path / file_name}', sess_options, providers=providers
             )
 
             elapse_times = self.__inference(ort_sess)
@@ -318,18 +375,30 @@ def __inference(self, ort_sess):
         """
         precision = np.float16 if self._args.precision == Precision.FLOAT16 else np.float32
 
-        # Get input names from the ONNX session to determine input format
-        input_names = [input.name for input in ort_sess.get_inputs()]
+        # Get input metadata from the ONNX session to determine input format and shapes
+        ort_inputs = ort_sess.get_inputs()
+        input_names = [inp.name for inp in ort_inputs]
 
         # Determine input format based on what the model expects
         if 'pixel_values' in input_names:
-            # Vision model: use pixel_values (batch_size, 3, 224, 224)
-            pixel_values = np.random.randn(self._args.batch_size, 3, 224, 224).astype(dtype=precision)
+            # Vision model: derive (C, H, W) from the exported ONNX graph so that models
+            # with non-default shapes (e.g. 384x384 ViT, 1-channel medical models) work.
+            # Fall back to (3, 224, 224) only for dynamic / unknown axes.
+            meta = next(inp for inp in ort_inputs if inp.name == 'pixel_values')
+            dims = [d if isinstance(d, int) else None for d in (meta.shape or [])]
+            # Expected layout is (N, C, H, W); pad to length 4 if shorter.
+            dims = (dims + [None] * 4)[:4]
+            _, c, h, w = dims
+            c, h, w = c or 3, h or 224, w or 224
+            pixel_values = np.random.randn(self._args.batch_size, c, h, w).astype(dtype=precision)
             inputs = {'pixel_values': pixel_values}
         elif 'input_ids' in input_names:
-            # NLP model: use input_ids and attention_mask
+            # NLP model: use input_ids and attention_mask. Cap token IDs at the model's
+            # actual vocab_size to avoid out-of-range embedding lookups (undefined behavior
+            # on CUDA — silent NaNs / device-side asserts).
             seq_len = getattr(self._args, 'seq_length', 512)
-            input_ids = np.random.randint(0, 30000, (self._args.batch_size, seq_len)).astype(np.int64)
+            vocab_size = getattr(self._hf_config, 'vocab_size', None) or 30000
+            input_ids = np.random.randint(0, vocab_size, (self._args.batch_size, seq_len)).astype(np.int64)
             attention_mask = np.ones((self._args.batch_size, seq_len), dtype=np.int64)
             inputs = {'input_ids': input_ids, 'attention_mask': attention_mask}
         else: