From 4703f929d5f37734183ab0b7a5ea781165213d3e Mon Sep 17 00:00:00 2001 From: sharmis-tp Date: Tue, 10 Nov 2026 02:08:41 +0000 Subject: [PATCH 1/2] atomic checkpoint writes: write to temp, fsync, then rename (fixes #13) --- src/serialize/atomic.rs | 42 ++++++++++++++ src/serialize/checkpoint.rs | 45 +++++++++------ src/serialize/mod.rs | 1 + src/serialize/safetensors.rs | 20 +++++-- tests/robustness.rs | 108 ++++++++++++++++++++++++++++++++++- 5 files changed, 189 insertions(+), 27 deletions(-) create mode 100644 src/serialize/atomic.rs diff --git a/src/serialize/atomic.rs b/src/serialize/atomic.rs new file mode 100644 index 0000000..27253d5 --- /dev/null +++ b/src/serialize/atomic.rs @@ -0,0 +1,42 @@ +//! Atomic file writes: write to a per-process temp, `sync_all`, then rename. + +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; + +use crate::error::Result; + +/// Write to `path` atomically. +/// +/// Writes into `{path}.partial.{pid}`, calls [`File::sync_all`], then renames +/// the temp to `path`. A crash mid-write leaves the previous file intact. On +/// error the temp is removed so it never accumulates. +pub(crate) fn write_atomic( + path: &Path, + write_fn: impl FnOnce(&mut File) -> Result<()>, +) -> Result<()> { + let tmp = tmp_path(path); + let result = write_and_sync(&tmp, write_fn); + if result.is_err() { + let _ = std::fs::remove_file(&tmp); + return result; + } + std::fs::rename(&tmp, path)?; + Ok(()) +} + +fn write_and_sync( + path: &Path, + write_fn: impl FnOnce(&mut File) -> Result<()>, +) -> Result<()> { + let mut file = File::create(path)?; + write_fn(&mut file)?; + file.sync_all()?; + Ok(()) +} + +fn tmp_path(path: &Path) -> PathBuf { + let mut name = OsString::from(path.as_os_str()); + name.push(format!(".partial.{}", std::process::id())); + PathBuf::from(name) +} diff --git a/src/serialize/checkpoint.rs b/src/serialize/checkpoint.rs index 5f76b6b..2e4b336 100644 --- a/src/serialize/checkpoint.rs +++ b/src/serialize/checkpoint.rs @@ -21,6 +21,8 @@ use crate::error::{Error, Result}; use crate::nn::Module; use crate::tensor::Tensor; +use super::atomic::write_atomic; + const MAGIC: u32 = 0x4644_4C00; const VERSION: u32 = 1; @@ -66,28 +68,33 @@ pub fn load(module: &dyn Module, path: impl AsRef) -> Result<()> { } /// Write named tensors to `path`. +/// +/// The write is atomic: data lands in a per-process temp file, is fsynced, then +/// renamed over `path`, so a crash mid-write never destroys an existing file. pub fn save_tensors(tensors: &Checkpoint, path: impl AsRef) -> Result<()> { - let mut out = BufWriter::new(File::create(path)?); - - write_u32(&mut out, MAGIC)?; - write_u32(&mut out, VERSION)?; - write_u32(&mut out, tensors.len() as u32)?; - - for (name, tensor) in tensors { - write_u32(&mut out, name.len() as u32)?; - out.write_all(name.as_bytes())?; - - write_u32(&mut out, tensor.ndim() as u32)?; - for &dim in tensor.shape() { - write_u32(&mut out, dim as u32)?; - } - for value in tensor.to_vec() { - out.write_all(&value.to_le_bytes())?; + write_atomic(path.as_ref(), |file| { + let mut out = BufWriter::new(&mut *file); + + write_u32(&mut out, MAGIC)?; + write_u32(&mut out, VERSION)?; + write_u32(&mut out, tensors.len() as u32)?; + + for (name, tensor) in tensors { + write_u32(&mut out, name.len() as u32)?; + out.write_all(name.as_bytes())?; + + write_u32(&mut out, tensor.ndim() as u32)?; + for &dim in tensor.shape() { + write_u32(&mut out, dim as u32)?; + } + for value in tensor.to_vec() { + out.write_all(&value.to_le_bytes())?; + } } - } - out.flush()?; - Ok(()) + out.flush()?; + Ok(()) + }) } /// Read named tensors from `path`. diff --git a/src/serialize/mod.rs b/src/serialize/mod.rs index 4670d81..bf54c31 100644 --- a/src/serialize/mod.rs +++ b/src/serialize/mod.rs @@ -25,6 +25,7 @@ //! # Ok::<(), fastnn::Error>(()) //! ``` +pub(crate) mod atomic; pub mod checkpoint; pub mod half; pub mod safetensors; diff --git a/src/serialize/safetensors.rs b/src/serialize/safetensors.rs index 3f59946..c6ecf72 100644 --- a/src/serialize/safetensors.rs +++ b/src/serialize/safetensors.rs @@ -8,12 +8,14 @@ //! ``` use std::collections::{BTreeMap, HashMap}; +use std::io::Write; use std::path::Path; use crate::error::{Error, Result}; use crate::nn::Module; use crate::tensor::Tensor; +use super::atomic::write_atomic; use super::half::{f32_from_bf16, f32_from_f16}; /// What a partial load did. Empty `missing` and `unexpected` mean an exact match. @@ -28,6 +30,9 @@ pub struct LoadReport { } /// Write a module's parameters and buffers as `F32` safetensors. +/// +/// The write is atomic: data lands in a per-process temp file, is fsynced, then +/// renamed over `path`, so a crash mid-write never destroys an existing file. pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<()> { let mut tensors: BTreeMap = BTreeMap::new(); for (name, param) in module.named_parameters() { @@ -56,16 +61,19 @@ pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<( } header.push('}'); - let mut out = Vec::with_capacity(8 + header.len() + offset); - out.extend_from_slice(&(header.len() as u64).to_le_bytes()); - out.extend_from_slice(header.as_bytes()); + let mut buf = Vec::with_capacity(8 + header.len() + offset); + buf.extend_from_slice(&(header.len() as u64).to_le_bytes()); + buf.extend_from_slice(header.as_bytes()); for tensor in tensors.values() { for value in tensor.to_vec() { - out.extend_from_slice(&value.to_le_bytes()); + buf.extend_from_slice(&value.to_le_bytes()); } } - std::fs::write(path, out)?; - Ok(()) + + write_atomic(path.as_ref(), |file| { + file.write_all(&buf)?; + Ok(()) + }) } /// Load matching tensors from a safetensors file into a module. diff --git a/tests/robustness.rs b/tests/robustness.rs index dec40e9..0173312 100644 --- a/tests/robustness.rs +++ b/tests/robustness.rs @@ -1,8 +1,8 @@ -//! Resuming, anomaly detection, freezing and the fallible ops. +//! Resuming, anomaly detection, atomic writes, freezing and the fallible ops. use fastnn::autograd::detect_anomaly; use fastnn::prelude::*; -use fastnn::serialize::{load_training, save_training}; +use fastnn::serialize::{load_tensors, load_training, save_tensors, save_training, Checkpoint}; /// Deleted on drop. struct TempFile(std::path::PathBuf); @@ -407,3 +407,107 @@ fn fallible_ops_still_build_the_graph() { x.try_matmul(&w.tensor()).unwrap().sum().backward(); assert!(w.grad().is_some(), "try_matmul dropped the graph"); } + +// ---- Atomic checkpoint writes -------------------------------------------------- + +#[test] +fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() { + let file = TempFile::new("atomic-fdl"); + let model = small_model(); + save(&model, file.path()).unwrap(); + + // Record the weights of the good checkpoint. + let good = model.parameters()[0].value().to_vec(); + + // Verify the good checkpoint loads. + let loaded = load_tensors(file.path()).unwrap(); + assert!(!loaded.is_empty(), "good checkpoint should have tensors"); + + // Simulate a crash: truncate the file to a few bytes. + std::fs::write(file.path(), &[0u8; 8]).unwrap(); + + // The truncated file must fail to load. + assert!( + load_tensors(file.path()).is_err(), + "a truncated fdl checkpoint should not parse" + ); + + // Overwrite the truncated file with a fresh save -- the atomic write must + // succeed even when the target already exists and is corrupt. + save(&model, file.path()).unwrap(); + let fresh = small_model(); + load(&fresh, file.path()).unwrap(); + assert_eq!( + fresh.parameters()[0].value().to_vec(), + good, + "re-saved checkpoint should match the original weights" + ); +} + +#[test] +fn truncated_safetensors_does_not_destroy_the_good_copy() { + let file = TempFile::new("atomic-st"); + let path = file.path().with_extension("safetensors"); + + manual_seed(42); + let model = Sequential::new().add(Linear::new(4, 2)); + save_safetensors(&model, &path).unwrap(); + + let good = model.parameters()[0].value().to_vec(); + + // Simulate a crash: truncate the file. + std::fs::write(&path, &[0u8; 4]).unwrap(); + + // The truncated file must fail to load. + let fresh = Sequential::new().add(Linear::new(4, 2)); + assert!( + load_safetensors(&fresh, &path).is_err(), + "a truncated safetensors file should not parse" + ); + + // Re-save over the corrupt file; the atomic write must succeed. + save_safetensors(&model, &path).unwrap(); + let reloaded = Sequential::new().add(Linear::new(4, 2)); + let report = load_safetensors(&reloaded, &path).unwrap(); + assert!(report.missing.is_empty()); + assert_eq!( + reloaded.parameters()[0].value().to_vec(), + good, + "re-saved safetensors should match the original weights" + ); + let _ = std::fs::remove_file(&path); +} + +#[test] +fn atomic_write_cleans_up_temp_on_fdl_save() { + // After a successful save the .partial.{pid} file must not linger. + let file = TempFile::new("atomic-cleanup"); + let model = small_model(); + save(&model, file.path()).unwrap(); + + let tmp = file + .path() + .with_file_name(format!( + "{}.partial.{}", + file.path().file_name().unwrap().to_str().unwrap(), + std::process::id() + )); + assert!( + !tmp.exists(), + "temp file should be renamed away after a successful save" + ); +} + +#[test] +fn save_tensors_round_trips_through_atomic_write() { + let file = TempFile::new("atomic-rt"); + let mut tensors = Checkpoint::new(); + tensors.insert("a".into(), Tensor::from_vec(vec![1.0, 2.0, 3.0], &[3])); + tensors.insert("b".into(), Tensor::from_vec(vec![4.0], &[1])); + + save_tensors(&tensors, file.path()).unwrap(); + let loaded = load_tensors(file.path()).unwrap(); + assert_eq!(loaded.len(), 2); + assert_eq!(loaded["a"].to_vec(), vec![1.0, 2.0, 3.0]); + assert_eq!(loaded["b"].to_vec(), vec![4.0]); +} From 3adce3ea658867410593e7907a0fff8ef13fed0d Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Thu, 12 Nov 2026 11:56:56 +0300 Subject: [PATCH 2/2] satisfy clippy and rustfmt --- src/serialize/atomic.rs | 5 +---- tests/robustness.rs | 18 ++++++++---------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/src/serialize/atomic.rs b/src/serialize/atomic.rs index 27253d5..85a4e03 100644 --- a/src/serialize/atomic.rs +++ b/src/serialize/atomic.rs @@ -25,10 +25,7 @@ pub(crate) fn write_atomic( Ok(()) } -fn write_and_sync( - path: &Path, - write_fn: impl FnOnce(&mut File) -> Result<()>, -) -> Result<()> { +fn write_and_sync(path: &Path, write_fn: impl FnOnce(&mut File) -> Result<()>) -> Result<()> { let mut file = File::create(path)?; write_fn(&mut file)?; file.sync_all()?; diff --git a/tests/robustness.rs b/tests/robustness.rs index 0173312..e706eb7 100644 --- a/tests/robustness.rs +++ b/tests/robustness.rs @@ -2,7 +2,7 @@ use fastnn::autograd::detect_anomaly; use fastnn::prelude::*; -use fastnn::serialize::{load_tensors, load_training, save_tensors, save_training, Checkpoint}; +use fastnn::serialize::{Checkpoint, load_tensors, load_training, save_tensors, save_training}; /// Deleted on drop. struct TempFile(std::path::PathBuf); @@ -424,7 +424,7 @@ fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() { assert!(!loaded.is_empty(), "good checkpoint should have tensors"); // Simulate a crash: truncate the file to a few bytes. - std::fs::write(file.path(), &[0u8; 8]).unwrap(); + std::fs::write(file.path(), [0u8; 8]).unwrap(); // The truncated file must fail to load. assert!( @@ -456,7 +456,7 @@ fn truncated_safetensors_does_not_destroy_the_good_copy() { let good = model.parameters()[0].value().to_vec(); // Simulate a crash: truncate the file. - std::fs::write(&path, &[0u8; 4]).unwrap(); + std::fs::write(&path, [0u8; 4]).unwrap(); // The truncated file must fail to load. let fresh = Sequential::new().add(Linear::new(4, 2)); @@ -485,13 +485,11 @@ fn atomic_write_cleans_up_temp_on_fdl_save() { let model = small_model(); save(&model, file.path()).unwrap(); - let tmp = file - .path() - .with_file_name(format!( - "{}.partial.{}", - file.path().file_name().unwrap().to_str().unwrap(), - std::process::id() - )); + let tmp = file.path().with_file_name(format!( + "{}.partial.{}", + file.path().file_name().unwrap().to_str().unwrap(), + std::process::id() + )); assert!( !tmp.exists(), "temp file should be renamed away after a successful save"