geno_lewm.training.resume

resume

Closed state helpers for fresh-process Carbon training continuation.

capture_rng_state

capture_rng_state() -> dict[str, object]

Capture every global RNG domain used by the training process.

Source code in geno_lewm/training/resume.py

def capture_rng_state() -> dict[str, object]:
    """Capture every global RNG domain used by the training process."""
    numpy, torch = _runtime_modules()
    numpy_state = numpy.random.get_state()
    cuda_available = bool(torch.cuda.is_available())
    return {
        "python": _jsonable(random.getstate()),
        "numpy": {
            "algorithm": str(numpy_state[0]),
            "keys": [int(value) for value in numpy_state[1].tolist()],
            "position": int(numpy_state[2]),
            "has_gauss": int(numpy_state[3]),
            "cached_gaussian": float(numpy_state[4]),
        },
        "torch_cpu": torch.get_rng_state().cpu(),
        "torch_cuda": {
            "available": cuda_available,
            "device_count": int(torch.cuda.device_count()) if cuda_available else 0,
            "states": [state.cpu() for state in torch.cuda.get_rng_state_all()]
            if cuda_available
            else [],
        },
    }

restore_rng_state

restore_rng_state(state: Mapping[str, object]) -> None

Restore a closed RNG payload and reject device-domain drift.

Source code in geno_lewm/training/resume.py

def restore_rng_state(state: Mapping[str, object]) -> None:
    """Restore a closed RNG payload and reject device-domain drift."""
    if set(state) != _RNG_KEYS:
        raise InputError(
            "resume checkpoint RNG state set is incomplete",
            details={
                "missing": sorted(_RNG_KEYS - set(state)),
                "unexpected": sorted(set(state) - _RNG_KEYS),
            },
        )
    numpy, torch = _runtime_modules()
    numpy_state = _mapping(state["numpy"], "rng.numpy")
    if set(numpy_state) != {
        "algorithm",
        "keys",
        "position",
        "has_gauss",
        "cached_gaussian",
    }:
        raise InputError("resume checkpoint NumPy RNG state is not closed")
    cuda_state = _mapping(state["torch_cuda"], "rng.torch_cuda")
    if set(cuda_state) != {"available", "device_count", "states"}:
        raise InputError("resume checkpoint CUDA RNG state is not closed")
    keys = numpy_state.get("keys")
    cuda_states = cuda_state.get("states")
    if not isinstance(keys, list) or not isinstance(cuda_states, list):
        raise InputError("resume checkpoint RNG arrays must be lists")
    available = bool(torch.cuda.is_available())
    if cuda_state.get("available") is not available:
        raise InputError("CUDA RNG availability changed across checkpoint resume")
    expected_devices = int(torch.cuda.device_count()) if available else 0
    if cuda_state.get("device_count") != expected_devices or len(cuda_states) != expected_devices:
        raise InputError("CUDA RNG device count changed across checkpoint resume")
    python_state = _tuplify(state["python"])
    if not isinstance(python_state, tuple):
        raise InputError("resume checkpoint Python RNG state must resolve to a tuple")
    cached_gaussian = numpy_state.get("cached_gaussian")
    if isinstance(cached_gaussian, bool) or not isinstance(cached_gaussian, int | float):
        raise InputError("resume checkpoint NumPy cached Gaussian must be numeric")
    try:
        random.setstate(python_state)
        numpy.random.set_state(
            (
                _text(numpy_state, "algorithm"),
                numpy.asarray(keys, dtype=numpy.uint32),
                _integer(numpy_state, "position"),
                _integer(numpy_state, "has_gauss"),
                float(cached_gaussian),
            )
        )
        torch.set_rng_state(state["torch_cpu"])
        if available:
            torch.cuda.set_rng_state_all(cuda_states)
    except (TypeError, ValueError, RuntimeError) as exc:
        raise InputError("resume checkpoint RNG states could not be restored") from exc

write_resume_checkpoint

write_resume_checkpoint(path: Path, *, source: Mapping[str, object], training_contract: Mapping[str, object], identities: Mapping[str, object], progress: Mapping[str, object], states: Mapping[str, object], trainer_state: Mapping[str, object], rng_state: Mapping[str, object], metric_history: list[dict[str, object]]) -> dict[str, object]

Atomically write one closed checkpoint and return its in-memory payload.

Source code in geno_lewm/training/resume.py

def write_resume_checkpoint(
    path: Path,
    *,
    source: Mapping[str, object],
    training_contract: Mapping[str, object],
    identities: Mapping[str, object],
    progress: Mapping[str, object],
    states: Mapping[str, object],
    trainer_state: Mapping[str, object],
    rng_state: Mapping[str, object],
    metric_history: list[dict[str, object]],
) -> dict[str, object]:
    """Atomically write one closed checkpoint and return its in-memory payload."""
    if set(states) != _STATE_KEYS:
        raise InputError("resume checkpoint trainer-state set is incomplete")
    if set(rng_state) != _RNG_KEYS:
        raise InputError("resume checkpoint RNG state set is incomplete")
    consumed = progress.get("consumed_window_ids")
    if not isinstance(consumed, list) or any(not isinstance(item, str) for item in consumed):
        raise InputError("resume checkpoint consumed window order must be a list of strings")
    normalized_progress = dict(progress)
    normalized_progress["consumed_order_digest"] = canonical_json_sha256(consumed)
    if set(normalized_progress) != _PROGRESS_KEYS:
        raise InputError("resume checkpoint progress fields do not match the closed contract")
    _validate_progress(normalized_progress)
    payload: dict[str, object] = {
        "schema_version": CHECKPOINT_SCHEMA_VERSION,
        "source": dict(source),
        "training_contract": dict(training_contract),
        "identities": dict(identities),
        "progress": normalized_progress,
        "states": dict(states),
        "trainer_state": dict(trainer_state),
        "state_digests": {
            **{name: state_digest(value) for name, value in states.items()},
            "trainer": state_digest(trainer_state),
        },
        "rng_state": dict(rng_state),
        "rng_state_digests": {name: state_digest(value) for name, value in rng_state.items()},
        "metric_history": [dict(row) for row in metric_history],
    }
    payload["payload_digest"] = _payload_digest(payload)
    _, torch = _runtime_modules()
    with atomic_binary_writer(path) as stream:
        torch.save(payload, stream)
    return payload

load_resume_checkpoint

load_resume_checkpoint(path: Path) -> dict[str, Any]

Safely load and validate a closed production checkpoint.

Source code in geno_lewm/training/resume.py

def load_resume_checkpoint(path: Path) -> dict[str, Any]:
    """Safely load and validate a closed production checkpoint."""
    if not path.is_file():
        raise InputError("resume checkpoint is missing", details={"path": str(path)})
    _, torch = _runtime_modules()
    try:
        payload = torch.load(path, map_location="cpu", weights_only=True)
    except Exception as exc:
        raise InputError(
            "resume checkpoint could not be loaded safely",
            details={"path": str(path), "error": str(exc)},
        ) from exc
    if not isinstance(payload, dict):
        raise InputError("resume checkpoint must contain a mapping")
    keys = set(payload)
    if keys != _CHECKPOINT_KEYS:
        raise InputError(
            "resume checkpoint fields do not match the closed contract",
            details={
                "missing": sorted(_CHECKPOINT_KEYS - keys),
                "unexpected": sorted(keys - _CHECKPOINT_KEYS),
            },
        )
    if payload.get("schema_version") != CHECKPOINT_SCHEMA_VERSION:
        raise InputError("resume checkpoint schema version is unsupported")
    if payload.get("payload_digest") != _payload_digest(payload):
        raise InputError("resume checkpoint payload digest does not match its contents")
    states = _mapping(payload.get("states"), "checkpoint.states")
    if set(states) != _STATE_KEYS:
        raise InputError("resume checkpoint trainer-state set is incomplete")
    trainer_state = _mapping(payload.get("trainer_state"), "checkpoint.trainer_state")
    expected_state_digests = {
        **{name: state_digest(value) for name, value in states.items()},
        "trainer": state_digest(trainer_state),
    }
    if payload.get("state_digests") != expected_state_digests:
        raise InputError("resume checkpoint trainer-state digests do not match")
    rng_state = _mapping(payload.get("rng_state"), "checkpoint.rng_state")
    if set(rng_state) != _RNG_KEYS:
        raise InputError("resume checkpoint RNG state set is incomplete")
    if payload.get("rng_state_digests") != {
        name: state_digest(value) for name, value in rng_state.items()
    }:
        raise InputError("resume checkpoint RNG-state digests do not match")
    progress = _mapping(payload.get("progress"), "checkpoint.progress")
    if set(progress) != _PROGRESS_KEYS:
        raise InputError("resume checkpoint progress fields do not match the closed contract")
    _validate_progress(progress)
    return payload

state_digest

state_digest(value: object) -> str

Return a canonical digest for nested primitive and tensor state.

Source code in geno_lewm/training/resume.py

def state_digest(value: object) -> str:
    """Return a canonical digest for nested primitive and tensor state."""
    return canonical_json_sha256(_digest_view(value))