364 lines
13 KiB
Python
364 lines
13 KiB
Python
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)
|
|
def storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
|
for implementation in (LargeFileBase, LargeFileLocal):
|
|
monkeypatch.setattr(implementation, "cache_path", tmp_path / "cache")
|
|
monkeypatch.setattr(implementation, "max_cache_size", None)
|
|
yield
|
|
|
|
|
|
def test_configured_cache_directory_is_created(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
monkeypatch.setattr(LargeFileLocal, "cache_path", tmp_path / "custom-cache")
|
|
with LargeFileLocal("model", "w") as file:
|
|
file.write("weights")
|
|
assert LargeFileLocal("model").get().read_text() == "weights"
|
|
|
|
|
|
def test_unrelated_cache_entries_are_ignored() -> None:
|
|
LargeFileLocal.cache_path.mkdir(parents=True)
|
|
(LargeFileLocal.cache_path / "model-unrelated").write_text("unrelated")
|
|
(LargeFileLocal.cache_path / "model-0").write_text("weights")
|
|
assert LargeFileLocal("model").get().read_text() == "weights"
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["delete", "prune"])
|
|
def test_cached_directories_can_be_removed(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str
|
|
) -> None:
|
|
source = tmp_path / "source"
|
|
source.mkdir()
|
|
(source / "weights").write_text("weights")
|
|
LargeFileLocal("model", "w").push(source)
|
|
reader = LargeFileLocal("model")
|
|
cached = reader.get()
|
|
if operation == "delete":
|
|
reader.delete()
|
|
else:
|
|
monkeypatch.setattr(LargeFileLocal, "max_cache_size", "0")
|
|
reader.clean_up()
|
|
assert not cached.exists()
|
|
|
|
|
|
def test_cache_quota_counts_directory_contents(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
source = tmp_path / "source"
|
|
source.mkdir()
|
|
(source / "weights").write_bytes(b"x" * 100_000)
|
|
monkeypatch.setattr(LargeFileLocal, "max_cache_size", "1 KB")
|
|
LargeFileLocal("model", "w").push(source)
|
|
assert not (LargeFileLocal.cache_path / "model-0").exists()
|
|
|
|
|
|
def test_failed_write_removes_temporary_file() -> None:
|
|
temporary = None
|
|
try:
|
|
with pytest.raises(ValueError, match="failed"):
|
|
with LargeFileLocal("model", "w") as file:
|
|
temporary = Path(file.name)
|
|
file.write("unfinished")
|
|
raise ValueError("failed")
|
|
assert temporary is not None
|
|
assert not temporary.exists()
|
|
finally:
|
|
if temporary is not None:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
def test_evicted_remote_model_is_downloaded_again() -> None:
|
|
backend = MagicMock()
|
|
backend.list_objects_v2.return_value = {"Contents": [{"Key": "model/0"}]}
|
|
downloads = []
|
|
|
|
def download(
|
|
self: LargeFileS3, remote_path: Any, local_path: Path, hide_progress: bool
|
|
) -> None:
|
|
downloads.append(remote_path)
|
|
with tarfile.open(local_path, "w:gz") as archive:
|
|
payload = b"remote weights"
|
|
member = tarfile.TarInfo("model-0")
|
|
member.size = len(payload)
|
|
archive.addfile(member, io.BytesIO(payload))
|
|
|
|
with (
|
|
patch.object(LargeFileS3, "_client", backend),
|
|
patch.object(LargeFileS3, "_download", download),
|
|
):
|
|
reader = LargeFileS3("model")
|
|
cached = reader.get()
|
|
cached.unlink()
|
|
assert reader.get().read_text() == "remote weights"
|
|
assert downloads == ["model/0", "model/0"]
|
|
|
|
|
|
def test_upload_refreshes_available_versions() -> None:
|
|
writer = LargeFileLocal("model", "w")
|
|
with writer as file:
|
|
file.write("weights")
|
|
assert writer.versions_pretty == "0"
|
|
assert writer.get().read_text() == "weights"
|
|
|
|
|
|
def test_retention_counts_the_new_version() -> None:
|
|
for version in range(3):
|
|
with LargeFileLocal("model", "w", keep_last_n=1) as file:
|
|
file.write(str(version))
|
|
reader = LargeFileLocal("model")
|
|
assert reader.versions_pretty == "2"
|
|
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()
|
|
(source / "weights").write_text("model weights")
|
|
LargeFileLocal("model", "w").push(source)
|
|
reader = LargeFileLocal("model")
|
|
cached = reader.get()
|
|
reader.push(cached)
|
|
assert reader.version == 1
|
|
assert (cached / "weights").read_text() == "model weights"
|
|
assert (
|
|
LargeFileLocal("model", version=1).get() / "weights"
|
|
).read_text() == "model weights"
|
|
|
|
|
|
def test_gridfs_matches_only_exact_model_names() -> None:
|
|
files = [
|
|
SimpleNamespace(name=name, _id=name, length=1)
|
|
for name in ("model_0", "model_other_9", "other_model_10", "model_1_backup")
|
|
]
|
|
backend = MagicMock()
|
|
backend.find.side_effect = lambda query: [
|
|
file for file in files if query["filename"].search(file.name)
|
|
]
|
|
with patch.object(LargeFileMongo, "_client", backend):
|
|
reader = LargeFileMongo("model")
|
|
assert reader.version == 0
|
|
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 = [
|
|
{
|
|
"Contents": [{"Key": "model/0"}, {"Key": "model-other/99"}],
|
|
"IsTruncated": True,
|
|
"NextContinuationToken": "next",
|
|
},
|
|
{
|
|
"Contents": [
|
|
{"Key": "model/9"},
|
|
{"Key": "model/backup"},
|
|
{"Key": "model/9/extra"},
|
|
]
|
|
},
|
|
]
|
|
with patch.object(LargeFileS3, "_client", backend):
|
|
reader = LargeFileS3("model")
|
|
assert reader.version == 9
|
|
assert [instance.version for instance in reader._instances] == [0, 9]
|
|
assert backend.list_objects_v2.call_args_list[0].kwargs["Prefix"] == "model/"
|
|
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()
|