diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py index fb4d688..4ca161d 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/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]: diff --git a/great_ai/large_file/large_file/large_file_local.py b/great_ai/large_file/large_file/large_file_local.py index c1c5a12..27b5bf6 100644 --- a/great_ai/large_file/large_file/large_file_local.py +++ b/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() diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index 9ace902..ce4190b 100644 --- a/tests/large_file/test_storage_safety.py +++ b/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()