Count model directory contents against the cache quota

great_ai/large_file/large_file/large_file_base.py

@@ -327,7 +327,12 @@ class LargeFileBase(ABC):
             [f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime
         )
 
-        while sum(os.path.getsize(f) for f in least_recently_read) > allowed_size:
+        def size(path: Path) -> int:
+            if path.is_dir():
+                return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
+            return path.stat().st_size
+
+        while sum(size(f) for f in least_recently_read) > allowed_size:
             file = least_recently_read.pop(0)
             logger.info(
                 f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"

tests/large_file/test_storage_safety.py

@@ -45,3 +45,14 @@ def test_cached_directories_can_be_removed(
         monkeypatch.setattr(LargeFileLocal, "max_cache_size", "0")
         reader.clean_up()
     assert not cached.exists()
+
+
+def test_cache_quota_counts_directory_contents(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    source = tmp_path / "source"
+    source.mkdir()
+    (source / "weights").write_bytes(b"x" * 100_000)
+    monkeypatch.setattr(LargeFileLocal, "max_cache_size", "1 KB")
+    LargeFileLocal("model", "w").push(source)
+    assert not (LargeFileLocal.cache_path / "model-0").exists()