From fb31b3ed83e21ef2d20ebf61df8469c51d278b1e Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 16:52:35 +0000 Subject: [PATCH] Count model directory contents against the cache quota --- great_ai/large_file/large_file/large_file_base.py | 7 ++++++- tests/large_file/test_storage_safety.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) 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 4ca161d..ce0631e 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/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}" diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index ce4190b..0c4ba32 100644 --- a/tests/large_file/test_storage_safety.py +++ b/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()