Validate cached models against remote generation identities

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,
         )

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(
                 {

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"):

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

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"