remote: Optimize traverse/no_traverse behavior (#3501)

* remote: Optimize traverse/no_traverse behavior

- estimate remote file count by fetching a single parent cache dir, then
  determine whether or not to use no_traverse method for checking
  remainder of cache entries in `cache_exists`
- thread requests when traversing full remote file lists (by fetching
  one parent cache dir per thread)

* Use arbitrary prefix length for traverse remote cache queries

- default to 3
- for remotes that only support per-directory query use 2 (gdrive, hdfs)

* Add unit tests for new traverse/no_traverse behavior

* Fix deepsource warning

* Fix review issues

* fix CI

* Fix review issues

* obsolete no_traverse config option

- remove RemoteBASE.DEFAULT_NO_TRAVERSE and replace it with
  RemoteBASE.CAN_TRAVERSE (True for all remotes except http)

* review issues - use 'object exists' instead of 'no_traverse'

* fix obsolete no_traverse test issues

* fix SSHMocked func tests

dvc/config.py

@@ -92,7 +92,7 @@ class RelPath(str):
 REMOTE_COMMON = {
     "url": str,
     "checksum_jobs": All(Coerce(int), Range(1)),
-    "no_traverse": Bool,
+    Optional("no_traverse"): Bool,  # obsoleted
     "verify": Bool,
 }
 LOCAL_COMMON = {

dvc/remote/azure.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import posixpath
 import re
 from datetime import datetime, timedelta
 from urllib.parse import urlparse
@@ -103,8 +104,14 @@ class RemoteAZURE(RemoteBASE):
 
             next_marker = blobs.next_marker
 
-    def list_cache_paths(self):
-        return self._list_paths(self.path_info.bucket, self.path_info.path)
+    def list_cache_paths(self, prefix=None):
+        if prefix:
+            prefix = posixpath.join(
+                self.path_info.path, prefix[:2], prefix[2:]
+            )
+        else:
+            prefix = self.path_info.path
+        return self._list_paths(self.path_info.bucket, prefix)
 
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs

dvc/remote/base.py

@@ -80,8 +80,12 @@ class RemoteBASE(object):
     CHECKSUM_DIR_SUFFIX = ".dir"
     CHECKSUM_JOBS = max(1, min(4, cpu_count() // 2))
     DEFAULT_CACHE_TYPES = ["copy"]
-    DEFAULT_NO_TRAVERSE = True
     DEFAULT_VERIFY = False
+    LIST_OBJECT_PAGE_SIZE = 1000
+    TRAVERSE_WEIGHT_MULTIPLIER = 20
+    TRAVERSE_PREFIX_LEN = 3
+    TRAVERSE_THRESHOLD_SIZE = 500000
+    CAN_TRAVERSE = True
 
     CACHE_MODE = None
     SHARED_MODE_MAP = {None: (None, None), "group": (None, None)}
@@ -101,7 +105,6 @@ class RemoteBASE(object):
             or (self.repo and self.repo.config["core"].get("checksum_jobs"))
             or self.CHECKSUM_JOBS
         )
-        self.no_traverse = config.get("no_traverse", self.DEFAULT_NO_TRAVERSE)
         self.verify = config.get("verify", self.DEFAULT_VERIFY)
         self._dir_info = {}
 
@@ -688,7 +691,7 @@ class RemoteBASE(object):
     def checksum_to_path_info(self, checksum):
         return self.path_info / checksum[0:2] / checksum[2:]
 
-    def list_cache_paths(self):
+    def list_cache_paths(self, prefix=None):
         raise NotImplementedError
 
     def all(self):
@@ -804,11 +807,13 @@ class RemoteBASE(object):
 
         There are two ways of performing this check:
 
-        - Traverse: Get a list of all the files in the remote
+        - Traverse method: Get a list of all the files in the remote
             (traversing the cache directory) and compare it with
-            the given checksums.
+            the given checksums. Cache entries will be retrieved in parallel
+            threads according to prefix (i.e. entries starting with, "00...",
+            "01...", and so on) and a progress bar will be displayed.
 
-        - No traverse: For each given checksum, run the `exists`
+        - Exists method: For each given checksum, run the `exists`
             method and filter the checksums that aren't on the remote.
             This is done in parallel threads.
             It also shows a progress bar when performing the check.
@@ -819,12 +824,120 @@ class RemoteBASE(object):
         check if particular file exists much quicker, use their own
         implementation of cache_exists (see ssh, local).
 
+        Which method to use will be automatically determined after estimating
+        the size of the remote cache, and comparing the estimated size with
+        len(checksums). To estimate the size of the remote cache, we fetch
+        a small subset of cache entries (i.e. entries starting with "00...").
+        Based on the number of entries in that subset, the size of the full
+        cache can be estimated, since the cache is evenly distributed according
+        to checksum.
+
         Returns:
             A list with checksums that were found in the remote
         """
-        if not self.no_traverse:
-            return list(set(checksums) & set(self.all()))
+        # Remotes which do not use traverse prefix should override
+        # cache_exists() (see ssh, local)
+        assert self.TRAVERSE_PREFIX_LEN >= 2
 
+        if len(checksums) == 1 or not self.CAN_TRAVERSE:
+            return self._cache_object_exists(checksums, jobs, name)
+
+        # Fetch cache entries beginning with "00..." prefix for estimating the
+        # size of entire remote cache
+        checksums = frozenset(checksums)
+        prefix = "0" * self.TRAVERSE_PREFIX_LEN
+        total_prefixes = pow(16, self.TRAVERSE_PREFIX_LEN)
+        remote_checksums = set(
+            map(self.path_to_checksum, self.list_cache_paths(prefix=prefix))
+        )
+        if remote_checksums:
+            remote_size = total_prefixes * len(remote_checksums)
+        else:
+            remote_size = total_prefixes
+        logger.debug("Estimated remote size: {} files".format(remote_size))
+
+        traverse_pages = remote_size / self.LIST_OBJECT_PAGE_SIZE
+        # For sufficiently large remotes, traverse must be weighted to account
+        # for performance overhead from large lists/sets.
+        # From testing with S3, for remotes with 1M+ files, object_exists is
+        # faster until len(checksums) is at least 10k~100k
+        if remote_size > self.TRAVERSE_THRESHOLD_SIZE:
+            traverse_weight = traverse_pages * self.TRAVERSE_WEIGHT_MULTIPLIER
+        else:
+            traverse_weight = traverse_pages
+        if len(checksums) < traverse_weight:
+            logger.debug(
+                "Large remote ('{}' checksums < '{}' traverse weight), "
+                "using object_exists for remaining checksums".format(
+                    len(checksums), traverse_weight
+                )
+            )
+            return list(
+                checksums & remote_checksums
+            ) + self._cache_object_exists(
+                checksums - remote_checksums, jobs, name
+            )
+
+        if traverse_pages < 256 / self.JOBS:
+            # Threaded traverse will require making at least 255 more requests
+            # to the remote, so for small enough remotes, fetching the entire
+            # list at once will require fewer requests (but also take into
+            # account that this must be done sequentially rather than in
+            # parallel)
+            logger.debug(
+                "Querying {} checksums via default traverse".format(
+                    len(checksums)
+                )
+            )
+            return list(checksums & set(self.all()))
+
+        return self._cache_exists_traverse(
+            checksums, remote_checksums, jobs, name
+        )
+
+    def _cache_exists_traverse(
+        self, checksums, remote_checksums, jobs=None, name=None
+    ):
+        logger.debug(
+            "Querying {} checksums via threaded traverse".format(
+                len(checksums)
+            )
+        )
+
+        traverse_prefixes = ["{:02x}".format(i) for i in range(1, 256)]
+        if self.TRAVERSE_PREFIX_LEN > 2:
+            traverse_prefixes += [
+                "{0:0{1}x}".format(i, self.TRAVERSE_PREFIX_LEN)
+                for i in range(1, pow(16, self.TRAVERSE_PREFIX_LEN - 2))
+            ]
+        with Tqdm(
+            desc="Querying "
+            + ("cache in " + name if name else "remote cache"),
+            total=len(traverse_prefixes),
+            unit="dir",
+        ) as pbar:
+
+            def list_with_update(prefix):
+                ret = map(
+                    self.path_to_checksum,
+                    list(self.list_cache_paths(prefix=prefix)),
+                )
+                pbar.update_desc(
+                    "Querying cache in '{}'".format(self.path_info / prefix)
+                )
+                return ret
+
+            with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
+                in_remote = executor.map(list_with_update, traverse_prefixes,)
+                remote_checksums.update(
+                    itertools.chain.from_iterable(in_remote)
+                )
+            return list(checksums & remote_checksums)
+
+    def _cache_object_exists(self, checksums, jobs=None, name=None):
+        logger.debug(
+            "Querying {} checksums via object_exists".format(len(checksums))
+        )
         with Tqdm(
             desc="Querying "
             + ("cache in " + name if name else "remote cache"),

dvc/remote/gdrive.py

@@ -73,8 +73,10 @@ class RemoteGDrive(RemoteBASE):
     scheme = Schemes.GDRIVE
     path_cls = GDriveURLInfo
     REQUIRES = {"pydrive2": "pydrive2"}
-    DEFAULT_NO_TRAVERSE = False
     DEFAULT_VERIFY = True
+    # Always prefer traverse for GDrive since API usage quotas are a concern.
+    TRAVERSE_WEIGHT_MULTIPLIER = 1
+    TRAVERSE_PREFIX_LEN = 2
 
     GDRIVE_CREDENTIALS_DATA = "GDRIVE_CREDENTIALS_DATA"
     DEFAULT_USER_CREDENTIALS_FILE = "gdrive-user-credentials.json"
@@ -414,12 +416,18 @@ class RemoteGDrive(RemoteBASE):
         file_id = self._get_remote_id(from_info)
         self.gdrive_download_file(file_id, to_file, name, no_progress_bar)
 
-    def list_cache_paths(self):
+    def list_cache_paths(self, prefix=None):
         if not self.cache["ids"]:
             return
 
+        if prefix:
+            dir_ids = self.cache["dirs"].get(prefix[:2])
+            if not dir_ids:
+                return
+        else:
+            dir_ids = self.cache["ids"]
         parents_query = " or ".join(
-            "'{}' in parents".format(dir_id) for dir_id in self.cache["ids"]
+            "'{}' in parents".format(dir_id) for dir_id in dir_ids
         )
         query = "({}) and trashed=false".format(parents_query)
 

dvc/remote/gs.py

@@ -3,6 +3,7 @@ from datetime import timedelta
 from functools import wraps
 import io
 import os.path
+import posixpath
 import threading
 
 from funcy import cached_property, wrap_prop
@@ -126,14 +127,18 @@ class RemoteGS(RemoteBASE):
 
         blob.delete()
 
-    def _list_paths(self, path_info, max_items=None):
+    def _list_paths(self, path_info, max_items=None, prefix=None):
+        if prefix:
+            prefix = posixpath.join(path_info.path, prefix[:2], prefix[2:])
+        else:
+            prefix = path_info.path
         for blob in self.gs.bucket(path_info.bucket).list_blobs(
             prefix=path_info.path, max_results=max_items
         ):
             yield blob.name
 
-    def list_cache_paths(self):
-        return self._list_paths(self.path_info)
+    def list_cache_paths(self, prefix=None):
+        return self._list_paths(self.path_info, prefix=prefix)
 
     def walk_files(self, path_info):
         for fname in self._list_paths(path_info / ""):

dvc/remote/hdfs.py

@@ -21,6 +21,7 @@ class RemoteHDFS(RemoteBASE):
     REGEX = r"^hdfs://((?P<user>.*)@)?.*$"
     PARAM_CHECKSUM = "checksum"
     REQUIRES = {"pyarrow": "pyarrow"}
+    TRAVERSE_PREFIX_LEN = 2
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
@@ -152,16 +153,26 @@ class RemoteHDFS(RemoteBASE):
                 raise FileNotFoundError(*e.args)
             raise
 
-    def list_cache_paths(self):
+    def list_cache_paths(self, prefix=None):
         if not self.exists(self.path_info):
             return
 
-        dirs = deque([self.path_info.path])
+        if prefix:
+            root = posixpath.join(self.path_info.path, prefix[:2])
+        else:
+            root = self.path_info.path
+        dirs = deque([root])
 
         with self.hdfs(self.path_info) as hdfs:
             while dirs:
-                for entry in hdfs.ls(dirs.pop(), detail=True):
-                    if entry["kind"] == "directory":
-                        dirs.append(urlparse(entry["name"]).path)
-                    elif entry["kind"] == "file":
-                        yield urlparse(entry["name"]).path
+                try:
+                    for entry in hdfs.ls(dirs.pop(), detail=True):
+                        if entry["kind"] == "directory":
+                            dirs.append(urlparse(entry["name"]).path)
+                        elif entry["kind"] == "file":
+                            yield urlparse(entry["name"]).path
+                except IOError as e:
+                    # When searching for a specific prefix pyarrow raises an
+                    # exception if the specified cache dir does not exist
+                    if not prefix:
+                        raise e

dvc/remote/http.py

@@ -6,7 +6,6 @@ from funcy import cached_property, memoize, wrap_prop, wrap_with
 
 from dvc.path_info import HTTPURLInfo
 import dvc.prompt as prompt
-from dvc.config import ConfigError
 from dvc.exceptions import DvcException, HTTPError
 from dvc.progress import Tqdm
 from dvc.remote.base import RemoteBASE
@@ -32,6 +31,7 @@ class RemoteHTTP(RemoteBASE):
     REQUEST_TIMEOUT = 10
     CHUNK_SIZE = 2 ** 16
     PARAM_CHECKSUM = "etag"
+    CAN_TRAVERSE = False
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
@@ -45,12 +45,6 @@ class RemoteHTTP(RemoteBASE):
         else:
             self.path_info = None
 
-        if not self.no_traverse:
-            raise ConfigError(
-                "HTTP doesn't support traversing the remote to list existing "
-                "files. Use: `dvc remote modify <name> no_traverse true`"
-            )
-
         self.auth = config.get("auth", None)
         self.custom_auth_header = config.get("custom_auth_header", None)
         self.password = config.get("password", None)

dvc/remote/local.py

@@ -56,7 +56,8 @@ class RemoteLOCAL(RemoteBASE):
     def supported(cls, config):
         return True
 
-    def list_cache_paths(self):
+    def list_cache_paths(self, prefix=None):
+        assert prefix is None
         assert self.path_info is not None
         return walk_files(self.path_info)
 

dvc/remote/oss.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import posixpath
 import threading
 
 from funcy import cached_property, wrap_prop
@@ -95,8 +96,14 @@ class RemoteOSS(RemoteBASE):
         for blob in oss2.ObjectIterator(self.oss_service, prefix=prefix):
             yield blob.key
 
-    def list_cache_paths(self):
-        return self._list_paths(self.path_info.path)
+    def list_cache_paths(self, prefix=None):
+        if prefix:
+            prefix = posixpath.join(
+                self.path_info.path, prefix[:2], prefix[2:]
+            )
+        else:
+            prefix = self.path_info.path
+        return self._list_paths(prefix)
 
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs

dvc/remote/s3.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import posixpath
 import threading
 
 from funcy import cached_property, wrap_prop
@@ -190,24 +191,27 @@ class RemoteS3(RemoteBASE):
         logger.debug("Removing {}".format(path_info))
         self.s3.delete_object(Bucket=path_info.bucket, Key=path_info.path)
 
-    def _list_objects(self, path_info, max_items=None):
+    def _list_objects(self, path_info, max_items=None, prefix=None):
         """ Read config for list object api, paginate through list objects."""
         kwargs = {
             "Bucket": path_info.bucket,
             "Prefix": path_info.path,
             "PaginationConfig": {"MaxItems": max_items},
         }
+        if prefix:
+            kwargs["Prefix"] = posixpath.join(path_info.path, prefix[:2])
         paginator = self.s3.get_paginator(self.list_objects_api)
         for page in paginator.paginate(**kwargs):
             yield from page.get("Contents", ())
 
-    def _list_paths(self, path_info, max_items=None):
+    def _list_paths(self, path_info, max_items=None, prefix=None):
         return (
-            item["Key"] for item in self._list_objects(path_info, max_items)
+            item["Key"]
+            for item in self._list_objects(path_info, max_items, prefix)
         )
 
-    def list_cache_paths(self):
-        return self._list_paths(self.path_info)
+    def list_cache_paths(self, prefix=None):
+        return self._list_paths(self.path_info, prefix=prefix)
 
     def isfile(self, path_info):
         from botocore.exceptions import ClientError

dvc/remote/ssh/__init__.py

@@ -249,7 +249,8 @@ class RemoteSSH(RemoteBASE):
             else:
                 yield io.TextIOWrapper(fd, encoding=encoding)
 
-    def list_cache_paths(self):
+    def list_cache_paths(self, prefix=None):
+        assert prefix is None
         with self.ssh(self.path_info) as ssh:
             # If we simply return an iterator then with above closes instantly
             yield from ssh.walk_files(self.path_info.path)
@@ -297,7 +298,7 @@ class RemoteSSH(RemoteBASE):
         faster than current approach (relying on exists(path_info)) applied in
         remote/base.
         """
-        if not self.no_traverse:
+        if not self.CAN_TRAVERSE:
             return list(set(checksums) & set(self.all()))
 
         # possibly prompt for credentials before "Querying" progress output

tests/func/test_data_cloud.py

@@ -265,11 +265,11 @@ class TestRemoteSSHMocked(SSHMocked, TestDataCloudBase):
         repo = self.get_url()
         keyfile = self._get_keyfile()
 
+        self._get_cloud_class().CAN_TRAVERSE = False
         config = copy.deepcopy(TEST_CONFIG)
         config["remote"][TEST_REMOTE] = {
             "url": repo,
             "keyfile": keyfile,
-            "no_traverse": False,
         }
         self.dvc.config = config
         self.cloud = DataCloud(self.dvc)

tests/unit/remote/test_base.py

@@ -2,6 +2,7 @@ from unittest import TestCase
 
 import mock
 
+from dvc.path_info import PathInfo
 from dvc.remote.base import RemoteBASE
 from dvc.remote.base import RemoteCmdError
 from dvc.remote.base import RemoteMissingDepsError
@@ -35,3 +36,77 @@ class TestCmdError(TestCase, TestRemoteBASE):
         ):
             with self.assertRaises(RemoteCmdError):
                 self.REMOTE_CLS(repo, config).remove("file")
+
+
+@mock.patch.object(RemoteBASE, "_cache_exists_traverse")
+@mock.patch.object(RemoteBASE, "_cache_object_exists")
+@mock.patch.object(
+    RemoteBASE, "path_to_checksum", side_effect=lambda x: x,
+)
+def test_cache_exists(path_to_checksum, object_exists, traverse):
+    remote = RemoteBASE(None, {})
+
+    # remote does not support traverse
+    remote.CAN_TRAVERSE = False
+    with mock.patch.object(
+        remote, "list_cache_paths", return_value=list(range(256))
+    ):
+        checksums = list(range(1000))
+        remote.cache_exists(checksums)
+        object_exists.assert_called_with(checksums, None, None)
+        traverse.assert_not_called()
+
+    remote.CAN_TRAVERSE = True
+
+    # large remote, small local
+    object_exists.reset_mock()
+    traverse.reset_mock()
+    with mock.patch.object(
+        remote, "list_cache_paths", return_value=list(range(256))
+    ):
+        checksums = list(range(1000))
+        remote.cache_exists(checksums)
+        object_exists.assert_called_with(
+            frozenset(range(256, 1000)), None, None
+        )
+        traverse.assert_not_called()
+
+    # large remote, large local
+    object_exists.reset_mock()
+    traverse.reset_mock()
+    remote.JOBS = 16
+    with mock.patch.object(
+        remote, "list_cache_paths", return_value=list(range(256))
+    ):
+        checksums = list(range(1000000))
+        remote.cache_exists(checksums)
+        object_exists.assert_not_called()
+        traverse.assert_called_with(
+            frozenset(checksums), set(range(256)), None, None
+        )
+
+    # default traverse
+    object_exists.reset_mock()
+    traverse.reset_mock()
+    remote.TRAVERSE_WEIGHT_MULTIPLIER = 1
+    with mock.patch.object(remote, "list_cache_paths", return_value=[0]):
+        checksums = set(range(1000000))
+        remote.cache_exists(checksums)
+        traverse.assert_not_called()
+        object_exists.assert_not_called()
+
+
+@mock.patch.object(
+    RemoteBASE, "list_cache_paths", return_value=[],
+)
+@mock.patch.object(
+    RemoteBASE, "path_to_checksum", side_effect=lambda x: x,
+)
+def test_cache_exists_traverse(path_to_checksum, list_cache_paths):
+    remote = RemoteBASE(None, {})
+    remote.path_info = PathInfo("foo")
+    remote._cache_exists_traverse({0}, set())
+    for i in range(1, 16):
+        list_cache_paths.assert_any_call(prefix="{:03x}".format(i))
+    for i in range(1, 256):
+        list_cache_paths.assert_any_call(prefix="{:02x}".format(i))

tests/unit/remote/test_gdrive.py

@@ -28,7 +28,6 @@ class TestRemoteGDrive(object):
     def test_init(self):
         remote = RemoteGDrive(Repo(), self.CONFIG)
         assert str(remote.path_info) == self.CONFIG["url"]
-        assert not remote.no_traverse
 
     def test_drive(self):
         remote = RemoteGDrive(Repo(), self.CONFIG)

tests/unit/remote/test_http.py

@@ -1,23 +1,11 @@
 import pytest
 
-from dvc.config import ConfigError
 from dvc.exceptions import HTTPError
 from dvc.path_info import URLInfo
 from dvc.remote.http import RemoteHTTP
 from tests.utils.httpd import StaticFileServer
 
 
-def test_no_traverse_compatibility(dvc):
-    config = {
-        "url": "http://example.com/",
-        "path_info": "file.html",
-        "no_traverse": False,
-    }
-
-    with pytest.raises(ConfigError):
-        RemoteHTTP(dvc, config)
-
-
 def test_download_fails_on_error_code(dvc):
     with StaticFileServer() as httpd:
         url = "http://localhost:{}/".format(httpd.server_port)

tests/unit/test_config.py

@@ -1,16 +0,0 @@
-from dvc.config import COMPILED_SCHEMA
-
-
-def test_remote_config_no_traverse():
-    d = COMPILED_SCHEMA({"remote": {"myremote": {"url": "url"}}})
-    assert "no_traverse" not in d["remote"]["myremote"]
-
-    d = COMPILED_SCHEMA(
-        {"remote": {"myremote": {"url": "url", "no_traverse": "fAlSe"}}}
-    )
-    assert not d["remote"]["myremote"]["no_traverse"]
-
-    d = COMPILED_SCHEMA(
-        {"remote": {"myremote": {"url": "url", "no_traverse": "tRuE"}}}
-    )
-    assert d["remote"]["myremote"]["no_traverse"]