Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

dvc.objects: add callbacks for listing and estimating hashes

Browse Source 
This PR gets rid of Tqdm progress bars inside `dvc.objects`.
Those are replaced with multi-phase callback in `hashes_exist`, and
simple callable that takes `completed` value in
`odb._estimate_remote_size()`.

The `odb.list_hashes_exists()` and `odb.all()` now return iterables,
hence the caller has the responsibility of handling progress.

This is now handled in `dvc.data` by using a custom `Tqdm` progressbar.
This is ugly and is temporary, but my motivation is to remove tqdm from dvc.objects
first. We can gradually get rid of these from `dvc.data` too, but
that requires more knowledge about `dvc.data` and involves more
design work. All in time, I guess. :)
This commit is contained in:
Saugat Pachhai (सौगात)
2022-04-26 07:38:24 +05:45
committed by Saugat Pachhai
parent f39e35445a
commit 112ea1d448
6 changed files with 135 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/data/_progress.py
+32
View File
@@ -0,0 +1,32 @@
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)
dvc/data/db/local.py
+5 -8
View File
@@ -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
from dvc.objects.db import ObjectDB, noop, wrap_iter
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,15 +60,12 @@ class LocalObjectDB(ObjectDB):
return f"{self.cache_path}{os.sep}{hash_[0:2]}{os.sep}{hash_[2:]}"
def hashes_exist(
self, hashes, jobs=None, name=None
self, hashes, jobs=None, progress=noop
): # pylint: disable=unused-argument
ret = []
progress = partial(progress, "querying", len(hashes))
for hash_ in Tqdm(
hashes,
unit="file",
desc="Querying " + ("cache in " + name if name else "local cache"),
):
for hash_ in wrap_iter(hashes, progress):
hash_info = HashInfo(self.fs.PARAM_CHECKSUM, hash_)
try:
self.check(hash_info)
dvc/data/gc.py
+5 -5
View File
@@ -2,6 +2,8 @@ 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:
@@ -22,11 +24,9 @@ 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
for hash_ in sorted(
odb.all(jobs, odb.fs_path),
key=_is_dir_hash,
reverse=True,
):
hashes = QueryingProgress(odb.all(jobs), name=odb.fs_path)
for hash_ in sorted(hashes, key=_is_dir_hash, reverse=True):
if hash_ in used_hashes:
continue
fs_path = odb.hash_to_path(hash_)
dvc/data/status.py
+20 -3
View File
@@ -30,11 +30,17 @@ 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:
indexed_dir_exists.update(odb.list_hashes_exists(indexed_dirs, jobs))
hashes = QueryingProgress(
odb.list_hashes_exists(indexed_dirs, jobs=jobs),
total=len(indexed_dirs),
)
indexed_dir_exists.update(hashes)
missing_dirs = indexed_dirs.difference(indexed_dir_exists)
if missing_dirs:
logger.debug(
@@ -46,7 +52,13 @@ 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_exists.update(odb.list_hashes_exists(dir_hashes - dir_exists, jobs))
dir_missing = dir_hashes - dir_exists
dir_exists.update(
QueryingProgress(
odb.list_hashes_exists(dir_missing, jobs=jobs),
total=len(dir_missing),
)
)
# If .dir hash exists in the ODB, assume directory contents
# also exists
@@ -132,7 +144,12 @@ def status(
hashes.difference_update(exists)
if hashes:
exists.update(odb.hashes_exist(hashes, name=odb.fs_path, jobs=jobs))
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)
)
return StatusResult(
{hash_infos[hash_] for hash_ in exists},
{hash_infos[hash_] for hash_ in (hashes - exists)},
dvc/objects/db.py
+71 -83
View File
@@ -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,6 +18,16 @@ 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
@@ -232,9 +242,11 @@ class ObjectDB:
* self.fs.LIST_OBJECT_PAGE_SIZE,
)
def _estimate_remote_size(self, hashes=None, name=None):
def _estimate_remote_size(self, hashes=None, progress=noop):
"""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)
@@ -243,35 +255,29 @@ class ObjectDB:
else:
max_hashes = None
with Tqdm(
desc="Estimating size of "
+ (f"cache in '{name}'" if name else "remote cache"),
unit="file",
) as pbar:
def iter_with_pbar(hashes):
total = 0
for hash_ in hashes:
total += total_prefixes
progress(total)
yield hash_
def iter_with_pbar(hashes):
for hash_ in hashes:
pbar.update(total_prefixes)
yield hash_
if max_hashes:
hashes = self._hashes_with_limit(
max_hashes / total_prefixes, prefix
)
else:
hashes = self._list_hashes(prefix)
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")
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, name=None
):
def _list_hashes_traverse(self, remote_size, remote_hashes, jobs=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.
@@ -286,6 +292,8 @@ 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
@@ -295,81 +303,48 @@ 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
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)
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)
def all(self, jobs=None, name=None):
def all(self, jobs=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(name=name)
remote_size, remote_hashes = self._estimate_remote_size()
return self._list_hashes_traverse(
remote_size, remote_hashes, jobs=jobs, name=name
remote_size, remote_hashes, jobs=jobs
)
def _remove_unpacked_dir(self, hash_):
pass
def list_hashes_exists(self, hashes, jobs=None, name=None):
def list_hashes_exists(self, hashes, jobs=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 Tqdm(
desc="Querying "
+ ("cache in " + name if name else "remote cache"),
total=len(hashes),
unit="file",
) as pbar:
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)
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):
def hashes_exist(self, hashes, jobs=None, progress=noop):
"""Check if the given hashes are stored in the remote.
There are two ways of performing this check:
@@ -399,6 +374,9 @@ 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
"""
@@ -415,11 +393,16 @@ class ObjectDB:
if (
len(hashes) == 1 or not self.fs.CAN_TRAVERSE
) and not always_traverse:
remote_hashes = self.list_hashes_exists(hashes, jobs, name)
return remote_hashes
remote_hashes = self.list_hashes_exists(hashes, jobs)
callback = partial(progress, "querying", len(hashes))
return list(wrap_iter(remote_hashes, callback))
# Max remote size allowed for us to use traverse method
remote_size, remote_hashes = self._estimate_remote_size(hashes, name)
estimator_cb = partial(progress, "estimating", None)
remote_size, remote_hashes = self._estimate_remote_size(
hashes, progress=estimator_cb
)
traverse_pages = remote_size / self.fs.LIST_OBJECT_PAGE_SIZE
# For sufficiently large remotes, traverse must be weighted to account
@@ -439,14 +422,19 @@ class ObjectDB:
len(hashes), traverse_weight
)
)
return list(hashes & remote_hashes) + self.list_hashes_exists(
hashes - remote_hashes, jobs, name
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 ret
logger.debug(f"Querying '{len(hashes)}' hashes via traverse")
remote_hashes = set(
self._list_hashes_traverse(
remote_size, remote_hashes, jobs=jobs, name=name
)
remote_hashes = self._list_hashes_traverse(
remote_size, remote_hashes, jobs=jobs
)
return list(hashes & set(remote_hashes))
callback = partial(progress, "querying", remote_size)
return list(hashes & set(wrap_iter(remote_hashes, callback)))
tests/unit/remote/test_base.py
+2 -3
View File
@@ -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, None)
object_exists.assert_called_with(hashes, 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, None
frozenset(range(max_hashes, 1000)), None
)
traverse.assert_not_called()
@@ -63,7 +63,6 @@ def test_hashes_exist(object_exists, traverse, dvc):
256 * pow(16, odb.fs.TRAVERSE_PREFIX_LEN),
set(range(256)),
jobs=None,
name=None,
)