Spaces:
Running
Running
File size: 4,880 Bytes
3d46076 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Versioned population checkpoint envelope (v4.1, Phase 2).
Separates simulation state from research/provenance history while storing
both, so save -> terminate -> restore -> continue is exactly equivalent to
uninterrupted execution:
state_hash - Population.population_hash() (organisms + world)
research_hash - event history + lineages + teaching sessions
combined_hash - sha256(state_hash + research_hash)
version - FlyBrain version (src/version.py)
schema_version - this envelope format (population_checkpoint_v1)
graph_identity - canonical graph identity (REAL_SUBGRAPH / ...)
experiment_identity - experiment seed + seed bundle + graph/circuit fingerprint
rng_bundle - full SeedBundle (deterministic continuation)
world/population state - complete Population.snapshot()
event history - population EventLog + teaching sessions (dropped by the
legacy Population.restore; preserved here)
Legacy Population.snapshot()/restore() are untouched (backward compatible).
"""
import dataclasses
import hashlib
import json
from typing import Any, Dict
from src.common.events import EventLog
CHECKPOINT_SCHEMA_VERSION = "population_checkpoint_v1"
def _canon(obj: Any) -> bytes:
return json.dumps(obj, sort_keys=True, default=str).encode()
def research_hash_of(pop) -> str:
h = hashlib.sha256()
h.update(pop.events.compute_hash().encode())
h.update(_canon({"genetic": pop.genetic_lineage, "cultural": pop.cultural_lineage}))
h.update(_canon([dataclasses.asdict(s) for s in pop.teaching_sessions]))
return h.hexdigest()
def experiment_identity_of(pop) -> str:
seeds = pop.seeds
return hashlib.sha256(_canon({
"experiment_seed": pop.experiment_seed,
"seeds": {k: getattr(seeds, k) for k in
("experiment_seed", "generation_seed", "organism_seed",
"development_seed", "mutation_seed", "world_seed", "teacher_seed")},
"graph_mode": pop.graph_mode.value,
"circuit_size": pop.circuit_size,
"genome_version": pop.genome_version,
"autonomy_mode": pop.autonomy_mode,
})).hexdigest()
def save_checkpoint(pop) -> Dict[str, Any]:
from src.version import VERSION
from src.connectome.types import GraphMode as _GM
state_hash = pop.population_hash()
research_hash = research_hash_of(pop)
combined = hashlib.sha256(f"{state_hash}|{research_hash}".encode()).hexdigest()
seeds = pop.seeds
return {
"version": VERSION,
"schema_version": CHECKPOINT_SCHEMA_VERSION,
"state_hash": state_hash,
"research_hash": research_hash,
"combined_hash": combined,
"graph_identity": _GM.canonical(pop.graph_mode),
"graph_mode": pop.graph_mode.value,
"experiment_identity": experiment_identity_of(pop),
"rng_bundle": {k: getattr(seeds, k) for k in
("experiment_seed", "generation_seed", "organism_seed",
"development_seed", "mutation_seed", "world_seed", "teacher_seed")},
"population": pop.snapshot(),
"event_history": list(pop.events.events),
"teaching_sessions": [dataclasses.asdict(s) for s in pop.teaching_sessions],
}
def restore_checkpoint(payload: Dict[str, Any]):
"""Restore a population from an envelope. Verifies hashes; raises on mismatch."""
from src.common.determinism import SeedBundle
from src.connectome.types import GraphMode as _GM
from src.culture.transmission import TeachingSession
from src.population.population import Population
if payload.get("schema_version") != CHECKPOINT_SCHEMA_VERSION:
raise ValueError(f"unsupported checkpoint schema {payload.get('schema_version')!r}")
rb = payload["rng_bundle"]
seeds = SeedBundle(experiment_seed=rb["experiment_seed"],
generation_seed=rb["generation_seed"],
organism_seed=rb["organism_seed"],
development_seed=rb["development_seed"],
mutation_seed=rb["mutation_seed"],
world_seed=rb["world_seed"],
teacher_seed=rb["teacher_seed"])
pop = Population.restore(payload["population"], seeds)
# Preserve research history the legacy path drops.
pop.events = EventLog()
pop.events.events = [dict(e) for e in payload.get("event_history", [])]
pop.teaching_sessions = [TeachingSession(**s)
for s in payload.get("teaching_sessions", [])]
# Verify: restored hashes must equal the recorded ones.
if pop.population_hash() != payload["state_hash"]:
raise ValueError("checkpoint state_hash mismatch after restore")
if research_hash_of(pop) != payload["research_hash"]:
raise ValueError("checkpoint research_hash mismatch after restore")
return pop
|