space-ml-sim 0.3.0__py3-none-any.whl

space_ml_sim/__init__.py

@@ -0,0 +1,3 @@
+"""space-ml-sim: Simulate AI inference on orbital satellite constellations under realistic space radiation."""
+
+__version__ = "0.3.0"

space_ml_sim/compute/__init__.py

@@ -0,0 +1,32 @@
+"""Compute module: fault injection, TMR, checkpointing, and scheduling."""
+
+from space_ml_sim.compute.fault_injector import FaultInjector, FaultReport
+from space_ml_sim.compute.transformer_fault import TransformerFaultInjector
+from space_ml_sim.compute.tmr import TMRWrapper
+from space_ml_sim.compute.checkpoint import CheckpointManager
+from space_ml_sim.compute.scheduler import InferenceScheduler
+from space_ml_sim.compute.quantization import (
+    quantize_model,
+    compare_quantization_resilience,
+    plot_quantization_comparison,
+)
+
+__all__ = [
+    "FaultInjector",
+    "FaultReport",
+    "TransformerFaultInjector",
+    "TMRWrapper",
+    "CheckpointManager",
+    "InferenceScheduler",
+    "quantize_model",
+    "compare_quantization_resilience",
+    "plot_quantization_comparison",
+]
+
+# Optional ONNX support — only available when 'onnx' and 'onnxruntime' are installed.
+try:
+    from space_ml_sim.compute.onnx_adapter import OnnxModel, load_onnx
+
+    __all__ += ["OnnxModel", "load_onnx"]
+except ImportError:
+    pass  # onnx extras not installed

space_ml_sim/compute/checkpoint.py

@@ -0,0 +1,78 @@
+"""Model checkpointing for fault recovery.
+
+Periodically saves model state so it can be restored after
+radiation-induced corruption is detected.
+"""
+
+from __future__ import annotations
+
+import copy
+from collections import deque
+from typing import Any
+
+import torch
+
+
+class CheckpointManager:
+    """Manages model checkpoints for radiation fault recovery.
+
+    Keeps a sliding window of recent checkpoints so the model
+    can be rolled back to a known-good state.
+    """
+
+    def __init__(self, max_checkpoints: int = 3) -> None:
+        """Initialize checkpoint manager.
+
+        Args:
+            max_checkpoints: Maximum number of checkpoints to retain.
+        """
+        self._checkpoints: deque[dict[str, Any]] = deque(maxlen=max_checkpoints)
+        self._max = max_checkpoints
+
+    def save(self, model: torch.nn.Module, metadata: dict[str, Any] | None = None) -> int:
+        """Save a model checkpoint.
+
+        Args:
+            model: Model to checkpoint.
+            metadata: Optional metadata (e.g., step count, accuracy).
+
+        Returns:
+            Index of the saved checkpoint.
+        """
+        entry = {
+            "state_dict": copy.deepcopy(model.state_dict()),
+            "metadata": metadata or {},
+        }
+        self._checkpoints.append(entry)
+        return len(self._checkpoints) - 1
+
+    def restore(self, model: torch.nn.Module, index: int = -1) -> dict[str, Any]:
+        """Restore a model from a checkpoint.
+
+        Args:
+            model: Model to restore into.
+            index: Checkpoint index (-1 for most recent).
+
+        Returns:
+            Metadata associated with the restored checkpoint.
+
+        Raises:
+            IndexError: If no checkpoints are available.
+        """
+        if not self._checkpoints:
+            raise IndexError("No checkpoints available")
+
+        entry = self._checkpoints[index]
+        # No deepcopy needed — load_state_dict copies tensor data into
+        # the model's existing buffers without mutating the source dict.
+        model.load_state_dict(entry["state_dict"])
+        return entry["metadata"]
+
+    @property
+    def count(self) -> int:
+        """Number of stored checkpoints."""
+        return len(self._checkpoints)
+
+    def clear(self) -> None:
+        """Remove all checkpoints."""
+        self._checkpoints.clear()

space_ml_sim/compute/fault_injector.py

@@ -0,0 +1,277 @@
+"""ML-aware radiation fault injection using PyTorch hooks.
+
+Injects radiation-modeled faults into neural network inference:
+    1. Weight SEU: Flip random bits in weight tensors (persistent memory fault)
+    2. Activation SET: Flip random bits in activations during forward pass (transient)
+    3. Stuck-at: Zero out weights in TID-degraded regions (permanent)
+
+Key insight: MSB flips in the IEEE 754 exponent field cause catastrophic errors,
+while LSB flips in the mantissa are often benign. Real radiation is uniform
+across bit positions.
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+from space_ml_sim.environment.radiation import RadiationEnvironment
+from space_ml_sim.models.chip_profiles import ChipProfile
+
+
+@dataclass(frozen=True)
+class FaultReport:
+    """Summary of faults injected during a simulation run."""
+
+    total_faults_injected: int = 0
+    weight_faults: int = 0
+    activation_faults: int = 0
+    layers_affected: tuple[str, ...] = ()
+    bit_positions_flipped: tuple[int, ...] = ()
+
+
+class FaultInjector:
+    """Inject radiation-modeled faults into PyTorch model inference.
+
+    Works independently of the orbital mechanics module — can be used
+    standalone for fault tolerance research on any PyTorch model.
+    """
+
+    def __init__(
+        self,
+        rad_env: RadiationEnvironment,
+        chip_profile: ChipProfile,
+        seed: int | None = None,
+    ) -> None:
+        self.rad_env = rad_env
+        self.chip = chip_profile
+        self._hooks: list[torch.utils.hooks.RemovableHook] = []
+        self._rng = np.random.default_rng(seed)
+
+    @staticmethod
+    def flip_random_bits(tensor: torch.Tensor, num_flips: int) -> list[int]:
+        """Flip random bits in a float32 tensor's IEEE 754 binary representation.
+
+        Args:
+            tensor: A float32 tensor (modified in-place).
+            num_flips: Number of random bit flips to inject.
+
+        Returns:
+            List of bit positions that were flipped (0-31).
+        """
+        if num_flips == 0:
+            return []
+
+        # Ensure contiguous memory — view() fails on non-contiguous tensors
+        # (common with transposed weights in transformers)
+        if not tensor.is_contiguous():
+            tensor.data = tensor.contiguous()
+        flat = tensor.view(-1)
+        n = flat.numel()
+        if n == 0:
+            return []
+
+        # Map float dtype to corresponding integer dtype for bit manipulation
+        dtype_map = {
+            torch.float32: (torch.int32, 32),
+            torch.float16: (torch.int16, 16),
+            torch.bfloat16: (torch.int16, 16),
+        }
+        int_dtype, num_bits = dtype_map.get(flat.dtype, (torch.int32, 32))
+
+        indices = torch.randint(0, n, (num_flips,))
+        bit_positions = torch.randint(0, num_bits, (num_flips,))
+        masks = (1 << bit_positions).to(int_dtype)
+
+        int_view = flat.view(int_dtype)
+        for i in range(num_flips):
+            int_view[indices[i]] ^= masks[i]
+
+        return bit_positions.tolist()
+
+    def inject_weight_faults(
+        self,
+        model: torch.nn.Module,
+        num_faults: int | None = None,
+        inference_time_seconds: float = 0.001,
+    ) -> FaultReport:
+        """Inject SEU bit flips into model weights.
+
+        If num_faults is None, the count is sampled from a Poisson distribution
+        based on the radiation environment, chip profile, and inference time.
+
+        Args:
+            model: PyTorch model (weights modified in-place).
+            num_faults: Exact number of faults, or None for radiation-sampled.
+            inference_time_seconds: Duration used for Poisson rate calculation.
+
+        Returns:
+            FaultReport summarizing injected faults.
+        """
+        if num_faults is None:
+            total_weight_bits = sum(p.numel() * 32 for p in model.parameters())
+            expected = self.rad_env.base_seu_rate * total_weight_bits * inference_time_seconds
+            num_faults = int(self._rng.poisson(expected))
+
+        if num_faults == 0:
+            return FaultReport()
+
+        # Distribute faults across layers proportional to parameter count
+        params = [(name, p) for name, p in model.named_parameters() if p.requires_grad]
+        sizes = [p.numel() for _, p in params]
+        total = sum(sizes)
+        if total == 0:
+            return FaultReport()
+
+        all_layers: list[str] = []
+        all_bits: list[int] = []
+        total_injected = 0
+
+        for name, param in params:
+            layer_faults = max(1, int(num_faults * param.numel() / total))
+            layer_faults = min(layer_faults, num_faults - total_injected)
+            if layer_faults <= 0:
+                continue
+
+            with torch.no_grad():
+                bits = self.flip_random_bits(param.data, layer_faults)
+
+            total_injected += layer_faults
+            all_layers.append(name)
+            all_bits.extend(bits)
+
+            if total_injected >= num_faults:
+                break
+
+        return FaultReport(
+            total_faults_injected=total_injected,
+            weight_faults=total_injected,
+            layers_affected=tuple(all_layers),
+            bit_positions_flipped=tuple(all_bits),
+        )
+
+    def register_activation_hooks(
+        self,
+        model: torch.nn.Module,
+        fault_probability: float = 0.001,
+    ) -> None:
+        """Register forward hooks for transient activation fault injection.
+
+        Args:
+            model: PyTorch model to hook.
+            fault_probability: Probability of a bit flip per tensor element per forward pass.
+        """
+        self.remove_hooks()
+        rng = self._rng
+
+        def make_hook(layer_name: str):
+            def hook(module, input, output):
+                if isinstance(output, torch.Tensor) and output.numel() > 0:
+                    num_faults = int(rng.binomial(output.numel(), fault_probability))
+                    if num_faults > 0:
+                        FaultInjector.flip_random_bits(output, num_faults)
+                return output
+
+            return hook
+
+        for name, module in model.named_modules():
+            # Hook leaf modules only
+            if len(list(module.children())) == 0:
+                h = module.register_forward_hook(make_hook(name))
+                self._hooks.append(h)
+
+    def remove_hooks(self) -> None:
+        """Remove all registered activation hooks."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+
+    def sweep(
+        self,
+        model: torch.nn.Module,
+        dataloader: torch.utils.data.DataLoader,
+        fault_counts: list[int],
+        num_trials: int = 5,
+    ) -> "pd.DataFrame":
+        """Run fault injection sweep across multiple fault counts and trials.
+
+        For each fault count, runs num_trials with independent fault injections
+        and measures accuracy degradation.
+
+        Args:
+            model: Base PyTorch model (will be deep-copied per trial).
+            dataloader: Evaluation data loader.
+            fault_counts: List of fault counts to sweep.
+            num_trials: Number of independent trials per fault count.
+
+        Returns:
+            DataFrame with columns: fault_count, trial, accuracy,
+            top5_accuracy, critical_failure, faults_injected, layers_affected.
+        """
+        import pandas as pd
+
+        results: list[dict] = []
+
+        # Pre-allocate one working model and a clean state dict
+        # to avoid deepcopy per trial (saves ~50% memory)
+        test_model = copy.deepcopy(model)
+        clean_state = copy.deepcopy(model.state_dict())
+
+        for fc in fault_counts:
+            for trial in range(num_trials):
+                test_model.load_state_dict(copy.deepcopy(clean_state))
+                test_model.eval()
+
+                report = self.inject_weight_faults(test_model, num_faults=fc)
+
+                acc, top5 = _evaluate_model(test_model, dataloader)
+                results.append(
+                    {
+                        "fault_count": fc,
+                        "trial": trial,
+                        "accuracy": acc,
+                        "top5_accuracy": top5,
+                        "critical_failure": acc < 0.1,
+                        "faults_injected": report.total_faults_injected,
+                        "layers_affected": len(report.layers_affected),
+                    }
+                )
+
+        return pd.DataFrame(results)
+
+
+def _evaluate_model(
+    model: torch.nn.Module,
+    dataloader: torch.utils.data.DataLoader,
+) -> tuple[float, float]:
+    """Evaluate a model on a dataloader, returning (top1_accuracy, top5_accuracy)."""
+    correct = 0
+    correct_top5 = 0
+    total = 0
+
+    with torch.no_grad():
+        for images, labels in dataloader:
+            outputs = model(images)
+            _, predicted = outputs.max(1)
+            correct += predicted.eq(labels).sum().item()
+
+            if outputs.size(1) >= 5:
+                _, top5_pred = outputs.topk(5, dim=1)
+                for i in range(len(labels)):
+                    if labels[i] in top5_pred[i]:
+                        correct_top5 += 1
+            else:
+                correct_top5 += predicted.eq(labels).sum().item()
+
+            total += labels.size(0)
+
+    acc = correct / total if total > 0 else 0.0
+    top5 = correct_top5 / total if total > 0 else 0.0
+    return acc, top5

space_ml_sim/compute/inference_node.py

@@ -0,0 +1,88 @@
+"""Inference node abstraction for running ML models on a satellite."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+import torch
+from pydantic import BaseModel, Field
+
+from space_ml_sim.models.chip_profiles import ChipProfile
+
+
+class NodeStatus(str, Enum):
+    """Inference node operational status."""
+
+    IDLE = "idle"
+    RUNNING = "running"
+    FAULTED = "faulted"
+
+
+class InferenceResult(BaseModel):
+    """Result of a single inference run."""
+
+    predictions: list[int] = Field(default_factory=list)
+    latency_ms: float = 0.0
+    faults_during_inference: int = 0
+    status: NodeStatus = NodeStatus.IDLE
+
+
+class InferenceNode:
+    """Wraps a PyTorch model as a compute node on a satellite.
+
+    Tracks inference count, fault history, and thermal/power constraints.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        chip_profile: ChipProfile,
+        device: str = "cpu",
+    ) -> None:
+        self.model = model.to(device).eval()
+        self.chip_profile = chip_profile
+        self.device = device
+        self.status = NodeStatus.IDLE
+        self.inference_count = 0
+        self.total_faults = 0
+
+    def run_inference(self, inputs: torch.Tensor) -> InferenceResult:
+        """Run inference on a batch of inputs.
+
+        Args:
+            inputs: Input tensor batch.
+
+        Returns:
+            InferenceResult with predictions and metadata.
+        """
+        self.status = NodeStatus.RUNNING
+        try:
+            with torch.no_grad():
+                outputs = self.model(inputs.to(self.device))
+                predictions = outputs.argmax(dim=1).tolist()
+        except Exception:
+            self.status = NodeStatus.FAULTED
+            raise
+
+        self.inference_count += 1
+        self.status = NodeStatus.IDLE
+
+        return InferenceResult(
+            predictions=predictions,
+            status=NodeStatus.IDLE,
+        )
+
+    def can_run(self, power_available_w: float, temperature_c: float) -> bool:
+        """Check if the node has enough power and is within thermal limits.
+
+        Args:
+            power_available_w: Available power in watts.
+            temperature_c: Current temperature in Celsius.
+
+        Returns:
+            True if the node can safely run inference.
+        """
+        return (
+            power_available_w >= self.chip_profile.tdp_watts
+            and temperature_c <= self.chip_profile.max_temp_c
+        )