invarlock 0.2.0__py3-none-any.whl

invarlock/security.py

@@ -0,0 +1,176 @@
+"""
+InvarLock Security Utilities
+========================
+
+Runtime hardening helpers used by the CLI and automation surfaces.
+
+- Network guard: disables outbound socket connections unless explicitly allowed.
+- Secure temporary directory helper ensuring 0o700 permissions and cleanup.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import shutil
+import socket
+import stat
+import tempfile
+from collections.abc import Callable, Iterator
+from pathlib import Path
+from typing import Any
+
+__all__ = [
+    "NetworkGuard",
+    "enforce_network_policy",
+    "enforce_default_security",
+    "temporarily_allow_network",
+    "network_policy_allows",
+    "secure_tempdir",
+    "is_secure_path",
+]
+
+_NETWORK_DISABLED_ERROR = RuntimeError(
+    "Network access disabled by InvarLock security policy. "
+    "Set INVARLOCK_ALLOW_NETWORK=1 if connectivity is required."
+)
+
+
+class NetworkGuard:
+    """Installable network guard that blocks outbound socket connections."""
+
+    def __init__(self) -> None:
+        self._installed = False
+        self._original_socket_cls: type[socket.socket] | None = None
+        self._original_create_connection: Callable[..., socket.socket] | None = None
+
+    @property
+    def installed(self) -> bool:
+        """Whether the guard is currently blocking network access."""
+        return self._installed
+
+    def install(self) -> None:
+        """Install the guard if not already active."""
+        if self._installed:
+            return
+
+        self._original_socket_cls = socket.socket
+        self._original_create_connection = socket.create_connection
+
+        guard_error = _NETWORK_DISABLED_ERROR
+
+        class GuardedSocket(socket.socket):
+            """Socket subclass that blocks connect calls."""
+
+            def connect(self_inner, address: Any) -> None:
+                raise guard_error
+
+        def guarded_create_connection(
+            address: Any,
+            timeout: float | None = None,
+            source_address: Any | None = None,
+        ) -> socket.socket:
+            raise guard_error
+
+        setattr(socket, "socket", GuardedSocket)  # noqa: B010
+        setattr(socket, "create_connection", guarded_create_connection)  # noqa: B010
+        self._installed = True
+
+    def restore(self) -> None:
+        """Restore the original socket implementations."""
+        if not self._installed:
+            return
+
+        if self._original_socket_cls is not None:
+            setattr(socket, "socket", self._original_socket_cls)  # noqa: B010
+        if self._original_create_connection is not None:
+            setattr(socket, "create_connection", self._original_create_connection)  # noqa: B010
+        self._installed = False
+
+
+_GUARD = NetworkGuard()
+
+
+def enforce_network_policy(allow: bool) -> None:
+    """
+    Apply the global network policy.
+
+    Args:
+        allow: When True the guard is removed, otherwise network access is blocked.
+    """
+    if allow:
+        _GUARD.restore()
+    else:
+        _GUARD.install()
+
+
+def network_policy_allows() -> bool:
+    """Return True if outbound connections are currently permitted."""
+    return not _GUARD.installed
+
+
+def enforce_default_security() -> None:
+    """
+    Enforce default runtime security posture.
+
+    Network access is denied unless INVARLOCK_ALLOW_NETWORK is set to a truthy value.
+    """
+    allow_env = os.environ.get("INVARLOCK_ALLOW_NETWORK", "")
+    allow_network = allow_env.strip().lower() in {"1", "true", "yes", "on"}
+    enforce_network_policy(allow_network)
+
+
+@contextlib.contextmanager
+def temporarily_allow_network() -> Iterator[None]:
+    """
+    Temporarily allow network access inside the context block.
+
+    Restores the previous policy when exiting the context.
+    """
+    was_installed = _GUARD.installed
+    if was_installed:
+        _GUARD.restore()
+    try:
+        yield
+    finally:
+        if was_installed:
+            _GUARD.install()
+
+
+@contextlib.contextmanager
+def secure_tempdir(
+    prefix: str = "invarlock-", base_dir: str | os.PathLike[str] | None = None
+) -> Iterator[Path]:
+    """
+    Create a temporary directory with 0o700 permissions that is removed on exit.
+
+    Args:
+        prefix: Directory name prefix.
+        base_dir: Optional base directory.
+
+    Yields:
+        Path to the secure temporary directory.
+    """
+    path = Path(tempfile.mkdtemp(prefix=prefix, dir=base_dir))
+    os.chmod(path, 0o700)
+    try:
+        yield path
+    finally:
+        shutil.rmtree(path, ignore_errors=True)
+
+
+def is_secure_path(path: Path) -> bool:
+    """
+    Check whether a path has secure (0o700) permissions.
+
+    Args:
+        path: Path to validate.
+
+    Returns:
+        True if the path exists and has 0o700 permissions, False otherwise.
+    """
+    try:
+        mode = path.stat().st_mode
+    except FileNotFoundError:
+        return False
+    return stat.S_IMODE(mode) == 0o700

invarlock/sparsity_utils.py

@@ -0,0 +1,323 @@
+"""
+InvarLock Sparsity Utilities
+========================
+
+Utilities for working with sparse models and sparsity patterns.
+Helper functions for analyzing and manipulating model sparsity.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+try:
+    import torch
+    import torch.nn as nn
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+__all__ = [
+    "calculate_sparsity",
+    "get_zero_mask",
+    "apply_mask",
+    "count_parameters",
+    "get_sparsity_stats",
+    "create_structured_mask",
+    "validate_sparsity_target",
+]
+
+
+def calculate_sparsity(tensor) -> float:
+    """
+    Calculate sparsity ratio of a tensor.
+
+    Args:
+        tensor: PyTorch tensor
+
+    Returns:
+        Sparsity ratio (fraction of zero elements)
+    """
+    if not TORCH_AVAILABLE:
+        return 0.0
+
+    if not isinstance(tensor, torch.Tensor):
+        return 0.0
+
+    total_elements = tensor.numel()
+    if total_elements == 0:
+        return 0.0
+
+    zero_elements = (tensor == 0).sum().item()
+    return zero_elements / total_elements
+
+
+def get_zero_mask(tensor, threshold: float = 1e-8):
+    """
+    Get boolean mask of effectively zero elements.
+
+    Args:
+        tensor: PyTorch tensor
+        threshold: Threshold below which values are considered zero
+
+    Returns:
+        Boolean mask where True indicates zero/near-zero elements
+    """
+    if not TORCH_AVAILABLE:
+        return None
+
+    if not isinstance(tensor, torch.Tensor):
+        return None
+
+    return torch.abs(tensor) <= threshold
+
+
+def apply_mask(tensor, mask, fill_value: float = 0.0):
+    """
+    Apply a boolean mask to zero out elements.
+
+    Args:
+        tensor: PyTorch tensor to modify
+        mask: Boolean mask (True = zero out)
+        fill_value: Value to fill masked positions
+
+    Returns:
+        Modified tensor
+    """
+    if not TORCH_AVAILABLE:
+        return tensor
+
+    if not isinstance(tensor, torch.Tensor) or not isinstance(mask, torch.Tensor):
+        return tensor
+
+    result = tensor.clone()
+    result[mask] = fill_value
+    return result
+
+
+def count_parameters(module, only_trainable: bool = True) -> dict[str, int]:
+    """
+    Count parameters in a module.
+
+    Args:
+        module: PyTorch module
+        only_trainable: Count only trainable parameters
+
+    Returns:
+        Dictionary with parameter counts
+    """
+    if not TORCH_AVAILABLE or not isinstance(module, nn.Module):
+        return {"total": 0, "trainable": 0, "non_trainable": 0}
+
+    total_params = 0
+    trainable_params = 0
+    non_trainable_params = 0
+
+    for param in module.parameters():
+        param_count = param.numel()
+        total_params += param_count
+
+        if param.requires_grad:
+            trainable_params += param_count
+        else:
+            non_trainable_params += param_count
+
+    return {
+        "total": total_params,
+        "trainable": trainable_params,
+        "non_trainable": non_trainable_params,
+    }
+
+
+def get_sparsity_stats(module) -> dict[str, Any]:
+    """
+    Get comprehensive sparsity statistics for a module.
+
+    Args:
+        module: PyTorch module
+
+    Returns:
+        Dictionary with sparsity statistics
+    """
+    if not TORCH_AVAILABLE or not isinstance(module, nn.Module):
+        return {}
+
+    stats = {
+        "overall_sparsity": 0.0,
+        "layer_sparsities": {},
+        "parameter_counts": count_parameters(module),
+        "sparse_layers": 0,
+        "dense_layers": 0,
+    }
+
+    total_elements = 0
+    total_zeros = 0
+
+    for name, param in module.named_parameters():
+        if param.numel() == 0:
+            continue
+
+        layer_sparsity = calculate_sparsity(param)
+        stats["layer_sparsities"][name] = layer_sparsity
+
+        # Count for overall statistics
+        elements = param.numel()
+        zeros = (param == 0).sum().item()
+
+        total_elements += elements
+        total_zeros += zeros
+
+        # Classify layer as sparse or dense
+        if layer_sparsity > 0.1:  # 10% threshold for "sparse"
+            stats["sparse_layers"] += 1
+        else:
+            stats["dense_layers"] += 1
+
+    # Calculate overall sparsity
+    if total_elements > 0:
+        stats["overall_sparsity"] = total_zeros / total_elements
+
+    return stats
+
+
+def create_structured_mask(
+    shape: tuple[int, ...], pattern: str = "block", block_size: int = 4
+) -> torch.Tensor | None:
+    """
+    Create block/group sparsity masks.
+
+    Args:
+        shape: Shape of the tensor
+        pattern: Sparsity pattern ('block', 'column', 'row')
+        block_size: Size of blocks for block sparsity
+
+    Returns:
+        Boolean mask tensor
+    """
+    if not TORCH_AVAILABLE:
+        return None
+
+    mask = torch.zeros(shape, dtype=torch.bool)
+
+    if pattern == "block" and len(shape) >= 2:
+        # Block sparsity pattern
+        for i in range(0, shape[0], block_size):
+            for j in range(0, shape[1], block_size):
+                # Create checkerboard pattern of blocks
+                if (i // block_size + j // block_size) % 2 == 0:
+                    end_i = min(i + block_size, shape[0])
+                    end_j = min(j + block_size, shape[1])
+                    mask[i:end_i, j:end_j] = True
+
+    elif pattern == "column" and len(shape) >= 2:
+        # Column sparsity - zero out every other column
+        mask[:, ::2] = True
+
+    elif pattern == "row" and len(shape) >= 2:
+        # Row sparsity - zero out every other row
+        mask[::2, :] = True
+
+    else:
+        warnings.warn(
+            f"Unsupported sparsity pattern '{pattern}' for shape {shape}", stacklevel=2
+        )
+
+    return mask
+
+
+def validate_sparsity_target(target_sparsity: Any) -> bool:
+    """
+    Validate sparsity target value.
+
+    Args:
+        target_sparsity: Target sparsity ratio (0.0 to 1.0)
+
+    Returns:
+        True if valid
+    """
+    if not isinstance(target_sparsity, int | float):
+        return False
+    return 0.0 <= target_sparsity <= 1.0
+
+
+def get_magnitude_mask(tensor, sparsity_ratio: float) -> torch.Tensor | None:
+    """
+    Create magnitude-based sparsity mask.
+
+    Args:
+        tensor: PyTorch tensor
+        sparsity_ratio: Target sparsity fraction
+
+    Returns:
+        Boolean mask where True indicates weights to zero
+    """
+    if not TORCH_AVAILABLE or not isinstance(tensor, torch.Tensor):
+        return None
+
+    if not validate_sparsity_target(sparsity_ratio):
+        raise ValueError(f"Invalid sparsity ratio: {sparsity_ratio}")
+
+    # Flatten tensor for magnitude ranking
+    flat_tensor = tensor.view(-1)
+
+    # Find threshold for zeroing
+    num_to_zero = int(sparsity_ratio * flat_tensor.numel())
+    if num_to_zero == 0:
+        return torch.zeros_like(tensor, dtype=torch.bool)
+
+    # Get magnitude and find threshold
+    magnitudes = torch.abs(flat_tensor)
+    threshold_value = torch.kthvalue(magnitudes, num_to_zero).values
+
+    # Create mask
+    mask = torch.abs(tensor) <= threshold_value
+
+    return mask
+
+
+def analyze_sparsity_impact(original_tensor, edited_tensor) -> dict[str, Any]:
+    """
+    Analyze the impact of applied sparsity on a tensor.
+
+    Args:
+        original_tensor: Original tensor before changes
+        edited_tensor: Tensor after applying sparsity
+
+    Returns:
+        Dictionary with impact analysis
+    """
+    if not TORCH_AVAILABLE:
+        return {}
+
+    if not (
+        isinstance(original_tensor, torch.Tensor)
+        and isinstance(edited_tensor, torch.Tensor)
+    ):
+        return {}
+
+    # Calculate basic statistics
+    original_sparsity = calculate_sparsity(original_tensor)
+    final_sparsity = calculate_sparsity(edited_tensor)
+
+    # Calculate magnitude changes
+    magnitude_change = torch.norm(edited_tensor - original_tensor).item()
+    relative_change = (
+        magnitude_change / torch.norm(original_tensor).item()
+        if torch.norm(original_tensor) > 0
+        else 0.0
+    )
+
+    analysis = {
+        "original_sparsity": original_sparsity,
+        "final_sparsity": final_sparsity,
+        "sparsity_increase": final_sparsity - original_sparsity,
+        "magnitude_change": magnitude_change,
+        "relative_change": relative_change,
+        "compression_ratio": final_sparsity / original_sparsity
+        if original_sparsity > 0
+        else float("inf"),
+    }
+
+    return analysis

invarlock/utils/__init__.py

@@ -0,0 +1,150 @@
+"""
+InvarLock Utilities
+===============
+
+Common utility functions used across InvarLock modules.
+
+This package also exposes submodules such as `invarlock.utils.digest` for
+hashing and provenance utilities.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import psutil
+
+try:  # Torch is optional; utils can be imported without it.
+    import torch  # type: ignore[import]
+except Exception:  # pragma: no cover - exercised when torch is missing
+    torch = None  # type: ignore[assignment]
+
+
+def extract_input_ids(
+    batch: Any, device: str | None = None, strict: bool = False
+) -> torch.Tensor:
+    """
+    Extract input_ids from various batch formats.
+
+    Args:
+        batch: Input batch (tensor, dict, or other format)
+        device: Target device for tensor
+        strict: Whether to raise errors on format issues
+
+    Returns:
+        Extracted input_ids tensor
+    """
+    if isinstance(batch, torch.Tensor):
+        input_ids = batch
+    elif isinstance(batch, dict):
+        if "input_ids" in batch:
+            input_ids = batch["input_ids"]
+        elif "inputs" in batch:
+            input_ids = batch["inputs"]
+        else:
+            if strict:
+                raise ValueError(
+                    f"Dict batch missing 'input_ids' or 'inputs' keys: {list(batch.keys())}"
+                )
+            # Try first tensor value
+            for value in batch.values():
+                if isinstance(value, torch.Tensor):
+                    input_ids = value
+                    break
+            else:
+                raise ValueError("No tensor found in batch dict")
+    elif hasattr(batch, "input_ids"):
+        input_ids = batch.input_ids
+    else:
+        if strict:
+            raise ValueError(f"Unsupported batch format: {type(batch)}")
+        # Try to convert directly
+        input_ids = torch.tensor(batch)
+
+    # Move to device if specified
+    if device is not None:
+        input_ids = input_ids.to(device)
+
+    return input_ids
+
+
+def get_model_device(model: torch.nn.Module) -> torch.device:
+    """Get the device of a model."""
+    return next(model.parameters()).device
+
+
+def ensure_tensor(data: Any, device: torch.device | None = None) -> torch.Tensor:
+    """Ensure data is a tensor on the correct device."""
+    if not isinstance(data, torch.Tensor):
+        data = torch.tensor(data)
+
+    if device is not None:
+        data = data.to(device)
+
+    return data
+
+
+def safe_divide(numerator: float, denominator: float, default: float = 0.0) -> float:
+    """Safely divide two numbers, returning default if denominator is zero."""
+    if abs(denominator) < 1e-12:
+        return default
+    return numerator / denominator
+
+
+def dict_to_device(
+    data: dict[str, torch.Tensor], device: torch.device
+) -> dict[str, torch.Tensor]:
+    """Move all tensors in a dictionary to the specified device."""
+    return {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in data.items()
+    }
+
+
+def format_number(num: float, precision: int = 3) -> str:
+    """Format a number for display."""
+    if abs(num) < 1e-3:
+        return f"{num:.2e}"
+    elif abs(num) < 1:
+        return f"{num:.{precision + 1}f}"
+    else:
+        return f"{num:.{precision}f}"
+
+
+def get_memory_usage() -> dict[str, float]:
+    """Get current memory usage in MB."""
+    import gc
+
+    # Force garbage collection
+    gc.collect()
+
+    # Get process memory
+    process = psutil.Process()
+    memory_info = process.memory_info()
+
+    result = {
+        "rss_mb": memory_info.rss / 1024 / 1024,  # Resident Set Size
+        "vms_mb": memory_info.vms / 1024 / 1024,  # Virtual Memory Size
+    }
+
+    # Add CUDA memory if available
+    try:
+        if torch is not None and hasattr(torch, "cuda") and torch.cuda.is_available():
+            result["cuda_allocated_mb"] = torch.cuda.memory_allocated() / 1024 / 1024
+            result["cuda_reserved_mb"] = torch.cuda.memory_reserved() / 1024 / 1024
+    except Exception:
+        # If torch is unavailable or querying CUDA fails, fall back to CPU-only stats.
+        pass
+
+    return result
+
+
+__all__ = [
+    "extract_input_ids",
+    "get_model_device",
+    "ensure_tensor",
+    "safe_divide",
+    "dict_to_device",
+    "format_number",
+    "get_memory_usage",
+]

invarlock/utils/digest.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any, Protocol
+
+_ENC = "utf-8"
+
+
+class _HashLike(Protocol):
+    def update(self, data: bytes, /) -> None: ...
+    def hexdigest(self, /) -> str: ...
+
+
+def _h() -> _HashLike:
+    return hashlib.blake2s(digest_size=32)
+
+
+def hash_bytes(b: bytes, *, salt: bytes | None = None) -> str:
+    h = _h()
+    if salt:
+        h.update(salt)
+    h.update(b)
+    return h.hexdigest()
+
+
+def hash_json(obj: Any, *, salt: str | None = None) -> str:
+    s = json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+    return hash_bytes(s.encode(_ENC), salt=salt.encode(_ENC) if salt else None)
+
+
+def hash_int_array(arr, *, salt: str | None = None, byteorder: str = "little") -> str:
+    # Accept numpy arrays to avoid extra copies
+    try:
+        import numpy as _np
+    except Exception:  # pragma: no cover - numpy always present in tests
+        # Fallback: best-effort conversion
+        b = bytes(int(x) & 0xFFFFFFFF for x in arr)
+        return hash_bytes(b, salt=salt.encode(_ENC) if salt else None)
+
+    a = _np.asarray(arr, dtype=_np.int32, order="C")
+    return hash_bytes(a.tobytes(order="C"), salt=salt.encode(_ENC) if salt else None)
+
+
+__all__ = ["hash_bytes", "hash_json", "hash_int_array"]