cli: improve startup speed

fsspec callbacks were making us lose 35ms (~207 ms vs 170ms now).
These are now extracted from dvc.progress to dvc.fs._callback.

dvc/commands/live.py

@@ -1,5 +1,4 @@
 import argparse
-from pathlib import Path
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import fix_subparsers
@@ -16,6 +15,8 @@ class CmdLive(CmdBase):
         metrics, plots = self.repo.live.show(target=target, revs=revs)
 
         if plots:
+            from pathlib import Path
+
             html_path = Path.cwd() / (self.args.target + "_html")
 
             renderers = match_renderers(plots, self.repo.plots.templates)

dvc/fs/_callback.py

@@ -0,0 +1,49 @@
+import fsspec
+
+from dvc.progress import Tqdm
+
+
+class FsspecCallback(fsspec.Callback):
+    def __init__(self, progress_bar):
+        self.progress_bar = progress_bar
+        super().__init__()
+
+    def set_size(self, size):
+        if size is not None:
+            self.progress_bar.total = size
+            self.progress_bar.refresh()
+            super().set_size(size)
+
+    def relative_update(self, inc=1):
+        self.progress_bar.update(inc)
+        super().relative_update(inc)
+
+    def absolute_update(self, value):
+        self.progress_bar.update_to(value)
+        super().absolute_update(value)
+
+    @staticmethod
+    def wrap_fn(cb, fn):
+        def wrapped(*args, **kwargs):
+            res = fn(*args, **kwargs)
+            cb.relative_update()
+            return res
+
+        return wrapped
+
+
+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)
+
+
+DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback()

dvc/fs/base.py

@@ -10,7 +10,7 @@ from funcy import cached_property
 from tqdm.utils import CallbackIOWrapper
 
 from dvc.exceptions import DvcException
-from dvc.progress import DEFAULT_CALLBACK, FsspecCallback
+from dvc.fs._callback import DEFAULT_CALLBACK, FsspecCallback
 from dvc.ui import ui
 from dvc.utils import tmp_fname
 from dvc.utils.fs import makedirs, move

dvc/fs/dvc.py

@@ -5,7 +5,7 @@ import typing
 from dvc.exceptions import OutputNotFoundError
 from dvc.utils import relpath
 
-from ..progress import DEFAULT_CALLBACK
+from ._callback import DEFAULT_CALLBACK
 from ._metadata import Metadata
 from .base import FileSystem
 

dvc/fs/fsspec_wrapper.py

@@ -5,8 +5,7 @@ from typing import IO, TYPE_CHECKING, Any, Dict, Iterator, Optional, overload
 from funcy import cached_property
 from tqdm.utils import CallbackIOWrapper
 
-from dvc.progress import DEFAULT_CALLBACK
-
+from ._callback import DEFAULT_CALLBACK
 from .base import FileSystem
 
 FSPath = str

dvc/fs/http.py

@@ -4,9 +4,9 @@ from typing import Any
 from funcy import cached_property, memoize, wrap_with
 
 from dvc import prompt
-from dvc.progress import DEFAULT_CALLBACK
 from dvc.scheme import Schemes
 
+from ._callback import DEFAULT_CALLBACK
 from .fsspec_wrapper import AnyFSPath, FSSpecWrapper, NoDirectoriesMixin
 
 

dvc/fs/local.py

@@ -6,7 +6,7 @@ from dvc.system import System
 from dvc.utils import is_exec, tmp_fname
 from dvc.utils.fs import copy_fobj_to_file, copyfile, makedirs, move, remove
 
-from ..progress import DEFAULT_CALLBACK
+from ._callback import DEFAULT_CALLBACK
 from .base import FileSystem
 
 logger = logging.getLogger(__name__)

dvc/fs/repo.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Tuple, Type, Union
 
 from funcy import lfilter, wrap_with
 
-from ..progress import DEFAULT_CALLBACK
+from ._callback import DEFAULT_CALLBACK
 from .base import FileSystem
 from .dvc import DvcFileSystem
 

dvc/fs/s3.py

@@ -5,9 +5,9 @@ from collections import defaultdict
 
 from funcy import cached_property, wrap_prop
 
-from dvc.progress import DEFAULT_CALLBACK
 from dvc.scheme import Schemes
 
+from ._callback import DEFAULT_CALLBACK
 from .fsspec_wrapper import ObjectFSWrapper
 
 _AWS_CONFIG_PATH = os.path.join(os.path.expanduser("~"), ".aws", "config")

dvc/fs/ssh.py

@@ -8,7 +8,7 @@ from dvc import prompt
 from dvc.scheme import Schemes
 from dvc.utils.fs import as_atomic
 
-from ..progress import DEFAULT_CALLBACK
+from ._callback import DEFAULT_CALLBACK
 from .fsspec_wrapper import FSSpecWrapper
 
 _SSH_TIMEOUT = 60 * 30

dvc/progress.py

@@ -3,7 +3,6 @@ import logging
 import sys
 from threading import RLock
 
-import fsspec
 from tqdm import tqdm
 
 from dvc.env import DVC_IGNORE_ISATTY
@@ -137,6 +136,8 @@ class Tqdm(tqdm):
         return wrapped
 
     def as_callback(self):
+        from dvc.fs._callback import FsspecCallback
+
         return FsspecCallback(self)
 
     def close(self):
@@ -164,49 +165,3 @@ class Tqdm(tqdm):
             d["ncols_desc"] = d["ncols_info"] = 1
             d["prefix"] = ""
         return d
-
-
-class FsspecCallback(fsspec.Callback):
-    def __init__(self, progress_bar):
-        self.progress_bar = progress_bar
-        super().__init__()
-
-    def set_size(self, size):
-        if size is not None:
-            self.progress_bar.total = size
-            self.progress_bar.refresh()
-            super().set_size(size)
-
-    def relative_update(self, inc=1):
-        self.progress_bar.update(inc)
-        super().relative_update(inc)
-
-    def absolute_update(self, value):
-        self.progress_bar.update_to(value)
-        super().absolute_update(value)
-
-    @staticmethod
-    def wrap_fn(cb, fn):
-        def wrapped(*args, **kwargs):
-            res = fn(*args, **kwargs)
-            cb.relative_update()
-            return res
-
-        return wrapped
-
-
-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)
-
-
-DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback()

dvc/utils/fs.py

@@ -200,7 +200,7 @@ def copyfile(src, dest, callback=None, no_progress_bar=False, name=None):
     try:
         System.reflink(src, dest)
     except OSError:
-        from dvc.progress import tdqm_or_callback_wrapped
+        from dvc.fs._callback import tdqm_or_callback_wrapped
 
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
             with tdqm_or_callback_wrapped(