Enforce download cache quotas without returning evicted paths

great_ai/large_file/large_file/large_file_base.py

@@ -222,6 +222,8 @@ class LargeFileBase(ABC):
         else:
             logger.info(f"File {self._local_name} found in cache")
 
+        os.utime(destination, None)
+        self._prune_cache(protected=destination)
         return destination
 
     def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
@@ -388,7 +390,7 @@ class LargeFileBase(ABC):
         else:
             raise ValueError("Unsupported file mode.")
 
-    def _prune_cache(self) -> None:
+    def _prune_cache(self, protected: Optional[Path] = None) -> None:
         self.cache_path.mkdir(parents=True, exist_ok=True)
 
         if self.max_cache_size is None:
@@ -411,8 +413,19 @@ class LargeFileBase(ABC):
                 return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
             return path.stat().st_size
 
+        if protected is not None and size(protected) > allowed_size:
+            self._remove_cached_path(protected)
+            self._cache_identity_path.unlink(missing_ok=True)
+            raise ValueError(
+                f"File {self._name} exceeds max_cache_size={self.max_cache_size}; "
+                "increase the cache quota to read this file"
+            )
+
         while sum(size(f) for f in least_recently_read) > allowed_size:
             file = least_recently_read.pop(0)
+            if file == protected:
+                least_recently_read.append(file)
+                continue
             logger.info(
                 f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
             )
@@ -420,6 +433,10 @@ class LargeFileBase(ABC):
                 shutil.rmtree(file)
             else:
                 file.unlink()
+            metadata = (
+                self.cache_path / ".metadata" / sha256(file.name.encode()).hexdigest()
+            )
+            metadata.unlink(missing_ok=True)
 
     @abstractmethod
     def _find_remote_instances(self) -> List[DataInstance]:

tests/large_file/test_download_quota.py

@@ -0,0 +1,40 @@
+import io
+import tarfile
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+from great_ai.large_file import LargeFileS3
+
+
+def test_download_quota_keeps_the_requested_file_and_evicts_older_models(
+    tmp_path: Path,
+) -> None:
+    backend = MagicMock()
+    backend.list_objects_v2.side_effect = lambda **kwargs: {
+        "Contents": [{"Key": kwargs["Prefix"] + "0", "ETag": "one"}]
+    }
+
+    def download(
+        self: LargeFileS3, remote_path: Any, path: Path, hide_progress: bool
+    ) -> None:
+        data = b"x" * (2000 if self._name == "oversized" else 600)
+        with tarfile.open(path, "w:gz") as archive:
+            member = tarfile.TarInfo(self._local_name)
+            member.size = len(data)
+            archive.addfile(member, io.BytesIO(data))
+
+    with (
+        patch.object(LargeFileS3, "_client", backend),
+        patch.object(LargeFileS3, "_download", download),
+        patch.object(LargeFileS3, "cache_path", tmp_path / "cache"),
+        patch.object(LargeFileS3, "max_cache_size", "1 KB"),
+    ):
+        first = LargeFileS3("first").get()
+        second = LargeFileS3("second").get()
+        assert not first.exists()
+        assert second.read_bytes() == b"x" * 600
+        with pytest.raises(ValueError, match="exceeds max_cache_size"):
+            LargeFileS3("oversized").get()
+        assert second.exists()