From 80ca3eedc847a840d516facb63f59d4db651c450 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Thu, 13 Jun 2019 22:21:38 +0700 Subject: [PATCH] dvc: unify remote upload/download code Bugs fixed along the way: - unable to copy between buckets in GS - unneeded connection init in RemoteGS.copy() - connection not closed on some errors in RemoteSSH.download() - inconsistent progress bar handling --- dvc/remote/azure.py | 85 ++++-------------------- dvc/remote/base.py | 102 +++++++++++++++++++++++++++-- dvc/remote/gs.py | 109 ++++++------------------------- dvc/remote/hdfs.py | 73 ++++++--------------- dvc/remote/https.py | 1 - dvc/remote/local/__init__.py | 86 +++++-------------------- dvc/remote/oss.py | 76 ++++------------------ dvc/remote/s3.py | 121 ++++++++++------------------------- dvc/remote/ssh/__init__.py | 101 +++++++++-------------------- dvc/remote/ssh/connection.py | 15 +---- 10 files changed, 246 insertions(+), 523 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index b31c65c5c..dfda92f6b 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -13,8 +13,7 @@ try: except ImportError: BlockBlobService = None -from dvc.utils import tmp_fname, move -from dvc.utils.compat import urlparse, makedirs, fspath_py35 +from dvc.utils.compat import urlparse from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBASE @@ -121,74 +120,18 @@ class RemoteAZURE(RemoteBASE): def list_cache_paths(self): return self._list_paths(self.path_info.bucket, self.path_info.path) - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) - - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != self.scheme: - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - logger.debug("Uploading '{}' to '{}'".format(from_info, to_info)) - if not name: - name = from_info.name - - cb = None if no_progress_bar else Callback(name) - - try: - self.blob_service.create_blob_from_path( - to_info.bucket, - to_info.path, - from_info.fspath, - progress_callback=cb, - ) - except Exception: - msg = "failed to upload '{}'".format(from_info) - logger.warning(msg) - else: - progress.finish_target(name) - - def download( - self, - from_infos, - to_infos, - no_progress_bar=False, - names=None, - resume=False, + def _upload( + self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs ): - names = self._verify_path_args(from_infos, to_infos, names) + cb = None if no_progress_bar else Callback(name) + self.blob_service.create_blob_from_path( + to_info.bucket, to_info.path, from_file, progress_callback=cb + ) - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != self.scheme: - raise NotImplementedError - - if to_info.scheme != "local": - raise NotImplementedError - - logger.debug("Downloading '{}' to '{}'".format(from_info, to_info)) - - tmp_file = tmp_fname(to_info) - if not name: - name = to_info.name - - cb = None if no_progress_bar else Callback(name) - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - - try: - self.blob_service.get_blob_to_path( - from_info.bucket, - from_info.path, - tmp_file, - progress_callback=cb, - ) - except Exception: - msg = "failed to download '{}'".format(from_info) - logger.warning(msg) - else: - move(tmp_file, to_info.fspath) - - if not no_progress_bar: - progress.finish_target(name) + def _download( + self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs + ): + cb = None if no_progress_bar else Callback(name) + self.blob_service.get_blob_to_path( + from_info.bucket, from_info.path, to_file, progress_callback=cb + ) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 8f1be8257..12ef420b2 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -1,12 +1,13 @@ from __future__ import unicode_literals -from dvc.utils.compat import str, basestring, urlparse +from dvc.utils.compat import str, basestring, urlparse, fspath_py35, makedirs import os import json import logging import tempfile import itertools +from contextlib import contextmanager from operator import itemgetter from multiprocessing import cpu_count from concurrent.futures import ThreadPoolExecutor @@ -15,7 +16,7 @@ import dvc.prompt as prompt from dvc.config import Config from dvc.exceptions import DvcException, ConfirmRemoveError from dvc.progress import progress -from dvc.utils import LARGE_DIR_SIZE, tmp_fname, to_chunks +from dvc.utils import LARGE_DIR_SIZE, tmp_fname, to_chunks, move from dvc.state import StateBase from dvc.path_info import PathInfo, URLInfo @@ -381,18 +382,105 @@ class RemoteBASE(object): return self._save_file(path_info, checksum) + @contextmanager + def transfer_context(self): + yield None + + def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): + if not hasattr(self, "_upload"): + raise RemoteActionNotImplemented("upload", self.scheme) + names = self._verify_path_args(to_infos, from_infos, names) + + with self.transfer_context() as ctx: + for from_info, to_info, name in zip(from_infos, to_infos, names): + if to_info.scheme != self.scheme: + raise NotImplementedError + + if from_info.scheme != "local": + raise NotImplementedError + + msg = "Uploading '{}' to '{}'" + logger.debug(msg.format(from_info, to_info)) + + if not name: + name = from_info.name + + if not no_progress_bar: + progress.update_target(name, 0, None) + + try: + self._upload( + from_info.fspath, + to_info, + name=name, + ctx=ctx, + no_progress_bar=no_progress_bar, + ) + except Exception: + msg = "failed to upload '{}' to '{}'" + logger.exception(msg.format(from_info, to_info)) + continue + + if not no_progress_bar: + progress.finish_target(name) + def download( self, from_infos, to_infos, no_progress_bar=False, - name=None, + names=None, resume=False, ): - raise RemoteActionNotImplemented("download", self.scheme) + if not hasattr(self, "_download"): + raise RemoteActionNotImplemented("download", self.scheme) - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - raise RemoteActionNotImplemented("upload", self.scheme) + names = self._verify_path_args(from_infos, to_infos, names) + + with self.transfer_context() as ctx: + for to_info, from_info, name in zip(to_infos, from_infos, names): + if from_info.scheme != self.scheme: + raise NotImplementedError + + if to_info.scheme == self.scheme != "local": + self.copy(from_info, to_info, ctx=ctx) + continue + + if to_info.scheme != "local": + raise NotImplementedError + + msg = "Downloading '{}' to '{}'".format(from_info, to_info) + logger.debug(msg) + + tmp_file = tmp_fname(to_info) + if not name: + name = to_info.name + + if not no_progress_bar: + # real progress is not always available, + # lets at least show start and finish + progress.update_target(name, 0, None) + + makedirs(fspath_py35(to_info.parent), exist_ok=True) + + try: + self._download( + from_info, + tmp_file, + name=name, + ctx=ctx, + resume=resume, + no_progress_bar=no_progress_bar, + ) + except Exception: + msg = "failed to download '{}' to '{}'" + logger.exception(msg.format(from_info, to_info)) + continue + + move(tmp_file, fspath_py35(to_info)) + + if not no_progress_bar: + progress.finish_target(name) def remove(self, path_info): raise RemoteActionNotImplemented("remove", self.scheme) @@ -401,7 +489,7 @@ class RemoteBASE(object): self.copy(from_info, to_info) self.remove(from_info) - def copy(self, from_info, to_info): + def copy(self, from_info, to_info, ctx=None): raise RemoteActionNotImplemented("copy", self.scheme) def exists(self, path_infos): diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 519da7962..ac7f31835 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -2,21 +2,18 @@ from __future__ import unicode_literals import logging import itertools - -from dvc.scheme import Schemes +from contextlib import contextmanager try: from google.cloud import storage except ImportError: storage = None -from dvc.utils import tmp_fname, move -from dvc.utils.compat import makedirs, fspath_py35 from dvc.remote.base import RemoteBASE from dvc.config import Config -from dvc.progress import progress from dvc.exceptions import DvcException from dvc.path_info import CloudURLInfo +from dvc.scheme import Schemes logger = logging.getLogger(__name__) @@ -67,18 +64,17 @@ class RemoteGS(RemoteBASE): md5 = base64.b64decode(b64_md5) return codecs.getencoder("hex")(md5)[0].decode("utf-8") - def copy(self, from_info, to_info, gs=None): - gs = gs if gs else self.gs + def copy(self, from_info, to_info, ctx=None): + gs = ctx or self.gs - blob = gs.bucket(from_info.bucket).get_blob(from_info.path) + from_bucket = gs.bucket(from_info.bucket) + blob = from_bucket.get_blob(from_info.path) if not blob: msg = "'{}' doesn't exist in the cloud".format(from_info.path) raise DvcException(msg) - bucket = self.gs.bucket(to_info.bucket) - bucket.copy_blob( - blob, self.gs.bucket(to_info.bucket), new_name=to_info.path - ) + to_bucket = gs.bucket(to_info.bucket) + from_bucket.copy_blob(blob, to_bucket, new_name=to_info.path) def remove(self, path_info): if path_info.scheme != "gs": @@ -123,83 +119,16 @@ class RemoteGS(RemoteBASE): return results - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) + @contextmanager + def transfer_context(self): + yield self.gs - gs = self.gs + def _upload(self, from_file, to_info, ctx=None, **_kwargs): + bucket = ctx.bucket(to_info.bucket) + blob = bucket.blob(to_info.path) + blob.upload_from_filename(from_file) - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != "gs": - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - logger.debug("Uploading '{}' to '{}'".format(from_info, to_info)) - - if not name: - name = from_info.name - - if not no_progress_bar: - progress.update_target(name, 0, None) - - try: - bucket = gs.bucket(to_info.bucket) - blob = bucket.blob(to_info.path) - blob.upload_from_filename(from_info.fspath) - except Exception: - msg = "failed to upload '{}' to '{}'" - logger.exception(msg.format(from_info, to_info)) - continue - - progress.finish_target(name) - - def download( - self, - from_infos, - to_infos, - no_progress_bar=False, - names=None, - resume=False, - ): - names = self._verify_path_args(from_infos, to_infos, names) - - gs = self.gs - - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != "gs": - raise NotImplementedError - - if to_info.scheme == "gs": - self.copy(from_info, to_info, gs=gs) - continue - - if to_info.scheme != "local": - raise NotImplementedError - - msg = "Downloading '{}' to '{}'".format(from_info, to_info) - logger.debug(msg) - - tmp_file = tmp_fname(to_info) - if not name: - name = to_info.name - - if not no_progress_bar: - # percent_cb is not available for download_to_filename, so - # lets at least update progress at pathpoints(start, finish) - progress.update_target(name, 0, None) - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - try: - bucket = gs.bucket(from_info.bucket) - blob = bucket.get_blob(from_info.path) - blob.download_to_filename(tmp_file) - except Exception: - msg = "failed to download '{}' to '{}'" - logger.exception(msg.format(from_info, to_info)) - continue - - move(tmp_file, fspath_py35(to_info)) - - if not no_progress_bar: - progress.finish_target(name) + def _download(self, from_info, to_file, ctx=None, **_kwargs): + bucket = ctx.bucket(from_info.bucket) + blob = bucket.get_blob(from_info.path) + blob.download_to_filename(to_file) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 206d61bda..23534c000 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -72,7 +72,7 @@ class RemoteHDFS(RemoteBASE): ) return self._group(regex, stdout, "checksum") - def copy(self, from_info, to_info): + def copy(self, from_info, to_info, **_kwargs): dname = posixpath.dirname(to_info.path) self.hadoop_fs("mkdir -p {}".format(dname), user=to_info.user) self.hadoop_fs( @@ -103,64 +103,27 @@ class RemoteHDFS(RemoteBASE): except RemoteCmdError: return False - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) + def _upload(self, from_file, to_info, **_kwargs): + self.hadoop_fs( + "mkdir -p {}".format(to_info.parent.url), user=to_info.user + ) - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != "hdfs": - raise NotImplementedError + tmp_file = tmp_fname(to_info.url) - if from_info.scheme != "local": - raise NotImplementedError + self.hadoop_fs( + "copyFromLocal {} {}".format(from_file, tmp_file), + user=to_info.user, + ) - self.hadoop_fs( - "mkdir -p {}".format(to_info.parent.url), user=to_info.user - ) + self.hadoop_fs( + "mv {} {}".format(tmp_file, to_info.url), user=to_info.user + ) - tmp_file = tmp_fname(to_info.url) - - self.hadoop_fs( - "copyFromLocal {} {}".format(from_info.fspath, tmp_file), - user=to_info.user, - ) - - self.hadoop_fs( - "mv {} {}".format(tmp_file, to_info.url), user=to_info.user - ) - - def download( - self, - from_infos, - to_infos, - no_progress_bar=False, - names=None, - resume=False, - ): - names = self._verify_path_args(from_infos, to_infos, names) - - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != "hdfs": - raise NotImplementedError - - if to_info.scheme == "hdfs": - self.copy(from_info, to_info) - continue - - if to_info.scheme != "local": - raise NotImplementedError - - dname = to_info.parent.fspath - if not os.path.exists(dname): - os.makedirs(dname) - - tmp_file = tmp_fname(to_info.fspath) - - self.hadoop_fs( - "copyToLocal {} {}".format(from_info.url, tmp_file), - user=from_info.user, - ) - - os.rename(tmp_file, to_info.fspath) + def _download(self, from_info, to_file, **_kwargs): + self.hadoop_fs( + "copyToLocal {} {}".format(from_info.url, to_file), + user=from_info.user, + ) def list_cache_paths(self): try: diff --git a/dvc/remote/https.py b/dvc/remote/https.py index e775a5161..2a111330d 100644 --- a/dvc/remote/https.py +++ b/dvc/remote/https.py @@ -7,4 +7,3 @@ from .http import RemoteHTTP class RemoteHTTPS(RemoteHTTP): scheme = Schemes.HTTPS - REGEX = r"^https://.*$" diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index e33a59851..cd6ffd67b 100644 --- a/dvc/remote/local/__init__.py +++ b/dvc/remote/local/__init__.py @@ -254,78 +254,26 @@ class RemoteLOCAL(RemoteBASE): if not self.changed_cache_file(checksum) ] - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) - - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != "local": - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - logger.debug("Uploading '{}' to '{}'".format(from_info, to_info)) - - if not name: - name = from_info.name - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - tmp_file = tmp_fname(to_info) - - try: - copyfile( - fspath_py35(from_info), - tmp_file, - name=name, - no_progress_bar=no_progress_bar, - ) - os.rename(tmp_file, fspath_py35(to_info)) - except Exception: - logger.exception( - "failed to upload '{}' to '{}'".format(from_info, to_info) - ) - - def download( - self, - from_infos, - to_infos, - no_progress_bar=False, - names=None, - resume=False, + def _upload( + self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs ): - names = self._verify_path_args(from_infos, to_infos, names) + makedirs(fspath_py35(to_info.parent), exist_ok=True) - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != "local": - raise NotImplementedError + tmp_file = tmp_fname(to_info) + copyfile( + from_file, tmp_file, name=name, no_progress_bar=no_progress_bar + ) + os.rename(tmp_file, fspath_py35(to_info)) - if to_info.scheme != "local": - raise NotImplementedError - - logger.debug("Downloading '{}' to '{}'".format(from_info, to_info)) - - if not name: - name = to_info.name - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - tmp_file = tmp_fname(to_info) - try: - copyfile( - fspath_py35(from_info), - tmp_file, - no_progress_bar=no_progress_bar, - name=name, - ) - - move(tmp_file, fspath_py35(to_info)) - except Exception: - logger.exception( - "failed to download '{}' to '{}'".format( - from_info, to_info - ) - ) - - continue + def _download( + self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs + ): + copyfile( + fspath_py35(from_info), + to_file, + no_progress_bar=no_progress_bar, + name=name, + ) def _group(self, checksum_infos, show_checksums=False): by_md5 = {} diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py index 4dc1818e8..12fd90cc1 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -11,9 +11,6 @@ try: except ImportError: oss2 = None -from dvc.utils import tmp_fname, move -from dvc.utils.compat import makedirs, fspath_py35 -from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBASE from dvc.remote.azure import Callback @@ -108,65 +105,18 @@ class RemoteOSS(RemoteBASE): def list_cache_paths(self): return self._list_paths(self.path_info.path) - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) - - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != self.scheme: - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - logger.debug("Uploading '{}' to '{}'".format(from_info, to_info)) - - if not name: - name = from_info.name - - cb = None if no_progress_bar else Callback(name) - - try: - self.oss_service.put_object_from_file( - to_info.path, from_info.fspath, progress_callback=cb - ) - except Exception: - msg = "failed to upload '{}'".format(from_info) - logger.warning(msg) - else: - progress.finish_target(name) - - def download( - self, - from_infos, - to_infos, - names=None, - no_progress_bar=False, - resume=False, + def _upload( + self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs ): - names = self._verify_path_args(from_infos, to_infos, names) - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != self.scheme: - raise NotImplementedError - if to_info.scheme != "local": - raise NotImplementedError + cb = None if no_progress_bar else Callback(name) + self.oss_service.put_object_from_file( + to_info.path, from_file, progress_callback=cb + ) - logger.debug("Downloading '{}' to '{}'".format(from_info, to_info)) - - tmp_file = tmp_fname(to_info) - if not name: - name = to_info.name - - cb = None if no_progress_bar else Callback(name) - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - try: - self.oss_service.get_object_to_file( - from_info.path, tmp_file, progress_callback=cb - ) - except Exception: - logger.warning("failed to download '{}'".format(from_info)) - else: - move(tmp_file, fspath_py35(to_info)) - finally: - if not no_progress_bar: - progress.finish_target(name) + def _download( + self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs + ): + cb = None if no_progress_bar else Callback(name) + self.oss_service.get_object_to_file( + from_info.path, to_file, progress_callback=cb + ) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index e8087a8a9..c8bc6c9bf 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -4,21 +4,19 @@ import os import threading import logging import itertools - -from dvc.scheme import Schemes +from contextlib import contextmanager try: import boto3 except ImportError: boto3 = None -from dvc.utils import tmp_fname, move -from dvc.utils.compat import makedirs, fspath_py35 from dvc.progress import progress from dvc.config import Config from dvc.remote.base import RemoteBASE from dvc.exceptions import DvcException, ETagMismatchError from dvc.path_info import CloudURLInfo +from dvc.scheme import Schemes logger = logging.getLogger(__name__) @@ -197,9 +195,8 @@ class RemoteS3(RemoteBASE): if etag != cached_etag: raise ETagMismatchError(etag, cached_etag) - def copy(self, from_info, to_info, s3=None): - s3 = s3 if s3 else self.s3 - self._copy(s3, from_info, to_info, self.extra_args) + def copy(self, from_info, to_info, ctx=None): + self._copy(ctx or self.s3, from_info, to_info, self.extra_args) def remove(self, path_info): if path_info.scheme != "s3": @@ -250,90 +247,42 @@ class RemoteS3(RemoteBASE): return results - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) + @contextmanager + def transfer_context(self): + yield self.s3 - s3 = self.s3 + def _upload( + self, from_file, to_info, name=None, ctx=None, no_progress_bar=False + ): + total = os.path.getsize(from_file) + cb = None if no_progress_bar else Callback(name, total) + ctx.upload_file( + from_file, + to_info.bucket, + to_info.path, + Callback=cb, + ExtraArgs=self.extra_args, + ) - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != "s3": - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - logger.debug("Uploading '{}' to '{}'".format(from_info, to_info)) - - if not name: - name = from_info.name - - total = os.path.getsize(fspath_py35(from_info)) - cb = None if no_progress_bar else Callback(name, total) - - try: - s3.upload_file( - from_info.fspath, - to_info.bucket, - to_info.path, - Callback=cb, - ExtraArgs=self.extra_args, - ) - except Exception: - msg = "failed to upload '{}'".format(from_info) - logger.exception(msg) - continue - - progress.finish_target(name) - - def download( + def _download( self, - from_infos, - to_infos, + from_info, + to_file, + name=None, + ctx=None, no_progress_bar=False, - names=None, resume=False, ): - names = self._verify_path_args(from_infos, to_infos, names) + s3 = ctx - s3 = self.s3 + if no_progress_bar: + cb = None + else: + total = s3.head_object( + Bucket=from_info.bucket, Key=from_info.path + )["ContentLength"] + cb = Callback(name, total) - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != "s3": - raise NotImplementedError - - if to_info.scheme == "s3": - self.copy(from_info, to_info, s3=s3) - continue - - if to_info.scheme != "local": - raise NotImplementedError - - msg = "Downloading '{}' to '{}'".format(from_info, to_info) - logger.debug(msg) - - tmp_file = tmp_fname(to_info) - if not name: - name = to_info.name - - makedirs(fspath_py35(to_info.parent), exist_ok=True) - try: - if no_progress_bar: - cb = None - else: - total = s3.head_object( - Bucket=from_info.bucket, Key=from_info.path - )["ContentLength"] - cb = Callback(name, total) - - s3.download_file( - from_info.bucket, from_info.path, tmp_file, Callback=cb - ) - except Exception: - msg = "failed to download '{}'".format(from_info) - logger.exception(msg) - continue - - move(tmp_file, fspath_py35(to_info)) - - if not no_progress_bar: - progress.finish_target(name) + s3.download_file( + from_info.bucket, from_info.path, to_file, Callback=cb + ) diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index a067db341..e9d17fdcd 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -4,8 +4,6 @@ import os import getpass import logging -from dvc.scheme import Schemes - try: import paramiko except ImportError: @@ -16,6 +14,7 @@ from dvc.remote.ssh.connection import SSHConnection from dvc.config import Config from dvc.utils.compat import urlparse, StringIO from dvc.remote.base import RemoteBASE +from dvc.scheme import Schemes logger = logging.getLogger(__name__) @@ -155,12 +154,15 @@ class RemoteSSH(RemoteBASE): with self.ssh(path_info) as ssh: return ssh.isdir(path_info.path) - def copy(self, from_info, to_info): + def copy(self, from_info, to_info, ctx=None): if from_info.scheme != self.scheme or to_info.scheme != self.scheme: raise NotImplementedError - with self.ssh(from_info) as ssh: - ssh.cp(from_info.path, to_info.path) + if ctx: + ctx.cp(from_info.path, to_info.path) + else: + with self.ssh(from_info) as ssh: + ssh.cp(from_info.path, to_info.path) def remove(self, path_info): if path_info.scheme != self.scheme: @@ -176,75 +178,36 @@ class RemoteSSH(RemoteBASE): with self.ssh(from_info) as ssh: ssh.move(from_info.path, to_info.path) - def download( + def transfer_context(self): + return self.ssh(self.path_info) + + def _download( self, - from_infos, - to_infos, + from_info, + to_file, + name=None, + ctx=None, no_progress_bar=False, - names=None, resume=False, ): - names = self._verify_path_args(from_infos, to_infos, names) - ssh = self.ssh(from_infos[0]) + assert from_info.isin(self.path_info) + ctx.download( + from_info.path, + to_file, + progress_title=name, + no_progress_bar=no_progress_bar, + ) - for to_info, from_info, name in zip(to_infos, from_infos, names): - if from_info.scheme != self.scheme: - raise NotImplementedError - - if to_info.scheme == self.scheme: - ssh.cp(from_info.path, to_info.path) - continue - - if to_info.scheme != "local": - raise NotImplementedError - - logger.debug( - "Downloading '{src}' to '{dest}'".format( - src=from_info, dest=to_info - ) - ) - - try: - ssh.download( - from_info.path, - to_info.fspath, - progress_title=name, - no_progress_bar=no_progress_bar, - ) - except Exception: - logger.exception( - "failed to download '{src}' to '{dest}'".format( - src=from_info, dest=to_info - ) - ) - continue - - ssh.close() - - def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): - names = self._verify_path_args(to_infos, from_infos, names) - - with self.ssh(to_infos[0]) as ssh: - for from_info, to_info, name in zip(from_infos, to_infos, names): - if to_info.scheme != self.scheme: - raise NotImplementedError - - if from_info.scheme != "local": - raise NotImplementedError - - try: - ssh.upload( - from_info.fspath, - to_info.path, - progress_title=name, - no_progress_bar=no_progress_bar, - ) - except Exception: - logger.exception( - "failed to upload '{}' to '{}'".format( - from_info, to_info - ) - ) + def _upload( + self, from_file, to_info, name=None, ctx=None, no_progress_bar=False + ): + assert to_info.isin(self.path_info) + ctx.upload( + from_file, + to_info.path, + progress_title=name, + no_progress_bar=no_progress_bar, + ) def list_cache_paths(self): with self.ssh(self.path_info) as ssh: diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py index 1241ef0ae..b5ea969d0 100644 --- a/dvc/remote/ssh/connection.py +++ b/dvc/remote/ssh/connection.py @@ -10,7 +10,6 @@ except ImportError: paramiko = None from dvc.utils import tmp_fname -from dvc.utils.compat import makedirs from dvc.progress import progress from dvc.exceptions import DvcException from dvc.remote.base import RemoteCmdError @@ -203,23 +202,15 @@ class SSHConnection: def download(self, src, dest, no_progress_bar=False, progress_title=None): self._sftp_connect() - makedirs(os.path.dirname(dest), exist_ok=True) - tmp_file = tmp_fname(dest) - if no_progress_bar: - self._sftp.get(src, tmp_file) + self._sftp.get(src, dest) else: if not progress_title: - progress_title = os.path.basename(dest) + progress_title = os.path.basename(src) - self._sftp.get(src, tmp_file, callback=create_cb(progress_title)) + self._sftp.get(src, dest, callback=create_cb(progress_title)) progress.finish_target(progress_title) - if os.path.exists(dest): - os.remove(dest) - - os.rename(tmp_file, dest) - def move(self, src, dst): self.makedirs(posixpath.dirname(dst)) self._sftp_connect()