Repository: iterative/dvc
Pull request: remote: gs/s3: remove batch_exists
* [x] Have you followed the guidelines in our
      [Contributing document](https://dvc.org/doc/user-guide/contributing)?

* [x] Does your PR affect documented changes or does it add new functionality
      that should be documented? If yes, have you created a PR for
      [dvc.org](https://github.com/iterative/dvc.org) documenting it or at
      least opened an issue for it? If so, please add a link to it.

-----
Fix #2373 

Problem statement:
Collecting information from remote cache very slow
**Please provide information about your setup**
DVC version(i.e. `dvc --version`), Platform and method of installation (pip, homebrew, pkg Mac, exe (Windows), DEB(Linux), RPM(Linux))

Running under Windows 10 wsl as well as debian stretch with an azure blob storage as remote

$ dvc --version
0.54.1

installed from pip (Python 3)

---

When doing operations involving the remote cache dvc is very slow. Even when the files affected is very small if the total dvc repo has many files.

In a repo with 247 directories, 2444700 files totaling 132GB.

I was running a dvc pull on a dvc file pointing out a 4 file dataset of 1 byte each and this took 70 minutes.

jmollevi@LNOR070124:~/Projects/dvctest$ time dvc pull data/test.dvc
Preparing to download data from 'azure://jmollevidvctest/data/'
Preparing to collect status from azure://jmollevidvctest/data/
Collecting information from local cache...
[##############################] 100%

Collecting information from remote cache...
[##############################] 100% Analysing status
[##############################] 100% data/test

Preparing to download data from 'azure://jmollevidvctest/data/'

Preparing to collect status from azure://jmollevidvctest/data/

Collecting information from local cache...
[##############################] 100%

Collecting information from remote cache...
[##############################] 100% Analysing status
[##############################] 100% data/test/c
[##############################] 100% Created unpacked dir
[##############################] 100% Checkout finished!

real    70m2.014s
user    19m34.031s
sys     0m58.359s
jmollevi@LNOR070124:\~/Projects/dvctest$ grep  -e ''  data/test/*
data/test/a:a
data/test/b:b
data/test/c:e
data/test/f:d
jmollevi@LNOR070124:\~/Projects/dvctest$

the only changed files where in the data/test folder
the large folder with 247 directories and 2444700 files was not in local cache.

For nearly all of the time dvc was on the steps marked
Collecting information from ~local~ remote cache...

Other operations are similarly slow.

This slow performance is a blocker for me for an otherwise very suitable tool. (I am creating a training set for machine learning using semi-automated methods to generate all those files.)

Hints from issue discussion:
@JohanMollevik Thanks for repoting that! That seems like bug, let us investigate.
@pared noticed a typo in my original issue, edited it and corrected it
@JohanMollevik well, its not a bug, rather lack of optimization for azure. During pull we need to verify that checksum of your dataset (stored in '*.dvc' file) exists in remote. Basic implementation gets list of all checksums on remote and compares it with local, so in your case that might take some time. There is already implementation of optimization  for `s3` and `gs`, so we need to "translate" it to azure. We will get to it right away.

Commit message: remote: ssh: reintroduce batch_exists

Patch to review:
diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
index 0b8330ae3c..6e37545edb 100644
--- a/dvc/remote/azure.py
+++ b/dvc/remote/azure.py
@@ -6,6 +6,8 @@
 import logging
 from datetime import datetime, timedelta
 
+from funcy import cached_property
+
 from dvc.scheme import Schemes
 
 try:
@@ -71,29 +73,27 @@ def __init__(self, repo, config):
         if not self.connection_string:
             raise ValueError("azure storage connection string missing")
 
-        self.__blob_service = None
         self.path_info = (
             self.path_cls(url)
             if path
             else self.path_cls.from_parts(scheme=self.scheme, netloc=bucket)
         )
 
-    @property
+    @cached_property
     def blob_service(self):
-        if self.__blob_service is None:
-            logger.debug("URL {}".format(self.path_info))
-            logger.debug("Connection string {}".format(self.connection_string))
-            self.__blob_service = BlockBlobService(
-                connection_string=self.connection_string
+        logger.debug("URL {}".format(self.path_info))
+        logger.debug("Connection string {}".format(self.connection_string))
+        blob_service = BlockBlobService(
+            connection_string=self.connection_string
+        )
+        logger.debug("Container name {}".format(self.path_info.bucket))
+        try:  # verify that container exists
+            blob_service.list_blobs(
+                self.path_info.bucket, delimiter="/", num_results=1
             )
-            logger.debug("Container name {}".format(self.path_info.bucket))
-            try:  # verify that container exists
-                self.__blob_service.list_blobs(
-                    self.path_info.bucket, delimiter="/", num_results=1
-                )
-            except AzureMissingResourceHttpError:
-                self.__blob_service.create_container(self.path_info.bucket)
-        return self.__blob_service
+        except AzureMissingResourceHttpError:
+            blob_service.create_container(self.path_info.bucket)
+        return blob_service
 
     def remove(self, path_info):
         if path_info.scheme != self.scheme:
@@ -137,6 +137,10 @@ def _download(
             from_info.bucket, from_info.path, to_file, progress_callback=cb
         )
 
+    def exists(self, path_info):
+        paths = self._list_paths(path_info.bucket, path_info.path)
+        return any(path_info.path == path for path in paths)
+
     def open(self, path_info, mode="r", encoding=None):
         get_url = lambda: self._generate_download_url(path_info)  # noqa: E731
         return open_url(get_url, mode=mode, encoding=encoding)
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index 9a387774af..2105372ae8 100644
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -20,14 +20,7 @@
     DvcIgnoreInCollectedDirError,
 )
 from dvc.progress import progress, ProgressCallback
-from dvc.utils import (
-    LARGE_DIR_SIZE,
-    tmp_fname,
-    to_chunks,
-    move,
-    relpath,
-    makedirs,
-)
+from dvc.utils import LARGE_DIR_SIZE, tmp_fname, move, relpath, makedirs
 from dvc.state import StateNoop
 from dvc.path_info import PathInfo, URLInfo
 
@@ -644,15 +637,10 @@ def cache_exists(self, checksums, jobs=None):
         """
         progress_callback = ProgressCallback(len(checksums))
 
-        def exists_with_progress(chunks):
-            return self.batch_exists(chunks, callback=progress_callback)
-
-        if self.no_traverse and hasattr(self, "batch_exists"):
+        if self.no_traverse:
             with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
                 path_infos = [self.checksum_to_path_info(x) for x in checksums]
-                chunks = to_chunks(path_infos, num_chunks=self.JOBS)
-                results = executor.map(exists_with_progress, chunks)
-                in_remote = itertools.chain.from_iterable(results)
+                in_remote = executor.map(self.exists, path_infos)
                 ret = list(itertools.compress(checksums, in_remote))
                 progress_callback.finish("")
                 return ret
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
index b171b9c4ee..4d07226cde 100644
--- a/dvc/remote/gs.py
+++ b/dvc/remote/gs.py
@@ -1,7 +1,6 @@
 from __future__ import unicode_literals
 
 import logging
-import itertools
 from funcy import cached_property
 
 try:
@@ -89,17 +88,6 @@ def exists(self, path_info):
         paths = set(self._list_paths(path_info.bucket, path_info.path))
         return any(path_info.path == path for path in paths)
 
-    def batch_exists(self, path_infos, callback):
-        paths = []
-
-        for path_info in path_infos:
-            paths.append(self._list_paths(path_info.bucket, path_info.path))
-            callback.update(str(path_info))
-
-        paths = set(itertools.chain.from_iterable(paths))
-
-        return [path_info.path in paths for path_info in path_infos]
-
     def _upload(self, from_file, to_info, **_kwargs):
         bucket = self.gs.bucket(to_info.bucket)
         blob = bucket.blob(to_info.path)
diff --git a/dvc/remote/http.py b/dvc/remote/http.py
index 715465293a..f377f0ddc2 100644
--- a/dvc/remote/http.py
+++ b/dvc/remote/http.py
@@ -65,15 +65,6 @@ def _download(self, from_info, to_file, name=None, no_progress_bar=False):
     def exists(self, path_info):
         return bool(self._request("HEAD", path_info.url))
 
-    def batch_exists(self, path_infos, callback):
-        results = []
-
-        for path_info in path_infos:
-            results.append(self.exists(path_info))
-            callback.update(str(path_info))
-
-        return results
-
     def _content_length(self, url):
         return self._request("HEAD", url).headers.get("Content-Length")
 
diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py
index 79912edc36..cbf1fe3a75 100644
--- a/dvc/remote/oss.py
+++ b/dvc/remote/oss.py
@@ -119,3 +119,7 @@ def _download(
         self.oss_service.get_object_to_file(
             from_info.path, to_file, progress_callback=cb
         )
+
+    def exists(self, path_info):
+        paths = self._list_paths(path_info.path)
+        return any(path_info.path == path for path in paths)
diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py
index d9adab21b4..757db89806 100644
--- a/dvc/remote/s3.py
+++ b/dvc/remote/s3.py
@@ -3,7 +3,6 @@
 import os
 import threading
 import logging
-import itertools
 from funcy import cached_property
 
 try:
@@ -216,16 +215,6 @@ def exists(self, path_info):
         paths = self._list_paths(path_info.bucket, path_info.path)
         return any(path_info.path == path for path in paths)
 
-    def batch_exists(self, path_infos, callback):
-        paths = []
-
-        for path_info in path_infos:
-            paths.append(self._list_paths(path_info.bucket, path_info.path))
-            callback.update(str(path_info))
-
-        paths = set(itertools.chain.from_iterable(paths))
-        return [path_info.path in paths for path_info in path_infos]
-
     def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
         total = os.path.getsize(from_file)
         cb = None if no_progress_bar else Callback(name, total)
diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index 22a7da5aad..4202ba61a5 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/dvc/remote/ssh/__init__.py
@@ -1,12 +1,15 @@
 from __future__ import unicode_literals
 
+import errno
+import itertools
 import os
 import getpass
 import logging
-import itertools
-import errno
-from concurrent.futures import ThreadPoolExecutor
 import threading
+from concurrent.futures import ThreadPoolExecutor
+
+from dvc.progress import ProgressCallback
+from dvc.utils import to_chunks
 
 try:
     import paramiko
@@ -15,7 +18,6 @@
 
 import dvc.prompt as prompt
 from dvc.config import Config
-from dvc.utils import to_chunks
 from dvc.utils.compat import urlparse, StringIO
 from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
@@ -148,34 +150,6 @@ def exists(self, path_info):
         with self.ssh(path_info) as ssh:
             return ssh.exists(path_info.path)
 
-    def batch_exists(self, path_infos, callback):
-        def _exists(chunk_and_channel):
-            chunk, channel = chunk_and_channel
-            ret = []
-            for path in chunk:
-                try:
-                    channel.stat(path)
-                    ret.append(True)
-                except IOError as exc:
-                    if exc.errno != errno.ENOENT:
-                        raise
-                    ret.append(False)
-                callback.update(path)
-            return ret
-
-        with self.ssh(path_infos[0]) as ssh:
-            channels = ssh.open_max_sftp_channels()
-            max_workers = len(channels)
-
-            with ThreadPoolExecutor(max_workers=max_workers) as executor:
-                paths = [path_info.path for path_info in path_infos]
-                chunks = to_chunks(paths, num_chunks=max_workers)
-                chunks_and_channels = zip(chunks, channels)
-                outcome = executor.map(_exists, chunks_and_channels)
-                results = list(itertools.chain.from_iterable(outcome))
-
-            return results
-
     def get_file_checksum(self, path_info):
         if path_info.scheme != self.scheme:
             raise NotImplementedError
@@ -240,3 +214,54 @@ def walk(self, path_info):
     def makedirs(self, path_info):
         with self.ssh(path_info) as ssh:
             ssh.makedirs(path_info.path)
+
+    def batch_exists(self, path_infos, callback):
+        def _exists(chunk_and_channel):
+            chunk, channel = chunk_and_channel
+            ret = []
+            for path in chunk:
+                try:
+                    channel.stat(path)
+                    ret.append(True)
+                except IOError as exc:
+                    if exc.errno != errno.ENOENT:
+                        raise
+                    ret.append(False)
+                callback.update(path)
+            return ret
+
+        with self.ssh(path_infos[0]) as ssh:
+            channels = ssh.open_max_sftp_channels()
+            max_workers = len(channels)
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                paths = [path_info.path for path_info in path_infos]
+                chunks = to_chunks(paths, num_chunks=max_workers)
+                chunks_and_channels = zip(chunks, channels)
+                outcome = executor.map(_exists, chunks_and_channels)
+                results = list(itertools.chain.from_iterable(outcome))
+
+            return results
+
+    def cache_exists(self, checksums, jobs=None):
+        """This is older implementation used in remote/base.py
+        We are reusing it in RemoteSSH, because SSH's batch_exists proved to be
+        faster than current approach (relying on exists(path_info)) applied in
+        remote/base.
+        """
+        progress_callback = ProgressCallback(len(checksums))
+
+        def exists_with_progress(chunks):
+            return self.batch_exists(chunks, callback=progress_callback)
+
+        if self.no_traverse:
+            with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
+                path_infos = [self.checksum_to_path_info(x) for x in checksums]
+                chunks = to_chunks(path_infos, num_chunks=self.JOBS)
+                results = executor.map(exists_with_progress, chunks)
+                in_remote = itertools.chain.from_iterable(results)
+                ret = list(itertools.compress(checksums, in_remote))
+                progress_callback.finish("")
+                return ret
+
+        return list(set(checksums) & set(self.all()))


Code review comments:
diff_hunk: @@ -644,15 +637,10 @@ def cache_exists(self, checksums, jobs=None):
         """
         progress_callback = ProgressCallback(len(checksums))
 
-        def exists_with_progress(chunks):
-            return self.batch_exists(chunks, callback=progress_callback)
path: dvc/remote/base.py
text: @user1:
We've lost the progress bar :)