Remove the file existence check in TorchCheckpointIO

src/lightning/fabric/CHANGELOG.md

@@ -41,6 +41,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - The columns in the `metrics.csv` file produced by `CSVLogger` are now sorted alphabetically ([#19159](https://github.com/Lightning-AI/lightning/pull/19159))
 
 
+- `TorchCheckpointIO.remove_checkpoint()` no longer silently passes if the given file does not exist ([#19344](https://github.com/Lightning-AI/lightning/pull/19344))
+
+
 ### Deprecated
 
 -

src/lightning/fabric/plugins/io/torch_io.py

@@ -90,7 +90,6 @@ def remove_checkpoint(self, path: _PATH) -> None:
             path: Path to checkpoint
 
         """
-        fs = get_filesystem(path)
-        if fs.exists(path):
-            fs.rm(path, recursive=True)
-            log.debug(f"Removed checkpoint: {path}")
+        fs = get_filesystem(str(path))
+        fs.rm(str(path), recursive=True)
+        log.debug(f"Removed checkpoint: {path}")

src/lightning/pytorch/CHANGELOG.md

@@ -53,6 +53,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Reverted back to creating a checkpoint copy when `ModelCheckpoint(save_last=True)` instead of creating a symbolic link ([#19191](https://github.com/Lightning-AI/lightning/pull/19191))
 
 
+- `TorchCheckpointIO.remove_checkpoint()` no longer silently passes if the given file does not exist ([#19344](https://github.com/Lightning-AI/lightning/pull/19344))
+
+
 ### Deprecated
 
 - Deprecated all precision plugin classes under `lightning.pytorch.plugins` with the suffix `Plugin` in the name ([#18840](https://github.com/Lightning-AI/lightning/pull/18840))

src/lightning/pytorch/callbacks/on_exception_checkpoint.py

@@ -67,4 +67,6 @@ def on_exception(self, trainer: "pl.Trainer", *_: Any, **__: Any) -> None:
 
     @override
     def teardown(self, trainer: "pl.Trainer", *_: Any, **__: Any) -> None:
-        trainer.strategy.remove_checkpoint(self.ckpt_path)
+        if os.path.exists(self.ckpt_path):
+            # only exists if there was an exception
+            trainer.strategy.remove_checkpoint(self.ckpt_path)

tests/tests_fabric/plugins/io/test_torch_io.py

@@ -0,0 +1,57 @@
+# Copyright The Lightning AI team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+from lightning.fabric.plugins.io import TorchCheckpointIO
+
+
+def test_remove_checkpoint(tmp_path):
+    """Test that the IO can remove folders, files, and symlinks."""
+    io = TorchCheckpointIO()
+
+    # Path does not exist
+    with pytest.raises(FileNotFoundError):
+        io.remove_checkpoint("does_not_exist.txt")
+
+    # Single file
+    file = tmp_path / "file.txt"
+    file.touch()
+    io.remove_checkpoint(file)
+    assert not file.exists()
+
+    # Symlink
+    file = tmp_path / "file.txt"
+    file.touch()
+    link = tmp_path / "link.txt"
+    link.symlink_to(file)
+    io.remove_checkpoint(link)
+    assert file.exists()
+    assert not link.is_symlink()
+    file.unlink()
+
+    # Broken Symlink
+    file_not_exists = tmp_path / "not_exist.txt"
+    link = tmp_path / "link.txt"
+    link.symlink_to(file_not_exists)
+    assert not file_not_exists.exists()
+    io.remove_checkpoint(link)
+    assert not link.is_symlink()
+
+    # Folder with contents
+    folder = tmp_path / "folder"
+    nested_folder = folder / "nested_folder"
+    nested_folder.mkdir(parents=True)
+    file = nested_folder / "file.txt"
+    file.touch()
+    io.remove_checkpoint(folder)
+    assert not folder.exists()

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -960,6 +960,7 @@ def on_validation_epoch_end(self):
         max_epochs=len(monitor),
     )
     trainer.save_checkpoint = Mock()
+    trainer.strategy.remove_checkpoint = Mock()
 
     trainer.fit(model)