From 71abef74f542ae5832c070089a3be132859f7c27 Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 7 Aug 2019 12:34:37 +0200 Subject: [PATCH 1/9] remote: azure: implement batch_exists --- dvc/remote/azure.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 0b8330ae3..b94a67884 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -1,6 +1,7 @@ from __future__ import absolute_import from __future__ import unicode_literals +import itertools import os import re import logging @@ -154,3 +155,13 @@ class RemoteAZURE(RemoteBASE): path_info.bucket, path_info.path, sas_token=sas_token ) return download_url + + 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] From 237f3e248b93c690710ff9f0b684c8cce6a3cb82 Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 7 Aug 2019 13:29:47 +0200 Subject: [PATCH 2/9] remote: base: modify cache_exists to use exists instead of batch_exists --- dvc/remote/azure.py | 6 +++++- dvc/remote/base.py | 10 +++------- dvc/remote/gs.py | 12 ------------ dvc/remote/http.py | 9 --------- dvc/remote/oss.py | 4 ++++ dvc/remote/s3.py | 11 ----------- dvc/remote/ssh/__init__.py | 32 -------------------------------- 7 files changed, 12 insertions(+), 72 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index b94a67884..ded3ad935 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -1,12 +1,12 @@ from __future__ import absolute_import from __future__ import unicode_literals -import itertools import os import re import logging from datetime import datetime, timedelta + from dvc.scheme import Schemes try: @@ -138,6 +138,10 @@ class RemoteAZURE(RemoteBASE): 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 9a387774a..30df8696d 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -20,6 +20,7 @@ from dvc.exceptions import ( DvcIgnoreInCollectedDirError, ) from dvc.progress import progress, ProgressCallback +from dvc.utils import LARGE_DIR_SIZE, tmp_fname, move, relpath from dvc.utils import ( LARGE_DIR_SIZE, tmp_fname, @@ -644,15 +645,10 @@ class RemoteBASE(object): """ 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 = list(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 b171b9c4e..4d07226cd 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 @@ class RemoteGS(RemoteBASE): 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 715465293..f377f0ddc 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -65,15 +65,6 @@ class RemoteHTTP(RemoteBASE): 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 79912edc3..7e1d2313d 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -119,3 +119,7 @@ class RemoteOSS(RemoteBASE): 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.bucket, 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 d9adab21b..757db8980 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -3,7 +3,6 @@ from __future__ import unicode_literals import os import threading import logging -import itertools from funcy import cached_property try: @@ -216,16 +215,6 @@ class RemoteS3(RemoteBASE): 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 22a7da5aa..8bb5c69f2 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -3,9 +3,6 @@ from __future__ import unicode_literals import os import getpass import logging -import itertools -import errno -from concurrent.futures import ThreadPoolExecutor import threading try: @@ -15,7 +12,6 @@ except ImportError: 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 +144,6 @@ class RemoteSSH(RemoteBASE): 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 From 98de99f5a94f24515b3faa36c9d6b71523aef92a Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 7 Aug 2019 13:31:27 +0200 Subject: [PATCH 3/9] remote: azure: make blob_service cached property --- dvc/remote/azure.py | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index ded3ad935..f51e0931d 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -6,6 +6,7 @@ import re import logging from datetime import datetime, timedelta +from funcy import cached_property from dvc.scheme import Schemes @@ -72,29 +73,27 @@ class RemoteAZURE(RemoteBASE): 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: From 132b571e7232f2857064ddb1da4b749908b5a305 Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 7 Aug 2019 15:22:56 +0200 Subject: [PATCH 4/9] remote: oss: fix _list_cache_args --- dvc/remote/oss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py index 7e1d2313d..cbf1fe3a7 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -121,5 +121,5 @@ class RemoteOSS(RemoteBASE): ) def exists(self, path_info): - paths = self._list_paths(path_info.bucket, path_info.path) + paths = self._list_paths(path_info.path) return any(path_info.path == path for path in paths) From 394d6ffeeb98e9b3a40b2c31c295e6db8cbf5725 Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 7 Aug 2019 15:49:56 +0200 Subject: [PATCH 5/9] remote: ssh: connection: remove open_max_sftp_channels --- dvc/remote/ssh/connection.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py index 8362857b8..9d64ee5b4 100644 --- a/dvc/remote/ssh/connection.py +++ b/dvc/remote/ssh/connection.py @@ -289,15 +289,3 @@ class SSHConnection: def cp(self, src, dest): self.makedirs(posixpath.dirname(dest)) self.execute("cp {} {}".format(src, dest)) - - def open_max_sftp_channels(self): - # If there are more than 1 it means we've already opened max amount - if len(self._sftp_channels) <= 1: - while True: - try: - self._sftp_channels.append(self._ssh.open_sftp()) - except paramiko.ssh_exception.ChannelException: - if not self._sftp_channels: - raise - break - return self._sftp_channels From f224790aababaa488d97e9cf666c4ff64a2a403b Mon Sep 17 00:00:00 2001 From: pawel Date: Fri, 9 Aug 2019 11:56:27 +0200 Subject: [PATCH 6/9] remote: base: remove unnecessary list casting --- dvc/remote/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 30df8696d..275481f01 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -648,7 +648,7 @@ class RemoteBASE(object): 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] - in_remote = list(executor.map(self.exists, path_infos)) + in_remote = executor.map(self.exists, path_infos) ret = list(itertools.compress(checksums, in_remote)) progress_callback.finish("") return ret From 8233f959c795ab6b0d2582dc1e8c45376f6364c8 Mon Sep 17 00:00:00 2001 From: pawel Date: Mon, 12 Aug 2019 11:22:51 +0200 Subject: [PATCH 7/9] remote: ssh: reintroduce batch_exists --- dvc/remote/azure.py | 10 ------- dvc/remote/base.py | 10 +------ dvc/remote/ssh/__init__.py | 57 ++++++++++++++++++++++++++++++++++++ dvc/remote/ssh/connection.py | 12 ++++++++ 4 files changed, 70 insertions(+), 19 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index f51e0931d..6e37545ed 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -158,13 +158,3 @@ class RemoteAZURE(RemoteBASE): path_info.bucket, path_info.path, sas_token=sas_token ) return download_url - - 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] diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 275481f01..2105372ae 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -20,15 +20,7 @@ from dvc.exceptions import ( DvcIgnoreInCollectedDirError, ) from dvc.progress import progress, ProgressCallback -from dvc.utils import LARGE_DIR_SIZE, tmp_fname, move, relpath -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 diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index 8bb5c69f2..4202ba61a 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -1,9 +1,15 @@ from __future__ import unicode_literals +import errno +import itertools import os import getpass import logging import threading +from concurrent.futures import ThreadPoolExecutor + +from dvc.progress import ProgressCallback +from dvc.utils import to_chunks try: import paramiko @@ -208,3 +214,54 @@ class RemoteSSH(RemoteBASE): 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())) diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py index 9d64ee5b4..8362857b8 100644 --- a/dvc/remote/ssh/connection.py +++ b/dvc/remote/ssh/connection.py @@ -289,3 +289,15 @@ class SSHConnection: def cp(self, src, dest): self.makedirs(posixpath.dirname(dest)) self.execute("cp {} {}".format(src, dest)) + + def open_max_sftp_channels(self): + # If there are more than 1 it means we've already opened max amount + if len(self._sftp_channels) <= 1: + while True: + try: + self._sftp_channels.append(self._ssh.open_sftp()) + except paramiko.ssh_exception.ChannelException: + if not self._sftp_channels: + raise + break + return self._sftp_channels From c308ad8a58dcce8e609b355d6bfcc82ed7a728c7 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Tue, 13 Aug 2019 02:08:29 +0300 Subject: [PATCH 8/9] remote: base: don't forget about the progress bar Signed-off-by: Ruslan Kuprieiev --- dvc/remote/base.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 2105372ae..fa6ab3b35 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -635,17 +635,21 @@ class RemoteBASE(object): Returns: A list with checksums that were found in the remote """ + if not self.no_traverse: + return list(set(checksums) & set(self.all())) + progress_callback = ProgressCallback(len(checksums)) + def exists_with_progress(path_info): + ret = self.exists(path_info) + progress_callback.update(str(path_info)) + return ret - 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] - in_remote = executor.map(self.exists, path_infos) - ret = list(itertools.compress(checksums, in_remote)) - progress_callback.finish("") - return ret - - return list(set(checksums) & set(self.all())) + with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor: + path_infos = [self.checksum_to_path_info(x) for x in checksums] + in_remote = executor.map(exists_with_progress, path_infos) + ret = list(itertools.compress(checksums, in_remote)) + progress_callback.finish("") + return ret def already_cached(self, path_info): current = self.get_checksum(path_info) From 06d347178f352f4b318d35b439b4325182c3b3d7 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Tue, 13 Aug 2019 02:25:15 +0300 Subject: [PATCH 9/9] remote: base: remove whitespace Signed-off-by: Ruslan Kuprieiev --- dvc/remote/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index fa6ab3b35..20f8b7f16 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -639,6 +639,7 @@ class RemoteBASE(object): return list(set(checksums) & set(self.all())) progress_callback = ProgressCallback(len(checksums)) + def exists_with_progress(path_info): ret = self.exists(path_info) progress_callback.update(str(path_info))