Isolate model caches by storage backend and configuration

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)
-                copy(str(path), str(self.cache_path / self._local_name))
+                shutil.rmtree(self.cache_path / self._cache_name, ignore_errors=True)
+                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,