From de4dc07e546b223b63ee4f16afa4dac69515acff Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 16:55:04 +0000 Subject: [PATCH] Isolate model caches by storage backend and configuration --- docs/how-to-guides/large-file.md | 5 ++ .../large_file/large_file/large_file_base.py | 29 ++++++++++-- .../large_file/large_file/large_file_local.py | 4 ++ .../large_file/large_file/large_file_mongo.py | 8 ++++ .../large_file/large_file/large_file_s3.py | 13 ++++++ tests/large_file/test_storage_safety.py | 46 +++++++++++++++++++ 6 files changed, 100 insertions(+), 5 deletions(-) diff --git a/docs/how-to-guides/large-file.md b/docs/how-to-guides/large-file.md index f7babef..b19c3b2 100644 --- a/docs/how-to-guides/large-file.md +++ b/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 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 f3f640a..d66becc 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/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 diff --git a/great_ai/large_file/large_file/large_file_local.py b/great_ai/large_file/large_file/large_file_local.py index 27b5bf6..125cffd 100644 --- a/great_ai/large_file/large_file/large_file_local.py +++ b/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: 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 6e9d46e..2ed0da5 100644 --- a/great_ai/large_file/large_file/large_file_mongo.py +++ b/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, 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 60efc62..c717ddd 100644 --- a/great_ai/large_file/large_file/large_file_s3.py +++ b/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, diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index 6c6aeb0..10c5abe 100644 --- a/tests/large_file/test_storage_safety.py +++ b/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 = [