Remove cached model directories recursively

great_ai/large_file/large_file/large_file_base.py

@@ -332,7 +332,10 @@ class LargeFileBase(ABC):
             logger.info(
                 f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
             )
-            os.unlink(file)
+            if file.is_dir():
+                shutil.rmtree(file)
+            else:
+                file.unlink()
 
     @abstractmethod
     def _find_remote_instances(self) -> List[DataInstance]:

great_ai/large_file/large_file/large_file_local.py

@@ -1,3 +1,4 @@
+import shutil
 from pathlib import Path
 from typing import Any, List, Optional
 
@@ -80,4 +81,7 @@ class LargeFileLocal(LargeFileBase):
                 logger.info(
                     f"Removing old version (keep_last_n={self._keep_last_n}): {i.remote_path}"
                 )
-                i.remote_path.unlink()
+                if i.remote_path.is_dir():
+                    shutil.rmtree(i.remote_path)
+                else:
+                    i.remote_path.unlink()

tests/large_file/test_storage_safety.py

@@ -27,3 +27,21 @@ def test_unrelated_cache_entries_are_ignored() -> None:
     (LargeFileLocal.cache_path / "model-unrelated").write_text("unrelated")
     (LargeFileLocal.cache_path / "model-0").write_text("weights")
     assert LargeFileLocal("model").get().read_text() == "weights"
+
+
+@pytest.mark.parametrize("operation", ["delete", "prune"])
+def test_cached_directories_can_be_removed(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str
+) -> None:
+    source = tmp_path / "source"
+    source.mkdir()
+    (source / "weights").write_text("weights")
+    LargeFileLocal("model", "w").push(source)
+    reader = LargeFileLocal("model")
+    cached = reader.get()
+    if operation == "delete":
+        reader.delete()
+    else:
+        monkeypatch.setattr(LargeFileLocal, "max_cache_size", "0")
+        reader.clean_up()
+    assert not cached.exists()