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

fs.upload/download: simplify, remove unused kwargs, callbacks

Browse Source 
1. This PR simplifies `fs.upload`/`fs.download`/`fs.download_file`.
2. Removes unused kwargs.
3. `download` uses `dvc.utils.threadpool.ThreadPoolExecutor` which
   supports canceling futures.
4. `download` for directories supports nested/branched callbacks
   and progress bars.
This commit is contained in:
Saugat Pachhai (सौगात)
2022-04-26 06:57:10 +05:45
committed by Saugat Pachhai
parent 7211bd02ed
commit f39e35445a
6 changed files with 136 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/data/stage.py
+9 -4
View File
@@ -25,7 +25,8 @@ logger = logging.getLogger(__name__)
_STAGING_MEMFS_PATH = "dvc-staging"
def _upload_file(from_fs_path, fs, odb, upload_odb):
def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
from dvc.fs._callback import FsspecCallback
from dvc.utils import tmp_fname
from dvc.utils.stream import HashedStreamReader
@@ -34,9 +35,13 @@ def _upload_file(from_fs_path, fs, odb, upload_odb):
with fs.open(from_fs_path, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream:
stream = HashedStreamReader(stream)
size = fs.size(from_fs_path)
upload_odb.fs.upload(
stream, tmp_info, desc=fs_path.name(from_fs_path), total=size
)
with FsspecCallback.as_tqdm_callback(
callback,
desc=fs_path.name(from_fs_path),
bytes=True,
total=size,
) as cb:
upload_odb.fs.upload(stream, tmp_info, size=size, callback=cb)
odb.add(tmp_info, upload_odb.fs, stream.hash_info)
meta = Meta(size=size)
dvc/fs/_callback.py
+58 -18
View File
@@ -1,15 +1,16 @@
from functools import wraps
from typing import IO, TYPE_CHECKING, Optional, TypeVar, cast
from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast
import fsspec
from dvc.progress import Tqdm
if TYPE_CHECKING:
from typing import Callable
from typing_extensions import ParamSpec
from dvc.fs.base import FileSystem
from dvc.progress import Tqdm
_P = ParamSpec("_P")
_R = TypeVar("_R")
@@ -32,6 +33,39 @@ class FsspecCallback(fsspec.Callback):
return wrapped
def wrap_and_branch(
self, fn: "Callable", fs: "FileSystem" = None
) -> "Callable":
"""
Wraps a function, and pass a new child callback to it.
When the function completes, we increment the parent callback by 1.
"""
from .local import localfs
fs = fs or localfs
wrapped = self.wrap_fn(fn)
def make_callback(path1, path2):
# pylint: disable=assignment-from-none
child = self.branch(fs.path.name(path1), path2, {})
return self.as_callback(child)
@wraps(fn)
def func(path1, path2, **kwargs):
with make_callback(path1, path2) as callback:
return wrapped(path1, path2, callback=callback, **kwargs)
return func
def __enter__(self):
return self
def __exit__(self, *exc_args):
self.close()
def close(self):
"""Handle here on exit."""
@classmethod
def as_callback(
cls, maybe_callback: Optional["FsspecCallback"] = None
@@ -40,16 +74,33 @@ class FsspecCallback(fsspec.Callback):
return DEFAULT_CALLBACK
return maybe_callback
@classmethod
def as_tqdm_callback(
cls,
callback: Optional["FsspecCallback"] = None,
**tqdm_kwargs: Any,
) -> "FsspecCallback":
return callback or TqdmCallback(**tqdm_kwargs)
class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback):
pass
class TqdmCallback(FsspecCallback):
def __init__(self, progress_bar):
self.progress_bar = progress_bar
def __init__(self, progress_bar: "Tqdm" = None, **tqdm_kwargs):
from dvc.progress import Tqdm
self.progress_bar = progress_bar or Tqdm(**tqdm_kwargs)
super().__init__()
def __enter__(self):
self.progress_bar.__enter__()
return self
def close(self):
self.progress_bar.close()
def set_size(self, size):
if size is not None:
self.progress_bar.total = size
@@ -64,19 +115,8 @@ class TqdmCallback(FsspecCallback):
self.progress_bar.update_to(value)
super().absolute_update(value)
def tdqm_or_callback_wrapped(
fobj, method, total, callback=None, **pbar_kwargs
):
if callback:
from funcy import nullcontext
from tqdm.utils import CallbackIOWrapper
callback.set_size(total)
wrapper = CallbackIOWrapper(callback.relative_update, fobj, method)
return nullcontext(wrapper)
return Tqdm.wrapattr(fobj, method, total=total, bytes=True, **pbar_kwargs)
def branch(self, path_1: str, path_2: str, kwargs):
return TqdmCallback(bytes=True, total=-1, desc=path_1, **kwargs)
DEFAULT_CALLBACK = NoOpCallback()
dvc/fs/base.py
+51 -114
View File
@@ -1,9 +1,6 @@
import contextlib
import logging
import os
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import partialmethod
from multiprocessing import cpu_count
from typing import (
IO,
@@ -22,9 +19,8 @@ from typing import (
from funcy import cached_property
from dvc.exceptions import DvcException
from dvc.ui import ui
from dvc.utils import tmp_fname
from dvc.utils.fs import makedirs, move
from dvc.utils.fs import as_atomic, makedirs
from dvc.utils.threadpool import ThreadPoolExecutor
from ._callback import DEFAULT_CALLBACK, FsspecCallback
@@ -325,145 +321,86 @@ class FileSystem:
self,
from_info: Union[AnyFSPath, IO],
to_info: AnyFSPath,
total: int = None,
desc: str = None,
callback=None,
no_progress_bar: bool = False,
**pbar_args: Any,
size: int = None,
callback: FsspecCallback = None,
):
if not hasattr(from_info, "read"):
from .local import localfs
desc = desc or localfs.path.name(from_info)
if not hasattr(from_info, "read"):
logger.debug("Uploading '%s' to '%s'", from_info, to_info)
desc = localfs.path.name(from_info)
else:
desc = self.path.name(to_info)
stack = contextlib.ExitStack()
if not callback:
pbar = ui.progress(
with FsspecCallback.as_tqdm_callback(
callback,
desc=desc,
disable=no_progress_bar,
bytes=True,
total=total or -1,
**pbar_args,
)
stack.enter_context(pbar)
callback = pbar.as_callback()
with stack:
if total:
callback.set_size(total)
return self.put_file(
from_info, to_info, callback=callback, size=total
)
total=size or -1,
) as cb:
return self.put_file(from_info, to_info, callback=cb, size=size)
def download(
self,
from_info: AnyFSPath,
to_info: AnyFSPath,
name: str = None,
callback=None,
no_progress_bar: bool = False,
callback: "FsspecCallback" = None,
jobs: int = None,
_only_file: bool = False,
**kwargs: Any,
):
if not self.isdir(from_info):
return self.download_file(from_info, to_info, callback=callback)
from .local import localfs
if not hasattr(self, "get_file"):
raise RemoteActionNotImplemented("get_file", self.scheme)
download_dir = not _only_file and self.isdir(from_info)
desc = name or localfs.path.name(to_info)
stack = contextlib.ExitStack()
if not callback:
pbar_kwargs = {"unit": "Files"} if download_dir else {}
pbar = ui.progress(
total=-1,
desc="Downloading directory" if download_dir else desc,
bytes=not download_dir,
disable=no_progress_bar,
**pbar_kwargs,
pairs = {
info: localfs.path.join(
to_info, *self.path.relparts(info, from_info)
)
stack.enter_context(pbar)
callback = pbar.as_callback()
with stack:
if download_dir:
return self._download_dir(
from_info, to_info, callback=callback, jobs=jobs, **kwargs
)
return self._download_file(from_info, to_info, callback=callback)
download_file = partialmethod(download, _only_file=True)
def _download_dir(
self,
from_info: AnyFSPath,
to_info: AnyFSPath,
callback: FsspecCallback = DEFAULT_CALLBACK,
jobs: int = None,
**kwargs,
):
from .local import localfs
from_infos = list(self.find(from_info, **kwargs))
if not from_infos:
for info in self.find(from_info)
}
if not pairs:
makedirs(to_info, exist_ok=True)
return None
return
to_infos = (
localfs.path.join(to_info, *self.path.relparts(info, from_info))
for info in from_infos
)
callback.set_size(len(from_infos))
download_files = callback.wrap_fn(self._download_file)
with FsspecCallback.as_tqdm_callback(
callback,
total=-1,
desc=f"Downloading directory {self.path.name(from_info)}",
unit="files",
) as cb:
cb.set_size(len(pairs))
download_files = cb.wrap_and_branch(self.download_file, fs=self)
max_workers = jobs or self.jobs
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(download_files, from_info, to_info)
for from_info, to_info in zip(from_infos, to_infos)
]
# NOTE: unlike pulling/fetching cache, where we need to
# download everything we can, not raising an error here might
# turn very ugly, as the user might think that he has
# downloaded a complete directory, while having a partial one,
# which might cause unexpected results in his pipeline.
for future in as_completed(futures):
# NOTE: executor won't let us raise until all futures that
# it has are finished, so we need to cancel them ourselves
# before re-raising.
exc = future.exception()
if exc:
for entry in futures:
entry.cancel()
raise exc
with ThreadPoolExecutor(
max_workers=max_workers, cancel_on_error=True
) as executor:
list(
executor.imap_unordered(
lambda args: download_files(*args), pairs.items()
)
)
def _download_file(
def download_file(
self,
from_info: AnyFSPath,
to_info: AnyFSPath,
callback=DEFAULT_CALLBACK,
) -> None:
callback: FsspecCallback = None,
):
from .local import localfs
makedirs(localfs.path.parent(to_info), exist_ok=True)
tmp_file = tmp_fname(to_info)
logger.debug("Downloading '%s' to '%s'", from_info, to_info)
try:
# noqa, pylint: disable=no-member
self.get_file(from_info, tmp_file, callback=callback)
except Exception: # pylint: disable=broad-except
# do we need to rollback makedirs for previously not-existing
# directories?
with contextlib.suppress(FileNotFoundError):
os.unlink(tmp_file)
raise
move(tmp_file, to_info)
with FsspecCallback.as_tqdm_callback(
callback,
total=-1,
desc=self.path.name(from_info),
bytes=True,
) as cb:
with as_atomic(localfs, to_info, create_parents=True) as tmp_file:
self.get_file(from_info, tmp_file, callback=cb)
class ObjectFileSystem(FileSystem): # pylint: disable=abstract-method
dvc/fs/utils.py
+1 -1
View File
@@ -48,7 +48,7 @@ def _copy(
with from_fs.open(from_path, mode="rb") as fobj:
size = from_fs.size(from_path)
return to_fs.upload(fobj, to_path, total=size)
return to_fs.upload(fobj, to_path, size=size)
def _try_links(
dvc/progress.py
-5
View File
@@ -135,11 +135,6 @@ class Tqdm(tqdm):
return wrapped
def as_callback(self):
from dvc.fs._callback import TqdmCallback
return TqdmCallback(self)
def close(self):
self.postfix["info"] = ""
# remove ETA (either unknown or zero); remove completed bar
dvc/utils/fs.py
+13 -9
View File
@@ -200,17 +200,17 @@ def copyfile(src, dest, callback=None, no_progress_bar=False, name=None):
try:
System.reflink(src, dest)
except OSError:
from dvc.fs._callback import tdqm_or_callback_wrapped
from dvc.fs._callback import FsspecCallback
with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
with tdqm_or_callback_wrapped(
fdest,
"write",
total,
callback=callback,
with FsspecCallback.as_tqdm_callback(
callback,
total=total,
bytes=True,
disable=no_progress_bar,
desc=name,
) as wrapped:
) as cb:
wrapped = cb.wrap_attr(fdest, "write")
while True:
buf = fsrc.read(LOCAL_CHUNK_SIZE)
if not buf:
@@ -234,10 +234,14 @@ def walk_files(directory):
@contextmanager
def as_atomic(fs, to_info):
def as_atomic(fs, to_info, create_parents=False):
from dvc.utils import tmp_fname
tmp_info = fs.path.join(fs.path.parent(to_info), tmp_fname())
parent = fs.path.parent(to_info)
if create_parents:
fs.makedirs(parent, exist_ok=True)
tmp_info = fs.path.join(parent, tmp_fname())
try:
yield tmp_info
except BaseException: