Prevent concurrent model uploads from overwriting versions

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:

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:

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: