From 8f5c9872329b90564ed262e1b811dd8760e37fc6 Mon Sep 17 00:00:00 2001
From: Bhumit Chaudhry <78070889+bhumitschaudhry@users.noreply.github.com>
Date: Fri, 7 Aug 2026 15:25:55 +0530
Subject: [PATCH 1/3] fix: handle missing files gracefully in
 remove_real_and_linked_file

Fixes #297

The remove_real_and_linked_file() helper had two bugs that caused crashes
when delete_original=True was used during model splitting:

1. targetpath was only assigned inside an if-block, causing
   UnboundLocalError when the input was a regular file (not a symlink).
2. A FileNotFoundError was not caught, so deleting an already-removed
   file would crash mid-split, leaving users with partial splits and
   missing source shards.

Changes:
- Normalize to_delete via os.fspath() so Path objects compare correctly
  against str realpath.
- Initialize targetpath = None before the conditional.
- Catch FileNotFoundError on os.remove() and return gracefully.
- Add os.path.exists() guard before removing the symlink target.
- Add docstring explaining the function's behavior.
- Add tests for regular files, Path inputs, missing files, and
  (Linux/macOS) symlinks.

Co-authored-by: Claude <noreply@anthropic.com>
---
 air_llm/airllm/utils.py                       |  31 +++-
 air_llm/tests/test_delete_original.py         | 112 +++++++++++++++
 .../tests/test_delete_original_standalone.py  | 132 ++++++++++++++++++
 3 files changed, 270 insertions(+), 5 deletions(-)
 create mode 100644 air_llm/tests/test_delete_original.py
 create mode 100644 air_llm/tests/test_delete_original_standalone.py

diff --git a/air_llm/airllm/utils.py b/air_llm/airllm/utils.py
index ad45f995..be796107 100644
--- a/air_llm/airllm/utils.py
+++ b/air_llm/airllm/utils.py
@@ -196,12 +196,33 @@ def compress_layer_state_dict(layer_state_dict, compression=None):
     return compressed_layer_state_dict if compressed_layer_state_dict is not None else layer_state_dict
 
 def remove_real_and_linked_file(to_delete):
-    if (os.path.realpath(to_delete) != to_delete):
-        targetpath = os.path.realpath(to_delete)
+    """Remove a file, following symlinks to also remove the target if present.
 
-    os.remove(to_delete)
-    if (targetpath):
-         os.remove(targetpath)
+    If *to_delete* is a symlink the resolved target is removed after the link
+    itself.  For regular files the function simply deletes the file.
+
+    Parameters
+    ----------
+    to_delete : str | Path
+        Path to the file (or symlink) to remove.
+    """
+    to_delete = os.fspath(to_delete)
+    targetpath = None
+
+    realpath = os.path.realpath(to_delete)
+    if realpath != to_delete:
+        targetpath = realpath
+
+    # Remove the primary path (symlink or regular file).
+    try:
+        os.remove(to_delete)
+    except FileNotFoundError:
+        return
+
+    # If it was a symlink, also remove the resolved target so we don't leave
+    # orphaned blobs on disk.
+    if targetpath is not None and os.path.exists(targetpath):
+        os.remove(targetpath)
 
 
 
diff --git a/air_llm/tests/test_delete_original.py b/air_llm/tests/test_delete_original.py
new file mode 100644
index 00000000..957adf56
--- /dev/null
+++ b/air_llm/tests/test_delete_original.py
@@ -0,0 +1,112 @@
+"""Tests for remove_real_and_linked_file (issue #297).
+
+Covers:
+- Regular file deletion
+- Symlink deletion (link removed, target removed)
+- Path object input (not just str)
+- Missing file does not crash
+"""
+
+import importlib.util
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+# Import utils directly to avoid pulling in torch via airllm.__init__
+_utils_path = os.path.join(os.path.dirname(__file__), "..", "airllm", "utils.py")
+_spec = importlib.util.spec_from_file_location("airllm_utils", _utils_path)
+_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_mod)
+remove_real_and_linked_file = _mod.remove_real_and_linked_file
+
+
+class TestRemoveRealAndLinkedFile(unittest.TestCase):
+    """Unit tests for remove_real_and_linked_file."""
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        for entry in os.listdir(self.tmpdir):
+            path = os.path.join(self.tmpdir, entry)
+            if os.path.islink(path) or os.path.isfile(path):
+                os.remove(path)
+            elif os.path.isdir(path):
+                import shutil
+                shutil.rmtree(path)
+        os.rmdir(self.tmpdir)
+
+    # --- regular files ---------------------------------------------------
+
+    def test_regular_file_str_path(self):
+        """A normal file (str input) is deleted exactly once."""
+        filepath = os.path.join(self.tmpdir, "model-00001-of-00002.safetensors")
+        with open(filepath, "w") as f:
+            f.write("fake weights")
+
+        remove_real_and_linked_file(filepath)
+
+        self.assertFalse(os.path.exists(filepath))
+
+    def test_regular_file_path_object(self):
+        """A normal file (Path input) is deleted without UnboundLocalError."""
+        filepath = Path(self.tmpdir) / "model-00002-of-00002.safetensors"
+        filepath.write_text("fake weights")
+
+        remove_real_and_linked_file(filepath)
+
+        self.assertFalse(filepath.exists())
+
+    # --- symlinks --------------------------------------------------------
+
+    def test_symlink_removes_link_and_target(self):
+        """When to_delete is a symlink, both link and target are removed."""
+        target = os.path.join(self.tmpdir, "blobs", "abc123")
+        os.makedirs(os.path.dirname(target))
+        with open(target, "w") as f:
+            f.write("real data")
+
+        link = os.path.join(self.tmpdir, "snapshots", "model.safetensors")
+        os.makedirs(os.path.dirname(link))
+        os.symlink(target, link)
+
+        remove_real_and_linked_file(link)
+
+        self.assertFalse(os.path.exists(link))
+        self.assertFalse(os.path.exists(target))
+
+    def test_symlink_with_path_object(self):
+        """Symlink removal works with Path input too."""
+        target = Path(self.tmpdir) / "blob.bin"
+        target.write_bytes(b"\x00" * 64)
+
+        link = Path(self.tmpdir) / "link.bin"
+        os.symlink(str(target), str(link))
+
+        remove_real_and_linked_file(link)
+
+        self.assertFalse(link.exists())
+        self.assertFalse(target.exists())
+
+    # --- edge cases ------------------------------------------------------
+
+    def test_missing_file_does_not_crash(self):
+        """Deleting a non-existent path should not raise."""
+        nonexistent = os.path.join(self.tmpdir, "does-not-exist.bin")
+        # Should not raise
+        remove_real_and_linked_file(nonexistent)
+
+    def test_broken_symlink_does_not_crash(self):
+        """Deleting a symlink whose target is already gone should not raise."""
+        link = os.path.join(self.tmpdir, "broken-link.bin")
+        os.symlink("/nonexistent/target", link)
+
+        remove_real_and_linked_file(link)
+
+        self.assertFalse(os.path.exists(link))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/air_llm/tests/test_delete_original_standalone.py b/air_llm/tests/test_delete_original_standalone.py
new file mode 100644
index 00000000..be9fe6df
--- /dev/null
+++ b/air_llm/tests/test_delete_original_standalone.py
@@ -0,0 +1,132 @@
+"""Standalone test for remove_real_and_linked_file (issue #297).
+
+This script tests the function in isolation, avoiding the torch/safetensors
+import chain. Symlink tests are skipped on Windows (requires admin or
+developer mode).
+"""
+
+import os
+import platform
+import shutil
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+
+# ── Copy of the fixed function (from airllm/utils.py) ──────────────────
+
+def remove_real_and_linked_file(to_delete):
+    """Remove a file, following symlinks to also remove the target if present."""
+    to_delete = os.fspath(to_delete)
+    targetpath = None
+
+    realpath = os.path.realpath(to_delete)
+    if realpath != to_delete:
+        targetpath = realpath
+
+    try:
+        os.remove(to_delete)
+    except FileNotFoundError:
+        return
+
+    if targetpath is not None and os.path.exists(targetpath):
+        os.remove(targetpath)
+
+
+# ── Tests ──────────────────────────────────────────────────────────────
+
+class TestRemoveRealAndLinkedFile(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        for entry in os.listdir(self.tmpdir):
+            path = os.path.join(self.tmpdir, entry)
+            if os.path.islink(path) or os.path.isfile(path):
+                os.remove(path)
+            elif os.path.isdir(path):
+                shutil.rmtree(path)
+        os.rmdir(self.tmpdir)
+
+    # --- regular files ---------------------------------------------------
+
+    def test_regular_file_str_path(self):
+        """A normal file (str input) is deleted exactly once."""
+        filepath = os.path.join(self.tmpdir, "model-00001.safetensors")
+        with open(filepath, "w") as f:
+            f.write("fake weights")
+        remove_real_and_linked_file(filepath)
+        self.assertFalse(os.path.exists(filepath))
+
+    def test_regular_file_path_object(self):
+        """A normal file (Path input) is deleted without UnboundLocalError."""
+        filepath = Path(self.tmpdir) / "model-00002.safetensors"
+        filepath.write_text("fake weights")
+        remove_real_and_linked_file(filepath)
+        self.assertFalse(filepath.exists())
+
+    def test_regular_file_no_double_delete(self):
+        """targetpath should be None for a regular file — no double-delete."""
+        filepath = os.path.join(self.tmpdir, "shard.bin")
+        with open(filepath, "w") as f:
+            f.write("data")
+        # The old code raised UnboundLocalError here
+        remove_real_and_linked_file(filepath)
+        self.assertFalse(os.path.exists(filepath))
+
+    # --- edge cases ------------------------------------------------------
+
+    def test_missing_file_does_not_crash(self):
+        """Deleting a non-existent path should not raise."""
+        nonexistent = os.path.join(self.tmpdir, "does-not-exist.bin")
+        remove_real_and_linked_file(nonexistent)  # should not raise
+
+    def test_multiple_files_independently(self):
+        """Deleting several files in sequence works."""
+        for name in ["a.bin", "b.bin", "c.bin"]:
+            with open(os.path.join(self.tmpdir, name), "w") as f:
+                f.write(name)
+        for name in ["a.bin", "b.bin", "c.bin"]:
+            remove_real_and_linked_file(os.path.join(self.tmpdir, name))
+        self.assertEqual(os.listdir(self.tmpdir), [])
+
+    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    def test_symlink_removes_link_and_target(self):
+        """When to_delete is a symlink, both link and target are removed."""
+        target = os.path.join(self.tmpdir, "blobs", "abc123")
+        os.makedirs(os.path.dirname(target))
+        with open(target, "w") as f:
+            f.write("real data")
+
+        link = os.path.join(self.tmpdir, "snapshots", "model.safetensors")
+        os.makedirs(os.path.dirname(link))
+        os.symlink(target, link)
+
+        remove_real_and_linked_file(link)
+        self.assertFalse(os.path.exists(link))
+        self.assertFalse(os.path.exists(target))
+
+    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    def test_symlink_with_path_object(self):
+        """Symlink removal works with Path input too."""
+        target = Path(self.tmpdir) / "blob.bin"
+        target.write_bytes(b"\x00" * 64)
+        link = Path(self.tmpdir) / "link.bin"
+        os.symlink(str(target), str(link))
+        remove_real_and_linked_file(link)
+        self.assertFalse(link.exists())
+        self.assertFalse(target.exists())
+
+    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    def test_broken_symlink_does_not_crash(self):
+        """Deleting a symlink whose target is already gone should not raise."""
+        link = os.path.join(self.tmpdir, "broken-link.bin")
+        os.symlink("/nonexistent/target", link)
+        remove_real_and_linked_file(link)
+        self.assertFalse(os.path.exists(link))
+
+
+if __name__ == "__main__":
+    unittest.main()

From eb401847bb5dc42372393764291bada6df493b36 Mon Sep 17 00:00:00 2001
From: Bhumit Chaudhry <78070889+bhumitschaudhry@users.noreply.github.com>
Date: Fri, 7 Aug 2026 16:17:22 +0530
Subject: [PATCH 2/3] Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 air_llm/airllm/utils.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/air_llm/airllm/utils.py b/air_llm/airllm/utils.py
index be796107..f655326a 100644
--- a/air_llm/airllm/utils.py
+++ b/air_llm/airllm/utils.py
@@ -209,9 +209,10 @@ def remove_real_and_linked_file(to_delete):
     to_delete = os.fspath(to_delete)
     targetpath = None
 
-    realpath = os.path.realpath(to_delete)
-    if realpath != to_delete:
-        targetpath = realpath
+    # Only treat the path as "linked" when it's actually a symlink; realpath()
+    # can also differ for non-symlink paths (e.g., ".." components).
+    if os.path.islink(to_delete):
+        targetpath = os.path.realpath(to_delete)
 
     # Remove the primary path (symlink or regular file).
     try:
@@ -219,10 +220,12 @@ def remove_real_and_linked_file(to_delete):
     except FileNotFoundError:
         return
 
-    # If it was a symlink, also remove the resolved target so we don't leave
-    # orphaned blobs on disk.
+    # If it was a symlink, also remove the resolved target (best-effort).
     if targetpath is not None and os.path.exists(targetpath):
-        os.remove(targetpath)
+        try:
+            os.remove(targetpath)
+        except FileNotFoundError:
+            pass
 
 
 

From 86af6c4505a192d06ca0c798da4d053a3254d083 Mon Sep 17 00:00:00 2001
From: Bhumit Chaudhry <78070889+bhumitschaudhry@users.noreply.github.com>
Date: Fri, 7 Aug 2026 16:23:42 +0530
Subject: [PATCH 3/3] address review: use islink(), TOCTOU safety, skipUnless,
 move standalone test

Address Copilot review suggestions on PR #337:

1. Use os.path.islink() instead of realpath != to_delete to properly
   detect symlinks (avoids false positives from path normalization).
2. Wrap target removal in try/except FileNotFoundError for TOCTOU race
   safety under concurrent deletions.
3. Add cross-snapshot warning to docstring about shared HF cache blobs.
4. Add @skipUnless(platform != Windows) to symlink tests in
   test_delete_original.py.
5. Move test_delete_original_standalone.py to scripts/ so unittest
   discovery does not pick up the copied function.
6. Fix macOS skip condition: == Linux -> != Windows.

Co-authored-by: Claude <noreply@anthropic.com>
---
 air_llm/airllm/utils.py                       | 11 ++++++--
 air_llm/tests/test_delete_original.py         |  4 +++
 .../test_delete_original_standalone.py        | 26 ++++++++++++-------
 3 files changed, 29 insertions(+), 12 deletions(-)
 rename {air_llm/tests => scripts}/test_delete_original_standalone.py (84%)

diff --git a/air_llm/airllm/utils.py b/air_llm/airllm/utils.py
index f655326a..da59e1f3 100644
--- a/air_llm/airllm/utils.py
+++ b/air_llm/airllm/utils.py
@@ -201,6 +201,11 @@ def remove_real_and_linked_file(to_delete):
     If *to_delete* is a symlink the resolved target is removed after the link
     itself.  For regular files the function simply deletes the file.
 
+    .. warning::
+       When removing symlinks in Hugging Face cache directories, the resolved
+       target may be shared across snapshots or revisions.  Deleting it will
+       break other snapshots that still reference the same blob.
+
     Parameters
     ----------
     to_delete : str | Path
@@ -220,8 +225,10 @@ def remove_real_and_linked_file(to_delete):
     except FileNotFoundError:
         return
 
-    # If it was a symlink, also remove the resolved target (best-effort).
-    if targetpath is not None and os.path.exists(targetpath):
+    # If it was a symlink, also remove the resolved target so we don't leave
+    # orphaned blobs on disk.  Wrap in try/except for TOCTOU safety — another
+    # process may have already removed the target.
+    if targetpath is not None:
         try:
             os.remove(targetpath)
         except FileNotFoundError:
diff --git a/air_llm/tests/test_delete_original.py b/air_llm/tests/test_delete_original.py
index 957adf56..6f0e947b 100644
--- a/air_llm/tests/test_delete_original.py
+++ b/air_llm/tests/test_delete_original.py
@@ -9,6 +9,7 @@
 
 import importlib.util
 import os
+import platform
 import sys
 import tempfile
 import unittest
@@ -61,6 +62,7 @@ def test_regular_file_path_object(self):
 
     # --- symlinks --------------------------------------------------------
 
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_symlink_removes_link_and_target(self):
         """When to_delete is a symlink, both link and target are removed."""
         target = os.path.join(self.tmpdir, "blobs", "abc123")
@@ -77,6 +79,7 @@ def test_symlink_removes_link_and_target(self):
         self.assertFalse(os.path.exists(link))
         self.assertFalse(os.path.exists(target))
 
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_symlink_with_path_object(self):
         """Symlink removal works with Path input too."""
         target = Path(self.tmpdir) / "blob.bin"
@@ -98,6 +101,7 @@ def test_missing_file_does_not_crash(self):
         # Should not raise
         remove_real_and_linked_file(nonexistent)
 
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_broken_symlink_does_not_crash(self):
         """Deleting a symlink whose target is already gone should not raise."""
         link = os.path.join(self.tmpdir, "broken-link.bin")
diff --git a/air_llm/tests/test_delete_original_standalone.py b/scripts/test_delete_original_standalone.py
similarity index 84%
rename from air_llm/tests/test_delete_original_standalone.py
rename to scripts/test_delete_original_standalone.py
index be9fe6df..4697751c 100644
--- a/air_llm/tests/test_delete_original_standalone.py
+++ b/scripts/test_delete_original_standalone.py
@@ -1,8 +1,10 @@
 """Standalone test for remove_real_and_linked_file (issue #297).
 
 This script tests the function in isolation, avoiding the torch/safetensors
-import chain. Symlink tests are skipped on Windows (requires admin or
-developer mode).
+import chain.  It lives in scripts/ rather than tests/ so unittest discovery
+does not pick it up (it contains a copied function, not a production import).
+
+Symlink tests are skipped on Windows (requires admin or developer mode).
 """
 
 import os
@@ -15,23 +17,27 @@
 
 
 # ── Copy of the fixed function (from airllm/utils.py) ──────────────────
+# NOTE: This is intentionally duplicated for portability.  Keep in sync
+# with the real implementation in air_llm/airllm/utils.py.
 
 def remove_real_and_linked_file(to_delete):
     """Remove a file, following symlinks to also remove the target if present."""
     to_delete = os.fspath(to_delete)
     targetpath = None
 
-    realpath = os.path.realpath(to_delete)
-    if realpath != to_delete:
-        targetpath = realpath
+    if os.path.islink(to_delete):
+        targetpath = os.path.realpath(to_delete)
 
     try:
         os.remove(to_delete)
     except FileNotFoundError:
         return
 
-    if targetpath is not None and os.path.exists(targetpath):
-        os.remove(targetpath)
+    if targetpath is not None:
+        try:
+            os.remove(targetpath)
+        except FileNotFoundError:
+            pass
 
 
 # ── Tests ──────────────────────────────────────────────────────────────
@@ -92,7 +98,7 @@ def test_multiple_files_independently(self):
             remove_real_and_linked_file(os.path.join(self.tmpdir, name))
         self.assertEqual(os.listdir(self.tmpdir), [])
 
-    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_symlink_removes_link_and_target(self):
         """When to_delete is a symlink, both link and target are removed."""
         target = os.path.join(self.tmpdir, "blobs", "abc123")
@@ -108,7 +114,7 @@ def test_symlink_removes_link_and_target(self):
         self.assertFalse(os.path.exists(link))
         self.assertFalse(os.path.exists(target))
 
-    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_symlink_with_path_object(self):
         """Symlink removal works with Path input too."""
         target = Path(self.tmpdir) / "blob.bin"
@@ -119,7 +125,7 @@ def test_symlink_with_path_object(self):
         self.assertFalse(link.exists())
         self.assertFalse(target.exists())
 
-    @unittest.skipUnless(platform.system() == "Linux", "symlink test requires Linux/macOS")
+    @unittest.skipUnless(platform.system() != "Windows", "symlink tests require admin on Windows")
     def test_broken_symlink_does_not_crash(self):
         """Deleting a symlink whose target is already gone should not raise."""
         link = os.path.join(self.tmpdir, "broken-link.bin")