Isolate model caches by storage backend and configuration

docs/how-to-guides/large-file.md

@@ -70,6 +70,11 @@ Often, especially when working with data-heavy applications, large files can pro
 
 ## From the command-line 
 
+Remote cache filenames now include a storage-specific namespace. Existing local
+model stores remain readable, while remote models cached by older releases are
+downloaded again into the new namespace. Configure the same bucket or database
+when preparing and later using an offline remote cache.
+
 The main reason for using the `large-file` or `python3 -m great_ai.large_file` commands is to upload or download models from the terminal. For example, when building a docker image, it is best practice to cache the referred models.
 
 ### Setup

great_ai/large_file/large_file/large_file_base.py

@@ -4,9 +4,10 @@ import shutil
 import sys
 import tempfile
 from abc import ABC, abstractmethod
+from hashlib import sha256
 from pathlib import Path
 from types import TracebackType
-from typing import IO, Any, List, Literal, Optional, Type, Union, cast
+from typing import IO, Any, List, Literal, Optional, Tuple, Type, Union, cast
 
 from great_ai.utilities import ConfigFile, get_logger
 
@@ -47,6 +48,7 @@ class LargeFileBase(ABC):
     initialized = False
     cache_path = Path(".cache")
     max_cache_size: Optional[str] = "30GB"
+    _configuration_attributes: Tuple[str, ...] = ()
 
     def __init__(
         self,
@@ -84,6 +86,10 @@ class LargeFileBase(ABC):
 
         self._newline = newline
 
+        # Existing instances keep the storage configuration they were opened with.
+        for attribute in self._configuration_attributes:
+            setattr(self, attribute, getattr(self, attribute))
+        self._cache_namespace = self._get_cache_namespace()
         self.cache_path.mkdir(parents=True, exist_ok=True)
 
         self._find_instances()
@@ -175,7 +181,7 @@ class LargeFileBase(ABC):
             i.remote_path for i in self._instances if i.version == self._version
         )
 
-        destination = self.cache_path / self._local_name
+        destination = self.cache_path / self._cache_name
         if not destination.exists():
             logger.info(f"File {self._local_name} does not exist locally")
 
@@ -218,8 +224,8 @@ class LargeFileBase(ABC):
 
             try:
                 # Make local copy in the cache
-                shutil.rmtree(self.cache_path / self._local_name, ignore_errors=True)
+                shutil.rmtree(self.cache_path / self._cache_name, ignore_errors=True)
-                copy(str(path), str(self.cache_path / self._local_name))
+                copy(str(path), str(self.cache_path / self._cache_name))
             except shutil.SameFileError:
                 pass  # No worries
 
@@ -262,6 +268,18 @@ class LargeFileBase(ABC):
     def _local_name(self) -> str:
         return f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}{self.version}"
 
+    @property
+    def _cache_name(self) -> str:
+        # '@' is forbidden in model names, so even local names that resemble
+        # namespace hashes cannot collide with remote cache entries.
+        prefix = f"@{self._cache_namespace}-" if self._cache_namespace else ""
+        return f"{prefix}{self._local_name}"
+
+    def _get_cache_namespace(self) -> str:
+        return sha256(
+            f"{type(self).__module__}.{type(self).__qualname__}".encode()
+        ).hexdigest()
+
     def _find_instances(self) -> None:
         if self._cache_only_mode:
             self._instances = self._find_instances_from_cache()
@@ -273,7 +291,8 @@ class LargeFileBase(ABC):
     def _find_instances_from_cache(self) -> List[DataInstance]:
         logger.info(f"Fetching cached versions of {self._name}")
 
-        prefix = f"{self._name}{CACHE_NAME_VERSION_SEPARATOR}"
+        namespace = f"@{self._cache_namespace}-" if self._cache_namespace else ""
+        prefix = f"{namespace}{self._name}{CACHE_NAME_VERSION_SEPARATOR}"
         return [
             DataInstance(
                 name=self._name, version=int(f.name[len(prefix) :]), remote_path=f

great_ai/large_file/large_file/large_file_local.py

@@ -62,6 +62,10 @@ class LargeFileLocal(LargeFileBase):
     def _find_remote_instances(self) -> List[DataInstance]:
         return []
 
+    def _get_cache_namespace(self) -> str:
+        # Preserve existing local model stores; remote caches are namespaced.
+        return ""
+
     def _download(
         self, remote_path: Any, local_path: Path, hide_progress: bool
     ) -> None:

great_ai/large_file/large_file/large_file_mongo.py

@@ -1,4 +1,5 @@
 import re
+from hashlib import sha256
 from pathlib import Path
 from typing import Any, List
 
@@ -29,6 +30,13 @@ class LargeFileMongo(LargeFileBase):
     mongo_connection_string = None
     mongo_database = None
 
+    _configuration_attributes = ("mongo_connection_string", "mongo_database")
+
+    def _get_cache_namespace(self) -> str:
+        return sha256(
+            repr(("gridfs", self.mongo_connection_string, self.mongo_database)).encode()
+        ).hexdigest()
+
     @classmethod
     def configure_credentials(  # type: ignore
         cls,

great_ai/large_file/large_file/large_file_s3.py

@@ -1,3 +1,4 @@
+from hashlib import sha256
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
@@ -29,6 +30,18 @@ class LargeFileS3(LargeFileBase):
     bucket_name = None
     endpoint_url = None
 
+    _configuration_attributes = (
+        "region_name",
+        "access_key_id",
+        "secret_access_key",
+        "bucket_name",
+        "endpoint_url",
+    )
+
+    def _get_cache_namespace(self) -> str:
+        identity = ("s3", self.endpoint_url, self.region_name, self.bucket_name)
+        return sha256(repr(identity).encode()).hexdigest()
+
     @classmethod
     def configure_credentials(  # type: ignore
         cls,

tests/large_file/test_storage_safety.py

@@ -123,6 +123,52 @@ def test_gridfs_matches_only_exact_model_names() -> None:
     assert [instance.name for instance in reader._instances] == ["model"]
 
 
+def test_remote_cache_is_separate_from_local_cache() -> None:
+    with LargeFileLocal("model", "w") as file:
+        file.write("local weights")
+    local_path = LargeFileLocal("model").get()
+    backend = MagicMock()
+    backend.list_objects_v2.return_value = {"Contents": [{"Key": "model/0"}]}
+
+    def download(
+        self: LargeFileS3, remote_path: Any, local_path: Path, hide_progress: bool
+    ) -> None:
+        with tarfile.open(local_path, "w:gz") as archive:
+            data = b"remote weights"
+            info = tarfile.TarInfo("model-0")
+            info.size = len(data)
+            archive.addfile(info, io.BytesIO(data))
+
+    with (
+        patch.object(LargeFileS3, "_client", backend),
+        patch.object(LargeFileS3, "_download", download),
+    ):
+        remote_path = LargeFileS3("model").get()
+    assert remote_path != local_path
+    assert remote_path.read_text() == "remote weights"
+    assert local_path.read_text() == "local weights"
+
+
+def test_cache_identity_includes_storage_configuration(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    backend = MagicMock()
+    backend.list_objects_v2.return_value = {"Contents": [{"Key": "model/0"}]}
+    monkeypatch.setattr(LargeFileS3, "bucket_name", "first")
+    with patch.object(LargeFileS3, "_client", backend):
+        first = LargeFileS3("model")
+        monkeypatch.setattr(LargeFileS3, "bucket_name", "second")
+        second = LargeFileS3("model")
+        monkeypatch.setattr(
+            LargeFileS3, "endpoint_url", "https://another-store.example"
+        )
+        third = LargeFileS3("model")
+    assert len({first._cache_name, second._cache_name, third._cache_name}) == 3
+    assert first.bucket_name == "first"
+    assert second.bucket_name == "second"
+    assert second.endpoint_url != third.endpoint_url
+
+
 def test_s3_follows_every_listing_page() -> None:
     backend = MagicMock()
     backend.list_objects_v2.side_effect = [