invarlock 0.2.0__py3-none-any.whl

invarlock/core/checkpoint.py

@@ -0,0 +1,221 @@
+"""
+InvarLock Core Checkpoint System
+===========================
+
+Checkpoint and rollback functionality for safe model editing.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+from contextlib import contextmanager
+from typing import Any
+
+from .types import GuardOutcome
+
+
+def _use_chunked_snapshot() -> bool:
+    """Return True when chunked snapshot mode is enabled."""
+    return os.environ.get("INVARLOCK_SNAPSHOT_MODE", "bytes").lower() == "chunked"
+
+
+class PolicyCheckpoint:
+    """
+    Checkpoint manager for policy-based rollback decisions.
+    """
+
+    def __init__(self, model: Any, adapter: Any, policy: Any):
+        """
+        Initialize checkpoint.
+
+        Args:
+            model: Model to checkpoint
+            adapter: ModelAdapter for model operations
+            policy: Policy configuration
+        """
+        self.model = model
+        self.adapter = adapter
+        self.policy = policy
+        self.checkpoint_data: dict[str, Any] | None = None
+        self.rollback_performed = False
+
+    def create_checkpoint(self) -> None:
+        """Create a checkpoint of the current model state."""
+        if _use_chunked_snapshot() and hasattr(self.adapter, "snapshot_chunked"):
+            snapshot_path = self.adapter.snapshot_chunked(self.model)
+            self.checkpoint_data = {"mode": "chunked", "path": snapshot_path}
+        else:
+            self.checkpoint_data = {
+                "mode": "bytes",
+                "blob": self.adapter.snapshot(self.model),
+            }
+
+    def should_rollback(self, outcomes: list[GuardOutcome]) -> tuple[bool, str]:
+        """
+        Determine if rollback should be performed based on guard outcomes.
+
+        Args:
+            outcomes: List of guard outcomes
+
+        Returns:
+            (should_rollback, reason) tuple
+        """
+        # Check for abort actions
+        for outcome in outcomes:
+            if hasattr(outcome, "action") and outcome.action == "abort":
+                return True, "guard_abort"
+
+        # Check for rollback actions
+        for outcome in outcomes:
+            if hasattr(outcome, "action") and outcome.action == "rollback":
+                return True, "guard_rollback"
+
+        # Check policy configuration
+        if (
+            hasattr(self.policy, "enable_auto_rollback")
+            and self.policy.enable_auto_rollback
+        ):
+            # Check if any guards failed
+            for outcome in outcomes:
+                if hasattr(outcome, "passed") and not outcome.passed:
+                    return True, "auto_rollback"
+
+        return False, ""
+
+    def rollback(self, reason: str) -> bool:
+        """
+        Perform rollback to checkpoint.
+
+        Args:
+            reason: Reason for rollback
+
+        Returns:
+            True if rollback was successful
+        """
+        if self.checkpoint_data is None:
+            return False
+        try:
+            mode = self.checkpoint_data.get("mode", "bytes")
+            if mode == "chunked":
+                path = self.checkpoint_data.get("path")
+                if not path or not hasattr(self.adapter, "restore_chunked"):
+                    return False
+                self.adapter.restore_chunked(self.model, path)
+            else:
+                blob = self.checkpoint_data.get("blob")
+                self.adapter.restore(self.model, blob)
+            self.rollback_performed = True
+            return True
+        except Exception:
+            return False
+
+    def cleanup(self) -> None:
+        """Clean up checkpoint resources."""
+        if self.checkpoint_data and self.checkpoint_data.get("mode") == "chunked":
+            path = self.checkpoint_data.get("path")
+            if path and os.path.isdir(path):
+                shutil.rmtree(path, ignore_errors=True)
+        self.checkpoint_data = None
+
+
+@contextmanager
+def create_policy_checkpoint(model: Any, adapter: Any, policy: Any):
+    """
+    Context manager for policy-based checkpointing.
+
+    Args:
+        model: Model to checkpoint
+        adapter: ModelAdapter for operations
+        policy: Policy configuration
+
+    Yields:
+        PolicyCheckpoint instance
+    """
+    checkpoint = PolicyCheckpoint(model, adapter, policy)
+    checkpoint.create_checkpoint()
+
+    try:
+        yield checkpoint
+    finally:
+        checkpoint.cleanup()
+
+
+class CheckpointManager:
+    """
+    Manager for model checkpoints during pipeline execution.
+    """
+
+    def __init__(self):
+        """Initialize checkpoint manager."""
+        self.checkpoints: dict[str, dict[str, Any]] = {}
+        self.next_id = 1
+
+    def create_checkpoint(self, model: Any, adapter: Any) -> str:
+        """
+        Create a checkpoint of the model.
+
+        Args:
+            model: Model to checkpoint
+            adapter: ModelAdapter for serialization
+
+        Returns:
+            Checkpoint ID
+        """
+        checkpoint_id = f"checkpoint_{self.next_id}"
+        self.next_id += 1
+
+        try:
+            if _use_chunked_snapshot() and hasattr(adapter, "snapshot_chunked"):
+                snapshot_path = adapter.snapshot_chunked(model)
+                checkpoint_data = {"mode": "chunked", "path": snapshot_path}
+            else:
+                checkpoint_data = {
+                    "mode": "bytes",
+                    "blob": adapter.snapshot(model),
+                }
+            self.checkpoints[checkpoint_id] = checkpoint_data
+            return checkpoint_id
+        except Exception as e:
+            raise RuntimeError(f"Failed to create checkpoint: {e}") from e
+
+    def restore_checkpoint(self, model: Any, adapter: Any, checkpoint_id: str) -> bool:
+        """
+        Restore model from checkpoint.
+
+        Args:
+            model: Model to restore
+            adapter: ModelAdapter for deserialization
+            checkpoint_id: ID of checkpoint to restore
+
+        Returns:
+            True if restoration was successful
+        """
+        if checkpoint_id not in self.checkpoints:
+            return False
+
+        try:
+            checkpoint_data = self.checkpoints[checkpoint_id]
+            mode = checkpoint_data.get("mode", "bytes")
+            if mode == "chunked":
+                if not hasattr(adapter, "restore_chunked"):
+                    return False
+                adapter.restore_chunked(model, checkpoint_data.get("path"))
+            else:
+                adapter.restore(model, checkpoint_data.get("blob"))
+            return True
+        except Exception:
+            return False
+
+    def cleanup(self) -> None:
+        """Clean up all checkpoints."""
+        for data in self.checkpoints.values():
+            if data.get("mode") == "chunked":
+                path = data.get("path")
+                if path and os.path.isdir(path):
+                    shutil.rmtree(path, ignore_errors=True)
+        self.checkpoints.clear()
+        self.next_id = 1
+
+
+__all__ = ["PolicyCheckpoint", "create_policy_checkpoint", "CheckpointManager"]

invarlock/core/contracts.py

@@ -0,0 +1,73 @@
+"""
+InvarLock Contracts
+===============
+
+Lightweight runtime assertions for monotonic behaviour of guard/edit operations.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def enforce_relative_spectral_cap(
+    weight: torch.Tensor, baseline_sigma: float | torch.Tensor, cap_ratio: float
+) -> torch.Tensor:
+    """Clamp the spectral norm of ``weight`` to ``cap_ratio * baseline_sigma``."""
+    baseline_value = float(baseline_sigma)
+    if not torch.isfinite(torch.tensor(baseline_value)) or baseline_value <= 0:
+        return weight
+    with torch.no_grad():
+        sigma = _spectral_norm(weight)
+        limit = baseline_value * cap_ratio
+        if sigma > limit and sigma > 0:
+            weight.mul_(limit / sigma)
+    return weight
+
+
+def enforce_weight_energy_bound(
+    approx: torch.Tensor, exact: torch.Tensor, max_relative_error: float
+) -> torch.Tensor:
+    """Return ``approx`` if the relative error against ``exact`` is within bounds."""
+    denom = torch.norm(exact).clamp_min(1e-12)
+    rel_err = torch.norm(approx - exact) / denom
+    if rel_err <= max_relative_error:
+        return approx
+    return exact
+
+
+def rmt_correction_is_monotone(
+    corrected_sigma: float,
+    baseline_sigma: float,
+    max_ratio: float,
+    deadband: float,
+) -> bool:
+    """
+    Validate monotonicity for RMT correction.
+
+    ``corrected_sigma`` should not exceed ``baseline_sigma * (1 + deadband)``
+    and must remain ≤ ``max_ratio``.
+    """
+    if corrected_sigma < 0 or baseline_sigma <= 0 or max_ratio <= 0:
+        return False
+    if corrected_sigma > max_ratio:
+        return False
+    return corrected_sigma <= baseline_sigma * (1.0 + deadband)
+
+
+def _spectral_norm(weight: torch.Tensor) -> float:
+    """Compute the spectral norm (largest singular value) of ``weight``."""
+    if weight.ndim != 2:
+        weight = weight.view(weight.shape[0], -1)
+    try:
+        s = torch.linalg.svdvals(weight)
+    except RuntimeError:
+        s = torch.linalg.svdvals(weight.cpu()).to(weight.device)
+    return float(s.max().item())
+
+
+__all__ = [
+    "enforce_relative_spectral_cap",
+    "enforce_weight_energy_bound",
+    "rmt_correction_is_monotone",
+]

invarlock/core/error_utils.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from contextlib import ContextDecorator
+from dataclasses import dataclass
+from typing import Any, Generic, Literal, TypeVar
+
+from .exceptions import InvarlockError
+
+T = TypeVar("T", bound=InvarlockError)
+
+
+ContextFn = Callable[[BaseException], dict[str, Any] | None]
+
+
+@dataclass
+class _WrapErrors(ContextDecorator, Generic[T]):  # noqa: UP046
+    target_exc: type[T]
+    code: str
+    message: str
+    context_fn: ContextFn | None = None
+
+    # Context manager protocol
+    def __enter__(self) -> _WrapErrors:  # pragma: no cover - trivial
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: Any,
+    ) -> Literal[False]:
+        if exc is None:
+            return False
+        # If it's already a InvarlockError, do not double-wrap
+        if isinstance(exc, InvarlockError):
+            return False
+        ctx = self.context_fn(exc) if self.context_fn is not None else None
+        wrapped = self.target_exc(code=self.code, message=self.message, details=ctx)
+        raise wrapped from exc
+
+
+def wrap_errors(  # noqa: UP047
+    target_exc: type[T],
+    code: str,
+    message: str,
+    context_fn: ContextFn | None = None,
+) -> _WrapErrors[T]:
+    """Return a context manager/decorator that wraps arbitrary exceptions.
+
+    Usage as context manager:
+        with wrap_errors(AdapterError, "E202", "ADAPTER-LOAD-FAILED", ctx):
+            risky()
+
+    Usage as decorator:
+        @wrap_errors(ValidationError, "E301", "VALIDATION-FAILED")
+        def f(...): ...
+    """
+    return _WrapErrors(
+        target_exc=target_exc, code=code, message=message, context_fn=context_fn
+    )
+
+
+__all__ = ["wrap_errors"]

invarlock/core/events.py

@@ -0,0 +1,298 @@
+"""
+InvarLock Event Logger
+==================
+
+JSONL event logging for pipeline execution tracking.
+Provides structured logging for analysis and debugging.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import Mapping, Sequence
+from datetime import datetime
+from pathlib import Path
+from typing import Any, TextIO
+
+from .types import LogLevel
+
+__all__ = ["EventLogger"]
+
+
+class EventLogger:
+    """
+    JSONL event logger for InvarLock pipeline execution.
+
+    Logs structured events to a JSONL file for analysis,
+    debugging, and audit trails. Thread-safe and handles
+    file rotation if needed.
+    """
+
+    SENSITIVE_KEYWORDS: Sequence[str] = (
+        "token",
+        "secret",
+        "password",
+        "passphrase",
+        "api_key",
+        "credential",
+        "auth",
+        "email",
+    )
+
+    def __init__(
+        self,
+        log_path: Path,
+        auto_flush: bool = True,
+        *,
+        run_id: str | None = None,
+        redact_keywords: Sequence[str] | None = None,
+        max_string_length: int = 512,
+    ):
+        """
+        Initialize event logger.
+
+        Args:
+            log_path: Path to JSONL log file
+            auto_flush: Whether to flush after each write
+            run_id: Optional run identifier to include in every log entry
+            redact_keywords: Iterable of keywords whose values should be redacted
+            max_string_length: Maximum length for logged strings before truncation
+        """
+        self.log_path = Path(log_path)
+        self.auto_flush = auto_flush
+        self._file: TextIO | None = None
+        self._session_id = self._generate_session_id()
+        self._run_id = run_id
+        self._redact_keywords = tuple(
+            keyword.lower() for keyword in (redact_keywords or self.SENSITIVE_KEYWORDS)
+        )
+        # Honor caller-provided limit; clamp to a small positive minimum
+        self._max_string_length = max(1, int(max_string_length))
+
+        # Ensure parent directory exists
+        self.log_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Open file for writing
+        self._open_log_file()
+
+        # Log session start
+        session_start_payload = {
+            "session_id": self._session_id,
+            "log_path": str(self.log_path),
+        }
+        if self._run_id:
+            session_start_payload["run_id"] = self._run_id
+        self.log("logger", "session_start", LogLevel.INFO, session_start_payload)
+
+    def _generate_session_id(self) -> str:
+        """Generate unique session ID."""
+        return f"session_{int(time.time())}"
+
+    def _open_log_file(self) -> None:
+        """Open the log file for writing."""
+        try:
+            self._file = open(self.log_path, "a", encoding="utf-8")
+        except Exception as e:  # pragma: no cover - defensive guard
+            raise OSError(f"Failed to open log file {self.log_path}: {e}") from e
+
+    def log(
+        self,
+        component: str,
+        operation: str,
+        level: LogLevel,
+        data: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Log an event.
+
+        Args:
+            component: Component generating the event (e.g., "runner", "edit", "guard")
+            operation: Operation being performed (e.g., "start", "complete", "error")
+            level: Log level
+            data: Optional additional data
+        """
+        if not self._file:
+            return
+
+        event: dict[str, Any] = {
+            "timestamp": datetime.now().isoformat(),
+            "session_id": self._session_id,
+            "component": component,
+            "operation": operation,
+            "level": level.value,
+        }
+
+        if self._run_id:
+            event["run_id"] = self._run_id
+
+        if data:
+            event["data"] = self._sanitize_data(data)
+
+        try:
+            json_line = json.dumps(event, default=self._json_serializer)
+            self._file.write(json_line + "\n")
+
+            if self.auto_flush:
+                self._file.flush()
+
+        except Exception as e:  # pragma: no cover - fallback to stderr
+            import sys
+
+            print(f"Event logging failed: {e}", file=sys.stderr)
+
+    def _sanitize_data(self, data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Sanitize data for JSON serialization.
+
+        Removes non-serializable objects and large data structures,
+        and redacts potential secrets.
+        """
+
+        def sanitize_value(key: str | None, value: Any) -> Any:
+            key_lower = key.lower() if isinstance(key, str) else ""
+            if key_lower and any(word in key_lower for word in self._redact_keywords):
+                return "<redacted>"
+
+            if isinstance(value, str):
+                if any(word in value.lower() for word in self._redact_keywords):
+                    return "<redacted>"
+                if len(value) > self._max_string_length:
+                    return f"<str len={len(value)}>"
+                return value
+
+            if isinstance(value, Mapping):
+                return {
+                    inner_key: sanitize_value(str(inner_key), inner_value)
+                    for inner_key, inner_value in value.items()
+                }
+
+            # Common numeric array-like
+            if hasattr(value, "tolist"):
+                try:
+                    return value.tolist()
+                except Exception:
+                    pass
+
+            if isinstance(value, set | frozenset):
+                try:
+                    return list(value)
+                except Exception:
+                    pass
+
+            if isinstance(value, Sequence) and not isinstance(
+                value, str | bytes | bytearray
+            ):
+                return [sanitize_value(key, item) for item in value]
+
+            if isinstance(value, bytes):
+                return f"<bytes len={len(value)}>"
+            # Preserve JSON-native scalars; coerce others to a placeholder string
+            if value is None or isinstance(value, bool | int | float):
+                return value
+            return f"<{type(value).__name__}>"
+
+        return {key: sanitize_value(key, value) for key, value in data.items()}
+
+    def _json_serializer(self, obj: Any) -> Any:
+        """Custom JSON serializer for common non-serializable types."""
+        if hasattr(obj, "tolist"):  # numpy arrays
+            return obj.tolist()
+        if hasattr(obj, "__dict__"):  # custom objects
+            return str(obj)
+        if isinstance(obj, set | frozenset):
+            return list(obj)
+        if isinstance(obj, bytes):
+            return obj.decode("utf-8", errors="replace")
+        return str(obj)
+
+    def log_error(
+        self,
+        component: str,
+        operation: str,
+        error: Exception,
+        context: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Convenience method for logging errors.
+
+        Args:
+            component: Component where error occurred
+            operation: Operation that failed
+            error: The exception that occurred
+            context: Optional additional context
+        """
+        error_data: dict[str, Any] = {
+            "error_type": type(error).__name__,
+            "error_message": str(error),
+        }
+
+        if context:
+            error_data["context"] = context
+
+        self.log(component, operation, LogLevel.ERROR, error_data)
+
+    def log_metric(
+        self, component: str, metric_name: str, value: float, unit: str | None = None
+    ) -> None:
+        """
+        Convenience method for logging metrics.
+
+        Args:
+            component: Component reporting the metric
+            metric_name: Name of the metric
+            value: Metric value
+            unit: Optional unit (e.g., "seconds", "MB", "ratio")
+        """
+        metric_data = {"metric": metric_name, "value": value}
+
+        if unit:
+            metric_data["unit"] = unit
+
+        self.log(component, "metric", LogLevel.INFO, metric_data)
+
+    def log_checkpoint(
+        self, component: str, checkpoint_id: str, operation: str
+    ) -> None:
+        """
+        Convenience method for logging checkpoint operations.
+
+        Args:
+            component: Component handling the checkpoint
+            checkpoint_id: Unique checkpoint identifier
+            operation: Checkpoint operation ("create", "restore", "delete")
+        """
+        checkpoint_data = {
+            "checkpoint_id": checkpoint_id,
+            "checkpoint_operation": operation,
+        }
+
+        self.log(component, "checkpoint", LogLevel.INFO, checkpoint_data)
+
+    def close(self) -> None:
+        """Close the log file."""
+        if self._file:
+            session_end_payload = {"session_id": self._session_id}
+            if self._run_id:
+                session_end_payload["run_id"] = self._run_id
+            self.log("logger", "session_end", LogLevel.INFO, session_end_payload)
+
+            try:
+                self._file.close()
+            except Exception:  # pragma: no cover - best-effort cleanup
+                pass
+            finally:
+                self._file = None
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
+
+    def __del__(self):
+        """Destructor - ensure file is closed."""
+        if hasattr(self, "_file"):
+            self.close()