From 33b0eefec07690c6ad03f0110577e75039750fa7 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Mon, 21 Oct 2019 16:12:56 +0700 Subject: [PATCH] dvc: dedup used cache (#2630) * dvc: dedup used cache This is a preparation step to make garbage collector work with more commits. It has a value of its own: - less exists calls to show cloud status or prepare push/pull - same names are not repeated in status and pbar texts * dvc: delay checksum names concatenation --- dvc/cache.py | 27 +++++++++++ dvc/data_cloud.py | 20 ++++---- dvc/output/base.py | 33 ++++++-------- dvc/remote/base.py | 15 ++---- dvc/remote/local.py | 43 +++++------------ dvc/repo/__init__.py | 19 ++------ dvc/repo/fetch.py | 6 +-- dvc/repo/gc.py | 81 ++++++++------------------------- dvc/repo/push.py | 2 +- dvc/repo/status.py | 2 +- tests/func/test_data_cloud.py | 36 +++++++-------- tests/unit/remote/test_local.py | 14 ++---- 12 files changed, 121 insertions(+), 177 deletions(-) diff --git a/dvc/cache.py b/dvc/cache.py index 5c3e90c6c..8379a3a0f 100644 --- a/dvc/cache.py +++ b/dvc/cache.py @@ -3,6 +3,7 @@ from __future__ import unicode_literals import os +from collections import defaultdict from funcy import cached_property from dvc.utils.compat import builtin_str @@ -103,3 +104,29 @@ class Cache(object): ssh = _make_remote_property(Config.SECTION_CACHE_SSH) hdfs = _make_remote_property(Config.SECTION_CACHE_HDFS) azure = _make_remote_property(Config.SECTION_CACHE_AZURE) + + +class NamedCache(object): + def __init__(self): + self._items = defaultdict(lambda: defaultdict(set)) + self.repo = [] + + @classmethod + def make(cls, scheme, checksum, name): + cache = cls() + cache.add(scheme, checksum, name) + return cache + + def __getitem__(self, key): + return self._items[key] + + def add(self, scheme, checksum, name): + self._items[scheme][checksum].add(name) + + def update(self, cache, suffix=""): + for scheme, src in cache._items.items(): + dst = self._items[scheme] + for checksum, names in src.items(): + dst[checksum].update(names) + + self.repo.extend(cache.repo) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index af86efd64..63669f4b2 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -45,11 +45,11 @@ class DataCloud(object): def _init_remote(self, remote): return Remote(self.repo, name=remote) - def push(self, targets, jobs=None, remote=None, show_checksums=False): + def push(self, cache, jobs=None, remote=None, show_checksums=False): """Push data items in a cloud-agnostic way. Args: - targets (list): list of targets to push to the cloud. + cache (NamedCache): named checksums to push to the cloud. jobs (int): number of jobs that can be running simultaneously. remote (dvc.remote.base.RemoteBASE): optional remote to push to. By default remote from core.remote config option is used. @@ -57,17 +57,17 @@ class DataCloud(object): information messages. """ return self.repo.cache.local.push( - targets, + cache, jobs=jobs, remote=self.get_remote(remote, "push"), show_checksums=show_checksums, ) - def pull(self, targets, jobs=None, remote=None, show_checksums=False): + def pull(self, cache, jobs=None, remote=None, show_checksums=False): """Pull data items in a cloud-agnostic way. Args: - targets (list): list of targets to pull from the cloud. + cache (NamedCache): named checksums to pull from the cloud. jobs (int): number of jobs that can be running simultaneously. remote (dvc.remote.base.RemoteBASE): optional remote to pull from. By default remote from core.remote config option is used. @@ -75,25 +75,25 @@ class DataCloud(object): information messages. """ return self.repo.cache.local.pull( - targets, + cache, jobs=jobs, remote=self.get_remote(remote, "pull"), show_checksums=show_checksums, ) - def status(self, targets, jobs=None, remote=None, show_checksums=False): + def status(self, cache, jobs=None, remote=None, show_checksums=False): """Check status of data items in a cloud-agnostic way. Args: - targets (list): list of targets to check status for. + cache (NamedCache): named checksums to check status for. jobs (int): number of jobs that can be running simultaneously. remote (dvc.remote.base.RemoteBASE): optional remote to compare - targets to. By default remote from core.remote config option + cache to. By default remote from core.remote config option is used. show_checksums (bool): show checksums instead of file names in information messages. """ remote = self.get_remote(remote, "status") return self.repo.cache.local.status( - targets, jobs=jobs, remote=remote, show_checksums=show_checksums + cache, jobs=jobs, remote=remote, show_checksums=show_checksums ) diff --git a/dvc/output/base.py b/dvc/output/base.py index 44220d476..699ceb3d3 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -5,6 +5,7 @@ from copy import copy from schema import Or, Optional +from dvc.cache import NamedCache import dvc.prompt as prompt from dvc.exceptions import DvcException from dvc.utils.compat import str, urlparse @@ -345,12 +346,11 @@ class OutputBase(object): $ echo "bar" > directory/bar $ dvc add directory - It will return something similar to the following list: + It will return a NamedCache like: - [ - { 'path': 'directory/foo', 'md5': 'c157a79031e1', ... }, - { 'path': 'directory/bar', 'md5': 'd3b07384d113', ... }, - ] + nc = NamedCache() + nc.add(self.scheme, 'c157a79031e1', 'directory/foo') + nc.add(self.scheme, 'd3b07384d113', 'directory/bar') """ ret = [] @@ -358,12 +358,7 @@ class OutputBase(object): if self.cache.changed_cache_file(self.checksum): try: self.repo.cloud.pull( - [ - { - self.remote.PARAM_CHECKSUM: self.checksum, - "name": str(self), - } - ], + NamedCache.make("local", self.checksum, str(self)), jobs=jobs, remote=remote, show_checksums=False, @@ -385,13 +380,13 @@ class OutputBase(object): else: return ret + cache = NamedCache() for entry in self.dir_cache: - info = copy(entry) + checksum = entry[self.remote.PARAM_CHECKSUM] path_info = self.path_info / entry[self.remote.PARAM_RELPATH] - info["name"] = str(path_info) - ret.append(info) + cache.add(self.scheme, checksum, str(path_info)) - return ret + return cache def get_used_cache(self, **kwargs): """Get a dumpd of the given `out`, with an entry including the branch. @@ -403,7 +398,7 @@ class OutputBase(object): """ if not self.use_cache: - return [] + return NamedCache() if not self.info: logger.warning( @@ -411,14 +406,14 @@ class OutputBase(object): "not be collected. Use dvc repro to get your pipeline up to " "date.".format(self, self.stage) ) - return [] + return NamedCache() - ret = [{self.remote.PARAM_CHECKSUM: self.checksum, "name": str(self)}] + ret = NamedCache.make(self.scheme, self.checksum, str(self)) if not self.is_dir_checksum: return ret - ret.extend(self._collect_used_dir_cache(**kwargs)) + ret.update(self._collect_used_dir_cache(**kwargs)) return ret diff --git a/dvc/remote/base.py b/dvc/remote/base.py index ca95a4e19..5dd2ca1b6 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -620,14 +620,11 @@ class RemoteBASE(object): # We ignore all the non-cache looking files pass - def gc(self, cinfos): - used = self.extract_used_local_checksums(cinfos) + def gc(self, named_cache): + used = self.extract_used_local_checksums(named_cache) if self.scheme != "": - used |= { - info[self.PARAM_CHECKSUM] - for info in cinfos.get(self.scheme, []) - } + used.update(named_cache[self.scheme]) removed = False for checksum in self.all(): @@ -909,10 +906,8 @@ class RemoteBASE(object): 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"]} + def extract_used_local_checksums(self, named_cache): + used = set(named_cache["local"]) unpacked = self._get_unpacked_dir_names(used) return used | unpacked diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 384077180..e187f33b3 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -245,31 +245,9 @@ class RemoteLOCAL(RemoteBASE): def open(path_info, mode="r", encoding=None): return open(fspath_py35(path_info), mode=mode, encoding=encoding) - def _group(self, checksum_infos, show_checksums=False): - by_md5 = {} - - for info in checksum_infos: - md5 = info[self.PARAM_CHECKSUM] - - if show_checksums: - by_md5[md5] = {"name": md5} - continue - - name = info["name"] - branch = info.get("branch") - if branch: - name += "({})".format(branch) - - if md5 not in by_md5.keys(): - by_md5[md5] = {"name": name} - else: - by_md5[md5]["name"] += " " + name - - return by_md5 - def status( self, - checksum_infos, + named_cache, remote, jobs=None, show_checksums=False, @@ -278,8 +256,7 @@ class RemoteLOCAL(RemoteBASE): logger.debug( "Preparing to collect status from {}".format(remote.path_info) ) - ret = self._group(checksum_infos, show_checksums=show_checksums) or {} - md5s = list(ret) + md5s = list(named_cache[self.scheme]) logger.debug("Collecting information from local cache...") local_exists = self.cache_exists(md5s, jobs=jobs, name=self.cache_dir) @@ -298,6 +275,10 @@ class RemoteLOCAL(RemoteBASE): ) ) + ret = { + checksum: {"name": checksum if show_checksums else " ".join(names)} + for checksum, names in named_cache[self.scheme].items() + } self._fill_statuses(ret, local_exists, remote_exists) self._log_missing_caches(ret) @@ -337,7 +318,7 @@ class RemoteLOCAL(RemoteBASE): def _process( self, - checksum_infos, + named_cache, remote, jobs=None, show_checksums=False, @@ -365,7 +346,7 @@ class RemoteLOCAL(RemoteBASE): jobs = remote.JOBS status_info = self.status( - checksum_infos, + named_cache, remote, jobs=jobs, show_checksums=show_checksums, @@ -390,18 +371,18 @@ class RemoteLOCAL(RemoteBASE): return len(plans[0]) - def push(self, checksum_infos, remote, jobs=None, show_checksums=False): + def push(self, named_cache, remote, jobs=None, show_checksums=False): return self._process( - checksum_infos, + named_cache, remote, jobs=jobs, show_checksums=show_checksums, download=False, ) - def pull(self, checksum_infos, remote, jobs=None, show_checksums=False): + def pull(self, named_cache, remote, jobs=None, show_checksums=False): return self._process( - checksum_infos, + named_cache, remote, jobs=jobs, show_checksums=show_checksums, diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 36fd97d59..c632fc1b7 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -231,15 +231,9 @@ class Repo(object): A dictionary with Schemes (representing output's location) as keys, and a list with the outputs' `dumpd` as values. """ + from dvc.cache import NamedCache - cache = {} - cache["local"] = [] - cache["s3"] = [] - cache["gs"] = [] - cache["hdfs"] = [] - cache["ssh"] = [] - cache["azure"] = [] - cache["repo"] = [] + cache = NamedCache() for branch in self.brancher( all_branches=all_branches, all_tags=all_tags @@ -256,18 +250,15 @@ class Repo(object): for stage in stages: if stage.is_repo_import: - cache["repo"] += stage.deps + cache.repo += stage.deps continue for out in stage.outs: - scheme = out.path_info.scheme used_cache = out.get_used_cache( remote=remote, force=force, jobs=jobs ) - - cache[scheme].extend( - dict(entry, branch=branch) for entry in used_cache - ) + suffix = "({})".format(branch) if branch else "" + cache.update(used_cache, suffix=suffix) return cache diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index b12ac23c4..3f0bce1eb 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -49,16 +49,16 @@ def _fetch( try: downloaded += self.cloud.pull( - used["local"], jobs, remote=remote, show_checksums=show_checksums + used, jobs, remote=remote, show_checksums=show_checksums ) except NoRemoteError: - if not used["repo"] and used["local"]: + if not used.repo and used["local"]: raise except DownloadError as exc: failed += exc.amount - for dep in used["repo"]: + for dep in used.repo: try: out = dep.fetch() downloaded += out.get_files_number() diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index ef15d65f9..4480232ab 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -1,53 +1,14 @@ from __future__ import unicode_literals -import collections import logging +from dvc.cache import NamedCache from . import locked logger = logging.getLogger(__name__) -def _merge_cache_lists(clists): - merged_cache = collections.defaultdict(list) - - for cache_list in clists: - for scheme, cache in cache_list.items(): - for item in cache: - if item not in merged_cache[scheme]: - merged_cache[scheme].append(item) - - return merged_cache - - -def _load_all_used_cache( - repos, - all_branches=False, - with_deps=False, - all_tags=False, - remote=None, - force=False, - jobs=None, -): - clists = [] - - for repo in repos: - repo_clist = repo.used_cache( - targets=None, - all_branches=all_branches, - with_deps=with_deps, - all_tags=all_tags, - remote=remote, - force=force, - jobs=jobs, - ) - - clists.append(repo_clist) - - return clists - - def _do_gc(typ, func, clist): removed = func(clist) if not removed: @@ -79,37 +40,35 @@ def gc( stack.enter_context(repo.lock) stack.enter_context(repo.state) - all_clists = _load_all_used_cache( - all_repos + [self], - all_branches=all_branches, - with_deps=with_deps, - all_tags=all_tags, - remote=remote, - force=force, - jobs=jobs, - ) + used = NamedCache() + for repo in all_repos + [self]: + used.update( + repo.used_cache( + all_branches=all_branches, + with_deps=with_deps, + all_tags=all_tags, + remote=remote, + force=force, + jobs=jobs, + ) + ) - if len(all_clists) > 1: - clist = _merge_cache_lists(all_clists) - else: - clist = all_clists[0] - - _do_gc("local", self.cache.local.gc, clist) + _do_gc("local", self.cache.local.gc, used) if self.cache.s3: - _do_gc("s3", self.cache.s3.gc, clist) + _do_gc("s3", self.cache.s3.gc, used) if self.cache.gs: - _do_gc("gs", self.cache.gs.gc, clist) + _do_gc("gs", self.cache.gs.gc, used) if self.cache.ssh: - _do_gc("ssh", self.cache.ssh.gc, clist) + _do_gc("ssh", self.cache.ssh.gc, used) if self.cache.hdfs: - _do_gc("hdfs", self.cache.hdfs.gc, clist) + _do_gc("hdfs", self.cache.hdfs.gc, used) if self.cache.azure: - _do_gc("azure", self.cache.azure.gc, clist) + _do_gc("azure", self.cache.azure.gc, used) if cloud: - _do_gc("remote", self.cloud.get_remote(remote, "gc -c").gc, clist) + _do_gc("remote", self.cloud.get_remote(remote, "gc -c").gc, used) diff --git a/dvc/repo/push.py b/dvc/repo/push.py index ec5eb0afb..7d3593f54 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -23,5 +23,5 @@ def push( remote=remote, jobs=jobs, recursive=recursive, - )["local"] + ) return self.cloud.push(used, jobs, remote=remote) diff --git a/dvc/repo/status.py b/dvc/repo/status.py index 422504404..9b6808b2e 100644 --- a/dvc/repo/status.py +++ b/dvc/repo/status.py @@ -76,7 +76,7 @@ def _cloud_status( force=True, remote=remote, jobs=jobs, - )["local"] + ) ret = {} status_info = self.cloud.status(used, jobs, remote=remote) diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 6633a4f81..6ceeb07bb 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -16,6 +16,7 @@ from dvc.utils.compat import str from dvc.utils import env2bool from dvc.main import main from dvc.config import Config +from dvc.cache import NamedCache from dvc.data_cloud import DataCloud from dvc.remote import ( RemoteS3, @@ -289,72 +290,71 @@ class TestDataCloudBase(TestDvc): self.assertTrue(stage is not None) out = stage.outs[0] cache = out.cache_path - name = str(out) md5 = out.checksum - info = {"name": name, out.remote.PARAM_CHECKSUM: md5} + info = out.get_used_cache() stages = self.dvc.add(self.DATA_DIR) self.assertEqual(len(stages), 1) stage_dir = stages[0] self.assertTrue(stage_dir is not None) out_dir = stage_dir.outs[0] - cache_dir = out.cache_path - name_dir = str(out) - md5_dir = out.checksum - info_dir = {"name": name_dir, out_dir.remote.PARAM_CHECKSUM: md5_dir} + cache_dir = out_dir.cache_path + name_dir = str(out_dir) + md5_dir = out_dir.checksum + info_dir = NamedCache.make(out_dir.scheme, md5_dir, name_dir) with self.cloud.repo.state: # Check status - status = self.cloud.status([info], show_checksums=True) + status = self.cloud.status(info, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_NEW}} self.assertEqual(status, expected) - status_dir = self.cloud.status([info_dir], show_checksums=True) + status_dir = self.cloud.status(info_dir, show_checksums=True) expected = {md5_dir: {"name": md5_dir, "status": STATUS_NEW}} self.assertEqual(status_dir, expected) # Push and check status - self.cloud.push([info]) + self.cloud.push(info) self.assertTrue(os.path.exists(cache)) self.assertTrue(os.path.isfile(cache)) - self.cloud.push([info_dir]) + self.cloud.push(info_dir) self.assertTrue(os.path.isfile(cache_dir)) - status = self.cloud.status([info], show_checksums=True) + status = self.cloud.status(info, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_OK}} self.assertEqual(status, expected) - status_dir = self.cloud.status([info_dir], show_checksums=True) + status_dir = self.cloud.status(info_dir, show_checksums=True) expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}} self.assertEqual(status_dir, expected) # Remove and check status shutil.rmtree(self.dvc.cache.local.cache_dir) - status = self.cloud.status([info], show_checksums=True) + status = self.cloud.status(info, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_DELETED}} self.assertEqual(status, expected) - status_dir = self.cloud.status([info_dir], show_checksums=True) + status_dir = self.cloud.status(info_dir, show_checksums=True) expected = {md5_dir: {"name": md5_dir, "status": STATUS_DELETED}} self.assertEqual(status_dir, expected) # Pull and check status - self.cloud.pull([info]) + self.cloud.pull(info) self.assertTrue(os.path.exists(cache)) self.assertTrue(os.path.isfile(cache)) with open(cache, "r") as fd: self.assertEqual(fd.read(), self.FOO_CONTENTS) - self.cloud.pull([info_dir]) + self.cloud.pull(info_dir) self.assertTrue(os.path.isfile(cache_dir)) - status = self.cloud.status([info], show_checksums=True) + status = self.cloud.status(info, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_OK}} self.assertEqual(status, expected) - status_dir = self.cloud.status([info_dir], show_checksums=True) + status_dir = self.cloud.status(info_dir, show_checksums=True) expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}} self.assertTrue(status_dir, expected) diff --git a/tests/unit/remote/test_local.py b/tests/unit/remote/test_local.py index cc4f579fa..c9681d50f 100644 --- a/tests/unit/remote/test_local.py +++ b/tests/unit/remote/test_local.py @@ -1,3 +1,4 @@ +from dvc.cache import NamedCache from dvc.remote.local import RemoteLOCAL @@ -8,16 +9,11 @@ def test_status_download_optimization(mocker): """ remote = RemoteLOCAL(None, {}) - infos = [ - {"name": "foo", "md5": "acbd18db4cc2f85cedef654fccc4a4d8"}, - {"name": "bar", "md5": "37b51d194a7513e45b56f6524f2d51f2"}, - ] - - local_exists = [ - "acbd18db4cc2f85cedef654fccc4a4d8", - "37b51d194a7513e45b56f6524f2d51f2", - ] + infos = NamedCache() + infos.add("local", "acbd18db4cc2f85cedef654fccc4a4d8", "foo") + infos.add("local", "37b51d194a7513e45b56f6524f2d51f2", "bar") + local_exists = list(infos["local"]) mocker.patch.object(remote, "cache_exists", return_value=local_exists) other_remote = mocker.Mock()