From c52b12ac2ade208988f6bbbaa026e85a81196efc Mon Sep 17 00:00:00 2001 From: skshetry <18718008+skshetry@users.noreply.github.com> Date: Fri, 17 Mar 2023 16:04:16 +0545 Subject: [PATCH] dvcfs: support caching remote streams; use that in plots (#9183) * dvcfs: support caching remote streams; use that in plots * plots: show progress bar, cap max_workers to 16 * bump dvc-data to 0.44.1 --- dvc/fs/dvc.py | 7 +++- dvc/repo/plots/__init__.py | 70 ++++++++++++++++++++++----------- dvc/utils/serialize/__init__.py | 7 ++++ pyproject.toml | 2 +- 4 files changed, 59 insertions(+), 27 deletions(-) diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index 7a78ec878..2ace0af7e 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -251,12 +251,15 @@ class _DVCFileSystem(AbstractFileSystem): # pylint:disable=abstract-method try: return self.repo.fs.open(fs_path, mode=mode) except FileNotFoundError: - _, dvc_fs, subkey = self._get_subrepo_info(key) + repo, dvc_fs, subkey = self._get_subrepo_info(key) if not dvc_fs: raise dvc_path = _get_dvc_path(dvc_fs, subkey) - return dvc_fs.open(dvc_path, mode=mode) + kw = {} + if kwargs.get("cache_remote_stream", False): + kw["cache_odb"] = repo.cache.local + return dvc_fs.open(dvc_path, mode=mode, **kw) def isdvc(self, path, **kwargs) -> bool: """Is this entry dvc-tracked?""" diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py index cf36d36c3..2706533b9 100644 --- a/dvc/repo/plots/__init__.py +++ b/dvc/repo/plots/__init__.py @@ -20,13 +20,14 @@ from typing import ( import dpath import dpath.options -from funcy import distinct, first, project +from funcy import distinct, first, project, reraise from dvc.exceptions import DvcException from dvc.utils import error_handler, errored_revisions, onerror_collect from dvc.utils.objects import cached_property -from dvc.utils.serialize import LOADERS +from dvc.utils.serialize import PARSERS, EncodingError from dvc.utils.threadpool import ThreadPoolExecutor +from dvc_render.image import ImageRenderer if TYPE_CHECKING: from dvc.fs import FileSystem @@ -39,6 +40,9 @@ dpath.options.ALLOW_EMPTY_STRING_KEYS = True logger = logging.getLogger(__name__) +SUPPORTED_IMAGE_EXTENSIONS = ImageRenderer.EXTENSIONS + + class PlotMetricTypeError(DvcException): def __init__(self, file): super().__init__( @@ -196,7 +200,10 @@ class Plots: onerror=onerror, props=props, ): - _resolve_data_sources(data) + short_rev = "workspace" + if rev := getattr(self.repo.fs, "rev", None): + short_rev = rev[:7] + _resolve_data_sources(data, short_rev, cache_remote_stream=True) result.update(data) errored = errored_revisions(result) @@ -265,7 +272,11 @@ def _is_plot(out: "Output") -> bool: return bool(out.plot) -def _resolve_data_sources(plots_data: Dict): +def _resolve_data_sources( + plots_data: Dict, rev: str, cache_remote_stream: bool = False +): + from dvc.progress import Tqdm + values = list(plots_data.values()) to_resolve = [] while values: @@ -278,16 +289,26 @@ def _resolve_data_sources(plots_data: Dict): def resolve(value): data_source = value.pop("data_source") assert callable(data_source) - value.update(data_source()) + value.update(data_source(cache_remote_stream=cache_remote_stream)) + + if not to_resolve: + return executor = ThreadPoolExecutor( - max_workers=4 * cpu_count(), + max_workers=min(16, 4 * cpu_count()), thread_name_prefix="resolve_data", cancel_on_error=True, ) with executor: - # imap_unordered is lazy, wrapping to trigger it - list(executor.imap_unordered(resolve, to_resolve)) + iterable = executor.imap_unordered(resolve, to_resolve) + with Tqdm( + iterable, + total=len(to_resolve), + desc=f"Reading plot's data from {rev}", + unit="files", + unit_scale=False, + ) as progress_iterable: + list(progress_iterable) def _collect_plots( @@ -524,20 +545,25 @@ def unpack_if_dir(fs, path, props: Dict[str, str], onerror: Optional[Callable] = @error_handler -def parse(fs, path, props=None, **kwargs): +def parse(fs, path, props=None, **fs_kwargs): props = props or {} _, extension = os.path.splitext(path) - if extension in (".tsv", ".csv"): - header = props.get("header", True) - if extension == ".csv": - return _load_sv(path=path, fs=fs, delimiter=",", header=header) - return _load_sv(path=path, fs=fs, delimiter="\t", header=header) - if extension in LOADERS or extension in (".yml", ".yaml"): - return LOADERS[extension](path=path, fs=fs) - if extension in (".jpeg", ".jpg", ".gif", ".png", ".svg"): - with fs.open(path, "rb") as fd: + if extension in SUPPORTED_IMAGE_EXTENSIONS: + with fs.open(path, mode="rb", **fs_kwargs) as fd: return fd.read() - raise PlotMetricTypeError(path) + + if extension not in PARSERS.keys() | {".yml", ".yaml", ".csv", ".tsv"}: + raise PlotMetricTypeError(path) + + with reraise(UnicodeDecodeError, EncodingError(path, "utf8")): + with fs.open(path, mode="r", encoding="utf8", **fs_kwargs) as fd: + contents = fd.read() + + if extension in (".csv", ".tsv"): + header = props.get("header", True) + delim = "\t" if extension == ".tsv" else "," + return _load_sv(contents, delimiter=delim, header=header) + return PARSERS[extension](contents, path) def _plot_props(out: "Output") -> Dict: @@ -553,10 +579,7 @@ def _plot_props(out: "Output") -> Dict: return project(out.plot, PLOT_PROPS) -def _load_sv(path, fs, delimiter=",", header=True): - with fs.open(path, "r") as fd: - content = fd.read() - +def _load_sv(content, delimiter=",", header=True): if header: reader = csv.DictReader(io.StringIO(content), delimiter=delimiter) else: @@ -566,5 +589,4 @@ def _load_sv(path, fs, delimiter=",", header=True): delimiter=delimiter, fieldnames=[str(i) for i in range(len(first_row))], ) - return list(reader) diff --git a/dvc/utils/serialize/__init__.py b/dvc/utils/serialize/__init__.py index bd8bfa75a..c5f7f9759 100644 --- a/dvc/utils/serialize/__init__.py +++ b/dvc/utils/serialize/__init__.py @@ -12,6 +12,13 @@ LOADERS: DefaultDict[str, LoaderFn] = defaultdict( # noqa: F405 ) LOADERS.update({".toml": load_toml, ".json": load_json, ".py": load_py}) # noqa: F405 +PARSERS: DefaultDict[str, ParserFn] = defaultdict( # noqa: F405 + lambda: parse_yaml # noqa: F405 +) +PARSERS.update( + {".toml": parse_toml, ".json": parse_json, ".py": parse_py} # noqa: F405 +) + def load_path(fs_path, fs): suffix = fs.path.suffix(fs_path).lower() diff --git a/pyproject.toml b/pyproject.toml index da5f7c485..bb9a78239 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ dependencies = [ "configobj>=5.0.6", "distro>=1.3", "dpath<3,>=2.1.0", - "dvc-data>=0.44.0,<0.45", + "dvc-data>=0.44.1,<0.45", "dvc-http", "dvc-render>=0.3.1,<0.4.0", "dvc-studio-client>=0.5.0,<1",