Atomic checkpoint writes (fixes #13)

src/serialize/atomic.rs

@@ -0,0 +1,39 @@
+//! 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)
+}

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<Path>) -> 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<Path>) -> 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`.

src/serialize/mod.rs

@@ -25,6 +25,7 @@
 //! # Ok::<(), fastnn::Error>(())
 //! ```
 
+pub(crate) mod atomic;
 pub mod checkpoint;
 pub mod half;
 pub mod safetensors;

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<Path>) -> Result<()> {
     let mut tensors: BTreeMap<String, Tensor> = BTreeMap::new();
     for (name, param) in module.named_parameters() {
@@ -56,16 +61,19 @@ pub fn save_safetensors(module: &dyn Module, path: impl AsRef<Path>) -> 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.

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::{Checkpoint, load_tensors, load_training, save_tensors, save_training};
 
 /// Deleted on drop.
 struct TempFile(std::path::PathBuf);
@@ -407,3 +407,105 @@ 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]);
+}