Serialize cache downloads and uploads with the same model lock

great_ai/large_file/large_file/large_file_base.py

@@ -178,7 +178,10 @@ class LargeFileBase(ABC):
         Args:
             hide_progress: Do not show a progress update after each 10% of progress.
         """
+        with self._model_lock():
+            return self._get_cached_file(hide_progress)
 
+    def _get_cached_file(self, hide_progress: bool) -> Path:
         if not self._cache_only_mode:
             self._find_instances()
         instance = next(
@@ -233,12 +236,9 @@ class LargeFileBase(ABC):
         if isinstance(path, str):
             path = Path(path)
 
-        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"):
+        with self._model_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
@@ -306,6 +306,12 @@ class LargeFileBase(ABC):
             f"{type(self).__module__}.{type(self).__qualname__}".encode()
         ).hexdigest()
 
+    def _model_lock(self) -> FileLock:
+        directory = self.cache_path / ".locks"
+        directory.mkdir(parents=True, exist_ok=True)
+        key = sha256(f"{self._cache_namespace}:{self._name}".encode()).hexdigest()
+        return FileLock(directory / f"{key}.lock")
+
     @property
     def _cache_identity_path(self) -> Path:
         key = sha256(self._cache_name.encode()).hexdigest()

tests/large_file/test_cache_generations.py

@@ -1,3 +1,7 @@
+import io
+import tarfile
+import time
+from concurrent.futures import ThreadPoolExecutor
 from hashlib import sha256
 from pathlib import Path
 from typing import Any, Dict
@@ -44,3 +48,39 @@ def test_recreated_remote_version_invalidates_other_clients(tmp_path: Path) -> N
             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"]