From ed37c8874f532d3244e3794544e256d0941b3624 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Thu, 12 Nov 2026 13:41:31 +0300 Subject: [PATCH] one atomic write helper for checkpoints and the dataset cache --- src/data/mnist.rs | 77 ++----------------- src/serialize/atomic.rs | 142 +++++++++++++++++++++++++++++------ src/serialize/checkpoint.rs | 3 +- src/serialize/safetensors.rs | 3 +- tests/robustness.rs | 42 +++-------- 5 files changed, 133 insertions(+), 134 deletions(-) diff --git a/src/data/mnist.rs b/src/data/mnist.rs index 54ddf08..49103fb 100644 --- a/src/data/mnist.rs +++ b/src/data/mnist.rs @@ -2,11 +2,11 @@ //! `~/.fastnn/datasets/mnist`. use std::fs; -use std::io::{Read, Write}; +use std::io::Read; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicU32, Ordering}; use crate::error::{Error, Result}; +use crate::serialize::atomic::write_atomic_bytes; use crate::tensor::Tensor; use super::dataset::Dataset; @@ -163,35 +163,10 @@ fn fetch(cache: &Path, name: &str) -> Result> { let mut raw = Vec::new(); flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?; - write_atomic(&cache.join(name), &raw)?; + write_atomic_bytes(&cache.join(name), &raw)?; Ok(raw) } -/// Write `bytes` to `target` so a crash leaves it either absent or complete: -/// a temporary file next to it, `sync_all`, then a rename into place. -fn write_atomic(target: &Path, bytes: &[u8]) -> Result<()> { - // The pid and a counter keep concurrent writers, in and across processes, - // off each other's temporary file. - static SERIAL: AtomicU32 = AtomicU32::new(0); - let name = target.file_name().unwrap_or_default().to_string_lossy(); - let partial = target.with_file_name(format!( - "{name}.{}.{}.partial", - std::process::id(), - SERIAL.fetch_add(1, Ordering::Relaxed) - )); - - let written = (|| { - let mut file = fs::File::create(&partial)?; - file.write_all(bytes)?; - file.sync_all()?; - fs::rename(&partial, target) - })(); - if written.is_err() { - let _ = fs::remove_file(&partial); - } - written.map_err(Error::from) -} - /// IDX images: magic, count, rows, cols, then one byte per pixel. fn parse_images(bytes: &[u8]) -> Result<(Vec, usize)> { let header = read_header(bytes, 4, IMAGE_MAGIC, "images")?; @@ -283,48 +258,6 @@ mod tests { bytes } - fn partials_in(dir: &Path) -> usize { - fs::read_dir(dir) - .unwrap() - .filter(|entry| { - entry - .as_ref() - .unwrap() - .file_name() - .to_string_lossy() - .ends_with(".partial") - }) - .count() - } - - #[test] - fn write_atomic_writes_the_bytes_and_leaves_no_partial() { - let dir = TempDir::new("write"); - let target = dir.0.join("data.bin"); - write_atomic(&target, b"hello").unwrap(); - assert_eq!(fs::read(&target).unwrap(), b"hello"); - assert_eq!(partials_in(&dir.0), 0); - } - - #[test] - fn write_atomic_replaces_an_existing_file() { - let dir = TempDir::new("replace"); - let target = dir.0.join("data.bin"); - fs::write(&target, b"old").unwrap(); - write_atomic(&target, b"new").unwrap(); - assert_eq!(fs::read(&target).unwrap(), b"new"); - } - - #[test] - fn write_atomic_cleans_up_when_the_rename_fails() { - let dir = TempDir::new("rename-fails"); - // A file cannot replace a directory, so the rename fails after the write. - let target = dir.0.join("taken"); - fs::create_dir(&target).unwrap(); - assert!(write_atomic(&target, b"x").is_err()); - assert_eq!(partials_in(&dir.0), 0); - } - #[test] fn a_valid_cached_file_is_used_without_fetching() { let dir = TempDir::new("valid"); @@ -341,7 +274,7 @@ mod tests { let dir = TempDir::new("missing"); let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| { let bytes = image_file(2); - write_atomic(&cache.join(name), &bytes)?; + fs::write(cache.join(name), &bytes)?; Ok(bytes) }) .unwrap(); @@ -361,7 +294,7 @@ mod tests { let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| { fetches.set(fetches.get() + 1); let bytes = image_file(5); - write_atomic(&cache.join(name), &bytes)?; + fs::write(cache.join(name), &bytes)?; Ok(bytes) }) .unwrap(); diff --git a/src/serialize/atomic.rs b/src/serialize/atomic.rs index 85a4e03..9df4607 100644 --- a/src/serialize/atomic.rs +++ b/src/serialize/atomic.rs @@ -1,39 +1,131 @@ -//! Atomic file writes: write to a per-process temp, `sync_all`, then rename. +//! One atomic file write for checkpoints and the dataset cache: a temporary +//! file next to the target, `sync_all`, then a rename into place. use std::ffi::OsString; -use std::fs::File; +use std::fs::{self, File}; +use std::io::Write; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU32, Ordering}; use crate::error::Result; -/// Write to `path` atomically. +/// Write to `path` so a crash leaves it either absent or complete. /// -/// 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(()) +/// `write` fills a temporary file next to `path`, which is synced and renamed +/// into place only when `write` succeeds, and removed when anything fails. The +/// previous file is never touched until the rename. +pub(crate) fn write_atomic(path: &Path, write: impl FnOnce(&mut File) -> Result<()>) -> Result<()> { + let partial = partial_path(path); + let written: Result<()> = (|| { + let mut file = File::create(&partial)?; + write(&mut file)?; + file.sync_all()?; + fs::rename(&partial, path)?; + Ok(()) + })(); + if written.is_err() { + let _ = fs::remove_file(&partial); + } + written } -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(()) +/// [`write_atomic`] for bytes already in memory. +pub(crate) fn write_atomic_bytes(path: &Path, bytes: &[u8]) -> Result<()> { + write_atomic(path, |file| Ok(file.write_all(bytes)?)) } -fn tmp_path(path: &Path) -> PathBuf { +/// `{path}.{pid}.{serial}.partial`. The pid and the counter keep concurrent +/// writers, in and across processes, off each other's temporary file. +fn partial_path(path: &Path) -> PathBuf { + static SERIAL: AtomicU32 = AtomicU32::new(0); let mut name = OsString::from(path.as_os_str()); - name.push(format!(".partial.{}", std::process::id())); + name.push(format!( + ".{}.{}.partial", + std::process::id(), + SERIAL.fetch_add(1, Ordering::Relaxed) + )); PathBuf::from(name) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error; + + /// A fresh directory under the system temp dir, removed on drop. + struct TempDir(PathBuf); + + impl TempDir { + fn new(name: &str) -> TempDir { + let dir = + std::env::temp_dir().join(format!("fastnn-atomic-{}-{name}", std::process::id())); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + TempDir(dir) + } + } + + impl Drop for TempDir { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn partials_in(dir: &Path) -> usize { + fs::read_dir(dir) + .unwrap() + .filter(|entry| { + entry + .as_ref() + .unwrap() + .file_name() + .to_string_lossy() + .ends_with(".partial") + }) + .count() + } + + #[test] + fn writes_the_bytes_and_leaves_no_partial() { + let dir = TempDir::new("write"); + let target = dir.0.join("data.bin"); + write_atomic_bytes(&target, b"hello").unwrap(); + assert_eq!(fs::read(&target).unwrap(), b"hello"); + assert_eq!(partials_in(&dir.0), 0); + } + + #[test] + fn replaces_an_existing_file() { + let dir = TempDir::new("replace"); + let target = dir.0.join("data.bin"); + fs::write(&target, b"old").unwrap(); + write_atomic_bytes(&target, b"new").unwrap(); + assert_eq!(fs::read(&target).unwrap(), b"new"); + } + + #[test] + fn a_failed_write_leaves_the_old_file_and_no_partial() { + let dir = TempDir::new("write-fails"); + let target = dir.0.join("data.bin"); + fs::write(&target, b"old").unwrap(); + + let result = write_atomic(&target, |file| { + file.write_all(b"half")?; + Err(Error::Checkpoint("disk full".into())) + }); + + assert!(result.is_err()); + assert_eq!(fs::read(&target).unwrap(), b"old"); + assert_eq!(partials_in(&dir.0), 0); + } + + #[test] + fn a_failed_rename_leaves_no_partial() { + let dir = TempDir::new("rename-fails"); + // A file cannot replace a directory, so the rename fails after the write. + let target = dir.0.join("taken"); + fs::create_dir(&target).unwrap(); + assert!(write_atomic_bytes(&target, b"x").is_err()); + assert_eq!(partials_in(&dir.0), 0); + } +} diff --git a/src/serialize/checkpoint.rs b/src/serialize/checkpoint.rs index 4d076c1..7385a19 100644 --- a/src/serialize/checkpoint.rs +++ b/src/serialize/checkpoint.rs @@ -69,8 +69,7 @@ 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. +/// Written atomically: a crash mid-save leaves the previous file intact. pub fn save_tensors(tensors: &Checkpoint, path: impl AsRef) -> Result<()> { write_atomic(path.as_ref(), |file| { let mut out = BufWriter::new(&mut *file); diff --git a/src/serialize/safetensors.rs b/src/serialize/safetensors.rs index c6ecf72..c16ae43 100644 --- a/src/serialize/safetensors.rs +++ b/src/serialize/safetensors.rs @@ -31,8 +31,7 @@ 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. +/// Written atomically: a crash mid-save leaves the previous file intact. pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<()> { let mut tensors: BTreeMap = BTreeMap::new(); for (name, param) in module.named_parameters() { diff --git a/tests/robustness.rs b/tests/robustness.rs index e706eb7..0a5aa0a 100644 --- a/tests/robustness.rs +++ b/tests/robustness.rs @@ -408,10 +408,8 @@ fn fallible_ops_still_build_the_graph() { assert!(w.grad().is_some(), "try_matmul dropped the graph"); } -// ---- Atomic checkpoint writes -------------------------------------------------- - #[test] -fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() { +fn a_corrupt_checkpoint_can_be_saved_over() { let file = TempFile::new("atomic-fdl"); let model = small_model(); save(&model, file.path()).unwrap(); @@ -432,8 +430,6 @@ fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() { "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(); @@ -445,59 +441,39 @@ fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() { } #[test] -fn truncated_safetensors_does_not_destroy_the_good_copy() { +fn a_corrupt_safetensors_file_can_be_saved_over() { let file = TempFile::new("atomic-st"); - let path = file.path().with_extension("safetensors"); + let path = file.path(); manual_seed(42); let model = Sequential::new().add(Linear::new(4, 2)); - save_safetensors(&model, &path).unwrap(); + 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(); + 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(), + 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(); + save_safetensors(&model, path).unwrap(); let reloaded = Sequential::new().add(Linear::new(4, 2)); - let report = load_safetensors(&reloaded, &path).unwrap(); + 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() { +fn save_tensors_round_trips() { 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]));