Merge pull request #2090 from pared/hardlink_cache_check

remote: local: add cache status check hardlink optimization

dvc/remote/base.py

@@ -333,7 +333,7 @@ class RemoteBASE(object):
         logger.debug("'{}' hasn't changed.".format(path_info))
         return False
 
-    def link(self, from_info, to_info):
+    def link(self, from_info, to_info, link_type=None):
         self.copy(from_info, to_info)
 
     def _save_file(self, path_info, checksum, save_link=True):
@@ -385,7 +385,8 @@ class RemoteBASE(object):
     def walk(self, path_info):
         raise NotImplementedError
 
-    def protect(self, path_info):
+    @staticmethod
+    def protect(path_info):
         pass
 
     def save(self, path_info, checksum_info):
@@ -483,9 +484,7 @@ class RemoteBASE(object):
         )
 
     def gc(self, cinfos):
-        from dvc.remote.local import RemoteLOCAL
-
-        used = {info[RemoteLOCAL.PARAM_CHECKSUM] for info in cinfos["local"]}
+        used = self.extract_used_local_checksums(cinfos)
 
         if self.scheme != "":
             used |= {
@@ -539,11 +538,15 @@ class RemoteBASE(object):
         if self.changed_cache_file(checksum):
             return True
 
+        if not self._changed_unpacked_dir(checksum):
+            return False
+
         for entry in self.get_dir_cache(checksum):
-            checksum = entry[self.PARAM_CHECKSUM]
-            if self.changed_cache_file(checksum):
+            entry_checksum = entry[self.PARAM_CHECKSUM]
+            if self.changed_cache_file(entry_checksum):
                 return True
 
+        self._update_unpacked_dir(checksum)
         return False
 
     def changed_cache(self, checksum):
@@ -701,3 +704,19 @@ class RemoteBASE(object):
     @staticmethod
     def unprotect(path_info):
         pass
+
+    def _get_unpacked_dir_names(self, checksums):
+        return set()
+
+    def extract_used_local_checksums(self, cinfos):
+        from dvc.remote import RemoteLOCAL
+
+        used = {info[RemoteLOCAL.PARAM_CHECKSUM] for info in cinfos["local"]}
+        unpacked = self._get_unpacked_dir_names(used)
+        return used | unpacked
+
+    def _changed_unpacked_dir(self, checksum):
+        return True
+
+    def _update_unpacked_dir(self, checksum, progress_callback=None):
+        pass

dvc/remote/local/__init__.py

@@ -35,7 +35,7 @@ from dvc.utils import (
 )
 from dvc.config import Config
 from dvc.exceptions import DvcException
-from dvc.progress import progress
+from dvc.progress import progress, ProgressCallback
 from concurrent.futures import ThreadPoolExecutor
 
 
@@ -48,6 +48,8 @@ class RemoteLOCAL(RemoteBASE):
     PARAM_CHECKSUM = "md5"
     PARAM_PATH = "path"
 
+    UNPACKED_DIR_SUFFIX = ".unpacked"
+
     DEFAULT_CACHE_TYPES = ["reflink", "copy"]
     CACHE_TYPE_MAP = {
         "copy": shutil.copyfile,
@@ -123,47 +125,73 @@ class RemoteLOCAL(RemoteBASE):
         if not self.exists(path_info):
             os.makedirs(path_info.path)
 
-    @slow_link_guard
-    def link(self, cache_info, path_info):
-        cache = cache_info.path
-        path = path_info.path
+    def link(self, from_info, to_info, link_type=None):
+        from_path = from_info.path
+        to_path = to_info.path
 
-        assert os.path.isfile(cache)
+        assert os.path.isfile(from_path)
 
-        dname = os.path.dirname(path)
+        dname = os.path.dirname(to_path)
         if not os.path.exists(dname):
             os.makedirs(dname)
 
         # NOTE: just create an empty file for an empty cache
-        if os.path.getsize(cache) == 0:
-            open(path, "w+").close()
+        if os.path.getsize(from_path) == 0:
+            open(to_path, "w+").close()
 
-            msg = "Created empty file: {} -> {}".format(cache, path)
+            msg = "Created empty file: {} -> {}".format(from_path, to_path)
             logger.debug(msg)
             return
 
-        i = len(self.cache_types)
+        if not link_type:
+            link_types = self.cache_types
+        else:
+            link_types = [link_type]
+
+        self._try_links(from_info, to_info, link_types)
+
+    @classmethod
+    def _get_link_method(cls, link_type):
+        try:
+            return cls.CACHE_TYPE_MAP[link_type]
+        except KeyError:
+            raise DvcException(
+                "Cache type: '{}' not supported!".format(link_type)
+            )
+
+    def _link(self, from_info, to_info, link_method):
+        if os.path.lexists(to_info.path):
+            raise DvcException(
+                "Link '{}' already exists!".format(to_info.path)
+            )
+        else:
+            link_method(from_info.path, to_info.path)
+
+        if self.protected:
+            self.protect(to_info)
+
+        msg = "Created {}'{}': {} -> {}".format(
+            "protected " if self.protected else "",
+            self.cache_types[0],
+            from_info.path,
+            to_info.path,
+        )
+
+        logger.debug(msg)
+
+    @slow_link_guard
+    def _try_links(self, from_info, to_info, link_types):
+        i = len(link_types)
         while i > 0:
+            link_method = self._get_link_method(link_types[0])
             try:
-                self.CACHE_TYPE_MAP[self.cache_types[0]](cache, path)
-
-                if self.protected:
-                    self.protect(path_info)
-
-                msg = "Created {}'{}': {} -> {}".format(
-                    "protected " if self.protected else "",
-                    self.cache_types[0],
-                    cache,
-                    path,
-                )
-
-                logger.debug(msg)
+                self._link(from_info, to_info, link_method)
                 return
 
             except DvcException as exc:
                 msg = "Cache type '{}' is not supported: {}"
-                logger.debug(msg.format(self.cache_types[0], str(exc)))
-                del self.cache_types[0]
+                logger.debug(msg.format(link_types[0], str(exc)))
+                del link_types[0]
                 i -= 1
 
         raise DvcException("no possible cache types left to try out.")
@@ -211,7 +239,8 @@ class RemoteLOCAL(RemoteBASE):
         if path_info.scheme != "local":
             raise NotImplementedError
 
-        remove(path_info.path)
+        if self.exists(path_info):
+            remove(path_info.path)
 
     def move(self, from_info, to_info):
         if from_info.scheme != "local" or to_info.scheme != "local":
@@ -566,3 +595,68 @@ class RemoteLOCAL(RemoteBASE):
     @staticmethod
     def protect(path_info):
         os.chmod(path_info.path, stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH)
+
+    @classmethod
+    def _append_unpacked_suffix(cls, string):
+        return string + cls.UNPACKED_DIR_SUFFIX
+
+    def _get_unpacked_dir_path_info(self, checksum):
+        unpacked_dir_info = self.checksum_to_path_info(checksum)
+        unpacked_dir_info.path = self._append_unpacked_suffix(
+            unpacked_dir_info.path
+        )
+
+        return unpacked_dir_info
+
+    def _path_info_changed(self, path_info):
+        if self.exists(path_info) and self.state.get(path_info):
+            return False
+        return True
+
+    def _update_unpacked_dir(self, checksum):
+        unpacked_dir_info = self._get_unpacked_dir_path_info(checksum)
+
+        if not self._path_info_changed(unpacked_dir_info):
+            return
+
+        self.remove(unpacked_dir_info)
+
+        try:
+            dir_info = self.get_dir_cache(checksum)
+            self._create_unpacked_dir(checksum, dir_info, unpacked_dir_info)
+        except DvcException:
+            logger.warning(
+                "Could not create '{}'".format(unpacked_dir_info.path)
+            )
+
+            self.remove(unpacked_dir_info)
+
+    def _create_unpacked_dir(self, checksum, dir_info, unpacked_dir_info):
+        progress_callback = ProgressCallback(len(dir_info))
+        self.makedirs(unpacked_dir_info)
+        for index, entry in enumerate(dir_info):
+            entry_cache_info = self.checksum_to_path_info(
+                entry[self.PARAM_CHECKSUM]
+            )
+            relpath = entry[self.PARAM_RELPATH]
+
+            unpacked_entry_info = copy(entry_cache_info)
+            unpacked_entry_info.path = self.ospath.join(
+                unpacked_dir_info.path, relpath
+            )
+            self.link(entry_cache_info, unpacked_entry_info, "hardlink")
+
+            progress_callback.update("Creating unpacked dir.")
+        self.state.save(unpacked_dir_info, checksum)
+
+    def _changed_unpacked_dir(self, checksum):
+        status_unpacked_dir_info = self._get_unpacked_dir_path_info(checksum)
+
+        return not self.state.get(status_unpacked_dir_info)
+
+    def _get_unpacked_dir_names(self, checksums):
+        unpacked = set()
+        for c in checksums:
+            if self.is_dir_checksum(c):
+                unpacked.add(self._append_unpacked_suffix(c))
+        return unpacked

tests/func/remote/test_local.py

@@ -0,0 +1,84 @@
+import os
+
+import mock
+import pytest
+
+import dvc
+from dvc.exceptions import DvcException
+from dvc.remote.local import RemoteLOCAL
+from tests.utils import trees_equal
+
+
+class TestUnpackedDirStatusOptimization(object):
+    @pytest.fixture(autouse=True)
+    def setUp(self, dvc_repo, repo_dir):
+        stages = dvc_repo.add(repo_dir.DATA_DIR)
+        assert len(stages) == 1
+        self.dir_out = stages[0].outs[0]
+        self.unpacked_dir = RemoteLOCAL._append_unpacked_suffix(
+            self.dir_out.cache_path
+        )
+
+    def test_should_not_fail_add_on_unpacked_dir_creation_exception(
+        self, dvc_repo
+    ):
+        with mock.patch.object(
+            dvc.remote.local.RemoteLOCAL,
+            "_create_unpacked_dir",
+            side_effect=self._raise_dvc_exception,
+        ) as unpacked_create_spy, dvc_repo.state:
+            assert not dvc_repo.cache.local.changed_cache(
+                self.dir_out.checksum
+            )
+        assert unpacked_create_spy.call_count == 1
+
+    def test_should_remove_unpacked_dir_on_create_exception(self, dvc_repo):
+        # artificial unpacked dir for test purpose
+        os.makedirs(self.unpacked_dir)
+        assert os.path.exists(self.unpacked_dir)
+
+        with mock.patch.object(
+            dvc.remote.local.RemoteLOCAL,
+            "_create_unpacked_dir",
+            side_effect=self._raise_dvc_exception,
+        ), dvc_repo.state:
+            assert not dvc_repo.cache.local.changed_cache(
+                self.dir_out.checksum
+            )
+
+        assert not os.path.exists(self.unpacked_dir)
+
+    def test_should_create_unpacked_dir_on_status_check(
+        self, dvc_repo, repo_dir
+    ):
+        assert not os.path.exists(self.unpacked_dir)
+
+        with dvc_repo.state:
+            assert not dvc_repo.cache.local.changed_cache(
+                self.dir_out.checksum
+            )
+
+        assert os.path.exists(self.unpacked_dir)
+
+        trees_equal(repo_dir.DATA_DIR, self.unpacked_dir)
+
+    def test_should_proceed_with_full_check_on_modified_cache(self, dvc_repo):
+
+        some_file_md5 = self.dir_out.dir_cache[0]["md5"]
+        assert not os.path.exists(self.unpacked_dir)
+
+        with dvc_repo.state:
+            assert not dvc_repo.cache.local.changed_cache(
+                self.dir_out.checksum
+            )
+        assert os.path.exists(self.unpacked_dir)
+
+        with open(dvc_repo.cache.local.get(some_file_md5), "a") as fobj:
+            fobj.write("string")
+
+        with dvc_repo.state:
+            assert dvc_repo.cache.local.changed_cache(self.dir_out.checksum)
+
+    @staticmethod
+    def _raise_dvc_exception(*args):
+        raise DvcException("msg")

tests/func/test_checkout.py

@@ -422,7 +422,11 @@ class TestCheckoutShouldHaveSelfClearingProgressBar(TestDvc):
         self.write_args = [w_c[1][0] for w_c in write_calls]
 
         pattern = re.compile(".*\\[.{30}\\].*%.*")
-        progress_bars = [arg for arg in self.write_args if pattern.match(arg)]
+        progress_bars = [
+            arg
+            for arg in self.write_args
+            if pattern.match(arg) and "unpacked" not in arg
+        ]
 
         update_bars = progress_bars[:-1]
         finish_bar = progress_bars[-1]

tests/utils/__init__.py

@@ -1,4 +1,5 @@
 import os
+from filecmp import dircmp
 
 from dvc.scm import Git
 from mock import MagicMock
@@ -29,3 +30,13 @@ def cd(newdir):
         yield
     finally:
         os.chdir(prevdir)
+
+
+def trees_equal(dir_path_1, dir_path_2):
+
+    comparison = dircmp(dir_path_1, dir_path_2)
+
+    assert set(comparison.left_only) == set(comparison.right_only) == set()
+
+    for d in comparison.common_dirs:
+        trees_equal(os.path.join(dir_path_1, d), os.path.join(dir_path_2, d))