Fix bug when removing last checkpoint with deepspeed (#18793)

Co-authored-by: awaelchli <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

src/lightning/pytorch/CHANGELOG.md

@@ -29,7 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- Fixed an issue when replacing an existing `last.ckpt` file with a symlink ([#18793](https://github.com/Lightning-AI/lightning/pull/18793))
 
 
 

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -20,6 +20,7 @@
 import logging
 import os
 import re
+import shutil
 import time
 import warnings
 from copy import deepcopy
@@ -383,8 +384,10 @@ def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
     @staticmethod
     def _link_checkpoint(trainer: "pl.Trainer", filepath: str, linkpath: str) -> None:
         if trainer.is_global_zero:
-            if os.path.lexists(linkpath):
+            if os.path.islink(linkpath) or os.path.isfile(linkpath):
                 os.remove(linkpath)
+            elif os.path.isdir(linkpath):
+                shutil.rmtree(linkpath)
             os.symlink(filepath, linkpath)
         trainer.strategy.barrier()
 

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -514,6 +514,49 @@ def test_model_checkpoint_save_last(tmpdir):
     ModelCheckpoint.CHECKPOINT_NAME_LAST = "last"
 
 
+def test_model_checkpoint_link_checkpoint(tmp_path):
+    """Test that linking a checkpoint works and overwrites an existing link if present."""
+    trainer = Mock()
+
+    # link doesn't exist
+    file = tmp_path / "file"
+    file.touch()
+    link = tmp_path / "link"
+    ModelCheckpoint._link_checkpoint(trainer, filepath=str(file), linkpath=str(link))
+    assert os.path.islink(link)
+    assert os.path.realpath(link) == str(file)
+
+    # link exists (is a file)
+    new_file1 = tmp_path / "new_file1"
+    new_file1.touch()
+    ModelCheckpoint._link_checkpoint(trainer, filepath=str(new_file1), linkpath=str(link))
+    assert os.path.islink(link)
+    assert os.path.realpath(link) == str(new_file1)
+
+    # link exists (is a link)
+    new_file2 = tmp_path / "new_file2"
+    new_file2.touch()
+    ModelCheckpoint._link_checkpoint(trainer, filepath=str(new_file2), linkpath=str(link))
+    assert os.path.islink(link)
+    assert os.path.realpath(link) == str(new_file2)
+
+    # link exists (is a folder)
+    folder = tmp_path / "folder"
+    folder.mkdir()
+    folder_link = tmp_path / "folder_link"
+    folder_link.mkdir()
+    ModelCheckpoint._link_checkpoint(trainer, filepath=str(folder), linkpath=str(folder_link))
+    assert os.path.islink(folder_link)
+    assert os.path.realpath(folder_link) == str(folder)
+
+    # link exists (is a link to a folder)
+    new_folder = tmp_path / "new_folder"
+    new_folder.mkdir()
+    ModelCheckpoint._link_checkpoint(trainer, filepath=str(new_folder), linkpath=str(folder_link))
+    assert os.path.islink(folder_link)
+    assert os.path.realpath(folder_link) == str(new_folder)
+
+
 def test_invalid_top_k(tmpdir):
     """Make sure that a MisconfigurationException is raised for a negative save_top_k argument."""
     with pytest.raises(MisconfigurationException, match=r".*Must be >= -1"):