diff --git a/docs/how-to-guides/large-file.md b/docs/how-to-guides/large-file.md index b19c3b2..c0209ef 100644 --- a/docs/how-to-guides/large-file.md +++ b/docs/how-to-guides/large-file.md @@ -70,6 +70,13 @@ Often, especially when working with data-heavy applications, large files can pro ## From the command-line +Versions are allocated when an upload commits. Writers sharing a cache directory +serialize their uploads with a file lock. S3 and GridFS also reject conflicting +version commits from other hosts with `FileExistsError`; retry the upload to +allocate a fresh version. S3-compatible services must support conditional writes +(`If-None-Match`), including multipart completion. GridFS writers need permission +to create the unique index used for GreatAI's versioned filenames. + 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 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 49c7928..9a5aef9 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -9,6 +9,8 @@ from pathlib import Path from types import TracebackType from typing import IO, Any, List, Literal, Optional, Tuple, Type, Union, cast +from filelock import FileLock + from great_ai.utilities import ConfigFile, get_logger from ..helper import human_readable_to_byte @@ -213,34 +215,41 @@ class LargeFileBase(ABC): if isinstance(path, str): path = Path(path) - if "w" not in self._mode: + + lock_directory = self.cache_path / ".locks" + lock_directory.mkdir(parents=True, exist_ok=True) + lock_name = sha256(f"{self._cache_namespace}:{self._name}".encode()).hexdigest() + staging_directory = self.cache_path / ".staging" + staging_directory.mkdir(parents=True, exist_ok=True) + with FileLock(lock_directory / f"{lock_name}.lock"): + # Allocate at commit time, after earlier local writers have finished. self._find_instances() self._version = self._instances[-1].version + 1 if self._instances else 0 + with tempfile.TemporaryDirectory(dir=staging_directory) as tmp: + root = Path(tmp) + staged = root / self._local_name + if path.is_dir(): + shutil.copytree(path, staged) + else: + shutil.copy2(path, staged) - with tempfile.TemporaryDirectory() as tmp: - root = Path(tmp) - staged = root / self._local_name - if path.is_dir(): - shutil.copytree(path, staged) - else: - shutil.copy2(path, staged) + with tempfile.TemporaryDirectory() as archive_directory: + archive = shutil.make_archive( + str(Path(archive_directory) / self._local_name), + COMPRESSION_ALGORITHM, + root, + ) + # Remote backends must reject version conflicts atomically. + self._upload(Path(archive), hide_progress=hide_progress) - with tempfile.TemporaryDirectory() as archive_directory: - archive = shutil.make_archive( - str(Path(archive_directory) / self._local_name), - COMPRESSION_ALGORITHM, - root, - ) - self._upload(Path(archive), hide_progress=hide_progress) + # 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}") + shutil.move(str(staged), str(destination)) - # 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}") - shutil.move(str(staged), str(destination)) - - self._find_instances() - self.clean_up() + self._find_instances() + self.clean_up() def delete(self) -> None: """Delete all versions of the files under this `key`.""" @@ -338,7 +347,12 @@ class LargeFileBase(ABC): assert allowed_size >= 0 least_recently_read = sorted( - [f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime + [ + f + for f in self.cache_path.glob("*") + if f.name not in {".locks", ".staging"} + ], + key=lambda f: f.stat().st_atime, ) def size(path: Path) -> int: 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 2ed0da5..75a2285 100644 --- a/great_ai/large_file/large_file/large_file_mongo.py +++ b/great_ai/large_file/large_file/large_file_mongo.py @@ -6,6 +6,7 @@ from typing import Any, List from gridfs import DEFAULT_CHUNK_SIZE, GridFSBucket from pymongo import MongoClient from pymongo.database import Database +from pymongo.errors import DuplicateKeyError from ...utilities import get_logger from ..helper import DownloadProgressBar, UploadProgressBar, cached_property @@ -29,7 +30,6 @@ class LargeFileMongo(LargeFileBase): mongo_connection_string = None mongo_database = None - _configuration_attributes = ("mongo_connection_string", "mongo_database") def _get_cache_namespace(self) -> str: @@ -50,14 +50,17 @@ class LargeFileMongo(LargeFileBase): super().configure_credentials() @cached_property - def _client(self) -> GridFSBucket: + def _database(self) -> Database: if self.mongo_connection_string is None or self.mongo_database is None: raise ValueError( "Please configure the MongoDB access options by calling LargeFileMongo.configure_credentials or set cache_only_mode=True in the constructor." ) - db: Database = MongoClient(self.mongo_connection_string)[self.mongo_database] - return GridFSBucket(db) + return MongoClient(self.mongo_connection_string)[self.mongo_database] + + @cached_property + def _client(self) -> GridFSBucket: + return GridFSBucket(self._database) def _find_remote_instances(self) -> List[DataInstance]: logger.debug(f"Fetching Mongo (GridFS) versions of {self._name}") @@ -112,18 +115,33 @@ class LargeFileMongo(LargeFileBase): if not hide_progress else None ) - with self._client.open_upload_stream( - f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}" - ) as stream: - with open(local_path, "rb") as f: - while True: - content = f.read(DEFAULT_CHUNK_SIZE) - stream.write(content) - - if progress: - progress(len(content)) - if len(content) < DEFAULT_CHUNK_SIZE: - break + # Only GreatAI's versioned uploads participate in this unique index; + # unrelated GridFS users may still store duplicate filenames. + self._database.fs.files.create_index( + "filename", + name="great_ai_unique_version", + unique=True, + partialFilterExpression={"metadata.great_ai_version": True}, + ) + filename = f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}" + stream = self._client.open_upload_stream( + filename, metadata={"great_ai_version": True} + ) + try: + with stream: + with open(local_path, "rb") as f: + while True: + content = f.read(DEFAULT_CHUNK_SIZE) + stream.write(content) + if progress: + progress(len(content)) + if len(content) < DEFAULT_CHUNK_SIZE: + break + except DuplicateKeyError as error: + stream.abort() + raise FileExistsError( + f"Version {filename} was created by another writer; retry the upload" + ) from error def _delete_old_remote_versions(self) -> None: if self._keep_last_n is not None: 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 c717ddd..4c223f4 100644 --- a/great_ai/large_file/large_file/large_file_s3.py +++ b/great_ai/large_file/large_file/large_file_s3.py @@ -3,6 +3,9 @@ from pathlib import Path from typing import Any, Dict, List, Optional import boto3 +from boto3.exceptions import S3UploadFailedError +from boto3.s3.transfer import TransferConfig +from botocore.exceptions import ClientError from ...utilities import get_logger from ..helper import DownloadProgressBar, UploadProgressBar, cached_property @@ -29,7 +32,6 @@ class LargeFileS3(LargeFileBase): secret_access_key = None bucket_name = None endpoint_url = None - _configuration_attributes = ( "region_name", "access_key_id", @@ -127,16 +129,45 @@ class LargeFileS3(LargeFileBase): key = f"{self._name}/{self.version}" logger.info(f"Uploading {self._local_name} to S3 as {key}") - self._client.upload_file( - Filename=str(local_path), - Bucket=self.bucket_name, - Key=key, - Callback=( - None - if hide_progress - else UploadProgressBar(path=local_path, logger=logger) - ), - ) + def prevent_overwrite(params: Dict[str, Any], **_: Any) -> None: + if params.get("Bucket") == self.bucket_name and params.get("Key") == key: + params["IfNoneMatch"] = "*" + + # S3Transfer does not accept IfNoneMatch in ExtraArgs. These botocore + # hooks protect both single-part and multipart commits while retaining + # the managed transfer's streaming, retries, and progress reporting. + events = self._client.meta.events + operations = ("PutObject", "CompleteMultipartUpload") + for operation in operations: + events.register(f"before-parameter-build.s3.{operation}", prevent_overwrite) + try: + self._client.upload_file( + Filename=str(local_path), + Bucket=self.bucket_name, + Key=key, + Callback=( + None + if hide_progress + else UploadProgressBar(path=local_path, logger=logger) + ), + Config=TransferConfig(preferred_transfer_client="classic"), + ) + except (ClientError, S3UploadFailedError) as error: + cause: Optional[BaseException] = error + while cause is not None: + if isinstance(cause, ClientError) and cause.response["Error"][ + "Code" + ] in {"PreconditionFailed", "ConditionalRequestConflict", "412", "409"}: + raise FileExistsError( + f"Version {key} was created by another writer; retry the upload" + ) from error + cause = cause.__cause__ or cause.__context__ + raise + finally: + for operation in operations: + events.unregister( + f"before-parameter-build.s3.{operation}", prevent_overwrite + ) def _delete_old_remote_versions(self) -> None: if self._keep_last_n is not None: diff --git a/pyproject.toml b/pyproject.toml index dbe07b8..bbda82e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dependencies = [ "dill >= 0.4.1, < 1", "tqdm >= 4.70.1, < 5", "httpx >= 0.28.1, < 1", + "filelock >= 4.2.4, < 5", ] [project.optional-dependencies] diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index 5bc5193..0df110e 100644 --- a/tests/large_file/test_storage_safety.py +++ b/tests/large_file/test_storage_safety.py @@ -1,17 +1,22 @@ import io +import multiprocessing as mp import tarfile from pathlib import Path from types import SimpleNamespace from typing import Any, Iterator from unittest.mock import MagicMock, patch +import boto3 import pytest +from boto3.s3.transfer import TransferConfig +from botocore.stub import Stubber from great_ai.large_file import ( LargeFileBase, LargeFileLocal, LargeFileMongo, LargeFileS3, ) +from pymongo.errors import DuplicateKeyError @pytest.fixture(autouse=True) @@ -125,6 +130,52 @@ def test_retention_counts_the_new_version() -> None: assert reader.get().read_text() == "2" +def write_model(root: str, value: str, barrier: Any) -> None: + LargeFileBase.cache_path = LargeFileLocal.cache_path = Path(root) + LargeFileBase.max_cache_size = LargeFileLocal.max_cache_size = None + writer = LargeFileLocal("model", "w") + barrier.wait(timeout=10) + with writer as file: + file.write(value) + + +def test_overlapping_local_writers_preserve_both_versions() -> None: + first = LargeFileLocal("model", "w") + second = LargeFileLocal("model", "w") + for writer, value in [(first, "first"), (second, "second")]: + with writer as file: + file.write(value) + assert first.version == 0 + assert second.version == 1 + assert LargeFileLocal("model", version=0).get().read_text() == "first" + assert LargeFileLocal("model", version=1).get().read_text() == "second" + + +def test_concurrent_processes_preserve_both_versions(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + barrier = ctx.Barrier(2) + processes = [ + ctx.Process(target=write_model, args=(str(tmp_path / "cache"), value, barrier)) + for value in ("first", "second") + ] + try: + for process in processes: + process.start() + for process in processes: + process.join(timeout=15) + assert process.exitcode == 0 + finally: + for process in processes: + if process.is_alive(): + process.kill() + process.join() + process.close() + assert {LargeFileLocal("model", version=v).get().read_text() for v in (0, 1)} == { + "first", + "second", + } + + def test_uploading_cached_directory_preserves_source(tmp_path: Path) -> None: source = tmp_path / "source" source.mkdir() @@ -225,3 +276,89 @@ def test_s3_follows_every_listing_page() -> None: assert ( backend.list_objects_v2.call_args_list[1].kwargs["ContinuationToken"] == "next" ) + + +@pytest.mark.parametrize("multipart", [False, True]) +@pytest.mark.parametrize("conflict", [False, True]) +def test_s3_commits_are_conditional( + tmp_path: Path, multipart: bool, conflict: bool +) -> None: + client = boto3.client( + "s3", + region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + writer = LargeFileS3.__new__(LargeFileS3) + writer._client = client + writer.bucket_name = "models" + writer._name = "model" + writer._version = 0 + archive = tmp_path / "archive" + archive.write_bytes(b"model weights") + operation = "CompleteMultipartUpload" if multipart else "PutObject" + headers = [] + + def capture(model: Any, params: Any, **kwargs: Any) -> None: + if model.name == operation: + headers.append(params["headers"].get("If-None-Match")) + + client.meta.events.register_first("before-call.*.*", capture) + config = TransferConfig( + multipart_threshold=1 if multipart else 1024, + use_threads=False, + preferred_transfer_client="classic", + ) + with ( + Stubber(client) as stubber, + patch( + "great_ai.large_file.large_file.large_file_s3.TransferConfig", + return_value=config, + ), + ): + if multipart: + stubber.add_response("create_multipart_upload", {"UploadId": "upload"}) + stubber.add_response("upload_part", {"ETag": '"part"'}) + commit = "complete_multipart_upload" if multipart else "put_object" + if conflict: + stubber.add_client_error( + commit, service_error_code="PreconditionFailed", http_status_code=412 + ) + if multipart: + stubber.add_response("abort_multipart_upload", {}) + with pytest.raises(FileExistsError, match="another writer"): + writer._upload(archive, hide_progress=True) + else: + stubber.add_response(commit, {"ETag": '"model"'}) + writer._upload(archive, hide_progress=True) + stubber.assert_no_pending_responses() + assert headers == ["*"] + client.close() + + +def test_gridfs_rejects_conflicting_versions_and_aborts(tmp_path: Path) -> None: + backend = MagicMock() + stream = backend.open_upload_stream.return_value + stream.__enter__.return_value = stream + stream.__exit__.side_effect = DuplicateKeyError("duplicate filename") + database = MagicMock() + source = tmp_path / "weights" + source.write_text("weights") + with ( + patch.object(LargeFileMongo, "_client", backend), + patch.object(LargeFileMongo, "_database", database), + ): + writer = LargeFileMongo("model", "w") + with pytest.raises(FileExistsError, match="another writer"): + writer.push(source) + database.fs.files.create_index.assert_called_once_with( + "filename", + name="great_ai_unique_version", + unique=True, + partialFilterExpression={"metadata.great_ai_version": True}, + ) + backend.open_upload_stream.assert_called_once_with( + "model_0", metadata={"great_ai_version": True} + ) + stream.abort.assert_called_once() + assert not (writer.cache_path / writer._cache_name).exists()