Prevent concurrent model uploads from overwriting versions

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

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:
+                with tempfile.TemporaryDirectory() as archive_directory:
-            root = Path(tmp)
+                    archive = shutil.make_archive(
-            staged = root / self._local_name
+                        str(Path(archive_directory) / self._local_name),
-            if path.is_dir():
+                        COMPRESSION_ALGORITHM,
-                shutil.copytree(path, staged)
+                        root,
-            else:
+                    )
-                shutil.copy2(path, staged)
+                    # Remote backends must reject version conflicts atomically.
+                    self._upload(Path(archive), hide_progress=hide_progress)
 
-            with tempfile.TemporaryDirectory() as archive_directory:
+                # Stage the source before publishing: it may itself be cached.
-                archive = shutil.make_archive(
+                destination = self.cache_path / self._cache_name
-                    str(Path(archive_directory) / self._local_name),
+                if destination.exists():
-                    COMPRESSION_ALGORITHM,
+                    raise FileExistsError(f"Version already exists: {destination}")
-                    root,
+                shutil.move(str(staged), str(destination))
-                )
-                self._upload(Path(archive), hide_progress=hide_progress)
 
-            # Stage the source before publishing: it may itself be cached.
+            self._find_instances()
-            destination = self.cache_path / self._cache_name
+            self.clean_up()
-            if destination.exists():
-                raise FileExistsError(f"Version already exists: {destination}")
-            shutil.move(str(staged), str(destination))
-
-        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:

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 MongoClient(self.mongo_connection_string)[self.mongo_database]
-        return GridFSBucket(db)
+
+    @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(
+        # Only GreatAI's versioned uploads participate in this unique index;
-            f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}"
+        # unrelated GridFS users may still store duplicate filenames.
-        ) as stream:
+        self._database.fs.files.create_index(
-            with open(local_path, "rb") as f:
+            "filename",
-                while True:
+            name="great_ai_unique_version",
-                    content = f.read(DEFAULT_CHUNK_SIZE)
+            unique=True,
-                    stream.write(content)
+            partialFilterExpression={"metadata.great_ai_version": True},
-
+        )
-                    if progress:
+        filename = f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}"
-                        progress(len(content))
+        stream = self._client.open_upload_stream(
-                    if len(content) < DEFAULT_CHUNK_SIZE:
+            filename, metadata={"great_ai_version": True}
-                        break
+        )
+        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:

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(
+        def prevent_overwrite(params: Dict[str, Any], **_: Any) -> None:
-            Filename=str(local_path),
+            if params.get("Bucket") == self.bucket_name and params.get("Key") == key:
-            Bucket=self.bucket_name,
+                params["IfNoneMatch"] = "*"
-            Key=key,
+
-            Callback=(
+        # S3Transfer does not accept IfNoneMatch in ExtraArgs. These botocore
-                None
+        # hooks protect both single-part and multipart commits while retaining
-                if hide_progress
+        # the managed transfer's streaming, retries, and progress reporting.
-                else UploadProgressBar(path=local_path, logger=logger)
+        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:

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]

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