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
This commit is contained in:
Alexander Schepanovski
2019-10-21 12:12:56 +03:00
committed by Ruslan Kuprieiev
parent 1f58fc6f42
commit 33b0eefec0
12 changed files with 121 additions and 177 deletions
+27
View File
@@ -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)
+10 -10
View File
@@ -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
)
+14 -19
View File
@@ -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
+5 -10
View File
@@ -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
+12 -31
View File
@@ -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,
+5 -14
View File
@@ -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
+3 -3
View File
@@ -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()
+20 -61
View File
@@ -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)
+1 -1
View File
@@ -23,5 +23,5 @@ def push(
remote=remote,
jobs=jobs,
recursive=recursive,
)["local"]
)
return self.cloud.push(used, jobs, remote=remote)
+1 -1
View File
@@ -76,7 +76,7 @@ def _cloud_status(
force=True,
remote=remote,
jobs=jobs,
)["local"]
)
ret = {}
status_info = self.cloud.status(used, jobs, remote=remote)
+18 -18
View File
@@ -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)
+5 -9
View File
@@ -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()