From c10a77c969fd70d1e9a060cb82f49315c96056dd Mon Sep 17 00:00:00 2001 From: mldangelo-oai Date: Mon, 8 Jun 2026 23:48:33 -0700 Subject: [PATCH] fix: avoid irrelevant weight budget failures --- CHANGELOG.md | 1 + .../scanners/weight_distribution_scanner.py | 1089 ++++++-- .../test_weight_distribution_scanner.py | 2204 +++++++++++++---- 3 files changed, 2589 insertions(+), 705 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 16e56aa07..80aac7e51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - preserve registrable network domains and redact delimiter-split credentials in bounded finding evidence - preserve executable text-sidecar network findings for f-string calls, standard command wrappers, port-qualified Docker registries, and bounded xargs downloads while keeping prose references informational. - stabilize cache identity capture for compressed wrapper scans on Darwin `/private` path aliases and during unrelated temporary-file churn. +- bound per-tensor and cumulative weight-distribution extraction before materializing PyTorch, HDF5, TensorFlow, and ONNX payloads ## [0.2.47](https://github.com/promptfoo/modelaudit/compare/v0.2.46...v0.2.47) (2026-06-05) diff --git a/modelaudit/scanners/weight_distribution_scanner.py b/modelaudit/scanners/weight_distribution_scanner.py index 538b41984..c93979483 100644 --- a/modelaudit/scanners/weight_distribution_scanner.py +++ b/modelaudit/scanners/weight_distribution_scanner.py @@ -1,6 +1,7 @@ """Scanner for detecting anomalous weight distributions in model files.""" import inspect +import math import numbers import os import zipfile @@ -11,13 +12,59 @@ from .base import BaseScanner, IssueSeverity, ScanResult, logger _ANALYSIS_INCONCLUSIVE_REASON = "weight_distribution_analysis_incomplete" +_FINAL_LAYER_NAME_PATTERNS = ("fc", "classifier", "head", "output", "final", "dense") +_ZIP_MEMBER_READ_CHUNK_SIZE = 64 * 1024 +_MAX_PICKLE_METADATA_BYTES = 10 * 1024 * 1024 +_DEFAULT_MAX_TENSOR_BYTES = 100 * 1024 * 1024 +_DEFAULT_MAX_TOTAL_TENSOR_BYTES = 512 * 1024 * 1024 +_ESTIMATED_PICKLE_OBJECT_BYTES = 64 _MAX_RESTRICTED_PICKLE_BYTES = 10 * 1024 * 1024 _MAX_RESTRICTED_PICKLE_NODES = 1_000_000 _MAX_RESTRICTED_PICKLE_OPCODES = 500_000 _MAX_RESTRICTED_PICKLE_MEMO_ENTRIES = 100_000 -_PICKLE_OPCODE_BUDGET_BYTES = 128 _MAX_NUMERIC_SCALAR_BYTES = 16 _MIN_NUMPY_ARRAY_BUDGET_BYTES = 256 +_PICKLE_OBJECT_OPCODES = frozenset( + { + "BINBYTES", + "BINBYTES8", + "BINFLOAT", + "BININT", + "BININT1", + "BININT2", + "BINSTRING", + "BINUNICODE", + "BINUNICODE8", + "BYTEARRAY8", + "DICT", + "EMPTY_DICT", + "EMPTY_LIST", + "EMPTY_SET", + "EMPTY_TUPLE", + "FLOAT", + "FROZENSET", + "INT", + "LIST", + "LONG", + "LONG1", + "LONG4", + "NEWFALSE", + "NEWTRUE", + "NONE", + "SHORT_BINBYTES", + "SHORT_BINSTRING", + "SHORT_BINUNICODE", + "STRING", + "TUPLE", + "TUPLE1", + "TUPLE2", + "TUPLE3", + "UNICODE", + } +) +_ADDITIVE_EXTRACTION_DETAIL_KEYS = frozenset( + {"external_reference_tensors", "oversized_tensors", "tensor_read_failures"} +) class WeightDistributionScanner(BaseScanner): @@ -51,7 +98,16 @@ def __init__(self, config: dict[str, Any] | None = None): self.llm_vocab_threshold = self.config.get("llm_vocab_threshold", 10000) self.enable_llm_checks = self.config.get("enable_llm_checks", False) # Use max_array_size for in-memory array size limits (default 100MB) - self.max_array_size = self.config.get("max_array_size", 100 * 1024 * 1024) + self.max_array_size = self.config.get("max_array_size", _DEFAULT_MAX_TENSOR_BYTES) + default_total_tensor_bytes = ( + 0 + if self._configured_byte_limit(self.max_array_size, fallback=_DEFAULT_MAX_TENSOR_BYTES) is None + else _DEFAULT_MAX_TOTAL_TENSOR_BYTES + ) + self.max_total_tensor_bytes = self.config.get( + "max_weight_distribution_total_bytes", + default_total_tensor_bytes, + ) # Direct torch.load on untrusted files can trigger pickle RCE. Keep opt-in. self.enable_unsafe_torch_load = self.config.get("enable_unsafe_torch_load") is True # Flag set when weight extraction would be unsafe @@ -60,6 +116,7 @@ def __init__(self, config: dict[str, Any] | None = None): self.extraction_incomplete = False self.extraction_incomplete_reasons: list[str] = [] self.extraction_incomplete_details: dict[str, Any] = {} + self.retained_tensor_bytes = 0 @classmethod def can_handle(cls, path: str) -> bool: @@ -130,6 +187,7 @@ def scan(self, path: str) -> ScanResult: self.extraction_incomplete = False self.extraction_incomplete_reasons = [] self.extraction_incomplete_details = {} + self.retained_tensor_bytes = 0 try: # Extract weights based on file format @@ -251,8 +309,10 @@ def _record_extraction_incomplete(self, reason: str, **details: Any) -> None: if value is None: continue existing = self.extraction_incomplete_details.get(key) - if isinstance(existing, int) and isinstance(value, int): + if key in _ADDITIVE_EXTRACTION_DETAIL_KEYS and isinstance(existing, int) and isinstance(value, int): self.extraction_incomplete_details[key] = existing + value + elif existing == value: + continue elif isinstance(existing, list): if isinstance(value, list): existing.extend(value) @@ -287,15 +347,373 @@ def _mark_analysis_incomplete( }, ) + @staticmethod + def _configured_byte_limit(value: Any, *, fallback: int | None = None) -> int | None: + if isinstance(value, bool) or not isinstance(value, numbers.Real): + return fallback + numeric_value = float(value) + if not math.isfinite(numeric_value): + return fallback + if numeric_value == 0: + return None + if numeric_value < 0: + return fallback + return max(int(numeric_value), 1) + + def _max_tensor_bytes(self) -> int | None: + return self._configured_byte_limit(self.max_array_size, fallback=_DEFAULT_MAX_TENSOR_BYTES) + + def _max_total_tensor_bytes(self) -> int | None: + return self._configured_byte_limit( + self.max_total_tensor_bytes, + fallback=self._max_tensor_bytes(), + ) + + def _remaining_tensor_bytes(self) -> int | None: + limits = [limit for limit in (self._max_tensor_bytes(),) if limit is not None] + max_total_bytes = self._max_total_tensor_bytes() + if max_total_bytes is not None: + limits.append(max(max_total_bytes - self.retained_tensor_bytes, 0)) + if not limits: + return None + return min(limits) + + def _tensor_fits_budget( + self, + reason: str, + name: str, + *, + tensor_nbytes: int | None, + retain: bool = False, + ) -> bool: + max_tensor_bytes = self._max_tensor_bytes() + if tensor_nbytes is None or tensor_nbytes < 0: + self._record_oversized_tensor(reason, name, tensor_nbytes=tensor_nbytes) + return False + if max_tensor_bytes is not None and tensor_nbytes > max_tensor_bytes: + self._record_oversized_tensor(reason, name, tensor_nbytes=tensor_nbytes) + return False + + max_total_bytes = self._max_total_tensor_bytes() + projected_total = self.retained_tensor_bytes + tensor_nbytes + if max_total_bytes is not None and projected_total > max_total_bytes: + self._record_extraction_incomplete( + f"{reason}_total", + failed_tensors=[name], + oversized_tensors=1, + tensor_nbytes=tensor_nbytes, + retained_tensor_bytes=self.retained_tensor_bytes, + max_total_tensor_bytes=max_total_bytes, + ) + return False + + if retain: + self.retained_tensor_bytes = projected_total + return True + + @staticmethod + def _logical_tensor_nbytes(shape: Any, dtype: Any) -> int | None: + if dtype is None: + return None + try: + dimensions = [int(dim) for dim in shape] + except Exception: + return None + if any(dim < 0 for dim in dimensions): + return None + + try: + itemsize = int(dtype.itemsize) + except Exception: + try: + import numpy as np + + itemsize = int(np.dtype(dtype).itemsize) + except Exception: + return None + + if itemsize <= 0: + return None + + return math.prod(dimensions) * itemsize + + def _record_oversized_tensor(self, reason: str, name: str, *, tensor_nbytes: int | None) -> None: + self._record_extraction_incomplete( + reason, + failed_tensors=[name], + oversized_tensors=1, + tensor_nbytes=tensor_nbytes, + max_array_size=self._max_tensor_bytes(), + ) + + def _select_pytorch_data_pickle( + self, + archive: zipfile.ZipFile, + ) -> tuple[str, zipfile.ZipInfo] | None: + members = [member for member in archive.infolist() if not member.is_dir()] + pickle_candidates: list[tuple[str, zipfile.ZipInfo]] = [] + for member in members: + parts = member.filename.strip("/").split("/") + if parts and parts[-1] == "data.pkl": + pickle_candidates.append(("/".join(parts[:-1]), member)) + + if not pickle_candidates: + return None + + candidate_names = [member.filename for _root, member in pickle_candidates] + if len(candidate_names) != len(set(candidate_names)): + self._record_extraction_incomplete( + "pytorch_pickle_member_ambiguous", + pickle_member_count=len(candidate_names), + pickle_members=candidate_names, + ) + return None + + if len(pickle_candidates) == 1: + return pickle_candidates[0] + + member_names = {member.filename.strip("/") for member in members} + ranked: list[tuple[int, int, str, zipfile.ZipInfo]] = [] + for root, member in pickle_candidates: + prefix = f"{root}/" if root else "" + has_root_marker = any(f"{prefix}{marker}" in member_names for marker in ("version", "byteorder")) + has_numeric_storage = any( + name.startswith(f"{prefix}data/") + and "/" not in name[len(f"{prefix}data/") :] + and name[len(f"{prefix}data/") :].isdigit() + for name in member_names + ) + credibility = 2 if has_root_marker else int(has_numeric_storage) + depth = len(root.split("/")) if root else 0 + ranked.append((credibility, depth, root, member)) + + credible = [candidate for candidate in ranked if candidate[0] > 0] + if len(credible) != 1: + self._record_extraction_incomplete( + "pytorch_pickle_member_ambiguous", + pickle_member_count=len(candidate_names), + pickle_members=candidate_names, + ) + return None + _credibility, _depth, root, member = credible[0] + return root, member + + def _pytorch_load_within_budget(self, path: str) -> bool: + max_total_bytes = self._max_total_tensor_bytes() + max_tensor_bytes = self._max_tensor_bytes() + + try: + if zipfile.is_zipfile(path): + with zipfile.ZipFile(path, "r") as archive: + selected = self._select_pytorch_data_pickle(archive) + if selected is None: + return not self.extraction_incomplete + root, data_pkl_info = selected + prefix_parts = root.split("/") if root else [] + selected_members = [data_pkl_info] + selected_names = {data_pkl_info.filename} + for member_info in archive.infolist(): + parts = member_info.filename.strip("/").split("/") + is_storage = ( + len(parts) == len(prefix_parts) + 2 + and parts[: len(prefix_parts)] == prefix_parts + and parts[-2] == "data" + and parts[-1].isdigit() + ) + if member_info.is_dir() or not is_storage: + continue + if member_info.filename in selected_names: + self._record_extraction_incomplete( + "pytorch_archive_member_ambiguous", + failed_tensors=[member_info.filename], + ) + return False + selected_names.add(member_info.filename) + selected_members.append(member_info) + if max_tensor_bytes is not None and member_info.file_size > max_tensor_bytes: + self._record_oversized_tensor( + "pytorch_tensor_storage_size_limit", + member_info.filename, + tensor_nbytes=member_info.file_size, + ) + return False + + load_bytes = sum(member.file_size for member in selected_members) + + if max_total_bytes is not None and load_bytes > max_total_bytes: + self._record_extraction_incomplete( + "pytorch_load_size_limit", + failed_tensors=[path], + oversized_tensors=1, + tensor_nbytes=load_bytes, + max_total_tensor_bytes=max_total_bytes, + ) + return False + + data = self._read_zip_member_bounded(archive, data_pkl_info) + if data is None or not self._pickle_object_budget_allows(data, data_pkl_info.filename): + return False + return True + else: + load_bytes = os.path.getsize(path) + except (OSError, zipfile.BadZipFile): + return True + + if max_tensor_bytes is not None and load_bytes > max_tensor_bytes: + self._record_oversized_tensor( + "pytorch_legacy_load_size_limit", + os.path.basename(path), + tensor_nbytes=load_bytes, + ) + return False + if max_total_bytes is None or load_bytes <= max_total_bytes: + return True + + self._record_extraction_incomplete( + "pytorch_load_size_limit", + failed_tensors=[path], + oversized_tensors=1, + tensor_nbytes=load_bytes, + max_total_tensor_bytes=max_total_bytes, + ) + return False + + def _pickle_object_budget_allows(self, data: bytes, name: str) -> bool: + remaining_bytes = self._remaining_tensor_bytes() + budget_bytes = min(limit for limit in (remaining_bytes, _MAX_PICKLE_METADATA_BYTES) if limit is not None) + max_objects = max(budget_bytes // _ESTIMATED_PICKLE_OBJECT_BYTES, 1024) + object_count = 0 + + import pickletools + + for opcode, _arg, _pos in pickletools.genops(data): + if opcode.name not in _PICKLE_OBJECT_OPCODES: + continue + object_count += 1 + if object_count > max_objects: + self._record_extraction_incomplete( + "pytorch_pickle_object_budget", + failed_tensors=[name], + oversized_tensors=1, + pickle_objects=object_count, + max_pickle_objects=max_objects, + ) + return False + return True + + @staticmethod + def _torch_tensor_nbytes(tensor: Any) -> int | None: + try: + numel = int(tensor.numel()) + element_size = int(tensor.element_size()) + except Exception: + return None + if numel < 0 or element_size <= 0: + return None + return numel * element_size + + def _convert_torch_tensor( + self, + tensor: Any, + name: str, + seen_storages: set[int] | None = None, + ) -> Any | None: + tensor_nbytes = self._torch_tensor_nbytes(tensor) + if not self._tensor_fits_budget("pytorch_tensor_size_limit", name, tensor_nbytes=tensor_nbytes): + return None + + storage_identity: int + storage_nbytes: int | None + try: + storage = tensor.untyped_storage() + storage_identity = int(getattr(storage, "_cdata", id(storage))) + storage_nbytes = int(storage.nbytes()) + except Exception: + storage_identity = id(tensor) + storage_nbytes = tensor_nbytes + retain_storage = seen_storages is None or storage_identity not in seen_storages + if retain_storage and not self._tensor_fits_budget( + "pytorch_tensor_storage_size_limit", + name, + tensor_nbytes=storage_nbytes, + ): + return None + try: + array = tensor.detach().cpu().numpy() + except Exception as exc: + self._record_extraction_incomplete( + "pytorch_tensor_read_failed", + failed_tensors=[name], + tensor_read_failures=1, + exception_type=type(exc).__name__, + ) + return None + if not self._tensor_fits_budget( + "pytorch_tensor_size_limit", + name, + tensor_nbytes=int(array.nbytes), + ): + return None + if retain_storage: + if not self._tensor_fits_budget( + "pytorch_tensor_storage_size_limit", + name, + tensor_nbytes=storage_nbytes, + retain=True, + ): + return None + if seen_storages is not None: + seen_storages.add(storage_identity) + return array.T + + def _python_tensor_nbytes(self, value: Any) -> int | None: + max_bytes = self._remaining_tensor_bytes() + stack: list[tuple[Any, bool]] = [(value, False)] + active_containers: set[int] = set() + scalar_count = 0 + max_itemsize = 1 + + while stack: + item, exiting = stack.pop() + if isinstance(item, (list, tuple)): + item_id = id(item) + if exiting: + active_containers.remove(item_id) + continue + if item_id in active_containers: + return None + active_containers.add(item_id) + stack.append((item, True)) + stack.extend((child, False) for child in reversed(item)) + continue + if isinstance(item, bool): + itemsize = 1 + elif isinstance(item, complex): + itemsize = 16 + elif isinstance(item, float) or (isinstance(item, int) and -(2**63) <= item < 2**63): + itemsize = 8 + else: + return None + scalar_count += 1 + max_itemsize = max(max_itemsize, itemsize) + estimated_nbytes = scalar_count * max_itemsize + if max_bytes is not None and estimated_nbytes > max_bytes: + return estimated_nbytes + + return scalar_count * max_itemsize + + @staticmethod + def _python_tensor_has_matrix_shape(value: Any) -> bool: + if not isinstance(value, (list, tuple)) or not value: + return False + return isinstance(value[0], (list, tuple)) + def _restricted_pickle_array_limit(self) -> int: """Return a hard-bounded array budget for the primitive pickle fallback.""" - try: - configured_limit = int(self.max_array_size) - except (TypeError, ValueError): - configured_limit = 0 - if configured_limit <= 0: + remaining_bytes = self._remaining_tensor_bytes() + if remaining_bytes is None: return 100 * 1024 * 1024 - return min(configured_limit, 100 * 1024 * 1024) + return min(remaining_bytes, 100 * 1024 * 1024) def _new_primitive_array_budget(self) -> dict[str, int]: byte_limit = self._restricted_pickle_array_limit() @@ -308,7 +726,7 @@ def _new_primitive_array_budget(self) -> dict[str, int]: def _restricted_pickle_opcode_limit(self) -> int: return min( _MAX_RESTRICTED_PICKLE_OPCODES, - max(256, self._restricted_pickle_array_limit() // _PICKLE_OPCODE_BUDGET_BYTES), + max(1024, self._restricted_pickle_array_limit() // 8), ) def _bounded_primitive_array(self, value: Any, np: Any, budget: dict[str, int]) -> Any | None: @@ -316,6 +734,7 @@ def _bounded_primitive_array(self, value: Any, np: Any, budget: dict[str, int]) active_container_ids: set[int] = set() stack: list[tuple[Any, bool]] = [(value, False)] numeric_items = 0 + max_numeric_item_bytes = 1 metadata_items = 0 while stack: @@ -342,6 +761,12 @@ def _bounded_primitive_array(self, value: Any, np: Any, budget: dict[str, int]) continue if isinstance(node, numbers.Number): numeric_items += 1 + if isinstance(node, complex): + max_numeric_item_bytes = max(max_numeric_item_bytes, 16) + elif isinstance(node, (int, float)): + max_numeric_item_bytes = max(max_numeric_item_bytes, 8) + else: + max_numeric_item_bytes = max(max_numeric_item_bytes, _MAX_NUMERIC_SCALAR_BYTES) continue if isinstance(node, (str, bytes)) or node is None: metadata_items += 1 @@ -354,7 +779,7 @@ def _bounded_primitive_array(self, value: Any, np: Any, budget: dict[str, int]) return None required_bytes = max( - numeric_items * _MAX_NUMERIC_SCALAR_BYTES, + numeric_items * max_numeric_item_bytes, _MIN_NUMPY_ARRAY_BUDGET_BYTES, ) if required_bytes > budget["remaining_bytes"]: @@ -373,6 +798,80 @@ def _bounded_primitive_array(self, value: Any, np: Any, budget: dict[str, int]) ) return array + def _read_zip_member_bounded(self, archive: zipfile.ZipFile, member_info: zipfile.ZipInfo) -> bytes | None: + max_tensor_bytes = self._max_tensor_bytes() + remaining_tensor_bytes = self._remaining_tensor_bytes() + max_bytes = min( + limit + for limit in (max_tensor_bytes, remaining_tensor_bytes, _MAX_PICKLE_METADATA_BYTES) + if limit is not None + ) + if member_info.file_size > max_bytes: + self._record_extraction_incomplete( + "pytorch_zip_data_pkl_size_limit", + failed_tensors=[member_info.filename], + oversized_tensors=1, + tensor_nbytes=member_info.file_size, + max_array_size=max_tensor_bytes, + max_pickle_metadata_bytes=max_bytes, + ) + return None + + data = bytearray() + with archive.open(member_info, "r") as member: + while True: + chunk = member.read(_ZIP_MEMBER_READ_CHUNK_SIZE) + if not chunk: + break + data.extend(chunk) + if len(data) > max_bytes: + self._record_extraction_incomplete( + "pytorch_zip_data_pkl_size_limit", + failed_tensors=[member_info.filename], + oversized_tensors=1, + tensor_nbytes=len(data), + max_array_size=max_tensor_bytes, + max_pickle_metadata_bytes=max_bytes, + ) + return None + return bytes(data) + + @staticmethod + def _is_weight_tensor_name(name: str) -> bool: + lowered_name = name.lower() + return "kernel" in lowered_name or "weight" in lowered_name + + @staticmethod + def _weight_tensor_name_priority(name: str) -> int: + lowered_name = name.lower() + return int(any(pattern in lowered_name for pattern in _FINAL_LAYER_NAME_PATTERNS)) + + @staticmethod + def _tensorflow_dtype_itemsize(dtype: Any) -> int | None: + try: + import numpy as np + + as_numpy_dtype = getattr(dtype, "as_numpy_dtype", dtype) + numpy_dtype = np.dtype(as_numpy_dtype) + if numpy_dtype.hasobject or numpy_dtype.kind in "SU" or numpy_dtype.itemsize <= 0: + return None + return int(numpy_dtype.itemsize) + except Exception: + return None + + def _tensorflow_checkpoint_variable_nbytes(self, shape: Any, dtype: Any | None) -> int | None: + try: + dimensions = [int(dim) for dim in shape] + except Exception: + return None + if any(dim < 0 for dim in dimensions): + return None + + itemsize = self._tensorflow_dtype_itemsize(dtype) if dtype is not None else None + if itemsize is None: + return None + return math.prod(dimensions) * itemsize + def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: """Extract weights from PyTorch model files""" try: @@ -386,6 +885,7 @@ def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: self.extraction_unsafe_reason = None weights_info: dict[str, Any] = {} + seen_storages: set[int] = set() try: if not self.enable_unsafe_torch_load: @@ -410,6 +910,9 @@ def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: if supports_weights_only: load_kwargs["weights_only"] = True + if not self._pytorch_load_within_budget(path): + return {} + # Load model with map_location to CPU to avoid GPU requirements model_data = torch.load(path, **load_kwargs) @@ -420,10 +923,11 @@ def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: # Find final layer weights (classification head) for key, value in state_dict.items(): + key_text = str(key) if isinstance(value, torch.Tensor) and ( ( any( - pattern in key.lower() + pattern in key_text.lower() for pattern in [ "fc", "classifier", @@ -432,68 +936,49 @@ def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: "final", ] ) - and "weight" in key.lower() + and "weight" in key_text.lower() ) - or ("weight" in key.lower() and len(value.shape) >= 2) + or "weight" in key_text.lower() ): - # PyTorch uses (out_features, in_features) but we expect (in_features, out_features) - weights_info[key] = value.detach().cpu().numpy().T + if len(value.shape) < 2: + continue + array = self._convert_torch_tensor(value, key_text, seen_storages) + if array is not None: + weights_info[key_text] = array elif hasattr(model_data, "state_dict"): # Full model format state_dict = model_data.state_dict() for key, value in state_dict.items(): - if "weight" in key.lower() and isinstance(value, torch.Tensor): - # PyTorch uses (out_features, in_features) but we expect (in_features, out_features) - weights_info[key] = value.detach().cpu().numpy().T + key_text = str(key) + if "weight" in key_text.lower() and isinstance(value, torch.Tensor): + if len(value.shape) < 2: + continue + array = self._convert_torch_tensor(value, key_text, seen_storages) + if array is not None: + weights_info[key_text] = array except Exception as e: logger.debug(f"Failed to extract weights from {path}: {e}") - # Try loading as a zip file (newer PyTorch format) safe_fallback_processed = False try: with zipfile.ZipFile(path, "r") as z: - data_pkl_infos = [ - info - for info in z.infolist() - if info.filename.endswith("/data.pkl") or info.filename == "data.pkl" - ] - if len(data_pkl_infos) > 1: - self._record_extraction_incomplete( - "pytorch_pickle_member_ambiguous", - pickle_member_count=len(data_pkl_infos), - pickle_members=[info.filename for info in data_pkl_infos], - ) - safe_fallback_processed = True - elif data_pkl_infos: + selected = self._select_pytorch_data_pickle(z) + if selected is not None: import io import pickle import pickletools - pickle_info = data_pkl_infos[0] - data_pkl_path = pickle_info.filename - pickle_limit = min(self._restricted_pickle_array_limit(), _MAX_RESTRICTED_PICKLE_BYTES) - if pickle_info.file_size > pickle_limit: - self._record_extraction_incomplete( - "pytorch_pickle_read_limit_exceeded", - pickle_member=data_pkl_path, - pickle_size=pickle_info.file_size, - pickle_read_limit=pickle_limit, - ) - safe_fallback_processed = True - data = b"" - else: - with z.open(pickle_info, "r") as pickle_file: - data = pickle_file.read(pickle_limit + 1) - if len(data) > pickle_limit: - self._record_extraction_incomplete( - "pytorch_pickle_read_limit_exceeded", - pickle_member=data_pkl_path, - pickle_size=len(data), - pickle_read_limit=pickle_limit, - ) - safe_fallback_processed = True - data = b"" + _root, data_pkl_info = selected + data = self._read_zip_member_bounded(z, data_pkl_info) + if data is None: + self.extraction_unsafe = False + self.extraction_unsafe_reason = None + return weights_info + if not self._pickle_object_budget_allows(data, data_pkl_info.filename): + self.extraction_unsafe = False + self.extraction_unsafe_reason = None + return weights_info # Look for disallowed opcodes that could trigger code execution disallowed = { @@ -507,84 +992,79 @@ def _extract_pytorch_weights(self, path: str) -> dict[str, Any]: "NEWOBJ_EX", } unsafe = False - if data: - try: - pickle_opcode_limit = self._restricted_pickle_opcode_limit() - pickle_opcode_count = 0 - pickle_graph_over_budget = False - pickle_memo_entries = 0 - pickle_memo_limit = min( - _MAX_RESTRICTED_PICKLE_MEMO_ENTRIES, - pickle_opcode_limit, + pickle_opcode_limit = self._restricted_pickle_opcode_limit() + pickle_graph_over_budget = False + pickle_memo_entries = 0 + pickle_memo_limit = min(_MAX_RESTRICTED_PICKLE_MEMO_ENTRIES, pickle_opcode_limit) + for pickle_opcode_count, (opcode, arg, _pos) in enumerate(pickletools.genops(data), start=1): + if pickle_opcode_count > pickle_opcode_limit: + self._record_extraction_incomplete( + "pytorch_pickle_graph_budget_exceeded", + pickle_opcode_count=pickle_opcode_count, + pickle_opcode_limit=pickle_opcode_limit, ) - for pickle_opcode_count, (opcode, _arg, _pos) in enumerate( - pickletools.genops(data), start=1 + pickle_graph_over_budget = True + break + if opcode.name in {"PUT", "BINPUT", "LONG_BINPUT"}: + if ( + not isinstance(arg, int) + or arg < 0 + or arg > pickle_memo_entries + or arg >= pickle_memo_limit ): - if pickle_opcode_count > pickle_opcode_limit: - self._record_extraction_incomplete( - "pytorch_pickle_graph_budget_exceeded", - pickle_opcode_count=pickle_opcode_count, - pickle_opcode_limit=pickle_opcode_limit, - ) - safe_fallback_processed = True - pickle_graph_over_budget = True - break - if opcode.name in {"PUT", "BINPUT", "LONG_BINPUT"}: - if ( - not isinstance(_arg, int) - or _arg < 0 - or _arg > pickle_memo_entries - or _arg >= pickle_memo_limit - ): - self._record_extraction_incomplete( - "pytorch_pickle_graph_budget_exceeded", - pickle_memo_index=_arg, - pickle_memo_entries=pickle_memo_entries, - pickle_memo_limit=pickle_memo_limit, - pickle_memo_opcode=opcode.name, - ) - safe_fallback_processed = True - pickle_graph_over_budget = True - break - if _arg == pickle_memo_entries: - pickle_memo_entries += 1 - elif opcode.name == "MEMOIZE": - if pickle_memo_entries >= pickle_memo_limit: - self._record_extraction_incomplete( - "pytorch_pickle_graph_budget_exceeded", - pickle_memo_entries=pickle_memo_entries, - pickle_memo_limit=pickle_memo_limit, - pickle_memo_opcode=opcode.name, - ) - safe_fallback_processed = True - pickle_graph_over_budget = True - break - pickle_memo_entries += 1 - if opcode.name in disallowed: - unsafe = True - break - - if pickle_graph_over_budget: - pass - elif unsafe: - self.extraction_unsafe = True - self.extraction_unsafe_reason = ( - "Unsafe to extract weights from data.pkl in PyTorch archive" + self._record_extraction_incomplete( + "pytorch_pickle_graph_budget_exceeded", + pickle_memo_index=arg, + pickle_memo_entries=pickle_memo_entries, + pickle_memo_limit=pickle_memo_limit, + pickle_memo_opcode=opcode.name, ) - else: + pickle_graph_over_budget = True + break + if arg == pickle_memo_entries: + pickle_memo_entries += 1 + elif opcode.name == "MEMOIZE": + if pickle_memo_entries >= pickle_memo_limit: + self._record_extraction_incomplete( + "pytorch_pickle_graph_budget_exceeded", + pickle_memo_entries=pickle_memo_entries, + pickle_memo_limit=pickle_memo_limit, + pickle_memo_opcode=opcode.name, + ) + pickle_graph_over_budget = True + break + pickle_memo_entries += 1 + if opcode.name in disallowed: + unsafe = True + break + + if pickle_graph_over_budget: + safe_fallback_processed = True + elif unsafe: + self.extraction_unsafe = True + self.extraction_unsafe_reason = "Unsafe to extract weights from data.pkl in PyTorch archive" + else: + try: - class RestrictedUnpickler(pickle.Unpickler): - def find_class(self, module: str, name: str) -> Any: - raise pickle.UnpicklingError("global lookup not allowed") - - obj = RestrictedUnpickler(io.BytesIO(data)).load() - if isinstance(obj, dict): - primitive_budget = self._new_primitive_array_budget() - for key, value in obj.items(): - if not isinstance(key, str) or not ( - "weight" in key.lower() or "kernel" in key.lower() - ): - continue + class RestrictedUnpickler(pickle.Unpickler): + def find_class(self, module: str, name: str) -> Any: + raise pickle.UnpicklingError("global lookup not allowed") + + obj = RestrictedUnpickler(io.BytesIO(data)).load() + if isinstance(obj, dict): + primitive_budget = self._new_primitive_array_budget() + primitive_arrays: dict[int, Any | None] = {} + retained_primitive_arrays: set[int] = set() + primitive_array_names: dict[int, str] = {} + for key, value in obj.items(): + if not isinstance(key, str) or not self._is_weight_tensor_name(key): + continue + if not self._python_tensor_has_matrix_shape(value): + continue + value_identity = id(value) + if value_identity in primitive_arrays: + array = primitive_arrays[value_identity] + else: try: array = self._bounded_primitive_array(value, np, primitive_budget) except Exception as array_error: @@ -594,27 +1074,45 @@ def find_class(self, module: str, name: str) -> Any: tensor_read_failures=1, exception_type=type(array_error).__name__, ) + primitive_arrays[value_identity] = None continue - if array is None: - continue - if len(array.shape) >= 2: + primitive_arrays[value_identity] = array + if array is None or len(array.shape) < 2: + continue + if value_identity in retained_primitive_arrays: + previous_name = primitive_array_names[value_identity] + if self._weight_tensor_name_priority( + key + ) > self._weight_tensor_name_priority(previous_name): + weights_info.pop(previous_name, None) weights_info[key] = array - safe_fallback_processed = True - else: - self._record_extraction_incomplete( - "pytorch_pickle_unsupported_root", - pickle_root_type=type(obj).__name__, - ) - safe_fallback_processed = True + primitive_array_names[value_identity] = key + continue + if self._tensor_fits_budget( + "pytorch_zip_tensor_size_limit", + key, + tensor_nbytes=int(array.nbytes), + retain=True, + ): + weights_info[key] = array + retained_primitive_arrays.add(value_identity) + primitive_array_names[value_identity] = key + safe_fallback_processed = True + else: + self._record_extraction_incomplete( + "pytorch_pickle_unsupported_root", + pickle_root_type=type(obj).__name__, + ) + safe_fallback_processed = True except Exception as e2: # pragma: no cover - defensive - logger.debug( - f"Failed restricted unpickle for {path}: {e2}", - ) + logger.debug(f"Failed restricted unpickle for {path}: {e2}") self._record_extraction_incomplete( "pytorch_pickle_parse_failed", exception_type=type(e2).__name__, ) safe_fallback_processed = True + elif self.extraction_incomplete: + safe_fallback_processed = True except Exception as e2: # pragma: no cover - defensive logger.debug(f"Failed to extract weights from {path}: {e2}") @@ -640,13 +1138,125 @@ def _extract_keras_weights(self, path: str) -> dict[str, Any]: try: with h5py.File(path, "r") as f: - # Navigate through the HDF5 structure to find weights - def extract_weights(name, obj): - if ( - isinstance(obj, h5py.Dataset) - and ("kernel" in name or "weight" in name) - and np.issubdtype(obj.dtype, np.number) - ): + visited_group_states: set[tuple[Any, bool, bool]] = set() + materialized_datasets: dict[Any, Any] = {} + materialized_dataset_names: dict[Any, str] = {} + skipped_dataset_ids: set[Any] = set() + groups_to_visit: list[tuple[Any, str]] = [(f, "")] + + while groups_to_visit: + group, prefix = groups_to_visit.pop() + lowered_prefix = prefix.lower() + group_state = ( + group.id, + self._is_weight_tensor_name(prefix), + any(pattern in lowered_prefix for pattern in _FINAL_LAYER_NAME_PATTERNS), + ) + if group_state in visited_group_states: + continue + visited_group_states.add(group_state) + + for child_name in group: + name = f"{prefix}/{child_name}" if prefix else str(child_name) + try: + link = group.get(child_name, getlink=True) + except Exception as exc: + if self._is_weight_tensor_name(name): + self._record_extraction_incomplete( + "keras_hdf5_link_read_failed", + failed_tensors=[name], + tensor_read_failures=1, + exception_type=type(exc).__name__, + ) + continue + + if isinstance(link, h5py.ExternalLink): + if not self._is_weight_tensor_name(name) and not self._is_weight_tensor_name(link.path): + continue + self._record_extraction_incomplete( + "keras_hdf5_external_link_skipped", + failed_tensors=[name], + external_reference_tensors=1, + link_type=type(link).__name__, + ) + continue + + try: + obj = group.get(child_name, getlink=False) + except Exception as exc: + if self._is_weight_tensor_name(name): + self._record_extraction_incomplete( + "keras_hdf5_object_read_failed", + failed_tensors=[name], + tensor_read_failures=1, + exception_type=type(exc).__name__, + ) + continue + + if isinstance(obj, h5py.Group): + groups_to_visit.append((obj, name)) + continue + + if obj is None: + if self._is_weight_tensor_name(name): + self._record_extraction_incomplete( + "keras_hdf5_object_read_failed", + failed_tensors=[name], + tensor_read_failures=1, + exception_type="DanglingLink", + ) + continue + + if not ( + isinstance(obj, h5py.Dataset) + and self._is_weight_tensor_name(name) + and np.issubdtype(obj.dtype, np.number) + ): + continue + + dataset_identity = obj.id + if dataset_identity in materialized_datasets: + previous_name = materialized_dataset_names[dataset_identity] + if self._weight_tensor_name_priority(name) > self._weight_tensor_name_priority( + previous_name + ): + weights_info.pop(previous_name, None) + weights_info[name] = materialized_datasets[dataset_identity] + materialized_dataset_names[dataset_identity] = name + continue + if dataset_identity in skipped_dataset_ids: + continue + skipped_dataset_ids.add(dataset_identity) + + storage_properties = obj.id.get_create_plist() + if storage_properties.get_external_count() > 0: + self._record_extraction_incomplete( + "keras_hdf5_external_storage_skipped", + failed_tensors=[name], + external_reference_tensors=1, + ) + continue + + try: + virtual_source_count = storage_properties.get_virtual_count() + except ValueError: + virtual_source_count = 0 + if virtual_source_count > 0: + self._record_extraction_incomplete( + "keras_hdf5_virtual_dataset_skipped", + failed_tensors=[name], + external_reference_tensors=1, + ) + continue + + tensor_nbytes = self._logical_tensor_nbytes(obj.shape, obj.dtype) + if not self._tensor_fits_budget( + "keras_tensor_size_limit", + name, + tensor_nbytes=tensor_nbytes, + ): + continue + try: array = np.array(obj) except Exception as exc: @@ -657,11 +1267,19 @@ def extract_weights(name, obj): tensor_read_failures=1, exception_type=type(exc).__name__, ) - return + continue + if not self._tensor_fits_budget( + "keras_tensor_size_limit", + name, + tensor_nbytes=int(array.nbytes), + retain=True, + ): + continue + materialized_datasets[dataset_identity] = array + materialized_dataset_names[dataset_identity] = name + skipped_dataset_ids.discard(dataset_identity) weights_info[name] = array - f.visititems(extract_weights) - except Exception as e: logger.debug(f"Failed to extract weights from {path}: {e}") if os.path.splitext(path)[1].lower() == ".keras" and zipfile.is_zipfile(path): @@ -691,12 +1309,30 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: ckpt_prefix = os.path.join(path, "variables", "variables") if os.path.exists(ckpt_prefix + ".index"): - for name, _shape in tf.train.list_variables(ckpt_prefix): + variable_dtype_map: dict[str, Any] = {} + with suppress(Exception): + checkpoint_reader = tf.train.load_checkpoint(ckpt_prefix) + if hasattr(checkpoint_reader, "get_variable_to_dtype_map"): + variable_dtype_map = dict(checkpoint_reader.get_variable_to_dtype_map()) + + for name, shape in tf.train.list_variables(ckpt_prefix): if "weight" not in name.lower() and "kernel" not in name.lower(): continue + if len(shape) < 2: + continue + tensor_nbytes = self._tensorflow_checkpoint_variable_nbytes( + shape, + variable_dtype_map.get(name), + ) + if not self._tensor_fits_budget( + "tensorflow_checkpoint_tensor_size_limit", + name, + tensor_nbytes=tensor_nbytes, + ): + continue try: tensor = tf.train.load_variable(ckpt_prefix, name) - array = np.array(tensor) + array = np.asarray(tensor) except Exception as exc: logger.warning( "TensorFlow weight variable '%s' could not be read from %s: %s", name, path, exc @@ -708,7 +1344,12 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: exception_type=type(exc).__name__, ) continue - if self.max_array_size and self.max_array_size > 0 and array.nbytes > self.max_array_size: + if not self._tensor_fits_budget( + "tensorflow_checkpoint_tensor_size_limit", + name, + tensor_nbytes=int(array.nbytes), + retain=len(array.shape) >= 2, + ): continue # Only include 2D+ tensors for consistency with .pb file handling if len(array.shape) >= 2: @@ -724,7 +1365,7 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: from tensorflow.core.framework import graph_pb2 from tensorflow.core.protobuf import saved_model_pb2 - from modelaudit.utils.tensorflow_compat import tensor_proto_to_ndarray + from modelaudit.utils.tensorflow_compat import DTYPE_MAP, tensor_proto_to_ndarray nodes: list[Any] = [] saved_model = saved_model_pb2.SavedModel() @@ -748,10 +1389,27 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: continue tensor_proto = node.attr["value"].tensor + tensor_dtype = DTYPE_MAP.get(int(tensor_proto.dtype)) + if tensor_dtype is None or tensor_dtype.hasobject or tensor_dtype.kind in "SU": + self._record_extraction_incomplete( + "tensorflow_const_tensor_dtype_unsupported", + failed_tensors=[node.name], + tensor_read_failures=1, + tensor_dtype=int(tensor_proto.dtype), + ) + continue + remaining_tensor_bytes = self._remaining_tensor_bytes() + if remaining_tensor_bytes == 0: + self._record_oversized_tensor( + "tensorflow_const_tensor_size_limit", + node.name, + tensor_nbytes=None, + ) + continue try: array = tensor_proto_to_ndarray( tensor_proto, - max_tensor_bytes=self.max_array_size, + max_tensor_bytes=remaining_tensor_bytes, ) except Exception as exc: logger.warning( @@ -765,7 +1423,12 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: ) continue - if self.max_array_size and self.max_array_size > 0 and array.nbytes > self.max_array_size: + if not self._tensor_fits_budget( + "tensorflow_const_tensor_size_limit", + node.name, + tensor_nbytes=int(array.nbytes), + retain=len(array.shape) >= 2, + ): continue if len(array.shape) >= 2: weights_info[node.name] = array @@ -775,6 +1438,35 @@ def _extract_tensorflow_weights(self, path: str) -> dict[str, Any]: return weights_info + @staticmethod + def _onnx_inline_storage_nbytes(onnx: Any, initializer: Any) -> int: + raw_data = getattr(initializer, "raw_data", b"") + raw_bytes = len(raw_data) + data_type = getattr(initializer, "data_type", None) + tensor_proto = getattr(onnx, "TensorProto", None) + + packed_multiplier = 1 + for dtype_name in ("FLOAT4E2M1", "INT4", "UINT4"): + if tensor_proto is not None and data_type == getattr(tensor_proto, dtype_name, None): + packed_multiplier = 2 + break + for dtype_name in ("INT2", "UINT2"): + if tensor_proto is not None and data_type == getattr(tensor_proto, dtype_name, None): + packed_multiplier = 4 + break + + typed_bytes = 0 + for field_name, itemsize in ( + ("float_data", 4), + ("int32_data", 4), + ("int64_data", 8), + ("double_data", 8), + ("uint64_data", 8), + ): + typed_bytes += len(getattr(initializer, field_name, ())) * itemsize + typed_bytes += sum(len(value) for value in getattr(initializer, "string_data", ())) + return raw_bytes * packed_multiplier + typed_bytes + def _extract_onnx_weights(self, path: str) -> dict[str, Any]: """Extract weights from ONNX model files. @@ -790,6 +1482,18 @@ def _extract_onnx_weights(self, path: str) -> dict[str, Any]: weights_info: dict[str, Any] = {} try: + max_total_bytes = self._max_total_tensor_bytes() + model_size = os.path.getsize(path) + if max_total_bytes is not None and model_size > max_total_bytes: + self._record_extraction_incomplete( + "onnx_model_size_limit", + failed_tensors=[path], + oversized_tensors=1, + tensor_nbytes=model_size, + max_total_tensor_bytes=max_total_bytes, + ) + return {} + # Use load_external_data=False to prevent ValidationError when # external data files (e.g. weights.pb) are missing. This is the # common case for models downloaded from HuggingFace or distributed @@ -806,15 +1510,46 @@ def _extract_onnx_weights(self, path: str) -> dict[str, Any]: if len(initializer.dims) < 2: continue + initializer_name = str(initializer.name) + external_location = getattr(getattr(onnx, "TensorProto", None), "EXTERNAL", 1) + if getattr(initializer, "data_location", None) == external_location or bool( + getattr(initializer, "external_data", ()) + ): + self._record_extraction_incomplete( + "onnx_external_initializer_skipped", + failed_tensors=[initializer_name], + external_reference_tensors=1, + ) + continue + # Pre-check estimated byte size before materializing the # full array — avoids memory exhaustion on huge tensors. + tensor_dtype: Any | None = None with suppress(Exception): + tensor_dtype = onnx.helper.tensor_dtype_to_np_dtype(initializer.data_type) + if tensor_dtype is None: _onnx_mapping = getattr(onnx, "mapping", None) - if _onnx_mapping is not None and hasattr(_onnx_mapping, "TENSOR_TYPE_TO_NP_TYPE"): - tensor_dtype = _onnx_mapping.TENSOR_TYPE_TO_NP_TYPE[initializer.data_type] - estimated_size = int(np.prod(initializer.dims)) * np.dtype(tensor_dtype).itemsize - if self.max_array_size and self.max_array_size > 0 and estimated_size > self.max_array_size: - continue + with suppress(Exception): + if _onnx_mapping is not None and hasattr(_onnx_mapping, "TENSOR_TYPE_TO_NP_TYPE"): + tensor_dtype = _onnx_mapping.TENSOR_TYPE_TO_NP_TYPE[initializer.data_type] + with suppress(Exception): + numpy_dtype = np.dtype(tensor_dtype) + if numpy_dtype.hasobject or numpy_dtype.kind in "SU" or numpy_dtype.itemsize <= 0: + tensor_dtype = None + inline_storage_nbytes = self._onnx_inline_storage_nbytes(onnx, initializer) + if not self._tensor_fits_budget( + "onnx_initializer_storage_size_limit", + initializer_name, + tensor_nbytes=inline_storage_nbytes, + ): + continue + estimated_size = self._logical_tensor_nbytes(initializer.dims, tensor_dtype) + if not self._tensor_fits_budget( + "onnx_initializer_size_limit", + initializer_name, + tensor_nbytes=estimated_size, + ): + continue try: arr = onnx.numpy_helper.to_array(initializer) # type: ignore[possibly-unresolved-reference] @@ -824,14 +1559,19 @@ def _extract_onnx_weights(self, path: str) -> dict[str, Any]: ) self._record_extraction_incomplete( "onnx_initializer_read_failed", - failed_tensors=[initializer.name], + failed_tensors=[initializer_name], tensor_read_failures=1, exception_type=type(exc).__name__, ) continue - if self.max_array_size and self.max_array_size > 0 and arr.nbytes > self.max_array_size: + if not self._tensor_fits_budget( + "onnx_initializer_size_limit", + initializer_name, + tensor_nbytes=int(arr.nbytes), + retain=True, + ): continue - weights_info[initializer.name] = arr + weights_info[initializer_name] = arr except Exception as e: logger.warning(f"Failed to extract ONNX weights from {path}: {e}") @@ -875,18 +1615,7 @@ def _analyze_weight_distributions( final_layer_candidates = {} for name, weights in weights_info.items(): if ( - any( - pattern in name.lower() - for pattern in [ - "fc", - "classifier", - "head", - "output", - "final", - "dense", - ] - ) - and "weight" in name.lower() + any(pattern in name.lower() for pattern in _FINAL_LAYER_NAME_PATTERNS) and "weight" in name.lower() ) and len(weights.shape) == 2: # Ensure it's a 2D weight matrix final_layer_candidates[name] = weights diff --git a/tests/scanners/test_weight_distribution_scanner.py b/tests/scanners/test_weight_distribution_scanner.py index 18546f57e..976cb291a 100644 --- a/tests/scanners/test_weight_distribution_scanner.py +++ b/tests/scanners/test_weight_distribution_scanner.py @@ -129,6 +129,44 @@ def make_const_node(name: str, shape: list[int], float_values: list[float]) -> A assert weights["dense/kernel"].tolist() == [[1.0, 1.0], [1.0, 1.0]] +def test_extract_tensorflow_weights_rejects_string_const(tmp_path: Path) -> None: + import importlib + + import modelaudit.protos + + assert modelaudit.protos._check_vendored_protos() + + attr_value_pb2 = importlib.import_module("tensorflow.core.framework.attr_value_pb2") + graph_pb2 = importlib.import_module("tensorflow.core.framework.graph_pb2") + node_def_pb2 = importlib.import_module("tensorflow.core.framework.node_def_pb2") + tensor_pb2 = importlib.import_module("tensorflow.core.framework.tensor_pb2") + types_pb2 = importlib.import_module("tensorflow.core.framework.types_pb2") + + tensor = tensor_pb2.TensorProto(dtype=types_pb2.DT_STRING) + for size in (2, 2): + tensor.tensor_shape.dim.add(size=size) + tensor.string_val.extend([b"x" * 1024] * 4) + node = node_def_pb2.NodeDef(name="dense/kernel", op="Const") + node.attr["value"].CopyFrom(attr_value_pb2.AttrValue(tensor=tensor)) + + graph = graph_pb2.GraphDef() + graph.node.append(node) + model_path = tmp_path / "string-model.pb" + model_path.write_bytes(graph.SerializeToString()) + + scanner = WeightDistributionScanner( + { + "enable_unsafe_torch_load": True, + "max_array_size": 64, + "max_weight_distribution_total_bytes": 64, + } + ) + weights = scanner._extract_tensorflow_weights(str(model_path)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["tensorflow_const_tensor_dtype_unsupported"] + + # Skip tests if required libraries are not available def has_numpy(): try: @@ -386,311 +424,1639 @@ def test_partial_hdf5_weight_extraction_preserves_analyzed_findings(tmp_path: Pa assert any(check.name == "Weight Distribution Anomaly Detection" for check in result.checks) analysis_check = next(check for check in result.checks if check.name == "Weight Distribution Analysis") assert analysis_check.details["analysis_incomplete"] is True - assert analysis_check.details["tensor_read_failures"] == 1 + assert analysis_check.details["external_reference_tensors"] == 1 assert analysis_check.details["failed_tensors"] == ["model_weights/z_dense/weight_external:0"] -@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not available") -class TestWeightDistributionScanner: - """Test suite for weight distribution anomaly detection""" +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_oversized_weight_dataset_skips_materialization( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import h5py + import numpy as np - def _create_mock_architecture_analysis(self, is_llm=False, is_transformer=False): - """Helper method to create mock architecture analysis for testing""" - return { - "is_likely_transformer": is_transformer, - "is_likely_llm": is_llm, - "confidence": 0.8 if is_llm else 0.5, - "evidence": ["Mock evidence for testing"], - "architectural_features": {}, - "total_parameters": 100_000_000 if is_llm else 1_000_000, - "layer_count": 24 if is_llm else 3, - } + path = tmp_path / "oversized_weights.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file.create_dataset("model_weights/dense/kernel:0", shape=(16, 16), dtype=np.float32) - def test_scanner_initialization(self): - """Test scanner initialization with default and custom config""" - # Default initialization - scanner = WeightDistributionScanner() - assert scanner.z_score_threshold == 3.0 - assert scanner.cosine_similarity_threshold == 0.7 - assert scanner.weight_magnitude_threshold == 3.0 - assert scanner.max_array_size == 100 * 1024 * 1024 # Default 100MB + original_array = np.array - # Custom config - config = { - "z_score_threshold": 2.5, - "cosine_similarity_threshold": 0.8, - "weight_magnitude_threshold": 2.0, - "max_array_size": 50 * 1024 * 1024, # 50MB - } - scanner = WeightDistributionScanner(config) - assert scanner.z_score_threshold == 2.5 - assert scanner.cosine_similarity_threshold == 0.8 - assert scanner.weight_magnitude_threshold == 2.0 - assert scanner.max_array_size == 50 * 1024 * 1024 + def fail_if_oversized_dataset_is_materialized(obj: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(obj, h5py.Dataset) and obj.name.endswith("model_weights/dense/kernel:0"): + raise AssertionError("oversized HDF5 weight dataset should not be materialized") + return original_array(obj, *args, **kwargs) - def test_can_handle(self): - """Test file type detection""" - # Create temporary files to test can_handle - with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: - pt_path = f.name - with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: - h5_path = f.name - with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f: - txt_path = f.name - tf_dir = tempfile.mkdtemp() - open(os.path.join(tf_dir, "saved_model.pb"), "wb").close() + monkeypatch.setattr(np, "array", fail_if_oversized_dataset_is_materialized) - try: - # Should handle PyTorch files if torch is available - if _has_torch_cached(): - assert WeightDistributionScanner.can_handle(pt_path) + result = WeightDistributionScanner({"max_array_size": 1}).scan(str(path)) - # Should handle Keras files if h5py is available - if _has_h5py_cached(): - assert WeightDistributionScanner.can_handle(h5_path) + assert result.success is False + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + analysis_check = next(check for check in result.checks if check.name == "Weight Distribution Analysis") + assert "keras_tensor_size_limit" in analysis_check.details["extraction_incomplete_reasons"] + assert analysis_check.details["failed_tensors"] == ["model_weights/dense/kernel:0"] - if _has_tensorflow_cached(): - assert WeightDistributionScanner.can_handle(tf_dir) - # Should not handle unsupported extensions - assert not WeightDistributionScanner.can_handle(txt_path) - assert not WeightDistributionScanner.can_handle("directory/") - finally: - os.unlink(pt_path) - os.unlink(h5_path) - os.unlink(txt_path) - os.unlink(os.path.join(tf_dir, "saved_model.pb")) - os.rmdir(tf_dir) +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_external_weight_link_skips_without_following_target(tmp_path: Path) -> None: + import h5py - def test_analyze_layer_weights_outlier_detection(self): - """Test detection of outlier weight vectors""" - import numpy as np + path = tmp_path / "external_link_weights.h5" + target_path = tmp_path / "missing_external_weights.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file["model_weights/dense/kernel:0"] = h5py.ExternalLink(str(target_path), "/kernel") - scanner = WeightDistributionScanner() + result = WeightDistributionScanner().scan(str(path)) - # Create normal weights with one outlier - np.random.seed(42) - normal_weights = np.random.randn(100, 10) * 0.1 # Small weights + assert result.success is False + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + analysis_check = next(check for check in result.checks if check.name == "Weight Distribution Analysis") + assert "keras_hdf5_external_link_skipped" in analysis_check.details["extraction_incomplete_reasons"] + assert analysis_check.details["external_reference_tensors"] == 1 + assert analysis_check.details["failed_tensors"] == ["model_weights/dense/kernel:0"] - # Make one neuron an outlier with large weights - make it even more extreme - normal_weights[:, 5] = np.random.randn(100) * 10.0 # Much larger weights - architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) - anomalies = scanner._analyze_layer_weights("test_layer", normal_weights, architecture_analysis) +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_unrelated_external_link_does_not_make_weight_analysis_incomplete(tmp_path: Path) -> None: + import h5py + import numpy as np - # Should detect the outlier neuron - assert len(anomalies) > 0 + path = tmp_path / "external_metadata_link.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file.create_dataset("model_weights/dense/kernel:0", data=np.ones((2, 2), dtype=np.float32)) + metadata = hdf5_file.create_group("metadata") + metadata["asset"] = h5py.ExternalLink("missing-assets.h5", "/asset") - # Check for any type of anomaly (could be outlier or extreme value) - has_outlier = any("abnormal weight magnitudes" in a["description"] for a in anomalies) - has_extreme = any("extremely large weight values" in a["description"] for a in anomalies) - assert has_outlier or has_extreme + scanner = WeightDistributionScanner() + weights = scanner._extract_keras_weights(str(path)) - # If outlier detection worked, check the details - outlier_anomaly = next( - (a for a in anomalies if "abnormal weight magnitudes" in a["description"]), - None, - ) - if outlier_anomaly: - assert 5 in outlier_anomaly["details"]["outlier_neurons"] + assert list(weights) == ["model_weights/dense/kernel:0"] + assert scanner.extraction_incomplete is False - def test_analyze_layer_weights_dissimilar_vectors(self): - """Test detection of dissimilar weight vectors""" - import numpy as np - scanner = WeightDistributionScanner() +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_internal_soft_link_is_resolved(tmp_path: Path) -> None: + import h5py + import numpy as np - # Create similar weight vectors - np.random.seed(42) - base_vector = np.random.randn(100) - weights = np.column_stack( - [base_vector + np.random.randn(100) * 0.1 for _ in range(9)], - ) + path = tmp_path / "soft_link_weights.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file.create_dataset("storage/dense_values", data=np.ones((2, 2), dtype=np.float32)) + hdf5_file["model_weights/dense/kernel:0"] = h5py.SoftLink("/storage/dense_values") - # Add one completely different vector (potential backdoor) - random_vector = np.random.randn(100) * 2 - weights = np.column_stack([weights, random_vector]) + scanner = WeightDistributionScanner() + weights = scanner._extract_keras_weights(str(path)) - architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) - anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) + assert list(weights) == ["model_weights/dense/kernel:0"] + assert scanner.extraction_incomplete is False - # Should detect the dissimilar vector - dissimilar_anomaly = next( - (a for a in anomalies if "dissimilar weights" in a["description"]), - None, - ) - assert dissimilar_anomaly is not None - assert dissimilar_anomaly["details"]["neuron_index"] == 9 - def test_analyze_layer_weights_extreme_values(self): - """Test detection of extreme weight values""" - import numpy as np +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_group_soft_link_preserves_weight_alias_path(tmp_path: Path) -> None: + import h5py + import numpy as np - scanner = WeightDistributionScanner() + path = tmp_path / "group_soft_link_weights.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file.create_dataset("z_storage/dense_values", data=np.ones((2, 2), dtype=np.float32)) + hdf5_file["a_model_weights"] = h5py.SoftLink("/z_storage") - # Create normal weights - np.random.seed(42) - weights = np.random.randn(100, 10) * 0.1 + scanner = WeightDistributionScanner() + weights = scanner._extract_keras_weights(str(path)) - # Add extreme values to one neuron - weights[50:55, 3] = 10.0 # Very large values + assert list(weights) == ["a_model_weights/dense_values"] + assert scanner.extraction_incomplete is False - architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) - anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) - # Should detect extreme weights - extreme_anomaly = next( - (a for a in anomalies if "extremely large weight values" in a["description"]), - None, - ) - assert extreme_anomaly is not None - assert 3 in extreme_anomaly["details"]["affected_neurons"] +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_hard_link_aliases_materialize_dataset_once( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import h5py + import numpy as np - @pytest.mark.skipif(False, reason="Dynamic skip - see test method") - def test_pytorch_model_scan(self, tmp_path: Path) -> None: - """Test scanning a PyTorch model with anomalous weights""" - if not has_torch(): - pytest.skip("PyTorch not installed") + path = tmp_path / "hard_link_weights.h5" + with h5py.File(path, "w") as hdf5_file: + dataset = hdf5_file.create_dataset("weights/original_weight", data=np.ones((2, 2), dtype=np.float32)) + hdf5_file["weights/alias_a_weight"] = dataset + hdf5_file["weights/alias_b_weight"] = dataset - import torch + materialized_datasets = 0 + original_array = np.array - scanner = WeightDistributionScanner({"enable_unsafe_torch_load": True}) + def count_dataset_materializations(obj: Any, *args: Any, **kwargs: Any) -> Any: + nonlocal materialized_datasets + if isinstance(obj, h5py.Dataset): + materialized_datasets += 1 + return original_array(obj, *args, **kwargs) - # Create a simple model with anomalous weights - class SimpleModel(torch.nn.Module): - def __init__(self): - super().__init__() - self.fc1 = torch.nn.Linear(100, 50) - self.fc2 = torch.nn.Linear(50, 10) + monkeypatch.setattr(np, "array", count_dataset_materializations) - # Make one output neuron in fc2 anomalous - with torch.no_grad(): - self.fc2.weight.data = torch.randn(10, 50) * 0.1 - self.fc2.weight.data[5] = torch.randn(50) * 10.0 # Backdoor class - more extreme + scanner = WeightDistributionScanner() + weights = scanner._extract_keras_weights(str(path)) - model = SimpleModel() + assert len(weights) == 1 + assert materialized_datasets == 1 + assert scanner.retained_tensor_bytes == 16 - model_path = tmp_path / "model.pt" - torch.save(model.state_dict(), model_path) - result = scanner.scan(str(model_path)) - assert result.success +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_cumulative_tensor_budget_stops_before_second_materialization( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import h5py + import numpy as np - # If no issues found, it might be because the scanner couldn't extract weights - # This test is more about integration than specific anomaly detection - # So we'll make it more lenient - if len(result.issues) == 0: - # Check if any layers were analyzed - assert result.metadata.get("layers_analyzed", 0) >= 0 - else: - # Check that anomaly was detected - could be either type - has_magnitude = any("abnormal weight magnitudes" in issue.message for issue in result.issues) - has_extreme = any("extremely large weight values" in issue.message for issue in result.issues) - assert has_magnitude or has_extreme + path = tmp_path / "cumulative_weights.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file.create_dataset("weights/a_weight", data=np.ones((2, 2), dtype=np.float32)) + hdf5_file.create_dataset("weights/b_weight", data=np.ones((2, 2), dtype=np.float32)) - @pytest.mark.skipif(False, reason="Dynamic skip - see test method") - def test_keras_model_scan(self): - """Test scanning a Keras model""" - if not has_h5py(): - pytest.skip("h5py not installed") + materialized_names: list[str] = [] + original_array = np.array - import h5py - import numpy as np + def track_dataset_materialization(obj: Any, *args: Any, **kwargs: Any) -> Any: + if isinstance(obj, h5py.Dataset): + materialized_names.append(obj.name) + return original_array(obj, *args, **kwargs) - scanner = WeightDistributionScanner() + monkeypatch.setattr(np, "array", track_dataset_materialization) - # Create a simple H5 file with weights - with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: - with h5py.File(f.name, "w") as hf: - # Create weight arrays - normal_weights = np.random.randn(100, 10) * 0.1 - normal_weights[:, 5] = np.random.randn(100) * 3.0 # Anomalous + scanner = WeightDistributionScanner({"max_array_size": 32, "max_weight_distribution_total_bytes": 20}) + weights = scanner._extract_keras_weights(str(path)) - # Store as Keras would - hf.create_dataset("model_weights/dense_1/kernel:0", data=normal_weights) + assert list(weights) == ["weights/a_weight"] + assert materialized_names == ["/weights/a_weight"] + assert scanner.extraction_incomplete is True + assert scanner.extraction_incomplete_reasons == ["keras_tensor_size_limit_total"] + assert scanner.extraction_incomplete_details["max_total_tensor_bytes"] == 20 - temp_path = f.name - try: - result = scanner.scan(temp_path) - assert result.success - # Should detect anomaly in the weights - assert len(result.issues) > 0 +@pytest.mark.skipif(not HAS_NUMPY or not has_h5py(), reason="numpy and h5py required") +def test_hdf5_external_group_link_is_inconclusive_and_not_cached(tmp_path: Path) -> None: + import h5py - finally: + path = tmp_path / "external_group.h5" + with h5py.File(path, "w") as hdf5_file: + hdf5_file["layers"] = h5py.ExternalLink("missing.h5", "/model_weights") + + cache_dir = tmp_path / "cache" + reset_cache_manager() + try: + for _ in range(2): + result = core.scan_model_directory_or_file( + str(path), + scanners=["weight_distribution"], + cache_enabled=True, + cache_dir=str(cache_dir), + min_cache_file_size=0, + ) + assert result.success is False + assert core.determine_exit_code(result) == 2 + assert result.file_metadata[str(path)]["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert get_cache_manager(str(cache_dir), enabled=True).get_stats()["total_entries"] == 0 + finally: + reset_cache_manager() + + +def test_pytorch_primary_load_is_blocked_by_archive_budget( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + load_called = False + + def fake_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> dict[str, object]: + del map_location, weights_only + nonlocal load_called + load_called = True + return {} + + fake_torch.load = fake_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "large.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps({})) + archive.writestr("data/0", b"x" * 64) + + scanner = WeightDistributionScanner( + { + "enable_unsafe_torch_load": True, + "max_array_size": 1024, + "max_weight_distribution_total_bytes": 32, + } + ) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert load_called is False + assert scanner.extraction_incomplete_reasons == ["pytorch_load_size_limit"] + + +def test_pytorch_primary_storage_is_checked_before_load( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> object: + del map_location, weights_only + raise AssertionError("oversized storage should be rejected before torch.load") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "oversized-storage.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("archive/data.pkl", pickle.dumps({})) + archive.writestr("archive/data/0", b"x" * 64) + + scanner = WeightDistributionScanner( + { + "enable_unsafe_torch_load": True, + "max_array_size": 8, + "max_weight_distribution_total_bytes": 128, + } + ) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["pytorch_tensor_storage_size_limit"] + + +def test_pytorch_load_budget_ignores_unrelated_archive_members( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + load_called = False + + def fake_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> dict[str, object]: + del map_location, weights_only + nonlocal load_called + load_called = True + return {} + + fake_torch.load = fake_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "metadata-heavy.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("archive/data.pkl", pickle.dumps({})) + archive.writestr("archive/data/0", b"x" * 26) + archive.writestr("archive/metadata.json", b"x" * 1024) + + scanner = WeightDistributionScanner( + { + "enable_unsafe_torch_load": True, + "max_array_size": 64, + "max_weight_distribution_total_bytes": 64, + } + ) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert load_called is True + assert scanner.extraction_incomplete is False + + +def test_pytorch_pickle_object_budget_blocks_compact_container_amplification( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> object: + del map_location, weights_only + raise AssertionError("object-heavy pickle should be rejected before torch.load") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "object-heavy.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("archive/data.pkl", pickle.dumps([[] for _ in range(2000)], protocol=4)) + + scanner = WeightDistributionScanner({"max_array_size": 8192, "max_weight_distribution_total_bytes": 8192}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["pytorch_pickle_object_budget"] + + +def test_pytorch_blocked_load_fallback_honors_remaining_metadata_budget( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.5.1" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fake_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> dict[str, object]: + del map_location, weights_only + raise AssertionError("unsafe torch.load must remain blocked") + + fake_torch.load = fake_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "unsafe-version.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps({"layer.weight": [[1.0, 2.0], [3.0, 4.0]]})) + + original_open = zipfile.ZipFile.open + + def fail_if_data_pkl_is_opened( + archive: zipfile.ZipFile, + name: str | zipfile.ZipInfo, + mode: str = "r", + pwd: bytes | None = None, + *, + force_zip64: bool = False, + ) -> Any: + member_name = name.filename if isinstance(name, zipfile.ZipInfo) else name + if member_name == "data.pkl": + raise AssertionError("over-budget data.pkl should not be opened") + return original_open(archive, name, mode=mode, pwd=pwd, force_zip64=force_zip64) + + monkeypatch.setattr(zipfile.ZipFile, "open", fail_if_data_pkl_is_opened) + + scanner = WeightDistributionScanner({"max_array_size": 1024, "max_weight_distribution_total_bytes": 32}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert scanner.extraction_unsafe is False + assert scanner.extraction_incomplete_reasons == ["pytorch_zip_data_pkl_size_limit"] + assert scanner.extraction_incomplete_details["max_pickle_metadata_bytes"] == 32 + + +def test_pytorch_primary_tensor_size_is_checked_before_numpy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + shape = (2, 2) + + @staticmethod + def numel() -> int: + return 4 + + @staticmethod + def element_size() -> int: + return 4 + + @staticmethod + def detach() -> object: + raise AssertionError("oversized tensor should not be converted") + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fake_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> dict[str, object]: + del map_location, weights_only + return {"layer.weight": FakeTensor()} + + fake_torch.load = fake_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "tensor.pt" + path.write_bytes(b"x") + scanner = WeightDistributionScanner( + { + "enable_unsafe_torch_load": True, + "max_array_size": 8, + "max_weight_distribution_total_bytes": 1024, + } + ) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["pytorch_tensor_size_limit"] + assert scanner.extraction_incomplete_details["tensor_nbytes"] == 16 + + +def test_pytorch_unsupported_tensor_does_not_hide_valid_weights( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import numpy as np + + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + shape = (2, 2) + + def __init__(self, *, supported: bool) -> None: + self.supported = supported + + @staticmethod + def numel() -> int: + return 4 + + @staticmethod + def element_size() -> int: + return 4 + + def detach(self) -> "FakeTensor": + return self + + def cpu(self) -> "FakeTensor": + return self + + def numpy(self) -> Any: + if not self.supported: + raise TypeError("unsupported tensor layout") + return np.ones((2, 2), dtype=np.float32) + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fake_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> dict[str, object]: + del map_location, weights_only + return { + "dense.weight": FakeTensor(supported=True), + "sparse.weight": FakeTensor(supported=False), + } + + fake_torch.load = fake_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + path = tmp_path / "mixed.pt" + path.write_bytes(b"x") + scanner = WeightDistributionScanner({"enable_unsafe_torch_load": True}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert list(weights) == ["dense.weight"] + assert scanner.extraction_incomplete is True + assert scanner.extraction_incomplete_reasons == ["pytorch_tensor_read_failed"] + assert scanner.extraction_incomplete_details["failed_tensors"] == ["sparse.weight"] + + +def test_pytorch_zip_alias_expansion_is_rejected_before_numpy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import numpy as np + + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> object: + del map_location, weights_only + raise RuntimeError("force restricted fallback") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + row = [1.0] * 100 + payload = {"layer.weight": [row] * 100} + path = tmp_path / "aliased.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps(payload, protocol=4)) + + array_called = False + original_array = np.array + + def track_array(value: Any, *args: Any, **kwargs: Any) -> Any: + nonlocal array_called + array_called = True + return original_array(value, *args, **kwargs) + + monkeypatch.setattr(np, "array", track_array) + + scanner = WeightDistributionScanner({"max_array_size": 4096, "max_weight_distribution_total_bytes": 8192}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert weights == {} + assert array_called is False + assert scanner.extraction_incomplete_reasons == ["pytorch_tensor_materialization_failed"] + + +def test_pytorch_zip_small_shared_sequence_is_not_treated_as_cycle( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load(*_args: object, **_kwargs: object) -> object: + raise RuntimeError("force restricted fallback") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + row = [1.0, 2.0] + path = tmp_path / "small-alias.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps({"layer.weight": [row, row]}, protocol=4)) + + scanner = WeightDistributionScanner({"max_array_size": 1024, "max_weight_distribution_total_bytes": 1024}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert list(weights) == ["layer.weight"] + assert weights["layer.weight"].shape == (2, 2) + assert scanner.extraction_incomplete is False + + +def test_pytorch_zip_discarded_vector_does_not_consume_retained_budget( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> object: + del map_location, weights_only + raise RuntimeError("force restricted fallback") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + payload = { + "weight_names": [1] * 100, + "layer.weight": [[1] * 10 for _ in range(10)], + } + path = tmp_path / "vector_metadata.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps(payload, protocol=4)) + + scanner = WeightDistributionScanner({"max_array_size": 1000, "max_weight_distribution_total_bytes": 1000}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert list(weights) == ["layer.weight"] + assert scanner.retained_tensor_bytes == 800 + assert scanner.extraction_incomplete is False + + +def test_pytorch_zip_ignores_large_non_weight_metadata( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load( + _path: str, + *, + map_location: object, + weights_only: bool = False, + ) -> object: + del map_location, weights_only + raise RuntimeError("force restricted fallback") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + row = [1.0] * 100 + payload = { + "metadata": [row] * 100, + "layer.weight": [[1.0, 2.0], [3.0, 4.0]], + } + path = tmp_path / "metadata.pt" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps(payload, protocol=4)) + + scanner = WeightDistributionScanner({"max_array_size": 4096, "max_weight_distribution_total_bytes": 8192}) + weights = scanner._extract_pytorch_weights(str(path)) + + assert list(weights) == ["layer.weight"] + assert weights["layer.weight"].shape == (2, 2) + assert scanner.extraction_incomplete is False + + +def test_tensorflow_unknown_dtype_is_skipped_before_load_variable( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + saved_model_dir = tmp_path / "saved_model" + variables_dir = saved_model_dir / "variables" + variables_dir.mkdir(parents=True) + (variables_dir / "variables.index").write_bytes(b"checkpoint index") + + class FakeCheckpointReader: + @staticmethod + def get_variable_to_dtype_map() -> dict[str, object]: + return {"dense/kernel": object} + + class FakeTrain: + @staticmethod + def load_checkpoint(_prefix: str) -> FakeCheckpointReader: + return FakeCheckpointReader() + + @staticmethod + def list_variables(_prefix: str) -> list[tuple[str, list[int]]]: + return [("dense/kernel", [2, 2])] + + @staticmethod + def load_variable(_prefix: str, _name: str) -> object: + raise AssertionError("unknown dtype variable should not be loaded") + + fake_tensorflow: Any = types.ModuleType("tensorflow") + fake_tensorflow.train = FakeTrain + monkeypatch.setitem(sys.modules, "tensorflow", fake_tensorflow) + + scanner = WeightDistributionScanner() + weights = scanner._extract_tensorflow_weights(str(saved_model_dir)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["tensorflow_checkpoint_tensor_size_limit"] + + +def test_fixed_width_custom_numeric_dtype_remains_bounded() -> None: + import numpy as np + + scanner = WeightDistributionScanner() + + assert scanner._tensorflow_dtype_itemsize(np.dtype("V2")) == 2 + assert scanner._tensorflow_checkpoint_variable_nbytes([2, 2], np.dtype("V2")) == 8 + + +def test_onnx_external_and_unknown_initializers_fail_closed_before_materialization( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import numpy as np + + initializers = [ + types.SimpleNamespace( + name="external_weight", + dims=[2, 2], + data_type=1, + data_location=1, + external_data=[object()], + ), + types.SimpleNamespace( + name="unknown_weight", + dims=[2, 2], + data_type=999, + data_location=0, + external_data=[], + ), + types.SimpleNamespace( + name="valid_weight", + dims=[2, 2], + data_type=1, + data_location=0, + external_data=[], + ), + types.SimpleNamespace( + name="custom_numeric_weight", + dims=[2, 2], + data_type=16, + data_location=0, + external_data=[], + ), + ] + materialized: list[str] = [] + + def tensor_dtype_to_np_dtype(data_type: int) -> Any: + dtypes = {1: np.dtype("float32"), 16: np.dtype("V2")} + return dtypes[data_type] + + def to_array(initializer: Any) -> Any: + materialized.append(initializer.name) + return np.ones((2, 2), dtype=np.float32) + + fake_onnx: Any = types.ModuleType("onnx") + fake_onnx.TensorProto = types.SimpleNamespace(EXTERNAL=1) + fake_onnx.helper = types.SimpleNamespace(tensor_dtype_to_np_dtype=tensor_dtype_to_np_dtype) + fake_onnx.mapping = types.SimpleNamespace(TENSOR_TYPE_TO_NP_TYPE={1: np.dtype("float32"), 16: np.dtype("V2")}) + fake_onnx.numpy_helper = types.SimpleNamespace(to_array=to_array) + fake_onnx.load = lambda _path, load_external_data: types.SimpleNamespace( + graph=types.SimpleNamespace(initializer=initializers) + ) + monkeypatch.setitem(sys.modules, "onnx", fake_onnx) + + path = tmp_path / "model.onnx" + path.write_bytes(b"x") + scanner = WeightDistributionScanner() + weights = scanner._extract_onnx_weights(str(path)) + + assert list(weights) == ["valid_weight", "custom_numeric_weight"] + assert materialized == ["valid_weight", "custom_numeric_weight"] + assert scanner.extraction_incomplete_reasons == [ + "onnx_external_initializer_skipped", + "onnx_initializer_size_limit", + ] + + +def test_onnx_inline_storage_is_bounded_before_materialization( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + import numpy as np + + initializers = [ + types.SimpleNamespace( + name="oversized_weight", + dims=[2, 2], + data_type=1, + data_location=0, + external_data=[], + raw_data=b"x" * 128, + ), + types.SimpleNamespace( + name="packed_weight", + dims=[2, 2], + data_type=21, + data_location=0, + external_data=[], + raw_data=b"x" * 40, + ), + ] + + def fail_to_array(_initializer: Any) -> Any: + raise AssertionError("oversized inline storage should be rejected before to_array") + + fake_onnx: Any = types.ModuleType("onnx") + fake_onnx.TensorProto = types.SimpleNamespace(EXTERNAL=1, UINT4=21) + fake_onnx.helper = types.SimpleNamespace( + tensor_dtype_to_np_dtype=lambda data_type: np.dtype("float32") if data_type == 1 else np.dtype("uint8") + ) + fake_onnx.mapping = types.SimpleNamespace(TENSOR_TYPE_TO_NP_TYPE={}) + fake_onnx.numpy_helper = types.SimpleNamespace(to_array=fail_to_array) + fake_onnx.load = lambda _path, load_external_data: types.SimpleNamespace( + graph=types.SimpleNamespace(initializer=initializers) + ) + monkeypatch.setitem(sys.modules, "onnx", fake_onnx) + + path = tmp_path / "inline-storage.onnx" + path.write_bytes(b"x") + scanner = WeightDistributionScanner({"max_array_size": 64}) + weights = scanner._extract_onnx_weights(str(path)) + + assert weights == {} + assert scanner.extraction_incomplete_reasons == ["onnx_initializer_storage_size_limit"] + assert scanner.extraction_incomplete_details["failed_tensors"] == [ + "oversized_weight", + "packed_weight", + ] + + +def test_repeated_oversized_tensors_do_not_multiply_configured_limit() -> None: + scanner = WeightDistributionScanner({"max_array_size": 10}) + + scanner._record_oversized_tensor("tensor_size_limit", "first", tensor_nbytes=11) + scanner._record_oversized_tensor("tensor_size_limit", "second", tensor_nbytes=12) + + assert scanner.extraction_incomplete_details["oversized_tensors"] == 2 + assert scanner.extraction_incomplete_details["max_array_size"] == 10 + assert scanner.extraction_incomplete_details["tensor_nbytes"] == [11, 12] + + +@pytest.mark.skipif(not HAS_NUMPY, reason="numpy not available") +class TestWeightDistributionScanner: + """Test suite for weight distribution anomaly detection""" + + def _create_mock_architecture_analysis(self, is_llm=False, is_transformer=False): + """Helper method to create mock architecture analysis for testing""" + return { + "is_likely_transformer": is_transformer, + "is_likely_llm": is_llm, + "confidence": 0.8 if is_llm else 0.5, + "evidence": ["Mock evidence for testing"], + "architectural_features": {}, + "total_parameters": 100_000_000 if is_llm else 1_000_000, + "layer_count": 24 if is_llm else 3, + } + + def test_scanner_initialization(self): + """Test scanner initialization with default and custom config""" + # Default initialization + scanner = WeightDistributionScanner() + assert scanner.z_score_threshold == 3.0 + assert scanner.cosine_similarity_threshold == 0.7 + assert scanner.weight_magnitude_threshold == 3.0 + assert scanner.max_array_size == 100 * 1024 * 1024 # Default 100MB + + # Custom config + config = { + "z_score_threshold": 2.5, + "cosine_similarity_threshold": 0.8, + "weight_magnitude_threshold": 2.0, + "max_array_size": 50 * 1024 * 1024, # 50MB + } + scanner = WeightDistributionScanner(config) + assert scanner.z_score_threshold == 2.5 + assert scanner.cosine_similarity_threshold == 0.8 + assert scanner.weight_magnitude_threshold == 2.0 + assert scanner.max_array_size == 50 * 1024 * 1024 + + def test_numeric_byte_limit_config_is_coerced(self) -> None: + import numpy as np + + scanner = WeightDistributionScanner( + { + "max_array_size": 1.5, + "max_weight_distribution_total_bytes": np.int64(32), + } + ) + + assert scanner._max_tensor_bytes() == 1 + assert scanner._max_total_tensor_bytes() == 32 + + @pytest.mark.parametrize("invalid_limit", [True, False, "unbounded", -1, float("inf"), float("nan")]) + def test_invalid_tensor_byte_limit_uses_secure_default(self, invalid_limit: Any) -> None: + scanner = WeightDistributionScanner({"max_array_size": invalid_limit}) + + assert scanner._max_tensor_bytes() == 100 * 1024 * 1024 + assert scanner._max_total_tensor_bytes() == 512 * 1024 * 1024 + + def test_zero_tensor_byte_limit_remains_explicitly_unlimited(self) -> None: + scanner = WeightDistributionScanner({"max_array_size": 0}) + + assert scanner._max_tensor_bytes() is None + assert scanner._max_total_tensor_bytes() is None + + bounded_total_scanner = WeightDistributionScanner( + {"max_array_size": 0, "max_weight_distribution_total_bytes": 64} + ) + assert bounded_total_scanner._max_tensor_bytes() is None + assert bounded_total_scanner._max_total_tensor_bytes() == 64 + + def test_can_handle(self): + """Test file type detection""" + # Create temporary files to test can_handle + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: + pt_path = f.name + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + h5_path = f.name + with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f: + txt_path = f.name + tf_dir = tempfile.mkdtemp() + open(os.path.join(tf_dir, "saved_model.pb"), "wb").close() + + try: + # Should handle PyTorch files if torch is available + if _has_torch_cached(): + assert WeightDistributionScanner.can_handle(pt_path) + + # Should handle Keras files if h5py is available + if _has_h5py_cached(): + assert WeightDistributionScanner.can_handle(h5_path) + + if _has_tensorflow_cached(): + assert WeightDistributionScanner.can_handle(tf_dir) + + # Should not handle unsupported extensions + assert not WeightDistributionScanner.can_handle(txt_path) + assert not WeightDistributionScanner.can_handle("directory/") + finally: + os.unlink(pt_path) + os.unlink(h5_path) + os.unlink(txt_path) + os.unlink(os.path.join(tf_dir, "saved_model.pb")) + os.rmdir(tf_dir) + + def test_analyze_layer_weights_outlier_detection(self): + """Test detection of outlier weight vectors""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create normal weights with one outlier + np.random.seed(42) + normal_weights = np.random.randn(100, 10) * 0.1 # Small weights + + # Make one neuron an outlier with large weights - make it even more extreme + normal_weights[:, 5] = np.random.randn(100) * 10.0 # Much larger weights + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) + anomalies = scanner._analyze_layer_weights("test_layer", normal_weights, architecture_analysis) + + # Should detect the outlier neuron + assert len(anomalies) > 0 + + # Check for any type of anomaly (could be outlier or extreme value) + has_outlier = any("abnormal weight magnitudes" in a["description"] for a in anomalies) + has_extreme = any("extremely large weight values" in a["description"] for a in anomalies) + assert has_outlier or has_extreme + + # If outlier detection worked, check the details + outlier_anomaly = next( + (a for a in anomalies if "abnormal weight magnitudes" in a["description"]), + None, + ) + if outlier_anomaly: + assert 5 in outlier_anomaly["details"]["outlier_neurons"] + + def test_analyze_layer_weights_dissimilar_vectors(self): + """Test detection of dissimilar weight vectors""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create similar weight vectors + np.random.seed(42) + base_vector = np.random.randn(100) + weights = np.column_stack( + [base_vector + np.random.randn(100) * 0.1 for _ in range(9)], + ) + + # Add one completely different vector (potential backdoor) + random_vector = np.random.randn(100) * 2 + weights = np.column_stack([weights, random_vector]) + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) + anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) + + # Should detect the dissimilar vector + dissimilar_anomaly = next( + (a for a in anomalies if "dissimilar weights" in a["description"]), + None, + ) + assert dissimilar_anomaly is not None + assert dissimilar_anomaly["details"]["neuron_index"] == 9 + + def test_analyze_layer_weights_extreme_values(self): + """Test detection of extreme weight values""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create normal weights + np.random.seed(42) + weights = np.random.randn(100, 10) * 0.1 + + # Add extreme values to one neuron + weights[50:55, 3] = 10.0 # Very large values + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) + anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) + + # Should detect extreme weights + extreme_anomaly = next( + (a for a in anomalies if "extremely large weight values" in a["description"]), + None, + ) + assert extreme_anomaly is not None + assert 3 in extreme_anomaly["details"]["affected_neurons"] + + @pytest.mark.skipif(False, reason="Dynamic skip - see test method") + def test_pytorch_model_scan(self, tmp_path: Path) -> None: + """Test scanning a PyTorch model with anomalous weights""" + if not has_torch(): + pytest.skip("PyTorch not installed") + + import torch + + scanner = WeightDistributionScanner({"enable_unsafe_torch_load": True}) + + # Create a simple model with anomalous weights + class SimpleModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc1 = torch.nn.Linear(100, 50) + self.fc2 = torch.nn.Linear(50, 10) + + # Make one output neuron in fc2 anomalous + with torch.no_grad(): + self.fc2.weight.data = torch.randn(10, 50) * 0.1 + self.fc2.weight.data[5] = torch.randn(50) * 10.0 # Backdoor class - more extreme + + model = SimpleModel() + + model_path = tmp_path / "model.pt" + torch.save(model.state_dict(), model_path) + + result = scanner.scan(str(model_path)) + assert result.success + + # If no issues found, it might be because the scanner couldn't extract weights + # This test is more about integration than specific anomaly detection + # So we'll make it more lenient + if len(result.issues) == 0: + # Check if any layers were analyzed + assert result.metadata.get("layers_analyzed", 0) >= 0 + else: + # Check that anomaly was detected - could be either type + has_magnitude = any("abnormal weight magnitudes" in issue.message for issue in result.issues) + has_extreme = any("extremely large weight values" in issue.message for issue in result.issues) + assert has_magnitude or has_extreme + + @pytest.mark.skipif(False, reason="Dynamic skip - see test method") + def test_keras_model_scan(self): + """Test scanning a Keras model""" + if not has_h5py(): + pytest.skip("h5py not installed") + + import h5py + import numpy as np + + scanner = WeightDistributionScanner() + + # Create a simple H5 file with weights + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + with h5py.File(f.name, "w") as hf: + # Create weight arrays + normal_weights = np.random.randn(100, 10) * 0.1 + normal_weights[:, 5] = np.random.randn(100) * 3.0 # Anomalous + + # Store as Keras would + hf.create_dataset("model_weights/dense_1/kernel:0", data=normal_weights) + + temp_path = f.name + + try: + result = scanner.scan(temp_path) + assert result.success + # Should detect anomaly in the weights + assert len(result.issues) > 0 + + finally: + os.unlink(temp_path) + + @pytest.mark.skipif(False, reason="Dynamic skip - see test method") + def test_tensorflow_savedmodel_scan(self, tmp_path): + """Test scanning a TensorFlow SavedModel directory.""" + import sys + + if not has_tensorflow(): + pytest.skip("TensorFlow not installed") + + # Skip on Python 3.12+ due to TensorFlow/typing compatibility issues + if sys.version_info >= (3, 12): + pytest.skip("TensorFlow SavedModel has compatibility issues with Python 3.12+") + + import tensorflow as tf + + # Skip if tf.keras is not available (newer TensorFlow versions separate Keras) + if not hasattr(tf, "keras"): + pytest.skip("tf.keras not available (Keras may be a separate package)") + + scanner = WeightDistributionScanner() + + model = tf.keras.Sequential([tf.keras.layers.Dense(2, input_shape=(3,))]) # type: ignore[call-arg] + saved_path = tmp_path / "tf_model" + tf.saved_model.save(model, str(saved_path)) + + result = scanner.scan(str(saved_path)) + assert result.success + assert result.metadata.get("layers_analyzed", 0) > 0 + + def test_empty_model_handling(self): + """Test handling of models with no extractable weights""" + scanner = WeightDistributionScanner() + + # Create an empty file + with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: + f.write(b"") + temp_path = f.name + + try: + result = scanner.scan(temp_path) + # Should handle gracefully + assert result.success or len(result.issues) > 0 + + finally: os.unlink(temp_path) - @pytest.mark.skipif(False, reason="Dynamic skip - see test method") - def test_tensorflow_savedmodel_scan(self, tmp_path): - """Test scanning a TensorFlow SavedModel directory.""" - import sys + def test_pytorch_zip_data_pkl_safe_extraction( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + """Ensure safe pickle in PyTorch ZIP can be parsed without code execution""" + load_called = False + data = {"layer.weight": [[1.0, 2.0], [3.0, 4.0]]} + data_bytes = pickle.dumps(data, protocol=4) + zip_path = tmp_path / "model.pt" + with zipfile.ZipFile(zip_path, "w") as z: + z.writestr("data.pkl", data_bytes) + + def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: + nonlocal load_called + load_called = True + return {} + + _install_fake_torch(monkeypatch, fake_load) + scanner = WeightDistributionScanner() + weights = scanner._extract_pytorch_weights(str(zip_path)) + assert not scanner.extraction_unsafe + assert not scanner.extraction_incomplete + assert load_called is False + assert "layer.weight" in weights + assert weights["layer.weight"].shape == (2, 2) + + def test_pytorch_zip_data_pkl_size_limit_skips_before_opening_member( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + fake_torch: Any = types.ModuleType("torch") + fake_torch.__version__ = "2.6.0" + + class FakeTensor: # pragma: no cover - simple test double + pass + + fake_torch.Tensor = FakeTensor + fake_torch.device = lambda value: value + + def fail_load(*_args: object, **_kwargs: object) -> object: + raise RuntimeError("force zip fallback") + + fake_torch.load = fail_load + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + zip_path = tmp_path / "model.pt" + with zipfile.ZipFile(zip_path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps({"layer.weight": [[1.0, 2.0], [3.0, 4.0]]})) + + original_open = zipfile.ZipFile.open + + def fail_if_data_pkl_is_opened( + archive: zipfile.ZipFile, + name: str | zipfile.ZipInfo, + mode: str = "r", + pwd: bytes | None = None, + *, + force_zip64: bool = False, + ) -> Any: + member_name = name.filename if isinstance(name, zipfile.ZipInfo) else name + if member_name == "data.pkl": + raise AssertionError("oversized data.pkl should not be opened") + return original_open(archive, name, mode=mode, pwd=pwd, force_zip64=force_zip64) + + monkeypatch.setattr(zipfile.ZipFile, "open", fail_if_data_pkl_is_opened) + + scanner = WeightDistributionScanner({"max_array_size": 1}) + weights = scanner._extract_pytorch_weights(str(zip_path)) + + assert weights == {} + assert scanner.extraction_incomplete is True + assert scanner.extraction_incomplete_reasons == ["pytorch_zip_data_pkl_size_limit"] + assert scanner.extraction_incomplete_details["failed_tensors"] == ["data.pkl"] + + @pytest.mark.skipif(False, reason="Dynamic skip - see test method") + def test_pytorch_zip_data_pkl_unsafe_extraction(self, monkeypatch, tmp_path): + """Unsafe pickle opcodes should be flagged""" + if not has_torch(): + pytest.skip("PyTorch not installed") + + import torch + + model = torch.nn.Linear(2, 2) + zip_path = tmp_path / "model.pt" + torch.save(model.state_dict(), zip_path) + + def fail_load(*_args, **_kwargs): + raise RuntimeError("fail") + + # Mock torch.load directly since torch is now a lazy import + import torch + + monkeypatch.setattr(torch, "load", fail_load) + scanner = WeightDistributionScanner() + weights = scanner._extract_pytorch_weights(str(zip_path)) + assert weights == {} + assert scanner.extraction_unsafe + + def test_tensorflow_checkpoint_size_limit_skips_before_load_variable( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + import numpy as np + + saved_model_dir = tmp_path / "saved_model" + variables_dir = saved_model_dir / "variables" + variables_dir.mkdir(parents=True) + (variables_dir / "variables.index").write_bytes(b"checkpoint index") + + loaded_variables: list[str] = [] + + class FakeCheckpointReader: + def get_variable_to_dtype_map(self) -> dict[str, object]: + return { + "dense/kernel": np.dtype("float32"), + "small/kernel": np.dtype("float32"), + } + + class FakeTrain: + @staticmethod + def load_checkpoint(_prefix: str) -> FakeCheckpointReader: + return FakeCheckpointReader() + + @staticmethod + def list_variables(_prefix: str) -> list[tuple[str, list[int]]]: + return [ + ("dense/kernel", [1024, 1024]), + ("small/kernel", [2, 2]), + ] + + @staticmethod + def load_variable(_prefix: str, name: str) -> object: + if name == "dense/kernel": + raise AssertionError("oversized checkpoint variable should not be loaded") + loaded_variables.append(name) + return np.ones((2, 2), dtype=np.float32) + + fake_tensorflow: Any = types.ModuleType("tensorflow") + fake_tensorflow.train = FakeTrain + monkeypatch.setitem(sys.modules, "tensorflow", fake_tensorflow) + + scanner = WeightDistributionScanner({"max_array_size": 128}) + weights = scanner._extract_tensorflow_weights(str(saved_model_dir)) + + assert loaded_variables == ["small/kernel"] + assert list(weights) == ["small/kernel"] + assert scanner.extraction_incomplete is True + assert "tensorflow_checkpoint_tensor_size_limit" in scanner.extraction_incomplete_reasons + assert scanner.extraction_incomplete_details["failed_tensors"] == ["dense/kernel"] + + def test_tensorflow_checkpoint_skips_rank_one_weight_before_budgeting( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + import numpy as np + + saved_model_dir = tmp_path / "saved_model" + variables_dir = saved_model_dir / "variables" + variables_dir.mkdir(parents=True) + (variables_dir / "variables.index").write_bytes(b"checkpoint index") + + loaded_variables: list[str] = [] + + class FakeCheckpointReader: + @staticmethod + def get_variable_to_dtype_map() -> dict[str, object]: + return { + "class_weight": np.dtype("float32"), + "dense/kernel": np.dtype("float32"), + } + + class FakeTrain: + @staticmethod + def load_checkpoint(_prefix: str) -> FakeCheckpointReader: + return FakeCheckpointReader() + + @staticmethod + def list_variables(_prefix: str) -> list[tuple[str, list[int]]]: + return [("class_weight", [1024]), ("dense/kernel", [2, 2])] + + @staticmethod + def load_variable(_prefix: str, name: str) -> object: + loaded_variables.append(name) + return np.ones((2, 2), dtype=np.float32) + + fake_tensorflow: Any = types.ModuleType("tensorflow") + fake_tensorflow.train = FakeTrain + monkeypatch.setitem(sys.modules, "tensorflow", fake_tensorflow) + + scanner = WeightDistributionScanner({"max_array_size": 128}) + weights = scanner._extract_tensorflow_weights(str(saved_model_dir)) + + assert loaded_variables == ["dense/kernel"] + assert list(weights) == ["dense/kernel"] + assert scanner.extraction_incomplete is False + + def test_multiple_anomalies(self): + """Test detection of multiple types of anomalies in one layer""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create weights with multiple issues + np.random.seed(42) + weights = np.random.randn(100, 10) * 0.1 + + # Neuron 3: Large magnitude outlier + weights[:, 3] = np.random.randn(100) * 15.0 # More extreme outlier + + # Neuron 7: Dissimilar to others + weights[:, 7] = np.random.randn(100) * 0.5 + 10.0 + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) + anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) + + # Should detect at least one anomaly + assert len(anomalies) >= 1 + + # Check for any type of anomaly + has_magnitude_anomaly = any("abnormal weight magnitudes" in a["description"] for a in anomalies) + has_dissimilar_anomaly = any("dissimilar weights" in a["description"] for a in anomalies) + has_extreme_anomaly = any("extremely large weight values" in a["description"] for a in anomalies) + + assert has_magnitude_anomaly or has_dissimilar_anomaly or has_extreme_anomaly + + def test_llm_vocabulary_layer_handling(self): + """Test that LLM vocabulary layers don't produce false positives""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create a large vocabulary layer like in LLMs (e.g., 32k vocab) + np.random.seed(42) + vocab_size = 32000 + hidden_dim = 4096 + weights = np.random.randn(hidden_dim, vocab_size) * 0.02 # Typical LLM init + + # Add some natural variation (not anomalous) + for i in range(100): + weights[:, i] *= 1.2 # Some tokens might have slightly different scales + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) + + # Should not flag many neurons in an LLM + # With our new thresholds, we expect very few or no anomalies + assert len(anomalies) <= 1 # At most 1 anomaly type + + # If there are anomalies, they should affect very few neurons + for anomaly in anomalies: + if "outlier_neurons" in anomaly["details"]: + # Should be less than 0.1% of neurons + assert anomaly["details"]["total_outliers"] < vocab_size * 0.001 + + def test_llm_checks_disabled_by_default(self): + """Test that LLM checks are disabled by default""" + import numpy as np + + scanner = WeightDistributionScanner() + + # Create LLM-like weights + weights = np.random.randn(4096, 32000) * 0.02 + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) + + # Should return no anomalies since LLM checks are disabled by default + assert len(anomalies) == 0 + + def test_llm_checks_can_be_enabled(self): + """Test that LLM checks can be explicitly enabled via config""" + import numpy as np + + config = {"enable_llm_checks": True} + scanner = WeightDistributionScanner(config) + + # Create LLM-like weights with some outliers + np.random.seed(42) + weights = np.random.randn(4096, 32000) * 0.02 + # Make a few neurons extreme outliers + weights[:, 0] = np.random.randn(4096) * 10.0 + weights[:, 1] = np.random.randn(4096) * 10.0 + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) - if not has_tensorflow(): - pytest.skip("TensorFlow not installed") + # With LLM checks enabled, might detect extreme outliers with strict thresholds + # We made 2 extreme neurons, so could get up to 2 anomaly types (outlier + extreme) + assert len(anomalies) <= 2 - # Skip on Python 3.12+ due to TensorFlow/typing compatibility issues - if sys.version_info >= (3, 12): - pytest.skip("TensorFlow SavedModel has compatibility issues with Python 3.12+") + # Should only flag the 2 neurons we made extreme + for anomaly in anomalies: + if "outlier_neurons" in anomaly["details"]: + assert anomaly["details"]["total_outliers"] <= 2 - import tensorflow as tf + def test_gpt2_layer_pattern_detection(self): + """Test that GPT-2 style layer patterns are detected as LLM layers""" + import numpy as np - # Skip if tf.keras is not available (newer TensorFlow versions separate Keras) - if not hasattr(tf, "keras"): - pytest.skip("tf.keras not available (Keras may be a separate package)") + scanner = WeightDistributionScanner() + + # Test GPT-2 style layer names + gpt2_layer_names = [ + "h.0.mlp.c_fc.weight", + "h.1.attn.c_attn.weight", + "h.11.mlp.c_proj.weight", + "transformer.h.5.mlp.c_fc.weight", + ] + + # Create typical GPT-2 MLP weights (3072 -> 768 for GPT-2 base) + np.random.seed(42) + weights = np.random.randn(3072, 768) * 0.02 + + # Add some natural variation + weights[:, :10] *= 1.5 # Some neurons have different scales + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + for layer_name in gpt2_layer_names: + anomalies = scanner._analyze_layer_weights(layer_name, weights, architecture_analysis) + + # Should return no anomalies due to LLM detection + assert len(anomalies) == 0, f"Layer {layer_name} should be detected as LLM" + + def test_transformer_layer_pattern_detection(self): + """Test that transformer-related layers use structural analysis instead of name-based detection.""" + import numpy as np scanner = WeightDistributionScanner() - model = tf.keras.Sequential([tf.keras.layers.Dense(2, input_shape=(3,))]) # type: ignore[call-arg] - saved_path = tmp_path / "tf_model" - tf.saved_model.save(model, str(saved_path)) + # Test transformer-related layer names + transformer_patterns = [ + "encoder.layers.0.mlp.dense_h_to_4h.weight", + "decoder.attention.dense.weight", + "transformer.mlp.fc_in.weight", + "model.layers.5.mlp.gate_proj.weight", + ] - result = scanner.scan(str(saved_path)) - assert result.success - assert result.metadata.get("layers_analyzed", 0) > 0 + # Create transformer-like weights with some natural variation + np.random.seed(42) + weights = np.random.randn(1024, 4096) * 0.02 # Typical transformer dimensions + + # Add moderate natural variation (not extreme anomalies) + weights[:, :10] *= 1.2 # Some neurons have slightly different scales + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True, is_transformer=True) + for layer_name in transformer_patterns: + anomalies = scanner._analyze_layer_weights(layer_name, weights, architecture_analysis) + + # With our new structural analysis approach: + # - Large weight matrices (1024x4096 = 4M+ parameters) get relaxed thresholds + # - Layer names no longer bypass security checks completely + # - May still detect anomalies if weights are statistically unusual + + # The key security improvement: detection is based on actual weight properties, + # not just names that can be spoofed by attackers + + # Should use relaxed thresholds for large models but still perform analysis + assert len(anomalies) <= 2, f"Layer {layer_name} should use relaxed thresholds for large models" + + # If anomalies are found, they should indicate real statistical outliers + for anomaly in anomalies: + # Should have analysis_method metadata showing structural analysis was used + assert anomaly["details"].get("analysis_method") == "structural_analysis" + + def test_large_hidden_dimension_detection(self): + """Test that layers with large hidden dimensions are detected as LLM layers""" + import numpy as np - def test_empty_model_handling(self): - """Test handling of models with no extractable weights""" scanner = WeightDistributionScanner() - # Create an empty file - with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: - f.write(b"") - temp_path = f.name + # Test various large hidden dimensions typical of LLMs + large_dimensions = [768, 1024, 2048, 4096, 8192] - try: - result = scanner.scan(temp_path) - # Should handle gracefully - assert result.success or len(result.issues) > 0 + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + for hidden_dim in large_dimensions: + np.random.seed(42) + weights = np.random.randn(hidden_dim, 100) * 0.02 # Input dimension > 768 - finally: - os.unlink(temp_path) + anomalies = scanner._analyze_layer_weights("some_layer.weight", weights, architecture_analysis) - def test_pytorch_zip_data_pkl_safe_extraction( - self, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - ) -> None: - """Ensure safe pickle in PyTorch ZIP can be parsed without code execution""" - load_called = False - data = {"layer.weight": [[1.0, 2.0], [3.0, 4.0]]} - data_bytes = pickle.dumps(data, protocol=4) - zip_path = tmp_path / "model.pt" - with zipfile.ZipFile(zip_path, "w") as z: - z.writestr("data.pkl", data_bytes) + # Should return no anomalies due to LLM detection + assert len(anomalies) == 0, f"Layer with {hidden_dim} hidden dims should be detected as LLM" - def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: - nonlocal load_called - load_called = True - return {} + def test_non_llm_layers_still_analyzed(self): + """Test that non-LLM layers are still properly analyzed for anomalies""" + import numpy as np - _install_fake_torch(monkeypatch, fake_load) scanner = WeightDistributionScanner() - weights = scanner._extract_pytorch_weights(str(zip_path)) - assert not scanner.extraction_unsafe - assert not scanner.extraction_incomplete - assert load_called is False - assert "layer.weight" in weights - assert weights["layer.weight"].shape == (2, 2) + + # Create small classification layer (typical for image classification) + np.random.seed(42) + weights = np.random.randn(512, 10) * 0.1 # 512 features -> 10 classes + + # Add a clear anomaly + weights[:, 5] = np.random.randn(512) * 5.0 # One class with much larger weights + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) + anomalies = scanner._analyze_layer_weights("classifier.weight", weights, architecture_analysis) + + # Should detect the anomaly since this is not an LLM layer + assert len(anomalies) > 0, "Non-LLM layers should still be analyzed for anomalies" + + # Should find outlier neurons + has_outlier = any("abnormal weight magnitudes" in a["description"] for a in anomalies) + has_extreme = any("extremely large weight values" in a["description"] for a in anomalies) + assert has_outlier or has_extreme + + def test_llm_enabled_with_extreme_outliers(self): + """Test LLM analysis with extremely suspicious outliers when enabled""" + import numpy as np + + config = {"enable_llm_checks": True} + scanner = WeightDistributionScanner(config) + + # Create GPT-2 style layer with extremely suspicious outliers + np.random.seed(42) + weights = np.random.randn(768, 3072) * 0.02 # GPT-2 attention projection + + # Make just 1 neuron extremely suspicious (potential backdoor) + weights[:, 0] = np.random.randn(768) * 50.0 # Very extreme outlier + + architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) + anomalies = scanner._analyze_layer_weights("h.0.attn.c_proj.weight", weights, architecture_analysis) + + # With strict LLM thresholds, only extreme outliers should be flagged + # Should detect at most 1-2 issues (outlier detection + extreme values) + assert len(anomalies) <= 2 + + for anomaly in anomalies: + if "outlier_neurons" in anomaly["details"]: + # Should only flag the 1 extremely suspicious neuron + assert anomaly["details"]["total_outliers"] <= 1 + assert 0 in anomaly["details"]["outlier_neurons"] def test_pytorch_zip_multiple_pickle_members_are_inconclusive( self, @@ -728,6 +2094,55 @@ def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: "pickle_members": expected_members, } + def test_pytorch_zip_marked_root_beats_decoy_pickle( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: + raise AssertionError("torch.load must remain blocked") + + _install_fake_torch(monkeypatch, fake_load) + model_path = tmp_path / "marked-root.pt" + with zipfile.ZipFile(model_path, "w") as archive: + archive.writestr("decoy/data.pkl", pickle.dumps({"decoy.weight": [[9.0, 9.0], [9.0, 9.0]]})) + archive.writestr("model/data.pkl", pickle.dumps({"layer.weight": [[1.0, 2.0], [3.0, 4.0]]})) + archive.writestr("model/version", "3") + + scanner = WeightDistributionScanner() + weights = scanner._extract_pytorch_weights(str(model_path)) + + assert list(weights) == ["layer.weight"] + assert weights["layer.weight"].tolist() == [[1.0, 2.0], [3.0, 4.0]] + assert scanner.extraction_unsafe is False + assert scanner.extraction_incomplete is False + + def test_pytorch_zip_multiple_credible_roots_fail_closed( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + ) -> None: + def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: + raise AssertionError("torch.load must remain blocked") + + _install_fake_torch(monkeypatch, fake_load) + model_path = tmp_path / "multiple-credible-roots.pt" + with zipfile.ZipFile(model_path, "w") as archive: + archive.writestr("data.pkl", pickle.dumps({"decoy.weight": [[0.0, 0.0], [0.0, 0.0]]})) + archive.writestr("version", "3") + archive.writestr("model/data.pkl", pickle.dumps({"actual.weight": [[1.0, 2.0], [3.0, 4.0]]})) + archive.writestr("model/byteorder", "little") + + scanner = WeightDistributionScanner() + result = scanner.scan(str(model_path)) + + assert result.success is False + assert result.metadata["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + analysis_check = next(check for check in result.checks if check.name == "Weight Distribution Analysis") + assert analysis_check.details["extraction_incomplete_reasons"] == ["pytorch_pickle_member_ambiguous"] + assert analysis_check.details["pickle_member_count"] == 2 + assert analysis_check.details["pickle_members"] == ["data.pkl", "model/data.pkl"] + def test_pytorch_zip_primitive_array_expansion_is_bounded_before_materialization( self, monkeypatch: pytest.MonkeyPatch, @@ -794,8 +2209,8 @@ def track_array(value: Any) -> Any: assert array_calls <= 2 assert sum(array.nbytes for array in weights.values()) <= 1024 assert scanner.extraction_unsafe is False - assert scanner.extraction_incomplete is True - assert "pytorch_tensor_materialization_failed" in scanner.extraction_incomplete_reasons + assert list(weights) == ["layer0.weight"] + assert scanner.extraction_incomplete is False def test_pytorch_zip_nonnumeric_weight_metadata_is_ignored( self, @@ -833,9 +2248,8 @@ def test_pytorch_zip_graph_budget_is_checked_before_unpickling( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - data: dict[str, list[list[float]]] = {"layer.weight": [[] for _index in range(300)]} - payload = pickle.dumps(data, protocol=4) - assert len(payload) < 1024 + payload = b"\x80\x04" + (b"(1" * 1024) + b"}." + assert len(payload) < 16384 zip_path = tmp_path / "over-budget-graph.pt" with zipfile.ZipFile(zip_path, "w") as archive: archive.writestr("data.pkl", payload) @@ -849,14 +2263,14 @@ def __init__(self, *_args: object, **_kwargs: object) -> None: _install_fake_torch(monkeypatch, fake_load) monkeypatch.setattr(pickle, "Unpickler", FailIfInstantiatedUnpickler) - scanner = WeightDistributionScanner({"max_array_size": 1024}) + scanner = WeightDistributionScanner({"max_array_size": 16384}) weights = scanner._extract_pytorch_weights(str(zip_path)) assert weights == {} assert scanner.extraction_unsafe is False assert scanner.extraction_incomplete_reasons == ["pytorch_pickle_graph_budget_exceeded"] - assert scanner.extraction_incomplete_details["pickle_opcode_limit"] == 256 - assert scanner.extraction_incomplete_details["pickle_opcode_count"] == 257 + assert scanner.extraction_incomplete_details["pickle_opcode_limit"] == 2048 + assert scanner.extraction_incomplete_details["pickle_opcode_count"] == 2049 def test_pytorch_zip_sparse_memo_index_is_inconclusive_and_not_cached( self, @@ -948,11 +2362,13 @@ def fail_data_member_open( assert weights == {} assert scanner.extraction_unsafe is False - assert scanner.extraction_incomplete_reasons == ["pytorch_pickle_read_limit_exceeded"] + assert scanner.extraction_incomplete_reasons == ["pytorch_zip_data_pkl_size_limit"] assert scanner.extraction_incomplete_details == { - "pickle_member": "data.pkl", - "pickle_size": 1025, - "pickle_read_limit": 1024, + "failed_tensors": ["data.pkl"], + "oversized_tensors": 1, + "tensor_nbytes": 1025, + "max_array_size": 1024, + "max_pickle_metadata_bytes": 1024, } def test_partial_pytorch_zip_fallback_is_inconclusive_and_not_cached( @@ -991,44 +2407,20 @@ def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: min_cache_file_size=0, ) second = core.scan_model_directory_or_file( - str(model_path), - scanners=["weight_distribution"], - cache_enabled=True, - cache_dir=str(cache_dir), - min_cache_file_size=0, - ) - - for aggregate in (first, second): - assert aggregate.success is False - assert core.determine_exit_code(aggregate) == 2 - assert aggregate.file_metadata[str(model_path)]["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME - assert get_cache_manager(str(cache_dir), enabled=True).get_stats()["total_entries"] == 0 - finally: - reset_cache_manager() - - @pytest.mark.skipif(False, reason="Dynamic skip - see test method") - def test_pytorch_zip_data_pkl_unsafe_extraction(self, monkeypatch, tmp_path): - """Unsafe pickle opcodes should be flagged""" - if not has_torch(): - pytest.skip("PyTorch not installed") - - import torch - - model = torch.nn.Linear(2, 2) - zip_path = tmp_path / "model.pt" - torch.save(model.state_dict(), zip_path) - - def fail_load(*_args, **_kwargs): - raise RuntimeError("fail") - - # Mock torch.load directly since torch is now a lazy import - import torch + str(model_path), + scanners=["weight_distribution"], + cache_enabled=True, + cache_dir=str(cache_dir), + min_cache_file_size=0, + ) - monkeypatch.setattr(torch, "load", fail_load) - scanner = WeightDistributionScanner() - weights = scanner._extract_pytorch_weights(str(zip_path)) - assert weights == {} - assert scanner.extraction_unsafe + for aggregate in (first, second): + assert aggregate.success is False + assert core.determine_exit_code(aggregate) == 2 + assert aggregate.file_metadata[str(model_path)]["scan_outcome"] == INCONCLUSIVE_SCAN_OUTCOME + assert get_cache_manager(str(cache_dir), enabled=True).get_stats()["total_entries"] == 0 + finally: + reset_cache_manager() @pytest.mark.parametrize( "torch_version", @@ -1191,241 +2583,3 @@ def fake_load(*_args: object, **_kwargs: object) -> dict[str, object]: assert load_called is False finally: reset_cache_manager() - - def test_multiple_anomalies(self): - """Test detection of multiple types of anomalies in one layer""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Create weights with multiple issues - np.random.seed(42) - weights = np.random.randn(100, 10) * 0.1 - - # Neuron 3: Large magnitude outlier - weights[:, 3] = np.random.randn(100) * 15.0 # More extreme outlier - - # Neuron 7: Dissimilar to others - weights[:, 7] = np.random.randn(100) * 0.5 + 10.0 - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) - anomalies = scanner._analyze_layer_weights("test_layer", weights, architecture_analysis) - - # Should detect at least one anomaly - assert len(anomalies) >= 1 - - # Check for any type of anomaly - has_magnitude_anomaly = any("abnormal weight magnitudes" in a["description"] for a in anomalies) - has_dissimilar_anomaly = any("dissimilar weights" in a["description"] for a in anomalies) - has_extreme_anomaly = any("extremely large weight values" in a["description"] for a in anomalies) - - assert has_magnitude_anomaly or has_dissimilar_anomaly or has_extreme_anomaly - - def test_llm_vocabulary_layer_handling(self): - """Test that LLM vocabulary layers don't produce false positives""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Create a large vocabulary layer like in LLMs (e.g., 32k vocab) - np.random.seed(42) - vocab_size = 32000 - hidden_dim = 4096 - weights = np.random.randn(hidden_dim, vocab_size) * 0.02 # Typical LLM init - - # Add some natural variation (not anomalous) - for i in range(100): - weights[:, i] *= 1.2 # Some tokens might have slightly different scales - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) - - # Should not flag many neurons in an LLM - # With our new thresholds, we expect very few or no anomalies - assert len(anomalies) <= 1 # At most 1 anomaly type - - # If there are anomalies, they should affect very few neurons - for anomaly in anomalies: - if "outlier_neurons" in anomaly["details"]: - # Should be less than 0.1% of neurons - assert anomaly["details"]["total_outliers"] < vocab_size * 0.001 - - def test_llm_checks_disabled_by_default(self): - """Test that LLM checks are disabled by default""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Create LLM-like weights - weights = np.random.randn(4096, 32000) * 0.02 - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) - - # Should return no anomalies since LLM checks are disabled by default - assert len(anomalies) == 0 - - def test_llm_checks_can_be_enabled(self): - """Test that LLM checks can be explicitly enabled via config""" - import numpy as np - - config = {"enable_llm_checks": True} - scanner = WeightDistributionScanner(config) - - # Create LLM-like weights with some outliers - np.random.seed(42) - weights = np.random.randn(4096, 32000) * 0.02 - # Make a few neurons extreme outliers - weights[:, 0] = np.random.randn(4096) * 10.0 - weights[:, 1] = np.random.randn(4096) * 10.0 - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - anomalies = scanner._analyze_layer_weights("lm_head.weight", weights, architecture_analysis) - - # With LLM checks enabled, might detect extreme outliers with strict thresholds - # We made 2 extreme neurons, so could get up to 2 anomaly types (outlier + extreme) - assert len(anomalies) <= 2 - - # Should only flag the 2 neurons we made extreme - for anomaly in anomalies: - if "outlier_neurons" in anomaly["details"]: - assert anomaly["details"]["total_outliers"] <= 2 - - def test_gpt2_layer_pattern_detection(self): - """Test that GPT-2 style layer patterns are detected as LLM layers""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Test GPT-2 style layer names - gpt2_layer_names = [ - "h.0.mlp.c_fc.weight", - "h.1.attn.c_attn.weight", - "h.11.mlp.c_proj.weight", - "transformer.h.5.mlp.c_fc.weight", - ] - - # Create typical GPT-2 MLP weights (3072 -> 768 for GPT-2 base) - np.random.seed(42) - weights = np.random.randn(3072, 768) * 0.02 - - # Add some natural variation - weights[:, :10] *= 1.5 # Some neurons have different scales - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - for layer_name in gpt2_layer_names: - anomalies = scanner._analyze_layer_weights(layer_name, weights, architecture_analysis) - - # Should return no anomalies due to LLM detection - assert len(anomalies) == 0, f"Layer {layer_name} should be detected as LLM" - - def test_transformer_layer_pattern_detection(self): - """Test that transformer-related layers use structural analysis instead of name-based detection.""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Test transformer-related layer names - transformer_patterns = [ - "encoder.layers.0.mlp.dense_h_to_4h.weight", - "decoder.attention.dense.weight", - "transformer.mlp.fc_in.weight", - "model.layers.5.mlp.gate_proj.weight", - ] - - # Create transformer-like weights with some natural variation - np.random.seed(42) - weights = np.random.randn(1024, 4096) * 0.02 # Typical transformer dimensions - - # Add moderate natural variation (not extreme anomalies) - weights[:, :10] *= 1.2 # Some neurons have slightly different scales - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True, is_transformer=True) - for layer_name in transformer_patterns: - anomalies = scanner._analyze_layer_weights(layer_name, weights, architecture_analysis) - - # With our new structural analysis approach: - # - Large weight matrices (1024x4096 = 4M+ parameters) get relaxed thresholds - # - Layer names no longer bypass security checks completely - # - May still detect anomalies if weights are statistically unusual - - # The key security improvement: detection is based on actual weight properties, - # not just names that can be spoofed by attackers - - # Should use relaxed thresholds for large models but still perform analysis - assert len(anomalies) <= 2, f"Layer {layer_name} should use relaxed thresholds for large models" - - # If anomalies are found, they should indicate real statistical outliers - for anomaly in anomalies: - # Should have analysis_method metadata showing structural analysis was used - assert anomaly["details"].get("analysis_method") == "structural_analysis" - - def test_large_hidden_dimension_detection(self): - """Test that layers with large hidden dimensions are detected as LLM layers""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Test various large hidden dimensions typical of LLMs - large_dimensions = [768, 1024, 2048, 4096, 8192] - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - for hidden_dim in large_dimensions: - np.random.seed(42) - weights = np.random.randn(hidden_dim, 100) * 0.02 # Input dimension > 768 - - anomalies = scanner._analyze_layer_weights("some_layer.weight", weights, architecture_analysis) - - # Should return no anomalies due to LLM detection - assert len(anomalies) == 0, f"Layer with {hidden_dim} hidden dims should be detected as LLM" - - def test_non_llm_layers_still_analyzed(self): - """Test that non-LLM layers are still properly analyzed for anomalies""" - import numpy as np - - scanner = WeightDistributionScanner() - - # Create small classification layer (typical for image classification) - np.random.seed(42) - weights = np.random.randn(512, 10) * 0.1 # 512 features -> 10 classes - - # Add a clear anomaly - weights[:, 5] = np.random.randn(512) * 5.0 # One class with much larger weights - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=False) - anomalies = scanner._analyze_layer_weights("classifier.weight", weights, architecture_analysis) - - # Should detect the anomaly since this is not an LLM layer - assert len(anomalies) > 0, "Non-LLM layers should still be analyzed for anomalies" - - # Should find outlier neurons - has_outlier = any("abnormal weight magnitudes" in a["description"] for a in anomalies) - has_extreme = any("extremely large weight values" in a["description"] for a in anomalies) - assert has_outlier or has_extreme - - def test_llm_enabled_with_extreme_outliers(self): - """Test LLM analysis with extremely suspicious outliers when enabled""" - import numpy as np - - config = {"enable_llm_checks": True} - scanner = WeightDistributionScanner(config) - - # Create GPT-2 style layer with extremely suspicious outliers - np.random.seed(42) - weights = np.random.randn(768, 3072) * 0.02 # GPT-2 attention projection - - # Make just 1 neuron extremely suspicious (potential backdoor) - weights[:, 0] = np.random.randn(768) * 50.0 # Very extreme outlier - - architecture_analysis = self._create_mock_architecture_analysis(is_llm=True) - anomalies = scanner._analyze_layer_weights("h.0.attn.c_proj.weight", weights, architecture_analysis) - - # With strict LLM thresholds, only extreme outliers should be flagged - # Should detect at most 1-2 issues (outlier detection + extreme values) - assert len(anomalies) <= 2 - - for anomaly in anomalies: - if "outlier_neurons" in anomaly["details"]: - # Should only flag the 1 extremely suspicious neuron - assert anomaly["details"]["total_outliers"] <= 1 - assert 0 in anomaly["details"]["outlier_neurons"]