Preserve cached directory sources during uploads

great_ai/large_file/large_file/large_file_base.py

@@ -213,37 +213,31 @@ class LargeFileBase(ABC):
 
         if isinstance(path, str):
             path = Path(path)
+        if "w" not in self._mode:
+            self._find_instances()
+            self._version = self._instances[-1].version + 1 if self._instances else 0
 
         with tempfile.TemporaryDirectory() as tmp:
-            if path.is_file():
+            root = Path(tmp)
-                logger.info(f"Copying file for {self._local_name}")
+            staged = root / self._local_name
-                copy: Any = shutil.copy
+            if path.is_dir():
+                shutil.copytree(path, staged)
             else:
-                logger.info(f"Copying directory for {self._local_name}")
+                shutil.copy2(path, staged)
-                copy = shutil.copytree
 
-            try:
+            with tempfile.TemporaryDirectory() as archive_directory:
-                # Make local copy in the cache
+                archive = shutil.make_archive(
-                shutil.rmtree(self.cache_path / self._cache_name, ignore_errors=True)
+                    str(Path(archive_directory) / self._local_name),
-                copy(str(path), str(self.cache_path / self._cache_name))
-            except shutil.SameFileError:
-                pass  # No worries
-
-            copy(str(path), str(Path(tmp) / self._local_name))
-
-            with tempfile.TemporaryDirectory() as tmp2:
-                # A directory has to be zipped and it cannot contain the output of the zipping
-                logger.info(f"Compressing {self._local_name}")
-                shutil.make_archive(
-                    str(Path(tmp2) / self._local_name),
                     COMPRESSION_ALGORITHM,
-                    tmp,
+                    root,
                 )
+                self._upload(Path(archive), hide_progress=hide_progress)
 
-                file_to_be_uploaded = (
+            # Stage the source before publishing: it may itself be cached.
-                    Path(tmp2) / f"{self._local_name}{ARCHIVE_EXTENSION}"
+            destination = self.cache_path / self._cache_name
-                )
+            if destination.exists():
-                self._upload(file_to_be_uploaded, hide_progress=hide_progress)
+                raise FileExistsError(f"Version already exists: {destination}")
+            shutil.move(str(staged), str(destination))
 
         self.clean_up()
 

tests/large_file/test_storage_safety.py

@@ -108,6 +108,21 @@ def test_evicted_remote_model_is_downloaded_again() -> None:
     assert downloads == ["model/0", "model/0"]
 
 
+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)