Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2129 from Suor/download-refactor

Browse Source 
dvc: unify remote upload/download code
This commit is contained in:
Ruslan Kuprieiev
2019-06-14 11:56:43 +03:00
committed by GitHub
parent 765c8570d2 80ca3eedc8
commit 2f02280483
10 changed files with 246 additions and 523 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/remote/azure.py
+14 -71
View File
@@ -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
)
dvc/remote/base.py
+95 -7
View File
@@ -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):
dvc/remote/gs.py
+19 -90
View File
@@ -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)
dvc/remote/hdfs.py
+18 -55
View File
@@ -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:
dvc/remote/https.py
-1
View File
@@ -7,4 +7,3 @@ from .http import RemoteHTTP
class RemoteHTTPS(RemoteHTTP):
scheme = Schemes.HTTPS
REGEX = r"^https://.*$"
dvc/remote/local/__init__.py
+17 -69
View File
@@ -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 = {}
dvc/remote/oss.py
+13 -63
View File
@@ -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
)
dvc/remote/s3.py
+35 -86
View File
@@ -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
)
dvc/remote/ssh/__init__.py
+32 -69
View File
@@ -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:
dvc/remote/ssh/connection.py
+3 -12
View File
@@ -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()