Add atomic save to checkpoint routine (#20011)

src/lightning/fabric/CHANGELOG.md

@@ -9,7 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
--
+- Made saving non-distributed checkpoints fully atomic ([#20011](https://github.com/Lightning-AI/pytorch-lightning/pull/20011))
 
 -
 

src/lightning/fabric/utilities/cloud_io.py

@@ -76,7 +76,10 @@ def _atomic_save(checkpoint: Dict[str, Any], filepath: Union[str, Path]) -> None
     bytesbuffer = io.BytesIO()
     log.debug(f"Saving checkpoint: {filepath}")
     torch.save(checkpoint, bytesbuffer)
-    with fsspec.open(filepath, "wb") as f:
+
+    # We use a transaction here to avoid file corruption if the save gets interrupted
+    fs, urlpath = fsspec.core.url_to_fs(str(filepath))
+    with fs.transaction, fs.open(urlpath, "wb") as f:
         f.write(bytesbuffer.getvalue())
 
 

src/lightning/pytorch/CHANGELOG.md

@@ -9,7 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
--
+- Made saving non-distributed checkpoints fully atomic ([#20011](https://github.com/Lightning-AI/pytorch-lightning/pull/20011))
 
 -