diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 30f13a8f4..0074351f8 100644 --- a/dvc/remote/base.py +++ b/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 diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index 3e314da4f..3e50591d2 100644 --- a/dvc/remote/local/__init__.py +++ b/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 diff --git a/tests/func/remote/__init__.py b/tests/func/remote/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/func/remote/test_local.py b/tests/func/remote/test_local.py new file mode 100644 index 000000000..e26b2061f --- /dev/null +++ b/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") diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py index b9ed9d264..2ac817a9f 100644 --- a/tests/func/test_checkout.py +++ b/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] diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index 3bbe0c3c3..b43af8d06 100644 --- a/tests/utils/__init__.py +++ b/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))