86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
import io
|
|
import tarfile
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import Any, Dict
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from great_ai.large_file import LargeFileS3
|
|
|
|
|
|
def test_recreated_remote_version_invalidates_other_clients(tmp_path: Path) -> None:
|
|
remote: Dict[str, bytes] = {}
|
|
backend = MagicMock()
|
|
backend.list_objects_v2.side_effect = lambda **kwargs: {
|
|
"Contents": [
|
|
{"Key": key, "ETag": sha256(value).hexdigest()}
|
|
for key, value in remote.items()
|
|
]
|
|
}
|
|
backend.delete_object.side_effect = lambda **kwargs: remote.pop(kwargs["Key"])
|
|
|
|
def upload(self: LargeFileS3, path: Path, hide_progress: bool) -> None:
|
|
remote[f"{self._name}/{self.version}"] = path.read_bytes()
|
|
|
|
def download(
|
|
self: LargeFileS3, remote_path: Any, path: Path, hide_progress: bool
|
|
) -> None:
|
|
path.write_bytes(remote[remote_path])
|
|
|
|
with (
|
|
patch.object(LargeFileS3, "_client", backend),
|
|
patch.object(LargeFileS3, "_upload", upload),
|
|
patch.object(LargeFileS3, "_download", download),
|
|
patch.object(LargeFileS3, "max_cache_size", None),
|
|
):
|
|
with patch.object(LargeFileS3, "cache_path", tmp_path / "writer"):
|
|
with LargeFileS3("model", "w") as file:
|
|
file.write("old")
|
|
with patch.object(LargeFileS3, "cache_path", tmp_path / "reader"):
|
|
reader = LargeFileS3("model")
|
|
assert reader.get().read_text() == "old"
|
|
with patch.object(LargeFileS3, "cache_path", tmp_path / "writer"):
|
|
LargeFileS3("model").delete()
|
|
with LargeFileS3("model", "w") as file:
|
|
file.write("new")
|
|
assert LargeFileS3("model").get().read_text() == "new"
|
|
with patch.object(LargeFileS3, "cache_path", tmp_path / "reader"):
|
|
assert reader.get().read_text() == "new"
|
|
|
|
|
|
def test_concurrent_downloads_publish_one_directory(tmp_path: Path) -> None:
|
|
backend = MagicMock()
|
|
backend.list_objects_v2.return_value = {
|
|
"Contents": [{"Key": "model/0", "ETag": "one"}]
|
|
}
|
|
downloads = []
|
|
|
|
def download(
|
|
self: LargeFileS3, remote_path: Any, path: Path, hide_progress: bool
|
|
) -> None:
|
|
downloads.append(remote_path)
|
|
time.sleep(0.1)
|
|
with tarfile.open(path, "w:gz") as archive:
|
|
folder = tarfile.TarInfo("model-0")
|
|
folder.type = tarfile.DIRTYPE
|
|
folder.mode = 0o755
|
|
archive.addfile(folder)
|
|
file = tarfile.TarInfo("model-0/weights")
|
|
file.size = 7
|
|
file.mode = 0o644
|
|
archive.addfile(file, io.BytesIO(b"weights"))
|
|
|
|
with (
|
|
patch.object(LargeFileS3, "_client", backend),
|
|
patch.object(LargeFileS3, "_download", download),
|
|
patch.object(LargeFileS3, "cache_path", tmp_path / "cache"),
|
|
ThreadPoolExecutor(max_workers=2) as pool,
|
|
):
|
|
readers = [LargeFileS3("model") for _ in range(2)]
|
|
results = list(pool.map(lambda reader: reader.get(), readers))
|
|
assert results[0] == results[1]
|
|
assert list(results[0].iterdir()) == [results[0] / "weights"]
|
|
assert (results[0] / "weights").read_text() == "weights"
|
|
assert downloads == ["model/0"]
|