andras/great-ai
andras/great-ai
Watch
1
0
Fork
You've already forked great-ai
 0
Code Issues Pull requests 2 Projects Releases 4 Packages 1 Wiki Activity Actions

Prevent concurrent model uploads from overwriting versions

Browse source
This commit is contained in:
Andras Schmelczer 2026-10-27 16:55:10 +00:00
commit df83ebc834
6 changed files with 261 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

7
docs/how-to-guides/large-file.md
View file

@ -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

62
great_ai/large_file/large_file/large_file_base.py
View file

@ -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:

50
great_ai/large_file/large_file/large_file_mongo.py
View file

@ -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:

53
great_ai/large_file/large_file/large_file_s3.py
View file

@ -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:

1
pyproject.toml
View file

@ -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]

137
tests/large_file/test_storage_safety.py
View file

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