From d1e82bacf5bde1abdbc2ff36ceebce97187bb075 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 19:50:38 +0000 Subject: [PATCH] Validate cached models against remote generation identities --- .../large_file/large_file/large_file_base.py | 53 +++++++++++++++++-- .../large_file/large_file/large_file_mongo.py | 1 + .../large_file/large_file/large_file_s3.py | 9 +++- great_ai/large_file/models/data_instance.py | 3 +- tests/large_file/test_cache_generations.py | 46 ++++++++++++++++ 5 files changed, 105 insertions(+), 7 deletions(-) create mode 100644 tests/large_file/test_cache_generations.py 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 32cdf11..53d0884 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -179,11 +179,25 @@ class LargeFileBase(ABC): hide_progress: Do not show a progress update after each 10% of progress. """ - remote_path = next( - i.remote_path for i in self._instances if i.version == self._version + if not self._cache_only_mode: + self._find_instances() + instance = next( + (i for i in self._instances if i.version == self._version), None ) + if instance is None: + raise FileNotFoundError( + f"File {self._name} version {self.version} not found" + ) destination = self.cache_path / self._cache_name + if destination.exists() and not self._cache_only_mode: + identity = instance.cache_identity + if ( + identity is None + or not self._cache_identity_path.exists() + or self._cache_identity_path.read_text() != identity + ): + self._remove_cached_path(destination) if not destination.exists(): logger.info(f"File {self._local_name} does not exist locally") @@ -193,7 +207,7 @@ class LargeFileBase(ABC): local_root_path / f"{self._local_name}{ARCHIVE_EXTENSION}" ) self._download( - remote_path, tmp_file_archive, hide_progress=hide_progress + instance.remote_path, tmp_file_archive, hide_progress=hide_progress ) logger.info(f"Decompressing {self._local_name}") @@ -201,6 +215,7 @@ class LargeFileBase(ABC): str(tmp_file_archive), tmp, COMPRESSION_ALGORITHM, filter="data" ) shutil.move(str(local_root_path / self._local_name), str(destination)) + self._write_cache_identity() else: logger.info(f"File {self._local_name} found in cache") @@ -247,10 +262,15 @@ class LargeFileBase(ABC): # Stage the source before publishing: it may itself be cached. destination = self.cache_path / self._cache_name if destination.exists(): - raise FileExistsError(f"Version already exists: {destination}") + if self._cache_only_mode: + raise FileExistsError(f"Version already exists: {destination}") + # The remote conditional commit succeeded. Any existing local + # entry belongs to a deleted generation of this version. + self._remove_cached_path(destination) shutil.move(str(staged), str(destination)) self._find_instances() + self._write_cache_identity() self.clean_up() def delete(self) -> None: @@ -286,6 +306,29 @@ class LargeFileBase(ABC): f"{type(self).__module__}.{type(self).__qualname__}".encode() ).hexdigest() + @property + def _cache_identity_path(self) -> Path: + key = sha256(self._cache_name.encode()).hexdigest() + return self.cache_path / ".metadata" / key + + def _write_cache_identity(self) -> None: + if self._cache_only_mode: + return + identity = next( + (i.cache_identity for i in self._instances if i.version == self.version), + None, + ) + if identity is not None: + self._cache_identity_path.parent.mkdir(parents=True, exist_ok=True) + self._cache_identity_path.write_text(identity) + + @staticmethod + def _remove_cached_path(path: Path) -> None: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink() + def _find_instances(self) -> None: if self._cache_only_mode: self._instances = self._find_instances_from_cache() @@ -352,7 +395,7 @@ class LargeFileBase(ABC): [ f for f in self.cache_path.glob("*") - if f.name not in {".locks", ".staging"} + if f.name not in {".locks", ".staging", ".metadata"} ], key=lambda f: f.stat().st_atime, ) diff --git a/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/large_file/large_file/large_file_mongo.py index 75a2285..4fd11fb 100644 --- a/great_ai/large_file/large_file/large_file_mongo.py +++ b/great_ai/large_file/large_file/large_file_mongo.py @@ -72,6 +72,7 @@ class LargeFileMongo(LargeFileBase): ), version=int(f.name.split(MONGO_NAME_VERSION_SEPARATOR)[-1]), remote_path=(f._id, f.length), + cache_identity=str(f._id), ) for f in self._client.find( { diff --git a/great_ai/large_file/large_file/large_file_s3.py b/great_ai/large_file/large_file/large_file_s3.py index 4c223f4..fb753cb 100644 --- a/great_ai/large_file/large_file/large_file_s3.py +++ b/great_ai/large_file/large_file/large_file_s3.py @@ -96,7 +96,14 @@ class LargeFileS3(LargeFileBase): if suffix.isdigit(): instances.append( DataInstance( - name=self._name, version=int(suffix), remote_path=key + name=self._name, + version=int(suffix), + remote_path=key, + cache_identity=( + repr((item["ETag"], item.get("LastModified"))) + if item.get("ETag") is not None + else None + ), ) ) if not page.get("IsTruncated"): diff --git a/great_ai/large_file/models/data_instance.py b/great_ai/large_file/models/data_instance.py index bfef26c..b260226 100644 --- a/great_ai/large_file/models/data_instance.py +++ b/great_ai/large_file/models/data_instance.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Optional from pydantic import BaseModel @@ -7,3 +7,4 @@ class DataInstance(BaseModel): name: str version: int remote_path: Any + cache_identity: Optional[str] = None diff --git a/tests/large_file/test_cache_generations.py b/tests/large_file/test_cache_generations.py new file mode 100644 index 0000000..eafdd20 --- /dev/null +++ b/tests/large_file/test_cache_generations.py @@ -0,0 +1,46 @@ +from hashlib import sha256 +from pathlib import Path +from typing import Any, Dict +from unittest.mock import MagicMock, patch + +from great_ai.large_file import LargeFileS3 + + +def test_recreated_remote_version_invalidates_other_clients(tmp_path: Path) -> None: + remote: Dict[str, bytes] = {} + backend = MagicMock() + backend.list_objects_v2.side_effect = lambda **kwargs: { + "Contents": [ + {"Key": key, "ETag": sha256(value).hexdigest()} + for key, value in remote.items() + ] + } + backend.delete_object.side_effect = lambda **kwargs: remote.pop(kwargs["Key"]) + + def upload(self: LargeFileS3, path: Path, hide_progress: bool) -> None: + remote[f"{self._name}/{self.version}"] = path.read_bytes() + + def download( + self: LargeFileS3, remote_path: Any, path: Path, hide_progress: bool + ) -> None: + path.write_bytes(remote[remote_path]) + + with ( + patch.object(LargeFileS3, "_client", backend), + patch.object(LargeFileS3, "_upload", upload), + patch.object(LargeFileS3, "_download", download), + patch.object(LargeFileS3, "max_cache_size", None), + ): + with patch.object(LargeFileS3, "cache_path", tmp_path / "writer"): + with LargeFileS3("model", "w") as file: + file.write("old") + with patch.object(LargeFileS3, "cache_path", tmp_path / "reader"): + reader = LargeFileS3("model") + assert reader.get().read_text() == "old" + with patch.object(LargeFileS3, "cache_path", tmp_path / "writer"): + LargeFileS3("model").delete() + with LargeFileS3("model", "w") as file: + file.write("new") + assert LargeFileS3("model").get().read_text() == "new" + with patch.object(LargeFileS3, "cache_path", tmp_path / "reader"): + assert reader.get().read_text() == "new"