lfs: internal `_filter_paths()` function is prohibitively slow

[sisp]

The internal helper function _filter_paths() in the lfs subpackage is prohibitively slow at the moment.

For instance, when I host the (rather small) MNIST dataset (70,000 image files) using Git LFS and import it using dvc import, then _filter_paths() is effectively called like this:

paths = ["data/train/0/1.png", "data/train/0/2.png", ...]  # list of all 70,000 image files

for path in _filter_paths(paths, include=paths, exclude=None):
    # yield LFS pointer data

So, in this simple case, the current implementation has quadratic complexity in the number of LFS files and materializes generators unnecessarily, which is already significant with only 70,000 files – imagine large datasets.

I've done a quick benchmark to quantify the current implementation performance and compare it against two alternatives:

• Current implementation:

  import timeit
  
  setup = """\
  import fnmatch
  from typing import Iterable, Iterator, Optional
  
  import fsspec
  
  
  def _filter_paths(
      paths: Iterable[str], include: Optional[list[str]], exclude: Optional[list[str]]
  ) -> Iterator[str]:
      filtered = set()
      if include:
          for pattern in include:
              filtered.update(fnmatch.filter(paths, pattern))
      else:
          filtered.update(paths)
      if exclude:
          for pattern in exclude:
              filtered.difference_update(fnmatch.filter(paths, pattern))
      yield from filtered
  
  
  fs = fsspec.filesystem("file")
  paths = fs.find("data/")
  """
  
  
  timeit.timeit("list(_filter_paths(paths, include=paths, exclude=None))", setup=setup, number=1)
  # >>> 632.4136102720004 seconds

• A more efficient implementation with identical semantics by manually unionizing the filename regex patterns and matching each path against the pre-compiled single regex:

  import timeit
  
  setup = """\
  import fnmatch
  import re
  from typing import Iterable, Iterator, Optional
  
  import fsspec
  
  
  def _filter_paths(
      paths: Iterable[str], include: Optional[list[str]], exclude: Optional[list[str]]
  ) -> Iterator[str]:
      if include:
          include_match = re.compile(r"|".join(map(fnmatch.translate, include))).match
          paths = (path for path in paths if include_match(path) is not None)
      if exclude:
          exclude_match = re.compile(r"|".join(map(fnmatch.translate, exclude))).match
          paths = (path for path in paths if exclude_match(path) is None)
      yield from paths
  
  
  fs = fsspec.filesystem("file")
  paths = fs.find("data/")
  """
  
  
  timeit.timeit("list(_filter_paths(paths, include=paths, exclude=None))", setup=setup, number=1)
  # >>> 118.97085808499833 seconds

• Assuming we actually don't need Unix filename pattern matching (fnmatch) – it doesn't seem to be needed at the moment by DVC –, an implementation using set lookups (O(1) complexity) would be orders of magnitude faster:

  import timeit
  
  setup = """\
  import fnmatch
  import re
  from typing import Iterable, Iterator, Optional
  
  import fsspec
  
  
  def _filter_paths(
      paths: Iterable[str], include: Optional[list[str]], exclude: Optional[list[str]]
  ) -> Iterator[str]:
      if include:
          include_set = set(include)
          paths = (path for path in paths if path in include_set)
      if exclude:
          exclude_set = set(exclude)
          paths = (path for path in paths if path not in exclude_set)
      yield from paths
  
  
  fs = fsspec.filesystem("file")
  paths = fs.find("data/")
  """
  
  
  timeit.timeit("list(_filter_paths(paths, include=paths, exclude=None))", setup=setup, number=1)
  # >>> 0.00669500499861897 seconds

WDYT about the two alternative implementations? Is Unix filename pattern matching really needed?

[shcheklein]

Really great research, and yes it doesn't look good. Before we make a decision, a few questions / things to explore:

• in the _filter_paths(fs.find("/"), include, exclude): the fs.find("/") does it come from a repo root? Does it mean potentially data that is not even relevant to the include? I wonder if we could analyze directories in include and use those instead of /
• the way I understand this - in DVC we expand the directory and pass all the files in include (?). Can we actually pass dirname/* ?
• can we detect if there are no patterns (fnmatch is not needed) and fall back to set() ?

[sisp]

• in the _filter_paths(fs.find("/"), include, exclude): the fs.find("/") does it come from a repo root? [...]

Correct, / is the repo root because fs is a scmrepo.fs.GitFileSystem instance.

[...] Does it mean potentially data that is not even relevant to the include? [...]

Correct.

[...] I wonder if we could analyze directories in include and use those instead of / [...]

We probably could, although it might get tricky when Unix filename patterns are involved. I think there's a limit to what can be optimized given the current scmrepo.git.lfs.fetch() signature and the semantics of include/exclude. For instance, unionizing the Unix filename pattern regexes as shown above would be more performant with identical behavior.

• the way I understand this - in DVC we expand the directory and pass all the files in include (?). Can we actually pass dirname/* ?

Yes, I believe so. As you discovered, the way DVC calls dvc.scm.lfs_prefetch() (and thus scmrepo.git.lfs.fetch()) seems suboptimal to me now:

include in scmrepo.git.lfs.fetch() corresponds to all files listed in from_infos in dvc.fs.download() – all files in the Git/DVC repo IIUC. The enumeration of files in from_infos is needed to copy/download them. But I believe it's not needed to filter files for LFS object fetching – from_infos files have not been filtered with some advanced filtering like gitignore-like include/exclude rules, so this is just an enumeration of all files IIUC. But scmrepo.git.lfs.fetch._collect_objects() also enumerates files with fs.find("/"). Thus, include and fs.find("/") are redundant (not necessarily identical though when fs_path in dvc.fs.download() is not the Git repo root).

I think a minimally invasive change in dvc.fs.download() with a huge performance boost (when using DVC with LFS fetching) would be this:

-lfs_prefetch(fs, from_infos)
+lfs_prefetch(fs, [f"{fs.normpath(fs_path)}/**" if fs.isdir(fs_path) else fs_path])

(Is fs.normpath() correct here? I'm trying to make sure that fs_path has no trailing path separator.)

With this change, include in scmrepo.git.lfs.fetch() is a single-item list with either a Unix filename pattern like <subdir>/** or a file path. scmrepo.git.lfs.fetch._collect_objects() would still enumerate all files in the repo with fs.find("/") (that's quite fast even with many files) and scmrepo.git.lfs.fetch._filter_paths() (especially with the optimized implementation that still supports Unix filename pattern matching) would only match all files against a single pattern.

With this minor change, the above benchmark becomes:

import timeit

setup = """\
import fnmatch
import re
from typing import Iterable, Iterator, Optional

import fsspec


def _filter_paths(
    paths: Iterable[str], include: Optional[list[str]], exclude: Optional[list[str]]
) -> Iterator[str]:
    if include:
        include_match = re.compile(r"|".join(map(fnmatch.translate, include))).match
        paths = (path for path in paths if include_match(path) is not None)
    if exclude:
        exclude_match = re.compile(r"|".join(map(fnmatch.translate, exclude))).match
        paths = (path for path in paths if exclude_match(path) is None)
    yield from paths


fs = fsspec.filesystem("file")
paths = fs.find("data/")
"""


timeit.timeit("list(_filter_paths(paths, include=['data/**'], exclude=None))", setup=setup, number=1)
#                                                ^^^^^^^^^^^
#                                                 note this
# >>> 0.007259630001499318 seconds

Note that fs.find("data/") in the benchmark is basically equivalent to fs.find("/") in scmrepo.git.lfs.fetch._filter_paths() where fs is an instance of scmrepo.git.GitFileSystem – using fs.find(".") with a simple fsspec.implementations.local.LocalFileSystem instance would include all files in the .git directory as well.

I think the above _filter_paths() implementation and the small change in DVC would be a good way forward – given that you agree that the change in DVC is semantically correct.

• can we detect if there are no patterns (fnmatch is not needed) and fall back to set() ?

I haven't found an official check for this. We could try to come up with some handcrafted check, but that might be error-prone.

[shcheklein]

hey @sisp , sorry for the delay. The implementation you described sounds reasonable to me. I mean the lfs_prefetch(fs, [f"{fs.normpath(fs_path)}/**" if fs.isdir(fs_path) else fs_path]) change. Btw, thanks a lot of all the links and the proper syntax highlight in the summary - saves a LOT of time 🙏 Excellent style of writing comments :)

We definitely need to check if this is semantically correct. But even if it's not, it seems that it should be fixed anyways.

A few more questions / suggestions:

• lfs_prefetch - if it's not used anywhere except this download. Can we make it accept a single path then (a file or a dir).
• add and explicit optional param to the _lfs_fetch to specify root that we could use in find() later? seem possible on the level of the lfs_prefetch now, since we are dealing with a single path this way?

This way we can get rid of bottlenecks - traversing the whole repo (can take unpredictable amount of time) and this O^2 that you identified.

WDYT?

[sisp]

Thanks for the praise, @shcheklein! 😊 Adding links actually doesn't only ease the readers' lives but also mine when trying to write a meaningful comment especially with somewhat complex content – so win-win 😉.

I think your suggestion will work and keep scmrepo backwards compatible when done right. This would be my plan:

I'd replace dvc.scm.lfs_prefetch()'s paths: list[str] argument by path: str (accepting a directory or file path)

-def lfs_prefetch(fs: "FileSystem", paths: list[str]):
+def lfs_prefetch(fs: "FileSystem", path: str):
...
         _lfs_fetch(
             scm,
             [git_fs.rev],
-            include=[(path if path.startswith("/") else f"/{path}") for path in paths],
             progress=pbar.update_git,
+            root=path,
         ) 

and then add an optional argument root: str = "/" to scmrepo.git.lfs.fetch():

 def fetch(
     scm: "Git",
     revs: Optional[list[str]] = None,
     remote: Optional[str] = None,
     include: Optional[list[str]] = None,
     exclude: Optional[list[str]] = None,
     progress: Optional[Callable[["GitProgressEvent"], None]] = None,
+    root: str = "/",
 ):
...
-            for pointer in _collect_objects(scm, rev, include, exclude)
+            for pointer in _collect_objects(scm, rev, include, exclude, root)
...

Unfortunately (but not too bad), we need to append the root argument after progress for backwards compatibility because there's no keyword-only arguments marker. root then gets passed down to scmrepo.git.lfs.fetch._collect_objects() where we can simply pass it to the fs.find() call which should be super cheap when root is a file path and also conveniently returns and empty generator when a passed file path doesn't exist. There's one more detail though: I think it would be intuitive if include/exclude throughout the call stack were paths or Unix filename patterns relative to root (when root is a directory path – otherwise include/exclude make no sense anyway). Since fs.find(root) produces paths including the root path prefix, include/exclude passed to _filter_paths() must be prepended with root:

 def _collect_objects(
     scm: "Git",
     rev: str,
     include: Optional[list[str]],
     exclude: Optional[list[str]],
+    root: str,
 ) -> Iterator[Pointer]:
     fs = scm.get_fs(rev)
-    for path in _filter_paths(fs.find("/"), include, exclude):
+    root_dir = root if fs.isdir(root) else fs.parent(root)
+    for path in _filter_paths(
+        fs.find(root),
+        [fs.join(root_dir, i.lstrip(fs.sep)) for i in (include or [])],
+        [fs.join(root_dir, e.lstrip(fs.sep)) for e in (exclude or [])],
+    ):
...

Finally, the dvc.scm.lfs_prefetch() call would look like this – nice and simple:

-            lfs_prefetch(fs, from_infos)
+            lfs_prefetch(fs, fs_path)

WDYT?

[shcheklein]

Sorry for the delay.

There's one more detail though: I think it would be intuitive if include/exclude throughout the call stack were paths or Unix filename patterns relative to root

yes, makes perfect sense to me!

Unfortunately (but not too bad), we need to append the root argument after progress for backwards compatibility

my 2cs - I would not worry to much - I doubt that there are many consumers, also we recently bumped dependency on DVC side - so, it's not like DVC can be very flexible with scmrepo versions. Your call on this, I'm fine with either approach here.

---

thinking more about this though. I would probably go with your initial suggestion:

lfs_prefetch(fs, [f"{fs.normpath(fs_path)}/**" if fs.isdir(fs_path) else fs_path])

Primarily, because this construct with root feels redundant. We need it only for optimization, right? But include, exclude provides enough information to do this kind of optimization (we just don't want to deal with this atm - and that's fine). So, changing public interface just to do a fast improvement for a specific situation probably not the right way to go.

If we pass a single path, I would though definitely try to short cut to make DVC faster. In the _collect_objects do a check if it's single path - get its prefix (before any wildcards) and pass it to find.

Just an idea. Probably if analyzing this path is still complicated, we can go just with your single-line suggestion.

WDYT?

[sisp]

Agreed. 👍 I've submitted #355 and treeverse/dvc#10414 to implement the discussed optimizations.