gam_solve/

persistent_warm_start.rs

use gam_runtime::warm_start::{
    ConfiguredWarmStartStore, EntryKind, Fingerprinter, StoreError, StoreOptions,
};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};

// v2 (#2615): `PersistentBlockInnerSummary` carries the smoothing state its
// cached inner mode was solved at, so a restored mode can be keyed correctly.
const CACHE_VERSION: u32 = 2;
const MAX_ENTRY_BYTES: u64 = 16 * 1024 * 1024;
const MAX_TOTAL_BYTES: u64 = 256 * 1024 * 1024;
const CACHE_TTL_SECS: u64 = 60 * 60 * 24 * 365 * 10;

/// String tag identifying the on-disk cache schema, embedded directly in
/// cache keys.
///
/// The leading `schema2-` prefix is bumped manually only when the
/// serialized cache layout changes in a way that makes prior entries
/// unsafe to consume (struct fields added/removed, optimization
/// invariants altered, payload semantics shift). This is **deliberately
/// separate** from `CARGO_PKG_VERSION` so a routine library version bump
/// does NOT invalidate every user's warm-start cache.
pub fn cache_schema_tag() -> String {
    // Bumped from `schema2-` → `schema3-` when the three hand-written
    // hashers (`Fingerprinter`, `StableHasher`, `CacheDigestBuilder`) were
    // unified onto `Fingerprinter`. Prior on-disk warm-start entries are
    // walled off into the `schema2-` keyspace and cold-start once; this
    // is the intentional consequence of the unification, documented in
    // the commit that performs it. See `src/warm_start/key.rs` for the new
    // canonical hasher API.
    // Bumped to `v2` when the persistent warm-start key stopped hashing the
    // θ-dependent, lazily-refreshed isometry Jacobian cache slots
    // (`jacobian_cache` / `jacobian_second_cache` / `third_decoder_derivative`).
    // Those snapshots made the key non-reproducible across identical repeat
    // fits, so the outer `skip-outer-validation` warm hit was lost (#1048). The
    // bump walls off any entries written under the old, drifting keys; they are
    // simply never matched (and TTL-evicted) rather than aliasing.
    // Bumped to `v3` when the descriptor-indexed cross-fit `FitArtifact`
    // keyspace ("fit-artifact-key") was introduced alongside the existing
    // inner/outer warm-start records. The bump walls off any entries written
    // under the old layouts so a mixed-schema store never aliases a legacy
    // payload into the new artifact reader (and vice versa).
    "schema3-unified-fingerprinter-v3".to_string()
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PersistentWarmStartRecord {
    pub version: u32,
    pub key: String,
    pub package_version: String,
    pub created_unix_secs: u64,
    pub updated_unix_secs: u64,
    pub n_rows: usize,
    pub n_cols: usize,
    pub rho: Vec<f64>,
    pub beta: Vec<f64>,
    pub prev_rho: Option<Vec<f64>>,
    pub prev_beta: Option<Vec<f64>>,
    pub last_inner_iters: usize,
    pub last_inner_converged: bool,
    pub last_pirls_lm_lambda: Option<f64>,
    pub last_ift_prediction_residual: Option<f64>,
    pub last_pirls_accept_rho: Option<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PersistentBlockInnerSummary {
    pub log_likelihood: f64,
    pub penalty_value: f64,
    pub cycles: usize,
    pub converged: bool,
    pub block_logdet_h: f64,
    pub block_logdet_s: f64,
    /// Per-block `log λ` this mode was solved at, in block order (#2615).
    ///
    /// The record's own `rho` is the OPTIMIZER's coordinate vector, which for a
    /// labeled / joint-penalty family is neither the per-block strengths nor
    /// split into them — so it cannot serve as the reuse key. A cached Laplace
    /// mode replayed at a different smoothing state is a criterion evaluated at
    /// one ρ and reported at another; the state therefore travels with the mode.
    pub block_log_lambdas: Vec<Vec<f64>>,
    /// Joint-bundle `log λ` this mode was solved at; empty for a family that
    /// declares no joint penalties (#2615).
    pub joint_log_lambdas: Vec<f64>,
}

impl PersistentBlockInnerSummary {
    fn is_valid(&self) -> bool {
        self.log_likelihood.is_finite()
            && self.penalty_value.is_finite()
            && self.block_logdet_h.is_finite()
            && self.block_logdet_s.is_finite()
            && self
                .block_log_lambdas
                .iter()
                .all(|block| block.iter().all(|v| v.is_finite()))
            && self.joint_log_lambdas.iter().all(|v| v.is_finite())
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PersistentBlockWarmStartRecord {
    pub version: u32,
    pub key: String,
    pub package_version: String,
    pub created_unix_secs: u64,
    pub updated_unix_secs: u64,
    pub n_rows: usize,
    pub block_names: Vec<String>,
    pub block_dims: Vec<usize>,
    pub rho: Vec<f64>,
    pub block_beta: Vec<Vec<f64>>,
    pub active_sets: Vec<Option<Vec<usize>>>,
    #[serde(default)]
    pub inner: Option<PersistentBlockInnerSummary>,
}

impl PersistentBlockWarmStartRecord {
    pub fn new(
        key: String,
        n_rows: usize,
        block_names: Vec<String>,
        block_dims: Vec<usize>,
    ) -> Self {
        let now = unix_secs_now();
        Self {
            version: CACHE_VERSION,
            key,
            package_version: env!("CARGO_PKG_VERSION").to_string(),
            created_unix_secs: now,
            updated_unix_secs: now,
            n_rows,
            block_names,
            block_dims,
            rho: Vec::new(),
            block_beta: Vec::new(),
            active_sets: Vec::new(),
            inner: None,
        }
    }

    pub fn is_compatible(
        &self,
        key: &str,
        n_rows: usize,
        block_names: &[String],
        block_dims: &[usize],
        rho_len: usize,
    ) -> bool {
        self.version == CACHE_VERSION
            && self.key == key
            // Note: `package_version` is no longer required to match. A
            // library version bump that doesn't change the cache schema
            // (the common case for patch / minor releases) should NOT
            // invalidate users' on-disk warm-start caches. Schema-breaking
            // changes bump the `schemaN-` prefix in `cache_schema_tag()`,
            // which is encoded in the cache key itself.
            && self.n_rows == n_rows
            && self.block_names == block_names
            && self.block_dims == block_dims
            && self.rho.len() == rho_len
            && self.rho.iter().all(|v| v.is_finite())
            && self.block_beta.len() == block_dims.len()
            && self
                .block_beta
                .iter()
                .zip(block_dims.iter())
                .all(|(beta, dim)| beta.len() == *dim && beta.iter().all(|v| v.is_finite()))
            && self.active_sets.len() == block_dims.len()
            && self.inner.as_ref().is_none_or(|inner| inner.is_valid())
    }
}

impl PersistentWarmStartRecord {
    pub fn new(key: String, n_rows: usize, n_cols: usize) -> Self {
        let now = unix_secs_now();
        Self {
            version: CACHE_VERSION,
            key,
            package_version: env!("CARGO_PKG_VERSION").to_string(),
            created_unix_secs: now,
            updated_unix_secs: now,
            n_rows,
            n_cols,
            rho: Vec::new(),
            beta: Vec::new(),
            prev_rho: None,
            prev_beta: None,
            last_inner_iters: 0,
            last_inner_converged: false,
            last_pirls_lm_lambda: None,
            last_ift_prediction_residual: None,
            last_pirls_accept_rho: None,
        }
    }

    pub fn is_compatible(&self, key: &str, n_rows: usize, n_cols: usize) -> bool {
        self.version == CACHE_VERSION
            && self.key == key
            // Note: `package_version` is no longer required to match. A
            // library version bump that doesn't change the cache schema
            // (the common case for patch / minor releases) should NOT
            // invalidate users' on-disk warm-start caches. Schema-breaking
            // changes bump the `schemaN-` prefix in `cache_schema_tag()`,
            // which is encoded in the cache key itself.
            && self.n_rows == n_rows
            && self.n_cols == n_cols
            && self.rho.iter().all(|v| v.is_finite())
            && self.beta.len() == n_cols
            && self.beta.iter().all(|v| v.is_finite())
            && self
                .prev_rho
                .as_ref()
                .is_none_or(|rho| rho.len() == self.rho.len() && rho.iter().all(|v| v.is_finite()))
            && self
                .prev_beta
                .as_ref()
                .is_none_or(|beta| beta.len() == n_cols && beta.iter().all(|v| v.is_finite()))
    }
}

/// Configure persistent warm starts at the exact caller-supplied root.
///
/// This is lazy: request parsing and structural validation do not touch the
/// filesystem. The first real persistence operation opens the root once, and
/// all clones share that handle or the single best-effort unavailable decision.
pub fn configured_store(root: PathBuf) -> ConfiguredWarmStartStore {
    ConfiguredWarmStartStore::new(
        root,
        StoreOptions {
            size_budget_bytes: MAX_TOTAL_BYTES,
            ttl: Duration::from_secs(CACHE_TTL_SECS),
        },
    )
}

pub fn load_record(
    store: &ConfiguredWarmStartStore,
    key: &str,
) -> Option<PersistentWarmStartRecord> {
    best_effort(
        store,
        "load warm-start record",
        load_json_record(store, key),
    )
}

pub fn load_block_record(
    store: &ConfiguredWarmStartStore,
    key: &str,
) -> Option<PersistentBlockWarmStartRecord> {
    best_effort(
        store,
        "load custom-family warm-start record",
        load_json_record(store, key),
    )
}

pub fn store_record(store: &ConfiguredWarmStartStore, record: &PersistentWarmStartRecord) {
    best_effort(
        store,
        "store warm-start record",
        store_json_record(store, &record.key, record),
    );
}

pub fn store_block_record(
    store: &ConfiguredWarmStartStore,
    record: &PersistentBlockWarmStartRecord,
) {
    best_effort(
        store,
        "store custom-family warm-start record",
        store_json_record(store, &record.key, record),
    );
}

#[derive(Debug)]
enum PersistentStoreError {
    Encode(String),
    Unavailable(String),
    Rejected(String),
}

impl std::fmt::Display for PersistentStoreError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Encode(detail) => write!(f, "encode failed: {detail}"),
            Self::Unavailable(detail) => write!(f, "filesystem unavailable: {detail}"),
            Self::Rejected(detail) => write!(f, "store rejected operation: {detail}"),
        }
    }
}

fn classify_store_error(error: StoreError) -> PersistentStoreError {
    match error {
        StoreError::Io(error) => PersistentStoreError::Unavailable(error.to_string()),
        StoreError::Json(error) => PersistentStoreError::Rejected(error.to_string()),
    }
}

/// Interpret every persistent-cache failure in one place.
///
/// Filesystem refusal permanently disables this explicitly configured store
/// for the fit and emits one diagnostic from the store capability. Encoding or
/// store-contract defects remain loud in diagnostics and typed in the internal
/// `try_*` functions, but persistence can never fail a statistical fit.
fn best_effort<T: Default>(
    store: &ConfiguredWarmStartStore,
    operation: &'static str,
    result: Result<T, PersistentStoreError>,
) -> T {
    match result {
        Ok(value) => value,
        Err(PersistentStoreError::Unavailable(detail)) => {
            store.mark_unavailable(operation, &detail);
            T::default()
        }
        Err(error @ (PersistentStoreError::Encode(_) | PersistentStoreError::Rejected(_))) => {
            log::warn!(
                "[warm-start-cache] persistence defect operation={} explicit_root={}: {}",
                operation,
                store.root().display(),
                error
            );
            T::default()
        }
    }
}

/// Encode a record for the store.
///
/// See [`load_json_record`] before changing the encoding: JSON's inability to
/// represent non-finite `f64` is relied upon on the load side (gam#2721).
fn store_json_record<T: Serialize>(
    configured: &ConfiguredWarmStartStore,
    key: &str,
    record: &T,
) -> Result<(), PersistentStoreError> {
    let bytes = serde_json::to_vec(record)
        .map_err(|error| PersistentStoreError::Encode(error.to_string()))?;
    if bytes.len() as u64 > MAX_ENTRY_BYTES {
        return Ok(());
    }
    let Some(store) = configured.store() else {
        return Ok(());
    };
    let mut fp = Fingerprinter::new();
    fp.absorb_str(b"warm-start-key", key);
    store
        .save(&fp.finalize(), &bytes, None, None, EntryKind::Checkpoint)
        .map_err(classify_store_error)?;
    Ok(())
}

/// Decode a stored record.
///
/// THE JSON ENCODING IS LOAD-BEARING, NOT INCIDENTAL (gam#2721). `serde_json`
/// cannot represent a non-finite `f64`: it writes `NaN` / `+-inf` as `null`, and
/// deserializing `null` into `f64` fails here as
/// [`PersistentStoreError::Rejected`]. That is the SECOND of exactly two
/// independent barriers keeping a non-finite coefficient out of a restored
/// warm start, and so out of `state.beta` at the cached-seed restore in
/// `inner_blockwise_fit`, which applies no finiteness check of its own:
///
///   1. the write side refuses to persist a record whose `block_beta` holds a
///      non-finite value (`store_persistent_custom_family_warm_start`);
///   2. this encoding cannot carry one, so a record written by an older binary,
///      hand-edited, or corrupted is REJECTED at parse rather than decoded into
///      a `NaN`.
///
/// Barrier 1 covers only records this code wrote. Barrier 2 covers every other
/// origin, and it is a property of the FORMAT rather than of any check — so
/// swapping this encoding for one that can carry non-finite values (bincode,
/// msgpack, arrow, CBOR) removes it with no diff to any guard and no failing
/// test. A format change here must therefore come with a load-side finiteness
/// check on the decoded coefficient vectors.
fn load_json_record<T: for<'de> Deserialize<'de>>(
    configured: &ConfiguredWarmStartStore,
    key: &str,
) -> Result<Option<T>, PersistentStoreError> {
    let Some(store) = configured.store() else {
        return Ok(None);
    };
    let mut fp = Fingerprinter::new();
    fp.absorb_str(b"warm-start-key", key);
    let Some(entry) = store.lookup(&fp.finalize()).map_err(classify_store_error)? else {
        return Ok(None);
    };
    if entry.payload.len() as u64 > MAX_ENTRY_BYTES {
        return Ok(None);
    }
    serde_json::from_slice(&entry.payload)
        .map(Some)
        .map_err(|error| PersistentStoreError::Rejected(error.to_string()))
}

/// Open a [`gam_runtime::warm_start::Session`] for outer-iterate (rho-axis) checkpoints.
///
/// Uses a different fingerprint tag than the inner `warm-start-key`
/// absorption (see `load_json_record`) so the outer-iterate keyspace
/// is disjoint from the inner beta-record keyspace —
/// the two layers persist different payload shapes and must not alias.
pub fn open_outer_session(
    configured: &ConfiguredWarmStartStore,
    key: &str,
) -> Option<std::sync::Arc<gam_runtime::warm_start::Session>> {
    let mut fp = Fingerprinter::new();
    fp.absorb_str(b"outer-iterate-key", key);
    let fp = fp.finalize();
    configured.open_session(fp)
}

/// Persist a descriptor-indexed cross-fit `FitArtifact` under the
/// `fit-artifact-key` keyspace, keyed by the descriptor's structural key (so
/// an LOSO fold of the same model retrieves a prior full-data fit). The
/// schema tag is folded into the key so legacy layouts are walled off.
///
/// Best-effort: the public function cannot fail a fit. Its internal `try_*`
/// counterpart retains typed failures for contract tests, and `best_effort`
/// owns the sole production interpretation.
pub fn store_fit_artifact(
    store: &ConfiguredWarmStartStore,
    artifact: &crate::warm_start_artifact::FitArtifact,
) {
    best_effort(
        store,
        "store cross-fit artifact",
        try_store_fit_artifact(store, artifact),
    );
}

fn try_store_fit_artifact(
    configured: &ConfiguredWarmStartStore,
    artifact: &crate::warm_start_artifact::FitArtifact,
) -> Result<(), PersistentStoreError> {
    if !artifact.is_usable() {
        // Never persist a non-finite / wrong-schema artifact: it could only
        // ever be rejected on load anyway.
        return Ok(());
    }
    let bytes = serde_json::to_vec(artifact)
        .map_err(|error| PersistentStoreError::Encode(error.to_string()))?;
    if bytes.len() as u64 > MAX_ENTRY_BYTES {
        return Ok(());
    }
    let Some(store) = configured.store() else {
        return Ok(());
    };
    let key = artifact.descriptor.descriptor_key().to_hex();
    let mut fp = Fingerprinter::new();
    fp.absorb_str(b"fit-artifact-key", &cache_schema_tag());
    fp.absorb_str(b"fit-artifact-descriptor", &key);
    store
        .save(&fp.finalize(), &bytes, None, None, EntryKind::Checkpoint)
        .map_err(classify_store_error)?;
    Ok(())
}

/// Load the newest valid cross-fit `FitArtifact` whose descriptor key
/// matches `descriptor_key_hex` (the hex of [`crate::warm_start_artifact::FitDescriptor::descriptor_key`]).
///
/// Uses `lookup_latest` (newest-valid) rather than objective-ranked lookup:
/// descriptor-key matches can be different folds / row sets whose objectives
/// are not on a common scale, so "lowest objective" is the wrong rule.
/// Returns `None` (cold fallback) on any miss or non-finite payload.
pub fn load_fit_artifact_by_descriptor(
    store: &ConfiguredWarmStartStore,
    descriptor_key_hex: &str,
) -> Option<crate::warm_start_artifact::FitArtifact> {
    best_effort(
        store,
        "load cross-fit artifact",
        try_load_fit_artifact_by_descriptor(store, descriptor_key_hex),
    )
}

fn try_load_fit_artifact_by_descriptor(
    configured: &ConfiguredWarmStartStore,
    descriptor_key_hex: &str,
) -> Result<Option<crate::warm_start_artifact::FitArtifact>, PersistentStoreError> {
    let Some(store) = configured.store() else {
        return Ok(None);
    };
    let mut fp = Fingerprinter::new();
    fp.absorb_str(b"fit-artifact-key", &cache_schema_tag());
    fp.absorb_str(b"fit-artifact-descriptor", descriptor_key_hex);
    let Some(entry) = store
        .lookup_latest(&fp.finalize())
        .map_err(classify_store_error)?
    else {
        return Ok(None);
    };
    if entry.payload.len() as u64 > MAX_ENTRY_BYTES {
        return Ok(None);
    }
    let artifact: crate::warm_start_artifact::FitArtifact = serde_json::from_slice(&entry.payload)
        .map_err(|error| PersistentStoreError::Rejected(error.to_string()))?;
    // Finite-guard on the way out: a corrupt payload must cold-fallback,
    // never poison a fit.
    Ok(artifact.is_usable().then_some(artifact))
}

fn unix_secs_now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

#[cfg(test)]
mod warm_start_artifact_tests {
    use super::*;
    use crate::warm_start_artifact::{
        FIT_ARTIFACT_SCHEMA, FitArtifact, FitDescriptor, GlobalFitSummary, ResponseSig,
        SerializableBasisMeta, TermArtifact, TermRole, term_identity_from_block,
    };
    use serde::ser::Error as _;

    fn isolated_store() -> (tempfile::TempDir, ConfiguredWarmStartStore) {
        let directory = tempfile::tempdir().expect("create isolated warm-start root");
        let store = configured_store(directory.path().join("warm"));
        (directory, store)
    }

    fn sample_artifact(family: &str, var: &str, rho: Vec<f64>) -> FitArtifact {
        // Block-layer identity (the surviving, fold-invariant identity API):
        // the block name carries the variable, with one unlabeled penalty.
        let block_name = format!("s({var})");
        let id = term_identity_from_block(TermRole::Mean, &block_name, &[None], &[1], 10);
        FitArtifact {
            schema: FIT_ARTIFACT_SCHEMA,
            created_unix_secs: unix_secs_now(),
            descriptor: FitDescriptor {
                family_kind: family.to_string(),
                term_identities: vec![id],
                response_signature: ResponseSig {
                    family_kind: family.to_string(),
                    n_response_channels: 1,
                },
                row_population: None,
            },
            terms: vec![TermArtifact {
                identity: id,
                role: TermRole::Mean,
                basis_meta: SerializableBasisMeta {
                    kind: "block-spec".to_string(),
                    degree: None,
                    num_knots: None,
                    n_centers: Some(8),
                    nullspace_order: None,
                    matern_nu: None,
                    periodic: false,
                },
                joint_null_rotation: None,
                raw_beta: vec![0.1, -0.2, 0.3, 0.4, -0.5, 0.6, -0.7, 0.8],
                rho_for_term: rho,
            }],
            global: GlobalFitSummary {
                outer_objective: -42.0,
                converged: true,
                n_rows: 500,
            },
        }
    }

    #[test]
    fn artifact_round_trips_on_disk_by_descriptor() {
        let (_directory, store) = isolated_store();
        let artifact = sample_artifact("test-roundtrip", "x", vec![2.5]);
        let key_hex = artifact.descriptor.descriptor_key().to_hex();

        // The test owns a writable scratch root, so its assertion and
        // precondition are identical. Environment refusal is exercised through
        // the typed `try_*` seam instead of weakening the round-trip contract.
        try_store_fit_artifact(&store, &artifact).expect("store fit artifact");
        let loaded = load_fit_artifact_by_descriptor(&store, &key_hex)
            .expect("artifact must be retrievable by descriptor key");
        assert_eq!(loaded.schema, artifact.schema);
        assert_eq!(loaded.terms.len(), 1);
        assert_eq!(loaded.terms[0].identity, artifact.terms[0].identity);
        assert_eq!(loaded.terms[0].rho_for_term, vec![2.5]);
        assert_eq!(loaded.terms[0].raw_beta, artifact.terms[0].raw_beta);
        assert_eq!(
            loaded.descriptor.descriptor_key(),
            artifact.descriptor.descriptor_key()
        );
    }

    #[test]
    fn loso_fold_descriptor_matches_full_data_artifact() {
        let (_directory, store) = isolated_store();
        let family = "test-loso";
        // Full-data fit on 1000 rows.
        let mut full = sample_artifact(family, "x", vec![1.7]);
        full.descriptor.row_population = Some(crate::warm_start_artifact::RowPopulationTag {
            n_rows: 1000,
            label: Some("full".to_string()),
        });
        full.global.n_rows = 1000;
        let full_key = full.descriptor.descriptor_key().to_hex();

        // LOSO fold: same term identities, fewer rows. Its descriptor key
        // must equal the full-data key, so the load hits the stored artifact.
        let fold = sample_artifact(family, "x", vec![1.7]);
        let fold_key = fold.descriptor.descriptor_key().to_hex();
        assert_eq!(
            full_key, fold_key,
            "fold and full descriptor keys must match"
        );

        try_store_fit_artifact(&store, &full).expect("store full-data artifact");
        let loaded = load_fit_artifact_by_descriptor(&store, &fold_key)
            .expect("LOSO fold must retrieve the full-data artifact");
        assert_eq!(loaded.terms[0].rho_for_term, vec![1.7]);
    }

    #[test]
    fn io_refusal_remains_typed_before_best_effort_interpretation() {
        let error = StoreError::Io(std::io::Error::from(std::io::ErrorKind::PermissionDenied));
        assert!(matches!(
            classify_store_error(error),
            PersistentStoreError::Unavailable(_)
        ));
    }

    #[test]
    fn persistence_defects_remain_typed_before_best_effort_interpretation() {
        struct RefusesSerialization;

        impl Serialize for RefusesSerialization {
            fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
            where
                S: serde::Serializer,
            {
                Err(S::Error::custom("intentional encoding refusal"))
            }
        }

        let (_directory, store) = isolated_store();
        assert!(matches!(
            store_json_record(&store, "encode-defect", &RefusesSerialization),
            Err(PersistentStoreError::Encode(_))
        ));

        let malformed_json =
            serde_json::from_str::<serde_json::Value>("{").expect_err("fixture must be malformed");
        assert!(matches!(
            classify_store_error(StoreError::Json(malformed_json)),
            PersistentStoreError::Rejected(_)
        ));
    }
}