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

@@ -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?"""

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)

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()

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",