Function prefetch_safetensors

cake-core/src/utils/mod.rs:213–248

Pre-read safetensor files into the OS page cache so that subsequent mmap access doesn't trigger per-tensor page faults during layer loading. Uses OnceLock to skip redundant calls (e.g. multi-GPU VarBuilder creation).

(filenames: &[PathBuf])

/// mmap access doesn't trigger per-tensor page faults during layer loading.
/// Uses OnceLock to skip redundant calls (e.g. multi-GPU VarBuilder creation).
fn prefetch_safetensors(filenames: &[PathBuf]) -> Result<()> {
    use std::sync::OnceLock;
    static DONE: OnceLock<()> = OnceLock::new();

    if DONE.get().is_some() {
        log::info!("safetensor files already in page cache, skipping prefetch");
        return Ok(());
    }

    use std::io::Read;
    let start = std::time::Instant::now();
    let mut total_bytes: u64 = 0;
    let mut buf = Vec::new();
    for (i, filename) in filenames.iter().enumerate() {
        log::info!(
            "caching shard {}/{} ({}) ...",
            i + 1,
            filenames.len(),
            filename.file_name().unwrap_or_default().to_string_lossy()
        );
        buf.clear();
        std::fs::File::open(filename)
            .map_err(|e| anyhow!("prefetch: can't open {}: {e}", filename.display()))?
            .read_to_end(&mut buf)
            .map_err(|e| anyhow!("prefetch: can't read {}: {e}", filename.display()))?;
        total_bytes += buf.len() as u64;
    }
    log::info!(
        "pre-cached {} in {:.1}s",
        human_bytes::human_bytes(total_bytes as f64),
        start.elapsed().as_secs_f64()
    );

    DONE.set(()).ok();
    Ok(())
}

/// Create a VarBuilder with the tensors loaded from the index.
pub fn load_var_builder_from_index<'a>(