From fecc81e951efeaa8130264f726c27e92876422ae Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Mon, 25 Apr 2022 18:36:54 +0545 Subject: [PATCH] Revert "dvc.objects: add callbacks for listing and estimating hashes" This reverts commit 7599c6e20d7a6cd28b3433ba3cfed54ffc8d8135. --- dvc/data/_progress.py | 32 ------- dvc/data/db/local.py | 13 +-- dvc/data/gc.py | 10 +-- dvc/data/status.py | 23 +---- dvc/objects/db.py | 154 ++++++++++++++++++--------------- tests/unit/remote/test_base.py | 5 +- 6 files changed, 102 insertions(+), 135 deletions(-) delete mode 100644 dvc/data/_progress.py diff --git a/dvc/data/_progress.py b/dvc/data/_progress.py deleted file mode 100644 index 66de97e59..000000000 --- a/dvc/data/_progress.py +++ /dev/null @@ -1,32 +0,0 @@ -from dvc.progress import Tqdm - - -class QueryingProgress(Tqdm): - def __init__(self, iterable=None, total=None, name=None, phase="Querying"): - msg_part = "cache in " + f"'{name}'" if name else "remote cache" - msg_fmt = "{phase} " + msg_part - - self._estimating_msg = msg_fmt.format(phase="Estimating size of") - self._listing_msg = msg_fmt.format(phase="Querying") - super().__init__( - iterable=iterable, - desc=msg_fmt.format(phase=phase), - total=total, - unit="files", - unit_scale=False, - bar_format=self.BAR_FMT_DEFAULT, - ) - - def callback(self, phase, *args): - total = args[0] if args else self.total - completed = args[1] if len(args) > 1 else self.n - desc = self.desc - if phase == "estimating": - desc = self._estimating_msg - elif phase == "querying": - desc = self._listing_msg - - if desc != self.desc: - # let Tqdm take care of when to update - self.set_description_str(desc, refresh=False) - self.update_to(completed, total) diff --git a/dvc/data/db/local.py b/dvc/data/db/local.py index 06ad2ba09..cb4004e28 100644 --- a/dvc/data/db/local.py +++ b/dvc/data/db/local.py @@ -1,14 +1,14 @@ import logging import os import stat -from functools import partial from funcy import cached_property from shortuuid import uuid from dvc.hash_info import HashInfo -from dvc.objects.db import ObjectDB, noop, wrap_iter +from dvc.objects.db import ObjectDB from dvc.objects.errors import ObjectFormatError +from dvc.progress import Tqdm from dvc.utils import relpath from dvc.utils.fs import copyfile, remove, umask, walk_files @@ -60,12 +60,15 @@ class LocalObjectDB(ObjectDB): return f"{self.cache_path}{os.sep}{hash_[0:2]}{os.sep}{hash_[2:]}" def hashes_exist( - self, hashes, jobs=None, progress=noop + self, hashes, jobs=None, name=None ): # pylint: disable=unused-argument ret = [] - progress = partial(progress, "querying", len(hashes)) - for hash_ in wrap_iter(hashes, progress): + for hash_ in Tqdm( + hashes, + unit="file", + desc="Querying " + ("cache in " + name if name else "local cache"), + ): hash_info = HashInfo(self.fs.PARAM_CHECKSUM, hash_) try: self.check(hash_info) diff --git a/dvc/data/gc.py b/dvc/data/gc.py index f8c6fded2..338f3ccfa 100644 --- a/dvc/data/gc.py +++ b/dvc/data/gc.py @@ -2,8 +2,6 @@ def gc(odb, used, jobs=None, cache_odb=None, shallow=True): from dvc.data.tree import Tree from dvc.objects.errors import ObjectDBPermissionError - from ._progress import QueryingProgress - if odb.read_only: raise ObjectDBPermissionError("Cannot gc read-only ODB") if not cache_odb: @@ -24,9 +22,11 @@ def gc(odb, used, jobs=None, cache_odb=None, shallow=True): removed = False # hashes must be sorted to ensure we always remove .dir files first - - hashes = QueryingProgress(odb.all(jobs), name=odb.fs_path) - for hash_ in sorted(hashes, key=_is_dir_hash, reverse=True): + for hash_ in sorted( + odb.all(jobs, odb.fs_path), + key=_is_dir_hash, + reverse=True, + ): if hash_ in used_hashes: continue fs_path = odb.hash_to_path(hash_) diff --git a/dvc/data/status.py b/dvc/data/status.py index ecbfc4b82..9ec4d0f3d 100644 --- a/dvc/data/status.py +++ b/dvc/data/status.py @@ -30,17 +30,11 @@ class CompareStatusResult(NamedTuple): def _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb, jobs=None): # Validate our index by verifying all indexed .dir hashes # still exist on the remote - from ._progress import QueryingProgress - dir_hashes = set(dir_objs.keys()) indexed_dirs = set(index.dir_hashes()) indexed_dir_exists = set() if indexed_dirs: - hashes = QueryingProgress( - odb.list_hashes_exists(indexed_dirs, jobs=jobs), - total=len(indexed_dirs), - ) - indexed_dir_exists.update(hashes) + indexed_dir_exists.update(odb.list_hashes_exists(indexed_dirs, jobs)) missing_dirs = indexed_dirs.difference(indexed_dir_exists) if missing_dirs: logger.debug( @@ -52,13 +46,7 @@ def _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb, jobs=None): # Check if non-indexed (new) dir hashes exist on remote dir_exists = dir_hashes.intersection(indexed_dir_exists) - dir_missing = dir_hashes - dir_exists - dir_exists.update( - QueryingProgress( - odb.list_hashes_exists(dir_missing, jobs=jobs), - total=len(dir_missing), - ) - ) + dir_exists.update(odb.list_hashes_exists(dir_hashes - dir_exists, jobs)) # If .dir hash exists in the ODB, assume directory contents # also exists @@ -144,12 +132,7 @@ def status( hashes.difference_update(exists) if hashes: - from ._progress import QueryingProgress - - with QueryingProgress(phase="Checking", name=odb.fs_path) as pbar: - exists.update( - odb.hashes_exist(hashes, jobs=jobs, progress=pbar.callback) - ) + exists.update(odb.hashes_exist(hashes, name=odb.fs_path, jobs=jobs)) return StatusResult( {hash_infos[hash_] for hash_ in exists}, {hash_infos[hash_] for hash_ in (hashes - exists)}, diff --git a/dvc/objects/db.py b/dvc/objects/db.py index 1c03a3e6f..feeadd585 100644 --- a/dvc/objects/db.py +++ b/dvc/objects/db.py @@ -3,11 +3,11 @@ import logging from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from copy import copy -from functools import partial from typing import TYPE_CHECKING, Optional from dvc.objects.errors import ObjectDBPermissionError, ObjectFormatError from dvc.objects.file import HashFile +from dvc.progress import Tqdm if TYPE_CHECKING: from typing import Tuple @@ -18,16 +18,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def noop(*args, **kwargs): - pass - - -def wrap_iter(iterable, callback): - for index, item in enumerate(iterable, start=1): - yield item - callback(index) - - class ObjectDB: DEFAULT_VERIFY = False @@ -242,11 +232,9 @@ class ObjectDB: * self.fs.LIST_OBJECT_PAGE_SIZE, ) - def _estimate_remote_size(self, hashes=None, progress=noop): + def _estimate_remote_size(self, hashes=None, name=None): """Estimate fs size based on number of entries beginning with "00..." prefix. - - Takes a progress callback that returns current_estimated_size. """ prefix = "0" * self.fs.TRAVERSE_PREFIX_LEN total_prefixes = pow(16, self.fs.TRAVERSE_PREFIX_LEN) @@ -255,29 +243,35 @@ class ObjectDB: else: max_hashes = None - def iter_with_pbar(hashes): - total = 0 - for hash_ in hashes: - total += total_prefixes - progress(total) - yield hash_ + with Tqdm( + desc="Estimating size of " + + (f"cache in '{name}'" if name else "remote cache"), + unit="file", + ) as pbar: - if max_hashes: - hashes = self._hashes_with_limit( - max_hashes / total_prefixes, prefix - ) - else: - hashes = self._list_hashes(prefix) + def iter_with_pbar(hashes): + for hash_ in hashes: + pbar.update(total_prefixes) + yield hash_ - remote_hashes = set(iter_with_pbar(hashes)) - if remote_hashes: - remote_size = total_prefixes * len(remote_hashes) - else: - remote_size = total_prefixes - logger.debug(f"Estimated remote size: {remote_size} files") + if max_hashes: + hashes = self._hashes_with_limit( + max_hashes / total_prefixes, prefix + ) + else: + hashes = self._list_hashes(prefix) + + remote_hashes = set(iter_with_pbar(hashes)) + if remote_hashes: + remote_size = total_prefixes * len(remote_hashes) + else: + remote_size = total_prefixes + logger.debug(f"Estimated remote size: {remote_size} files") return remote_size, remote_hashes - def _list_hashes_traverse(self, remote_size, remote_hashes, jobs=None): + def _list_hashes_traverse( + self, remote_size, remote_hashes, jobs=None, name=None + ): """Iterate over all hashes found in this fs. Hashes are fetched in parallel according to prefix, except in cases where the remote size is very small. @@ -292,8 +286,6 @@ class ObjectDB: sense to use a generator to gradually iterate over it, without keeping all of it in memory. """ - from funcy import collecting - num_pages = remote_size / self.fs.LIST_OBJECT_PAGE_SIZE if num_pages < 256 / self.fs.jobs: # Fetching prefixes in parallel requires at least 255 more @@ -303,48 +295,81 @@ class ObjectDB: # NOTE: this ends up re-fetching hashes that were already # fetched during remote size estimation traverse_prefixes = [None] + initial = 0 else: yield from remote_hashes + initial = len(remote_hashes) traverse_prefixes = [f"{i:02x}" for i in range(1, 256)] if self.fs.TRAVERSE_PREFIX_LEN > 2: traverse_prefixes += [ "{0:0{1}x}".format(i, self.fs.TRAVERSE_PREFIX_LEN) for i in range(1, pow(16, self.fs.TRAVERSE_PREFIX_LEN - 2)) ] + with Tqdm( + desc="Querying " + + (f"cache in '{name}'" if name else "remote cache"), + total=remote_size, + initial=initial, + unit="file", + ) as pbar: + from funcy import collecting - list_hashes = collecting(self._list_hashes) - with ThreadPoolExecutor(max_workers=jobs or self.fs.jobs) as executor: - in_remote = executor.map(list_hashes, traverse_prefixes) - yield from itertools.chain.from_iterable(in_remote) + with ThreadPoolExecutor( + max_workers=jobs or self.fs.jobs + ) as executor: + list_hashes = collecting(pbar.wrap_fn(self._list_hashes)) + in_remote = executor.map(list_hashes, traverse_prefixes) + yield from itertools.chain.from_iterable(in_remote) - def all(self, jobs=None): + def all(self, jobs=None, name=None): """Iterate over all hashes in this fs. Hashes will be fetched in parallel threads according to prefix (except for small remotes) and a progress bar will be displayed. """ + logger.debug( + "Fetching all hashes from '{}'".format( + name if name else "remote cache" + ) + ) + if not self.fs.CAN_TRAVERSE: return self._list_hashes() - remote_size, remote_hashes = self._estimate_remote_size() + remote_size, remote_hashes = self._estimate_remote_size(name=name) return self._list_hashes_traverse( - remote_size, remote_hashes, jobs=jobs + remote_size, remote_hashes, jobs=jobs, name=name ) def _remove_unpacked_dir(self, hash_): pass - def list_hashes_exists(self, hashes, jobs=None): + def list_hashes_exists(self, hashes, jobs=None, name=None): """Return list of the specified hashes which exist in this fs. Hashes will be queried individually. """ logger.debug(f"Querying {len(hashes)} hashes via object_exists") - with ThreadPoolExecutor(max_workers=jobs or self.fs.jobs) as executor: - fs_paths = map(self.hash_to_path, hashes) - in_remote = executor.map(self.fs.exists, fs_paths) - yield from itertools.compress(hashes, in_remote) + with Tqdm( + desc="Querying " + + ("cache in " + name if name else "remote cache"), + total=len(hashes), + unit="file", + ) as pbar: - def hashes_exist(self, hashes, jobs=None, progress=noop): + def exists_with_progress(fs_path): + ret = self.fs.exists(fs_path) + pbar.update_msg(fs_path) + return ret + + with ThreadPoolExecutor( + max_workers=jobs or self.fs.jobs + ) as executor: + fs_paths = map(self.hash_to_path, hashes) + in_remote = executor.map(exists_with_progress, fs_paths) + ret = list(itertools.compress(hashes, in_remote)) + return ret + + def hashes_exist(self, hashes, jobs=None, name=None): """Check if the given hashes are stored in the remote. There are two ways of performing this check: @@ -374,9 +399,6 @@ class ObjectDB: cache can be estimated, since the cache is evenly distributed according to hash. - Takes a callback that returns value in the format of: - (phase, total, current). The phase can be {"estimating, "querying"}. - Returns: A list with hashes that were found in the remote """ @@ -393,16 +415,11 @@ class ObjectDB: if ( len(hashes) == 1 or not self.fs.CAN_TRAVERSE ) and not always_traverse: - remote_hashes = self.list_hashes_exists(hashes, jobs) - callback = partial(progress, "querying", len(hashes)) - return list(wrap_iter(remote_hashes, callback)) + remote_hashes = self.list_hashes_exists(hashes, jobs, name) + return remote_hashes # Max remote size allowed for us to use traverse method - - estimator_cb = partial(progress, "estimating", None) - remote_size, remote_hashes = self._estimate_remote_size( - hashes, progress=estimator_cb - ) + remote_size, remote_hashes = self._estimate_remote_size(hashes, name) traverse_pages = remote_size / self.fs.LIST_OBJECT_PAGE_SIZE # For sufficiently large remotes, traverse must be weighted to account @@ -422,19 +439,14 @@ class ObjectDB: len(hashes), traverse_weight ) ) - remaining_hashes = hashes - remote_hashes - ret = list(hashes & remote_hashes) - callback = partial(progress, "querying", len(remaining_hashes)) - ret.extend( - wrap_iter( - self.list_hashes_exists(remaining_hashes, jobs), callback - ) + return list(hashes & remote_hashes) + self.list_hashes_exists( + hashes - remote_hashes, jobs, name ) - return ret logger.debug(f"Querying '{len(hashes)}' hashes via traverse") - remote_hashes = self._list_hashes_traverse( - remote_size, remote_hashes, jobs=jobs + remote_hashes = set( + self._list_hashes_traverse( + remote_size, remote_hashes, jobs=jobs, name=name + ) ) - callback = partial(progress, "querying", remote_size) - return list(hashes & set(wrap_iter(remote_hashes, callback))) + return list(hashes & set(remote_hashes)) diff --git a/tests/unit/remote/test_base.py b/tests/unit/remote/test_base.py index b7e2b0ed6..7c1aa62c0 100644 --- a/tests/unit/remote/test_base.py +++ b/tests/unit/remote/test_base.py @@ -26,7 +26,7 @@ def test_hashes_exist(object_exists, traverse, dvc): with mock.patch.object(odb, "_list_hashes", return_value=list(range(256))): hashes = set(range(1000)) odb.hashes_exist(hashes) - object_exists.assert_called_with(hashes, None) + object_exists.assert_called_with(hashes, None, None) traverse.assert_not_called() odb.fs.CAN_TRAVERSE = True @@ -47,7 +47,7 @@ def test_hashes_exist(object_exists, traverse, dvc): ) assert max_hashes < 2048 object_exists.assert_called_with( - frozenset(range(max_hashes, 1000)), None + frozenset(range(max_hashes, 1000)), None, None ) traverse.assert_not_called() @@ -63,6 +63,7 @@ def test_hashes_exist(object_exists, traverse, dvc): 256 * pow(16, odb.fs.TRAVERSE_PREFIX_LEN), set(range(256)), jobs=None, + name=None, )