Spaces:
Paused
Paused
| """Restore a real Brain-5D core from snapshot, journal, and runtime sidecar.""" | |
| from __future__ import annotations | |
| import random | |
| from pathlib import Path | |
| from typing import Any, Protocol, assert_never, cast | |
| from src.config.loader import ConfigDict | |
| from src.core.network import NeuralNetwork, SpikeEvent | |
| from src.core.spatial_index import unpack_coords | |
| from .b5d import B5DReader | |
| from .checkpoint import RuntimeCheckpoint, read_runtime_checkpoint | |
| from .recovery import RecoveryManager | |
| from .structural_journal import StructuralChangeKind, StructuralChangeRecord | |
| from .structural_recovery import StructuralRecoveryManager | |
| class StructuralRestoreManipulator(Protocol): | |
| def create_neuron(self, coord: tuple[int, int, int, int, int]) -> int: ... | |
| def delete_neuron(self, neuron_id: int) -> None: ... | |
| def create_synapse( | |
| self, source_id: int, target_id: int, weight: float, delay: int | |
| ) -> None: ... | |
| def delete_synapse(self, source_id: int, target_id: int) -> None: ... | |
| class _StructuralReplayTarget: | |
| def __init__(self, manipulator: StructuralRestoreManipulator) -> None: | |
| self._manipulator = manipulator | |
| def apply_structural_record(self, record: StructuralChangeRecord) -> bool: | |
| if record.kind is StructuralChangeKind.NEURON_ADD: | |
| if record.coord is None: | |
| return False | |
| neuron_id = self._manipulator.create_neuron(record.coord) | |
| return record.neuron_id is None or neuron_id == record.neuron_id | |
| if record.kind is StructuralChangeKind.NEURON_REMOVE: | |
| if record.neuron_id is None: | |
| return False | |
| self._manipulator.delete_neuron(record.neuron_id) | |
| return True | |
| if record.kind is StructuralChangeKind.SYNAPSE_ADD: | |
| if ( | |
| record.source_id is None | |
| or record.target_id is None | |
| or record.weight is None | |
| or record.delay is None | |
| ): | |
| return False | |
| self._manipulator.create_synapse( | |
| record.source_id, | |
| record.target_id, | |
| record.weight, | |
| record.delay, | |
| ) | |
| return True | |
| if record.kind is StructuralChangeKind.SYNAPSE_REMOVE: | |
| if record.source_id is None or record.target_id is None: | |
| return False | |
| self._manipulator.delete_synapse(record.source_id, record.target_id) | |
| return True | |
| assert_never(record.kind) | |
| class RestoredNeuralNetwork(NeuralNetwork): | |
| """NeuralNetwork subclass exposing a narrow restoration hook.""" | |
| def restore_queued_count(self, value: int) -> None: | |
| """Restore queue counter after event slots have been reconstructed.""" | |
| self._queued_event_count = value | |
| def restore_network( | |
| snapshot_path: Path, | |
| journal_path: Path, | |
| checkpoint_path: Path, | |
| config: ConfigDict, | |
| recovered_path: Path, | |
| *, | |
| structural_journal_path: Path | None = None, | |
| ) -> RestoredNeuralNetwork: | |
| """Recover persisted state and return a runnable deterministic network. | |
| Args: | |
| snapshot_path: Path to the base `.b5d` snapshot. | |
| journal_path: Path to the `.b5d.journal` delta journal. | |
| checkpoint_path: Path to the runtime checkpoint JSON. | |
| config: Network configuration (dict). | |
| recovered_path: Path to write the recovered snapshot (if needed). | |
| structural_journal_path: Optional structural journal path. | |
| Returns: | |
| A restored RestoredNeuralNetwork instance. | |
| Raises: | |
| RuntimeError: If recovery fails. | |
| ValueError: If the snapshot is not restart-capable or required fields are missing. | |
| """ | |
| result = RecoveryManager(snapshot_path, journal_path).recover(recovered_path) | |
| if not result.success: | |
| raise RuntimeError(result.error or "snapshot recovery failed") | |
| checkpoint = read_runtime_checkpoint(checkpoint_path) | |
| rng = _restore_rng(checkpoint) | |
| # Cast ConfigDict to dict[str, Any] for the NeuralNetwork constructor | |
| network = RestoredNeuralNetwork(cast(dict[str, Any], config), rng) | |
| with B5DReader(recovered_path) as reader: | |
| for neuron_record in reader.iter_neurons(): | |
| neuron_id = network.add_neuron(unpack_coords(neuron_record.neuron_id)) | |
| neuron = network.neurons[neuron_id] | |
| # Ensure restart-capable snapshot was used | |
| if neuron_record.a is None or neuron_record.b is None: | |
| raise ValueError("restart-capable snapshot required") | |
| if neuron_record.c is None or neuron_record.d is None: | |
| raise ValueError("restart-capable snapshot required") | |
| if neuron_record.spike_cost is None: | |
| raise ValueError("restart-capable snapshot required") | |
| if neuron_record.spike_counter is None: | |
| raise ValueError("restart-capable snapshot required") | |
| if neuron_record.last_spike_tick is None: | |
| raise ValueError("restart-capable snapshot required") | |
| # Restore core neuron parameters | |
| neuron.a = neuron_record.a | |
| neuron.b = neuron_record.b | |
| neuron.c = neuron_record.c | |
| neuron.d = neuron_record.d | |
| neuron.spike_cost = neuron_record.spike_cost | |
| neuron.spike_counter = neuron_record.spike_counter | |
| neuron.last_spike_tick = neuron_record.last_spike_tick | |
| # Restore dynamic state from optical record | |
| neuron.v = neuron_record.optical.membrane_v | |
| neuron.u = neuron_record.optical.recovery_u | |
| neuron.energy = neuron_record.optical.energy | |
| neuron.threshold_adaptation = neuron_record.optical.threshold_adaptation | |
| for synapse_record in reader.iter_synapses(): | |
| network.connect( | |
| synapse_record.source_id, | |
| synapse_record.target_id, | |
| synapse_record.weight, | |
| synapse_record.delay, | |
| ) | |
| synapse = network.synapses[synapse_record.source_id][-1] | |
| synapse.eligibility = synapse_record.eligibility | |
| synapse.last_pre_spike = synapse_record.last_pre_spike | |
| # Replay structural plasticity journal if provided | |
| if structural_journal_path is not None: | |
| from src.manipulation.manipulator import Brain5DManipulator | |
| structural_target = _StructuralReplayTarget(Brain5DManipulator(network)) | |
| StructuralRecoveryManager().replay( | |
| structural_target, | |
| structural_journal_path, | |
| ) | |
| # Restore runtime state and exact neuron/synapse values from checkpoint | |
| _restore_runtime(network, checkpoint) | |
| _restore_exact_neuron_state(network, checkpoint) | |
| _restore_exact_synapse_state(network, checkpoint) | |
| return network | |
| def _restore_rng(checkpoint: RuntimeCheckpoint) -> random.Random: | |
| """Restore the random number generator from checkpoint state.""" | |
| rng = random.Random() | |
| rng.setstate( | |
| ( | |
| checkpoint.rng.version, | |
| tuple(checkpoint.rng.state), | |
| checkpoint.rng.gauss_next, | |
| ) | |
| ) | |
| return rng | |
| def _restore_runtime( | |
| network: RestoredNeuralNetwork, | |
| checkpoint: RuntimeCheckpoint, | |
| ) -> None: | |
| """Restore runtime state from checkpoint.""" | |
| network.current_tick = checkpoint.current_tick | |
| network.total_spikes = checkpoint.total_spikes | |
| network.total_events_processed = checkpoint.total_events_processed | |
| network.pending_currents = dict(checkpoint.pending_currents) | |
| network.input_cells = set(checkpoint.input_cells) | |
| network.output_cells = set(checkpoint.output_cells) | |
| # Rebuild event queue | |
| network.event_slots = [[] for _ in range(network.max_delay + 1)] | |
| queued = 0 | |
| for event in checkpoint.queued_events: | |
| slot = event.delivery_tick % len(network.event_slots) | |
| network.event_slots[slot].append( | |
| SpikeEvent( | |
| source_id=event.source_id, | |
| target_id=event.target_id, | |
| weight=event.weight, | |
| delivery_tick=event.delivery_tick, | |
| ) | |
| ) | |
| queued += 1 | |
| network.restore_queued_count(queued) | |
| def _restore_exact_neuron_state( | |
| network: RestoredNeuralNetwork, | |
| checkpoint: RuntimeCheckpoint, | |
| ) -> None: | |
| """Overlay exact values lost by the compact snapshot representation.""" | |
| for state in checkpoint.neuron_states: | |
| neuron = network.neurons.get(state.neuron_id) | |
| if neuron is None: | |
| raise ValueError( | |
| f"runtime checkpoint references missing neuron {state.neuron_id}" | |
| ) | |
| neuron.a = state.a | |
| neuron.b = state.b | |
| neuron.c = state.c | |
| neuron.d = state.d | |
| neuron.v = state.membrane_v | |
| neuron.u = state.recovery_u | |
| neuron.energy = state.energy | |
| neuron.spike_cost = state.spike_cost | |
| neuron.threshold_adaptation = state.threshold_adaptation | |
| neuron.spike_counter = state.spike_counter | |
| neuron.last_spike_tick = state.last_spike_tick | |
| neuron.last_external_current = state.last_external_current | |
| neuron.last_synaptic_current = state.last_synaptic_current | |
| neuron.firing_rate_estimate = state.firing_rate_estimate | |
| neuron._spike_count_window = ( # pyright: ignore[reportPrivateUsage] | |
| state.spike_count_window | |
| ) | |
| neuron._last_update_tick = ( # pyright: ignore[reportPrivateUsage] | |
| state.last_update_tick | |
| ) | |
| neuron.pre_trace = state.pre_trace | |
| neuron.post_trace = state.post_trace | |
| def _restore_exact_synapse_state( | |
| network: RestoredNeuralNetwork, | |
| checkpoint: RuntimeCheckpoint, | |
| ) -> None: | |
| """Overlay exact synaptic float state lost by V1 float32 fields.""" | |
| if not checkpoint.synapse_states: | |
| return | |
| for state in checkpoint.synapse_states: | |
| candidates = network.synapses.get(state.source_id) | |
| if candidates is None: | |
| raise ValueError( | |
| f"runtime checkpoint references missing source {state.source_id}" | |
| ) | |
| synapse = next( | |
| (item for item in candidates if item.target_id == state.target_id), | |
| None, | |
| ) | |
| if synapse is None: | |
| raise ValueError( | |
| "runtime checkpoint references missing synapse " | |
| f"{state.source_id}->{state.target_id}" | |
| ) | |
| synapse.weight = state.weight | |
| synapse.delay = state.delay | |
| synapse.eligibility = state.eligibility | |
| synapse.last_pre_spike = state.last_pre_spike | |
| def restore_homeostasis_state( | |
| homeostasis_engine: Any, | |
| checkpoint: RuntimeCheckpoint, | |
| ) -> None: | |
| """Restore homeostasis engine state from checkpoint v4. | |
| Args: | |
| homeostasis_engine: A HomeostasisEngine instance (must have _rates_hz dict). | |
| checkpoint: The RuntimeCheckpoint with homeostasis_state. | |
| Raises: | |
| ValueError: If a neuron referenced in the checkpoint no longer exists. | |
| """ | |
| if not checkpoint.homeostasis_state: | |
| return | |
| rates = {} | |
| for record in checkpoint.homeostasis_state: | |
| rates[int(record.neuron_id)] = float(record.rate_hz) | |
| # Restore the internal rates dict | |
| if hasattr(homeostasis_engine, "_rates_hz"): | |
| homeostasis_engine._rates_hz = rates | |
| def restore_learning_state( | |
| learning_engine: Any, | |
| checkpoint: RuntimeCheckpoint, | |
| ) -> None: | |
| """Restore learning engine state from checkpoint v4. | |
| Restores per-synapse traces (last_pre_tick, last_post_tick, eligibility) | |
| and pending reward signals. | |
| Args: | |
| learning_engine: A LearningEngine instance (must have _states dict | |
| and _pending_rewards list). | |
| checkpoint: The RuntimeCheckpoint with learning_state and pending_rewards. | |
| Raises: | |
| ValueError: If a synapse referenced in the checkpoint no longer exists. | |
| """ | |
| if not hasattr(learning_engine, "_states"): | |
| return | |
| # Restore per-synapse traces | |
| # Uses stable (pre_id, target_id) key matching LearningEngine's | |
| # deterministic topology indexing. | |
| if checkpoint.learning_state: | |
| states = learning_engine._states | |
| for record in checkpoint.learning_state: | |
| pre_id = int(record.pre_id) | |
| target_id = int(record.target_id) | |
| key = (pre_id, target_id) | |
| if key in states: | |
| state = states[key] | |
| state.last_pre_tick = record.last_pre_tick | |
| state.last_post_tick = record.last_post_tick | |
| # EligibilityTrace uses .value, not ._trace | |
| state.eligibility.value = record.eligibility_value | |
| # Restore per-trace last_tick so decay is accurate | |
| state.eligibility.last_tick = record.eligibility_last_tick | |
| # Restore pending rewards (independent of learning_state presence) | |
| if checkpoint.pending_rewards and hasattr(learning_engine, "_pending_rewards"): | |
| from src.learning.reward import RewardSignal | |
| learning_engine._pending_rewards = [ | |
| RewardSignal(value=float(r.value), tick=int(r.tick)) | |
| for r in checkpoint.pending_rewards | |
| ] | |
| # ============================================================================ | |
| # Production Restore Bundle | |
| # ============================================================================ | |
| class RestoredBundle: | |
| """Complete production restore result: network + optional engines. | |
| This bundles the restored neural network together with its optional | |
| homeostasis and learning engines, providing a single restore-and-continue | |
| entry point for production use. | |
| Attributes: | |
| network: The restored RestoredNeuralNetwork instance. | |
| homeostasis_engine: Optional restored HomeostasisEngine, or None. | |
| learning_engine: Optional restored LearningEngine, or None. | |
| """ | |
| def __init__( | |
| self, | |
| network: RestoredNeuralNetwork, | |
| homeostasis_engine: Any | None = None, | |
| learning_engine: Any | None = None, | |
| ) -> None: | |
| self.network = network | |
| self.homeostasis_engine = homeostasis_engine | |
| self.learning_engine = learning_engine | |
| def restore_full( | |
| snapshot_path: Path, | |
| journal_path: Path, | |
| checkpoint_path: Path, | |
| config: ConfigDict, | |
| recovered_path: Path, | |
| *, | |
| structural_journal_path: Path | None = None, | |
| create_homeostasis_engine: bool = False, | |
| create_learning_engine: bool = False, | |
| ) -> RestoredBundle: | |
| """Restore network and optionally create and restore engines. | |
| This is the production entry point that replaces calling | |
| ``restore_network()`` followed by manual engine creation and restore. | |
| It ensures the full restore chain is executed atomically: | |
| 1. Recover snapshot and restore network (same as restore_network()) | |
| 2. Optionally create HomeostasisEngine and restore its state | |
| 3. Optionally create LearningEngine and restore its state | |
| Args: | |
| snapshot_path: Path to the base ``.b5d`` snapshot. | |
| journal_path: Path to the ``.b5d.journal`` delta journal. | |
| checkpoint_path: Path to the runtime checkpoint JSON. | |
| config: Network configuration (dict). | |
| recovered_path: Path to write the recovered snapshot. | |
| structural_journal_path: Optional structural journal path. | |
| create_homeostasis_engine: If True, create and restore | |
| HomeostasisEngine from checkpoint. | |
| create_learning_engine: If True, create and restore | |
| LearningEngine from checkpoint. | |
| Returns: | |
| A RestoredBundle containing the network and optionally the engines. | |
| Raises: | |
| RuntimeError: If recovery fails. | |
| ValueError: If the snapshot is not restart-capable. | |
| """ | |
| network = restore_network( | |
| snapshot_path=snapshot_path, | |
| journal_path=journal_path, | |
| checkpoint_path=checkpoint_path, | |
| config=config, | |
| recovered_path=recovered_path, | |
| structural_journal_path=structural_journal_path, | |
| ) | |
| checkpoint = read_runtime_checkpoint(checkpoint_path) | |
| homeostasis_engine: Any = None | |
| learning_engine: Any = None | |
| if create_homeostasis_engine: | |
| from src.homeostasis.engine import HomeostasisEngine | |
| homeostasis_engine = HomeostasisEngine(network, dict(config)) | |
| restore_homeostasis_state(homeostasis_engine, checkpoint) | |
| homeostasis_engine.attach() | |
| if create_learning_engine: | |
| from src.learning.learning_engine import LearningEngine | |
| learning_engine = LearningEngine(network, dict(config)) | |
| restore_learning_state(learning_engine, checkpoint) | |
| learning_engine.attach() | |
| return RestoredBundle( | |
| network=network, | |
| homeostasis_engine=homeostasis_engine, | |
| learning_engine=learning_engine, | |
| ) | |