From 9d04a5026bf630c4042c8a7145f391f5bd95b87d Mon Sep 17 00:00:00 2001
From: supermario_leo <leo.stack@outlook.com>
Date: Wed, 22 Jul 2026 20:15:12 +0800
Subject: [PATCH] Raise clear error on dangling step_latest symlink instead of
 opaque assert

---
 deepspec/trainer/ckpt_manager.py            | 15 ++++--
 tests/__init__.py                           |  0
 tests/trainer/__init__.py                   |  0
 tests/trainer/test_ckpt_manager_discover.py | 51 +++++++++++++++++++++
 4 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/trainer/__init__.py
 create mode 100644 tests/trainer/test_ckpt_manager_discover.py

diff --git a/deepspec/trainer/ckpt_manager.py b/deepspec/trainer/ckpt_manager.py
index 6b99f65a..a3090e30 100644
--- a/deepspec/trainer/ckpt_manager.py
+++ b/deepspec/trainer/ckpt_manager.py
@@ -24,9 +24,18 @@
 
 def discover_latest_checkpoint(checkpoint_dir):
     latest_link = os.path.join(checkpoint_dir, "step_latest")
-    if not (os.path.islink(latest_link) or os.path.isdir(latest_link)):
-        return None
-    return os.path.realpath(latest_link)
+    if os.path.islink(latest_link):
+        resolved = os.path.realpath(latest_link)
+        if not os.path.isdir(resolved):
+            raise RuntimeError(
+                f"Dangling 'step_latest' symlink at {latest_link!r} resolves to "
+                f"missing target {resolved!r}. Remove the symlink to start a "
+                f"fresh run, or restore the checkpoint it referenced."
+            )
+        return resolved
+    if os.path.isdir(latest_link):
+        return os.path.realpath(latest_link)
+    return None
 
 
 def save_train_config(*, train_config, checkpoint_dir: str) -> str:
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/trainer/__init__.py b/tests/trainer/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/trainer/test_ckpt_manager_discover.py b/tests/trainer/test_ckpt_manager_discover.py
new file mode 100644
index 00000000..27f6b7e5
--- /dev/null
+++ b/tests/trainer/test_ckpt_manager_discover.py
@@ -0,0 +1,51 @@
+import os
+import shutil
+import tempfile
+import unittest
+
+from deepspec.trainer.ckpt_manager import discover_latest_checkpoint
+
+
+class TestDiscoverLatestCheckpoint(unittest.TestCase):
+    def setUp(self):
+        self.checkpoint_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.checkpoint_dir, ignore_errors=True)
+
+    def test_no_link_returns_none(self):
+        self.assertIsNone(discover_latest_checkpoint(self.checkpoint_dir))
+
+    def test_healthy_symlink_returns_resolved_dir(self):
+        target = os.path.join(self.checkpoint_dir, "step_10")
+        os.makedirs(target)
+        link = os.path.join(self.checkpoint_dir, "step_latest")
+        os.symlink(target, link)
+        self.assertEqual(
+            os.path.realpath(link),
+            discover_latest_checkpoint(self.checkpoint_dir),
+        )
+
+    def test_real_directory_returns_realpath(self):
+        link = os.path.join(self.checkpoint_dir, "step_latest")
+        os.makedirs(link)
+        self.assertEqual(
+            os.path.realpath(link),
+            discover_latest_checkpoint(self.checkpoint_dir),
+        )
+
+    def test_dangling_symlink_raises(self):
+        target = os.path.join(self.checkpoint_dir, "step_99")
+        os.makedirs(target)
+        link = os.path.join(self.checkpoint_dir, "step_latest")
+        os.symlink(target, link)
+        shutil.rmtree(target)
+        with self.assertRaises(RuntimeError) as ctx:
+            discover_latest_checkpoint(self.checkpoint_dir)
+        message = str(ctx.exception)
+        self.assertIn(link, message)
+        self.assertIn(os.path.realpath(target), message)
+
+
+if __name__ == "__main__":
+    unittest.main()