mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-10 07:18:27 +01:00
+57
-33
@@ -2,8 +2,10 @@
|
||||
from __future__ import print_function
|
||||
import logging
|
||||
from tqdm import tqdm
|
||||
from copy import deepcopy
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from funcy import merge
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TqdmThreadPoolExecutor(ThreadPoolExecutor):
|
||||
@@ -33,45 +35,67 @@ class Tqdm(tqdm):
|
||||
maximum-compatibility tqdm-based progressbars
|
||||
"""
|
||||
|
||||
BAR_FMT_DEFAULT = (
|
||||
"{percentage:3.0f}%|{bar:10}|"
|
||||
"{desc:{ncols_desc}.{ncols_desc}}{n}/{total}"
|
||||
" [{elapsed}<{remaining}, {rate_fmt:>11}{postfix}]"
|
||||
)
|
||||
BAR_FMT_NOTOTAL = (
|
||||
"{desc:{ncols_desc}.{ncols_desc}}{n}"
|
||||
" [{elapsed}<??:??, {rate_fmt:>11}{postfix}]"
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
iterable=None,
|
||||
disable=None,
|
||||
level=logging.ERROR,
|
||||
desc=None,
|
||||
leave=False,
|
||||
bar_format=None,
|
||||
bytes=False, # pylint: disable=W0622
|
||||
desc_truncate=None,
|
||||
leave=None,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
bytes : shortcut for
|
||||
`unit='B', unit_scale=True, unit_divisor=1024, miniters=1`
|
||||
desc_truncate : like `desc` but will truncate to 10 chars
|
||||
desc : persists after `close()`
|
||||
level : effective logging level for determining `disable`;
|
||||
used only if `disable` is unspecified
|
||||
kwargs : anything accepted by `tqdm.tqdm()`
|
||||
"""
|
||||
kwargs = deepcopy(kwargs)
|
||||
kwargs = kwargs.copy()
|
||||
kwargs.setdefault("unit_scale", True)
|
||||
if bytes:
|
||||
for k, v in dict(
|
||||
bytes_defaults = dict(
|
||||
unit="B", unit_scale=True, unit_divisor=1024, miniters=1
|
||||
).items():
|
||||
kwargs.setdefault(k, v)
|
||||
if desc_truncate is not None:
|
||||
kwargs.setdefault("desc", self.truncate(desc_truncate))
|
||||
if disable is None:
|
||||
disable = (
|
||||
logging.getLogger(__name__).getEffectiveLevel()
|
||||
>= logging.CRITICAL
|
||||
)
|
||||
kwargs = merge(bytes_defaults, kwargs)
|
||||
self.desc_persist = desc
|
||||
if disable is None:
|
||||
disable = logger.getEffectiveLevel() > level
|
||||
super(Tqdm, self).__init__(
|
||||
iterable=iterable, disable=disable, leave=leave, **kwargs
|
||||
iterable=iterable,
|
||||
disable=disable,
|
||||
leave=leave,
|
||||
desc=desc,
|
||||
bar_format="!",
|
||||
**kwargs
|
||||
)
|
||||
if bar_format is None:
|
||||
if self.__len__():
|
||||
self.bar_format = self.BAR_FMT_DEFAULT
|
||||
else:
|
||||
self.bar_format = self.BAR_FMT_NOTOTAL
|
||||
else:
|
||||
self.bar_format = bar_format
|
||||
self.refresh()
|
||||
|
||||
def update_desc(self, desc, n=1, truncate=True):
|
||||
def update_desc(self, desc, n=1):
|
||||
"""
|
||||
Calls `set_description(truncate(desc))` and `update(n)`
|
||||
Calls `set_description_str(desc)` and `update(n)`
|
||||
"""
|
||||
self.set_description(
|
||||
self.truncate(desc) if truncate else desc, refresh=False
|
||||
)
|
||||
self.set_description_str(desc, refresh=False)
|
||||
self.update(n)
|
||||
|
||||
def update_to(self, current, total=None):
|
||||
@@ -79,16 +103,16 @@ class Tqdm(tqdm):
|
||||
self.total = total # pylint: disable=W0613,W0201
|
||||
self.update(current - self.n)
|
||||
|
||||
@classmethod
|
||||
def truncate(cls, s, max_len=25, end=True, fill="..."):
|
||||
"""
|
||||
Guarantee len(output) < max_lenself.
|
||||
>>> truncate("hello", 4)
|
||||
'...o'
|
||||
"""
|
||||
if len(s) <= max_len:
|
||||
return s
|
||||
if len(fill) > max_len:
|
||||
return fill[-max_len:] if end else fill[:max_len]
|
||||
i = max_len - len(fill)
|
||||
return (fill + s[-i:]) if end else (s[:i] + fill)
|
||||
def close(self):
|
||||
if self.desc_persist is not None:
|
||||
self.set_description_str(self.desc_persist, refresh=False)
|
||||
super(Tqdm, self).close()
|
||||
|
||||
@property
|
||||
def format_dict(self):
|
||||
"""inject `ncols_desc` to fill the display width (`ncols`)"""
|
||||
d = super(Tqdm, self).format_dict
|
||||
ncols = d["ncols"] or 80
|
||||
ncols_desc = ncols - len(self.format_meter(ncols_desc=1, **d)) + 1
|
||||
d["ncols_desc"] = max(ncols_desc, 0)
|
||||
return d
|
||||
|
||||
+2
-6
@@ -115,9 +115,7 @@ class RemoteAZURE(RemoteBASE):
|
||||
def _upload(
|
||||
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
|
||||
):
|
||||
with Tqdm(
|
||||
desc_truncate=name, disable=no_progress_bar, bytes=True
|
||||
) as pbar:
|
||||
with Tqdm(desc=name, disable=no_progress_bar, bytes=True) as pbar:
|
||||
self.blob_service.create_blob_from_path(
|
||||
to_info.bucket,
|
||||
to_info.path,
|
||||
@@ -128,9 +126,7 @@ class RemoteAZURE(RemoteBASE):
|
||||
def _download(
|
||||
self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
|
||||
):
|
||||
with Tqdm(
|
||||
desc_truncate=name, disable=no_progress_bar, bytes=True
|
||||
) as pbar:
|
||||
with Tqdm(desc=name, disable=no_progress_bar, bytes=True) as pbar:
|
||||
self.blob_service.get_blob_to_path(
|
||||
from_info.bucket,
|
||||
from_info.path,
|
||||
|
||||
+10
-5
@@ -612,7 +612,7 @@ class RemoteBASE(object):
|
||||
return self._changed_dir_cache(checksum)
|
||||
return self.changed_cache_file(checksum)
|
||||
|
||||
def cache_exists(self, checksums, jobs=None):
|
||||
def cache_exists(self, checksums, jobs=None, name=None):
|
||||
"""Check if the given checksums are stored in the remote.
|
||||
|
||||
There are two ways of performing this check:
|
||||
@@ -630,7 +630,7 @@ class RemoteBASE(object):
|
||||
take much shorter time to just retrieve everything they have under
|
||||
a certain prefix (e.g. s3, gs, ssh, hdfs). Other remotes that can
|
||||
check if particular file exists much quicker, use their own
|
||||
implementation of cache_exists (see http, local).
|
||||
implementation of cache_exists (see ssh, local).
|
||||
|
||||
Returns:
|
||||
A list with checksums that were found in the remote
|
||||
@@ -638,15 +638,20 @@ class RemoteBASE(object):
|
||||
if not self.no_traverse:
|
||||
return list(set(checksums) & set(self.all()))
|
||||
|
||||
with Tqdm(total=len(checksums), unit="md5") as pbar:
|
||||
with Tqdm(
|
||||
desc="Querying "
|
||||
+ ("cache in " + name if name else "remote cache"),
|
||||
total=len(checksums),
|
||||
unit="file",
|
||||
) as pbar:
|
||||
|
||||
def exists_with_progress(path_info):
|
||||
ret = self.exists(path_info)
|
||||
pbar.update()
|
||||
pbar.update_desc(str(path_info))
|
||||
return ret
|
||||
|
||||
with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
|
||||
path_infos = [self.checksum_to_path_info(x) for x in checksums]
|
||||
path_infos = map(self.checksum_to_path_info, checksums)
|
||||
in_remote = executor.map(exists_with_progress, path_infos)
|
||||
ret = list(itertools.compress(checksums, in_remote))
|
||||
return ret
|
||||
|
||||
+1
-1
@@ -38,7 +38,7 @@ class RemoteHTTP(RemoteBASE):
|
||||
total=None if no_progress_bar else self._content_length(from_info),
|
||||
leave=False,
|
||||
bytes=True,
|
||||
desc_truncate=from_info.url if name is None else name,
|
||||
desc=from_info.url if name is None else name,
|
||||
disable=no_progress_bar,
|
||||
) as pbar:
|
||||
with open(to_file, "wb") as fd:
|
||||
|
||||
@@ -251,10 +251,15 @@ class RemoteLOCAL(RemoteBASE):
|
||||
|
||||
move(from_info, to_info, mode=mode)
|
||||
|
||||
def cache_exists(self, checksums, jobs=None):
|
||||
def cache_exists(self, checksums, jobs=None, name=None):
|
||||
return [
|
||||
checksum
|
||||
for checksum in Tqdm(checksums, unit="md5")
|
||||
for checksum in Tqdm(
|
||||
checksums,
|
||||
unit="file",
|
||||
desc="Querying "
|
||||
+ ("cache in " + name if name else "local cache"),
|
||||
)
|
||||
if not self.changed_cache_file(checksum)
|
||||
]
|
||||
|
||||
@@ -313,14 +318,14 @@ class RemoteLOCAL(RemoteBASE):
|
||||
show_checksums=False,
|
||||
download=False,
|
||||
):
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"Preparing to collect status from {}".format(remote.path_info)
|
||||
)
|
||||
ret = self._group(checksum_infos, show_checksums=show_checksums) or {}
|
||||
md5s = list(ret)
|
||||
|
||||
logger.info("Collecting information from local cache...")
|
||||
local_exists = self.cache_exists(md5s, jobs=jobs)
|
||||
logger.debug("Collecting information from local cache...")
|
||||
local_exists = self.cache_exists(md5s, jobs=jobs, name=self.cache_dir)
|
||||
|
||||
# This is a performance optimization. We can safely assume that,
|
||||
# if the resources that we want to fetch are already cached,
|
||||
@@ -329,8 +334,12 @@ class RemoteLOCAL(RemoteBASE):
|
||||
if download and sorted(local_exists) == sorted(md5s):
|
||||
remote_exists = local_exists
|
||||
else:
|
||||
logger.info("Collecting information from remote cache...")
|
||||
remote_exists = list(remote.cache_exists(md5s, jobs=jobs))
|
||||
logger.debug("Collecting information from remote cache...")
|
||||
remote_exists = list(
|
||||
remote.cache_exists(
|
||||
md5s, jobs=jobs, name=str(remote.path_info)
|
||||
)
|
||||
)
|
||||
|
||||
self._fill_statuses(ret, local_exists, remote_exists)
|
||||
|
||||
@@ -377,7 +386,7 @@ class RemoteLOCAL(RemoteBASE):
|
||||
show_checksums=False,
|
||||
download=False,
|
||||
):
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"Preparing to {} '{}'".format(
|
||||
"download data from" if download else "upload data to",
|
||||
remote.path_info,
|
||||
|
||||
+2
-6
@@ -107,9 +107,7 @@ class RemoteOSS(RemoteBASE):
|
||||
def _upload(
|
||||
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
|
||||
):
|
||||
with Tqdm(
|
||||
desc_truncate=name, disable=no_progress_bar, bytes=True
|
||||
) as pbar:
|
||||
with Tqdm(desc=name, disable=no_progress_bar, bytes=True) as pbar:
|
||||
self.oss_service.put_object_from_file(
|
||||
to_info.path, from_file, progress_callback=pbar.update_to
|
||||
)
|
||||
@@ -117,9 +115,7 @@ class RemoteOSS(RemoteBASE):
|
||||
def _download(
|
||||
self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
|
||||
):
|
||||
with Tqdm(
|
||||
desc_truncate=name, disable=no_progress_bar, bytes=True
|
||||
) as pbar:
|
||||
with Tqdm(desc=name, disable=no_progress_bar, bytes=True) as pbar:
|
||||
self.oss_service.get_object_to_file(
|
||||
from_info.path, to_file, progress_callback=pbar.update_to
|
||||
)
|
||||
|
||||
+2
-8
@@ -208,10 +208,7 @@ class RemoteS3(RemoteBASE):
|
||||
def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
|
||||
total = os.path.getsize(from_file)
|
||||
with Tqdm(
|
||||
disable=no_progress_bar,
|
||||
total=total,
|
||||
bytes=True,
|
||||
desc_truncate=name,
|
||||
disable=no_progress_bar, total=total, bytes=True, desc=name
|
||||
) as pbar:
|
||||
self.s3.upload_file(
|
||||
from_file,
|
||||
@@ -229,10 +226,7 @@ class RemoteS3(RemoteBASE):
|
||||
Bucket=from_info.bucket, Key=from_info.path
|
||||
)["ContentLength"]
|
||||
with Tqdm(
|
||||
disable=no_progress_bar,
|
||||
total=total,
|
||||
bytes=True,
|
||||
desc_truncate=name,
|
||||
disable=no_progress_bar, total=total, bytes=True, desc=name
|
||||
) as pbar:
|
||||
self.s3.download_file(
|
||||
from_info.bucket, from_info.path, to_file, Callback=pbar.update
|
||||
|
||||
@@ -258,7 +258,7 @@ class RemoteSSH(RemoteBASE):
|
||||
|
||||
return results
|
||||
|
||||
def cache_exists(self, checksums, jobs=None):
|
||||
def cache_exists(self, checksums, jobs=None, name=None):
|
||||
"""This is older implementation used in remote/base.py
|
||||
We are reusing it in RemoteSSH, because SSH's batch_exists proved to be
|
||||
faster than current approach (relying on exists(path_info)) applied in
|
||||
@@ -267,7 +267,12 @@ class RemoteSSH(RemoteBASE):
|
||||
if not self.no_traverse:
|
||||
return list(set(checksums) & set(self.all()))
|
||||
|
||||
with Tqdm(total=len(checksums), unit="md5") as pbar:
|
||||
with Tqdm(
|
||||
desc="Querying "
|
||||
+ ("cache in " + name if name else "remote cache"),
|
||||
total=len(checksums),
|
||||
unit="file",
|
||||
) as pbar:
|
||||
|
||||
def exists_with_progress(chunks):
|
||||
return self.batch_exists(chunks, callback=pbar.update_desc)
|
||||
@@ -278,5 +283,4 @@ class RemoteSSH(RemoteBASE):
|
||||
results = executor.map(exists_with_progress, chunks)
|
||||
in_remote = itertools.chain.from_iterable(results)
|
||||
ret = list(itertools.compress(checksums, in_remote))
|
||||
pbar.update_desc("", 0) # clear path name description
|
||||
return ret
|
||||
|
||||
@@ -169,7 +169,7 @@ class SSHConnection:
|
||||
|
||||
def download(self, src, dest, no_progress_bar=False, progress_title=None):
|
||||
with Tqdm(
|
||||
desc_truncate=progress_title or os.path.basename(src),
|
||||
desc=progress_title or os.path.basename(src),
|
||||
disable=no_progress_bar,
|
||||
bytes=True,
|
||||
) as pbar:
|
||||
@@ -186,7 +186,7 @@ class SSHConnection:
|
||||
progress_title = posixpath.basename(dest)
|
||||
|
||||
with Tqdm(
|
||||
desc_truncate=progress_title, disable=no_progress_bar, bytes=True
|
||||
desc=progress_title, disable=no_progress_bar, bytes=True
|
||||
) as pbar:
|
||||
self.sftp.put(src, tmp_file, callback=pbar.update_to)
|
||||
|
||||
|
||||
@@ -38,6 +38,8 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False):
|
||||
with self.state:
|
||||
_cleanup_unused_links(self, all_stages)
|
||||
total = get_all_files_numbers(stages)
|
||||
if total == 0:
|
||||
logger.info("Nothing to do")
|
||||
with Tqdm(
|
||||
total=total, unit="file", desc="Checkout", disable=total == 0
|
||||
) as pbar:
|
||||
@@ -51,4 +53,3 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False):
|
||||
)
|
||||
|
||||
stage.checkout(force=force, progress_callback=pbar.update_desc)
|
||||
pbar.update_desc("Checkout", 0) # clear path name description
|
||||
|
||||
@@ -51,16 +51,16 @@ def file_md5(fname):
|
||||
hash_md5 = hashlib.md5()
|
||||
binary = not istextfile(fname)
|
||||
size = os.path.getsize(fname)
|
||||
bar = False
|
||||
no_progress_bar = True
|
||||
if size >= LARGE_FILE_SIZE:
|
||||
bar = True
|
||||
no_progress_bar = False
|
||||
msg = "Computing md5 for a large file {}. This is only done once."
|
||||
logger.info(msg.format(relpath(fname)))
|
||||
name = relpath(fname)
|
||||
|
||||
with Tqdm(
|
||||
desc_truncate=name,
|
||||
disable=not bar,
|
||||
desc=name,
|
||||
disable=no_progress_bar,
|
||||
total=size,
|
||||
bytes=True,
|
||||
leave=False,
|
||||
@@ -132,10 +132,7 @@ def copyfile(src, dest, no_progress_bar=False, name=None):
|
||||
System.reflink(src, dest)
|
||||
except DvcException:
|
||||
with Tqdm(
|
||||
desc_truncate=name,
|
||||
disable=no_progress_bar,
|
||||
total=total,
|
||||
bytes=True,
|
||||
desc=name, disable=no_progress_bar, total=total, bytes=True
|
||||
) as pbar:
|
||||
with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
|
||||
while True:
|
||||
|
||||
Reference in New Issue
Block a user