remote: introduce cache_exists() method

Making this method separate from exists() makes it easier to optimize
requests for large amounts of cache, since we can re-use the same prefix
for retrieving all the necessary information in a one single request.

Related #1430

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/dependency/base.py

@@ -60,7 +60,7 @@ class DependencyBase(object):
 
     @property
     def exists(self):
-        return self.remote.exists([self.path_info])[0]
+        return self.remote.exists(self.path_info)
 
     def changed(self):
         raise NotImplementedError

dvc/remote/azure.py

@@ -49,7 +49,7 @@ class RemoteAzure(RemoteBase):
             or match.group('container_name')  # backward compatibility
             or os.getenv('AZURE_STORAGE_CONTAINER_NAME'))
 
-        self.prefix = urlparse(self.url).path.lstrip('/') or '/'
+        self.prefix = urlparse(self.url).path.lstrip('/')
 
         self.connection_string = (
             config.get(Config.SECTION_AZURE_CONNECTION_STRING)
@@ -145,29 +145,39 @@ class RemoteAzure(RemoteBase):
         return [{
             'scheme': self.scheme,
             'bucket': self.bucket,
-            'key': '{}/{}/{}'.format(self.prefix, md5[0:2], md5[2:])
+            'key': posixpath.join(self.prefix, md5[0:2], md5[2:])
         } for md5 in md5s]
 
-    def _all_keys(self, bucket, prefix):
-        return {blob.name
-                for blob in self.blob_service.list_blobs(bucket,
-                                                         prefix=prefix)}
+    def _list_keys(self, bucket, prefix):
+        blob_service = self.blob_service
+        next_marker = None
+        while True:
+            blobs = blob_service.list_blobs(bucket,
+                                            prefix=prefix,
+                                            marker=next_marker)
 
-    def exists(self, path_infos):
-        ret = []
+            for blob in blobs:
+                yield blob.name
 
-        if len(path_infos) == 0:
-            return ret
+            if not blobs.next_marker:
+                break
 
-        bucket = path_infos[0]['bucket']
-        key = path_infos[0]['key']
+            next_marker = blobs.next_marker
 
-        for path_info in path_infos:
-            assert path_info['scheme'] == self.scheme
-            assert path_info['bucket'] == bucket
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
 
-            key = path_info['key']
-            ret.append(key in self._all_keys(bucket, key))
+        if len(md5s) == 0:
+            return []
+
+        ret = len(md5s) * [False]
+        keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s]
+        for key in self._list_keys(self.bucket, self.prefix):
+            for i, k in enumerate(keys):
+                if k == key:
+                    ret[i] = True
+
+        assert len(ret) == len(keys) == len(md5s)
 
         return ret
 
@@ -246,8 +256,14 @@ class RemoteAzure(RemoteBase):
         return posixpath.dirname(relpath) + posixpath.basename(relpath)
 
     def _all(self):
-        keys = self._all_keys(self.bucket, self.prefix)
-        return [self._path_to_etag(key) for key in keys]
+        # NOTE: The list might be way too big(e.g. 100M entries, md5 for each
+        # is 32 bytes, so ~3200Mb list) and we don't really need all of it at
+        # the same time, so it makes sense to use a generator to gradually
+        # iterate over it, without keeping all of it in memory.
+        return (
+            self._path_to_etag(key)
+            for key in self._list_keys(self.bucket, self.prefix)
+        )
 
     def gc(self, cinfos):
         from dvc.remote.local import RemoteLOCAL

dvc/remote/base.py

@@ -131,6 +131,9 @@ class RemoteBase(object):
     def md5s_to_path_infos(self, md5s):
         raise NotImplementedError
 
+    def cache_exists(self, md5s):
+        raise NotImplementedError
+
     def exists(self, path_infos):
         raise NotImplementedError
 

dvc/remote/gs.py

@@ -33,6 +33,10 @@ class RemoteGS(RemoteBase):
         self.url = config.get(Config.SECTION_REMOTE_URL, storagepath)
         self.projectname = config.get(Config.SECTION_GCP_PROJECTNAME, None)
 
+        parsed = urlparse(self.url)
+        self.bucket = parsed.netloc
+        self.prefix = parsed.path.lstrip('/')
+
     @staticmethod
     def compat_config(config):
         ret = config.copy()
@@ -40,14 +44,6 @@ class RemoteGS(RemoteBase):
         ret[Config.SECTION_REMOTE_URL] = url
         return ret
 
-    @property
-    def bucket(self):
-        return urlparse(self.url).netloc
-
-    @property
-    def prefix(self):
-        return urlparse(self.url).path.lstrip('/')
-
     @property
     def gs(self):
         return storage.Client()
@@ -82,7 +78,7 @@ class RemoteGS(RemoteBase):
         cache = {'scheme': 'gs', 'bucket': self.bucket, 'key': key}
 
         if {self.PARAM_MD5: md5} != self.save_info(cache):
-            if self.exists([cache])[0]:
+            if self.exists(cache):
                 msg = 'Corrupted cache file {}'
                 logger.warn(msg.format(self.to_string(cache)))
                 self.remove(cache)
@@ -91,7 +87,7 @@ class RemoteGS(RemoteBase):
         return False
 
     def changed(self, path_info, checksum_info):
-        if not self.exists([path_info])[0]:
+        if not self.exists(path_info):
             return True
 
         md5 = checksum_info.get(self.PARAM_MD5, None)
@@ -146,7 +142,7 @@ class RemoteGS(RemoteBase):
             logger.warn(msg.format(md5, self.to_string(path_info)))
             return
 
-        if self.exists([path_info])[0]:
+        if self.exists(path_info):
             msg = "Data '{}' exists. Removing before checkout."
             logger.warn(msg.format(self.to_string(path_info)))
             self.remove(path_info)
@@ -179,26 +175,32 @@ class RemoteGS(RemoteBase):
                  'key': posixpath.join(self.prefix,
                                        md5[0:2], md5[2:])} for md5 in md5s]
 
-    def exists(self, path_infos):
-        ret = []
+    def _list_keys(self, bucket, prefix):
+        for blob in self.gs.bucket(bucket).list_blobs(prefix=prefix):
+            yield blob.name
 
-        if len(path_infos) == 0:
-            return ret
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] == 'gs'
 
-        gs = self.gs
-        bucket_name = path_infos[0]['bucket']
-        bucket = gs.bucket(bucket_name)
+        keys = self._list_keys(path_info['bucket'], path_info['key'])
+        return any(path_info['key'] == key for key in keys)
 
-        for path_info in path_infos:
-            assert self.scheme == path_info['scheme']
-            assert bucket_name == path_info['bucket']
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
+
+        if len(md5s) == 0:
+            return []
+
+        ret = len(md5s) * [False]
+        keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s]
+        for key in self._list_keys(self.bucket, self.prefix):
+            for i, k in enumerate(keys):
+                if k == key:
+                    ret[i] = True
+
+        assert len(ret) == len(keys) == len(md5s)
 
-            exists = False
-            key = path_info['key']
-            keys = [blob.name for blob in bucket.list_blobs(prefix=key)]
-            if key in keys:
-                exists = True
-            ret.append(exists)
         return ret
 
     def upload(self, from_infos, to_infos, names=None):
@@ -292,9 +294,14 @@ class RemoteGS(RemoteBase):
         return posixpath.dirname(relpath) + posixpath.basename(relpath)
 
     def _all_md5s(self):
-        blobs = self.gs.bucket(self.bucket).list_blobs(prefix=self.prefix)
-        blobs = list(blobs)
-        return [self._path_to_md5(blob.name) for blob in blobs]
+        # NOTE: The list might be way too big(e.g. 100M entries, md5 for each
+        # is 32 bytes, so ~3200Mb list) and we don't really need all of it at
+        # the same time, so it makes sense to use a generator to gradually
+        # iterate over it, without keeping all of it in memory.
+        return (
+            self._path_to_md5(key)
+            for key in self._list_keys(self.bucket, self.prefix)
+        )
 
     def gc(self, cinfos):
         used = [info[self.PARAM_MD5] for info in cinfos['gs']]

dvc/remote/hdfs.py

@@ -93,7 +93,7 @@ class RemoteHDFS(RemoteBase):
         cache['url'] = posixpath.join(self.url, checksum[0:2], checksum[2:])
 
         if {self.PARAM_CHECKSUM: checksum} != self.save_info(cache):
-            if self.exists([cache])[0]:
+            if self.exists(cache):
                 msg = 'Corrupted cache file {}'
                 logger.warn(msg.format(self.to_string(cache)))
                 self.remove(cache)
@@ -102,7 +102,7 @@ class RemoteHDFS(RemoteBase):
         return False
 
     def changed(self, path_info, checksum_info):
-        if not self.exists([path_info])[0]:
+        if not self.exists(path_info):
             return True
 
         checksum = checksum_info.get(self.PARAM_CHECKSUM, None)
@@ -148,7 +148,7 @@ class RemoteHDFS(RemoteBase):
             logger.warn(msg.format(checksum, self.to_string(path_info)))
             return
 
-        if self.exists([path_info])[0]:
+        if self.exists(path_info):
             msg = "Data '{}' exists. Removing before checkout."
             logger.warn(msg.format(self.to_string(path_info)))
             self.remove(path_info)
@@ -178,20 +178,28 @@ class RemoteHDFS(RemoteBase):
                  'url': posixpath.join(self.url,
                                        md5[0:2], md5[2:])} for md5 in md5s]
 
-    def exists(self, path_infos):
-        ret = []
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] == 'hdfs'
 
-        if len(path_infos) == 0:
-            return ret
+        try:
+            self.hadoop_fs('test -e {}'.format(path_info['url']))
+            return True
+        except RemoteHDFSCmdError:
+            return False
 
-        for path_info in path_infos:
-            assert path_info['scheme'] == self.scheme
-            try:
-                self.hadoop_fs('test -e {}'.format(path_info['url']))
-                exists = True
-            except RemoteHDFSCmdError:
-                exists = False
-            ret.append(exists)
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
+
+        if len(md5s) == 0:
+            return []
+
+        existing_md5s = self._all_checksums()
+        ret = len(md5s) * [False]
+        for existing_md5 in existing_md5s:
+            for i, md5 in enumerate(md5s):
+                if md5 == existing_md5:
+                    ret[i] = True
 
         return ret
 
@@ -242,6 +250,11 @@ class RemoteHDFS(RemoteBase):
         return posixpath.dirname(relpath) + posixpath.basename(relpath)
 
     def _all_checksums(self):
+        try:
+            self.hadoop_fs('test -e {}'.format(self.url))
+        except RemoteHDFSCmdError:
+            return []
+
         stdout = self.hadoop_fs('ls -R {}'.format(self.url))
         lines = stdout.split('\n')
         flist = []

dvc/remote/http.py

@@ -82,11 +82,19 @@ class RemoteHTTP(RemoteBase):
             if not no_progress_bar:
                 progress.finish_target(name)
 
-    def exists(self, path_infos):
-        return [
-            bool(self._request('HEAD', path_info.get('url')))
-            for path_info in path_infos
-        ]
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] in ['http', 'https']
+        return bool(self._request('HEAD', path_info.get('url')))
+
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
+        ret = []
+        for md5 in md5s:
+            url = posixpath.join(self.prefix, md5[0:2], md5[2:])
+            exists = bool(self._request('HEAD', url))
+            ret.append(exists)
+        return ret
 
     def save_info(self, path_info):
         if path_info['scheme'] not in ['http', 'https']:

dvc/remote/local.py

@@ -101,6 +101,11 @@ class RemoteLOCAL(RemoteBase):
             return True
         return False
 
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] == 'local'
+        return os.path.exists(path_info['path'])
+
     def changed_cache(self, md5):
         cache = self.get(md5)
         clist = [(cache, md5)]
@@ -465,7 +470,7 @@ class RemoteLOCAL(RemoteBase):
             - The checkusm stored in the State is different from the given one
             - There's no file in the cache
         """
-        if not self.exists([path_info])[0]:
+        if not self.exists(path_info):
             return True
 
         md5 = checksum_info.get(self.PARAM_MD5, None)
@@ -494,11 +499,12 @@ class RemoteLOCAL(RemoteBase):
                  'path': os.path.join(self.prefix,
                                       md5[0:2], md5[2:])} for md5 in md5s]
 
-    def exists(self, path_infos):
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
         ret = []
-        for path_info in path_infos:
-            assert path_info['scheme'] == 'local'
-            ret.append(os.path.exists(path_info['path']))
+        for md5 in md5s:
+            path = os.path.join(self.prefix, md5[0:2], md5[2:])
+            ret.append(os.path.exists(path))
         return ret
 
     def upload(self, from_infos, to_infos, names=None):
@@ -610,13 +616,9 @@ class RemoteLOCAL(RemoteBase):
         md5s, names = self._group(checksum_infos,
                                   show_checksums=show_checksums)
 
-        progress.update_target(title, 20, 100)
-
-        path_infos = remote.md5s_to_path_infos(md5s)
-
         progress.update_target(title, 30, 100)
 
-        remote_exists = remote.exists(path_infos)
+        remote_exists = remote.cache_exists(md5s)
 
         progress.update_target(title, 90, 100)
 
@@ -663,7 +665,7 @@ class RemoteLOCAL(RemoteBase):
 
         # NOTE: dummy call to try to establish a connection
         # to see if we need to ask user for a password.
-        remote.exists(remote.md5s_to_path_infos(['000']))
+        remote.cache_exists(['000'])
 
         progress.update_target(title, 70, 100)
 
@@ -711,7 +713,7 @@ class RemoteLOCAL(RemoteBase):
 
         # NOTE: filter files that are already uploaded
         md5s = [i[self.PARAM_MD5] for i in checksum_infos]
-        exists = remote.exists(remote.md5s_to_path_infos(md5s))
+        exists = remote.cache_exists(md5s)
 
         progress.update_target(title, 30, 100)
 

dvc/remote/s3.py

@@ -66,6 +66,10 @@ class RemoteS3(RemoteBase):
         if shared_creds:
             os.environ.setdefault('AWS_SHARED_CREDENTIALS_FILE', shared_creds)
 
+        parsed = urlparse(self.url)
+        self.bucket = parsed.netloc
+        self.prefix = parsed.path.lstrip('/')
+
     @staticmethod
     def compat_config(config):
         ret = config.copy()
@@ -73,14 +77,6 @@ class RemoteS3(RemoteBase):
         ret[Config.SECTION_REMOTE_URL] = url
         return ret
 
-    @property
-    def bucket(self):
-        return urlparse(self.url).netloc
-
-    @property
-    def prefix(self):
-        return urlparse(self.url).path.lstrip('/')
-
     @property
     def s3(self):
         session = boto3.session.Session(profile_name=self.profile,
@@ -106,7 +102,7 @@ class RemoteS3(RemoteBase):
                                                path_info['key'])}
 
     def changed(self, path_info, checksum_info):
-        if not self.exists([path_info])[0]:
+        if not self.exists(path_info):
             return True
 
         etag = checksum_info.get(self.PARAM_ETAG, None)
@@ -146,7 +142,7 @@ class RemoteS3(RemoteBase):
         cache = {'scheme': 's3', 'bucket': self.bucket, 'key': key}
 
         if {self.PARAM_ETAG: etag} != self.save_info(cache):
-            if self.exists([cache])[0]:
+            if self.exists(cache):
                 msg = 'Corrupted cache file {}'
                 logger.warn(msg.format(self.to_string(cache)))
                 self.remove(cache)
@@ -172,7 +168,7 @@ class RemoteS3(RemoteBase):
             logger.warn(msg.format(etag, self.to_string(path_info)))
             return
 
-        if self.exists([path_info])[0]:
+        if self.exists(path_info):
             msg = "Data '{}' exists. Removing before checkout."
             logger.warn(msg.format(self.to_string(path_info)))
             self.remove(path_info)
@@ -202,10 +198,9 @@ class RemoteS3(RemoteBase):
                  'key': posixpath.join(self.prefix,
                                        md5[0:2], md5[2:])} for md5 in md5s]
 
-    def _all_keys(self, bucket, prefix):
+    def _list_keys(self, bucket, prefix):
         s3 = self.s3
 
-        keys = []
         kwargs = {'Bucket': bucket,
                   'Prefix': prefix}
         while True:
@@ -219,7 +214,7 @@ class RemoteS3(RemoteBase):
                 break
 
             for obj in contents:
-                keys.append(obj['Key'])
+                yield obj['Key']
 
             token = resp.get('NextContinuationToken', None)
             if not token:
@@ -227,24 +222,27 @@ class RemoteS3(RemoteBase):
 
             kwargs['ContinuationToken'] = token
 
-        return keys
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] == 's3'
 
-    def exists(self, path_infos):
-        ret = []
+        keys = self._list_keys(path_info['bucket'], path_info['key'])
+        return any(path_info['key'] == key for key in keys)
 
-        if len(path_infos) == 0:
-            return ret
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
 
-        bucket = path_infos[0]['bucket']
+        if len(md5s) == 0:
+            return []
 
-        for path_info in path_infos:
-            assert path_info['scheme'] == self.scheme
-            assert path_info['bucket'] == bucket
-            exists = False
-            key = path_info['key']
-            if key in self._all_keys(bucket, key):
-                exists = True
-            ret.append(exists)
+        ret = len(md5s) * [False]
+        keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s]
+        for key in self._list_keys(self.bucket, self.prefix):
+            for i, k in enumerate(keys):
+                if k == key:
+                    ret[i] = True
+
+        assert len(keys) == len(ret) == len(md5s)
 
         return ret
 
@@ -342,8 +340,14 @@ class RemoteS3(RemoteBase):
         return posixpath.dirname(relpath) + posixpath.basename(relpath)
 
     def _all(self):
-        keys = self._all_keys(self.bucket, self.prefix)
-        return [self._path_to_etag(key) for key in keys]
+        # NOTE: The list might be way too big(e.g. 100M entries, md5 for each
+        # is 32 bytes, so ~3200Mb list) and we don't really need all of it at
+        # the same time, so it makes sense to use a generator to gradually
+        # iterate over it, without keeping all of it in memory.
+        return (
+            self._path_to_etag(key)
+            for key in self._list_keys(self.bucket, self.prefix)
+        )
 
     def gc(self, cinfos):
         used_etags = [info[self.PARAM_ETAG] for info in cinfos['s3']]

dvc/remote/ssh.py

@@ -105,31 +105,33 @@ class RemoteSSH(RemoteBase):
 
         return ssh
 
-    def exists(self, path_infos):
-        ret = []
+    def exists(self, path_info):
+        assert not isinstance(path_info, list)
+        assert path_info['scheme'] == 'ssh'
 
-        if len(path_infos) == 0:
-            return ret
-
-        host = path_infos[0]['host']
-        user = path_infos[0]['user']
-        port = path_infos[0]['port']
-        ssh = self.ssh(host=host,
-                       user=user,
-                       port=port)
-
-        for path_info in path_infos:
-            assert host == path_info['host']
-            assert user == path_info['user']
-            assert port == path_info['port']
+        with self.ssh(path_info['host'],
+                      path_info['user'],
+                      path_info['port']) as ssh:
             try:
                 self._exec(ssh, 'test -e {}'.format(path_info['path']))
                 exists = True
             except RemoteSSHCmdError:
                 exists = False
-            ret.append(exists)
 
-        ssh.close()
+        return exists
+
+    def cache_exists(self, md5s):
+        assert isinstance(md5s, list)
+
+        if len(md5s) == 0:
+            return []
+
+        existing_md5s = self._all_md5s()
+        ret = len(md5s) * [False]
+        for existing_md5 in existing_md5s:
+            for i, md5 in enumerate(md5s):
+                if md5 == existing_md5:
+                    ret[i] = True
 
         return ret
 
@@ -245,7 +247,7 @@ class RemoteSSH(RemoteBase):
         cache['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:])
 
         if {self.PARAM_MD5: md5} != self.save_info(cache):
-            if self.exists([cache])[0]:
+            if self.exists(cache):
                 msg = 'Corrupted cache file {}'
                 logger.warn(msg.format(self.to_string(cache)))
                 self.remove(cache)
@@ -254,7 +256,7 @@ class RemoteSSH(RemoteBase):
         return False
 
     def changed(self, path_info, checksum_info):
-        if not self.exists([path_info])[0]:
+        if not self.exists(path_info):
             return True
 
         md5 = checksum_info.get(self.PARAM_MD5, None)
@@ -296,7 +298,7 @@ class RemoteSSH(RemoteBase):
             logger.warn(msg.format(md5, self.to_string(path_info)))
             return
 
-        if self.exists([path_info])[0]:
+        if self.exists(path_info):
             msg = "Data '{}' exists. Removing before checkout."
             logger.warn(msg.format(self.to_string(path_info)))
             self.remove(path_info)