experiments: preserve file permissions across tmpdir execution (#4890)

* tree: use file_mode in upload() for trees that support modes

* experiments: preserve file_mode for files in the workspace

* skip permissions tests on windows

* LocalTree: separate chmod from protect

* fix rebase issue

dvc/repo/experiments/__init__.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import re
+import stat
 import tempfile
 import threading
 from collections import namedtuple
@@ -795,7 +796,6 @@ class Experiments:
             func = partial(
                 _log_exceptions(src_tree.download, "download"),
                 dir_mode=dest_tree.dir_mode,
-                file_mode=dest_tree.file_mode,
             )
             desc = "Downloading"
         else:
@@ -807,7 +807,23 @@ class Experiments:
             # TODO: parallelize this, currently --jobs for repro applies to
             # number of repro executors not download threads
             with ThreadPoolExecutor(max_workers=1) as dl_executor:
-                fails = sum(dl_executor.map(func, from_infos, to_infos, names))
+                mode = None
+                stat_func = getattr(src_tree, "stat", None)
+                futures = []
+                for from_info, to_info, name in zip(
+                    from_infos, to_infos, names
+                ):
+                    if stat_func:
+                        mode = stat.S_IMODE(stat_func(from_info).st_mode)
+                    futures.append(
+                        dl_executor.submit(
+                            func, from_info, to_info, name, file_mode=mode
+                        )
+                    )
+
+                fails = sum(
+                    future.result() for future in as_completed(futures)
+                )
 
         if fails:
             if download:

dvc/tree/base.py

@@ -338,7 +338,14 @@ class BaseTree:
         hash_info.size = dir_info.size
         return hash_info
 
-    def upload(self, from_info, to_info, name=None, no_progress_bar=False):
+    def upload(
+        self,
+        from_info,
+        to_info,
+        name=None,
+        no_progress_bar=False,
+        file_mode=None,
+    ):
         if not hasattr(self, "_upload"):
             raise RemoteActionNotImplemented("upload", self.scheme)
 
@@ -357,6 +364,7 @@ class BaseTree:
             to_info,
             name=name,
             no_progress_bar=no_progress_bar,
+            file_mode=file_mode,
         )
 
     def download(

dvc/tree/gdrive.py

@@ -564,7 +564,9 @@ class GDriveTree(BaseTree):
     def get_file_hash(self, path_info):
         raise NotImplementedError
 
-    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
         dirname = to_info.parent
         assert dirname
         parent_id = self._get_item_id(dirname, True)

dvc/tree/gs.py

@@ -200,7 +200,9 @@ class GSTree(BaseTree):
             size=blob.size,
         )
 
-    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
         bucket = self.gs.bucket(to_info.bucket)
         _upload_to_bucket(
             bucket,

dvc/tree/http.py

@@ -181,7 +181,9 @@ class HTTPTree(BaseTree):  # pylint:disable=abstract-method
                 for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE):
                     fd_wrapped.write(chunk)
 
-    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
         def chunks():
             with open(from_file, "rb") as fd:
                 with Tqdm.wrapattr(

dvc/tree/local.py

@@ -239,6 +239,24 @@ class LocalTree(BaseTree):
         os.chmod(tmp_info, self.file_mode)
         os.rename(tmp_info, to_info)
 
+    def chmod(self, path_info, mode):
+        path = os.fspath(path_info)
+        try:
+            os.chmod(path, mode)
+        except OSError as exc:
+            # There is nothing we need to do in case of a read-only file system
+            if exc.errno == errno.EROFS:
+                return
+
+            # In shared cache scenario, we might not own the cache file, so we
+            # need to check if cache file is already protected.
+            if exc.errno not in [errno.EPERM, errno.EACCES]:
+                raise
+
+            actual = stat.S_IMODE(os.stat(path).st_mode)
+            if actual != mode:
+                raise
+
     def _unprotect_file(self, path):
         if System.is_symlink(path) or System.is_hardlink(path):
             logger.debug(f"Unprotecting '{path}'")
@@ -276,24 +294,7 @@ class LocalTree(BaseTree):
             self._unprotect_file(path)
 
     def protect(self, path_info):
-        path = os.fspath(path_info)
-        mode = self.CACHE_MODE
-
-        try:
-            os.chmod(path, mode)
-        except OSError as exc:
-            # There is nothing we need to do in case of a read-only file system
-            if exc.errno == errno.EROFS:
-                return
-
-            # In shared cache scenario, we might not own the cache file, so we
-            # need to check if cache file is already protected.
-            if exc.errno not in [errno.EPERM, errno.EACCES]:
-                raise
-
-            actual = stat.S_IMODE(os.stat(path).st_mode)
-            if actual != mode:
-                raise
+        self.chmod(path_info, self.CACHE_MODE)
 
     def is_protected(self, path_info):
         try:
@@ -316,7 +317,13 @@ class LocalTree(BaseTree):
         return os.path.getsize(path_info)
 
     def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+        self,
+        from_file,
+        to_info,
+        name=None,
+        no_progress_bar=False,
+        file_mode=None,
+        **_kwargs,
     ):
         makedirs(to_info.parent, exist_ok=True)
 
@@ -325,7 +332,10 @@ class LocalTree(BaseTree):
             from_file, tmp_file, name=name, no_progress_bar=no_progress_bar
         )
 
-        self.protect(tmp_file)
+        if file_mode is not None:
+            self.chmod(tmp_file, file_mode)
+        else:
+            self.protect(tmp_file)
         os.replace(tmp_file, to_info)
 
     @staticmethod

dvc/tree/s3.py

@@ -342,7 +342,9 @@ class S3Tree(BaseTree):
                 size=obj.content_length,
             )
 
-    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
         with self._get_obj(to_info) as obj:
             total = os.path.getsize(from_file)
             with Tqdm(

dvc/tree/ssh/__init__.py

@@ -260,7 +260,9 @@ class SSHTree(BaseTree):
                 no_progress_bar=no_progress_bar,
             )
 
-    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
+    def _upload(
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    ):
         with self.ssh(to_info) as ssh:
             ssh.upload(
                 from_file,

dvc/tree/webdav.py

@@ -211,7 +211,7 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
 
     # Uploads file to remote
     def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False,
+        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs,
     ):
         # First try to create parent directories
         self.makedirs(to_info.parent)

tests/func/experiments/test_experiments.py

@@ -1,4 +1,6 @@
 import logging
+import os
+import stat
 
 import pytest
 from funcy import first
@@ -19,6 +21,18 @@ def test_new_simple(tmp_dir, scm, dvc, exp_stage, mocker):
     ).read_text() == "foo: 2"
 
 
+@pytest.mark.skipif(os.name == "nt", reason="Not supported for Windows.")
+def test_file_permissions(tmp_dir, scm, dvc, exp_stage, mocker):
+    mode = 0o755
+    os.chmod(tmp_dir / "copy.py", mode)
+    scm.add(["copy.py"])
+    scm.commit("set exec")
+
+    tmp_dir.gen("params.yaml", "foo: 2")
+    dvc.experiments.run(exp_stage.addressing)
+    assert stat.S_IMODE(os.stat(tmp_dir / "copy.py").st_mode) == mode
+
+
 def test_failed_exp(tmp_dir, scm, dvc, exp_stage, mocker, caplog):
     from dvc.stage.exceptions import StageCmdFailedError
 

tests/unit/tree/test_tree.py

@@ -1,8 +1,11 @@
+import os
+import stat
+
 import pytest
 
 from dvc.config import ConfigError
-from dvc.path_info import CloudURLInfo
-from dvc.tree import get_cloud_tree
+from dvc.path_info import CloudURLInfo, PathInfo
+from dvc.tree import LocalTree, get_cloud_tree
 
 
 def test_get_cloud_tree(tmp_dir, dvc):
@@ -46,3 +49,18 @@ def test_get_cloud_tree_validate(tmp_dir, dvc):
 
     with pytest.raises(ConfigError):
         get_cloud_tree(dvc, name="second")
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Not supported for Windows.")
+@pytest.mark.parametrize(
+    "mode, expected", [(None, LocalTree.CACHE_MODE), (0o777, 0o777)],
+)
+def test_upload_file_mode(tmp_dir, mode, expected):
+    tmp_dir.gen("src", "foo")
+    src = PathInfo(tmp_dir / "src")
+    dest = PathInfo(tmp_dir / "dest")
+    tree = LocalTree(None, {"url": os.fspath(tmp_dir)})
+    tree.upload(src, dest, file_mode=mode)
+    assert (tmp_dir / "dest").exists()
+    assert (tmp_dir / "dest").read_text() == "foo"
+    assert stat.S_IMODE(os.stat(dest).st_mode) == expected