mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-05 23:09:32 +01:00
Ruslan Kuprieiev
+4
-1
@@ -1,13 +1,14 @@
|
||||
import os
|
||||
|
||||
from dvc.config import Config
|
||||
from dvc.remote import Remote
|
||||
|
||||
|
||||
class Cache(object):
|
||||
CACHE_DIR = 'cache'
|
||||
|
||||
def __init__(self, project):
|
||||
from dvc.remote import Remote
|
||||
|
||||
config = project.config._config[Config.SECTION_CACHE]
|
||||
|
||||
local = config.get(Config.SECTION_CACHE_LOCAL, None)
|
||||
@@ -31,6 +32,8 @@ class Cache(object):
|
||||
self.hdfs = self._get_remote(project, config, Config.SECTION_CACHE_HDFS)
|
||||
|
||||
def _get_remote(self, project, config, name):
|
||||
from dvc.remote import Remote
|
||||
|
||||
remote = config.get(name, None)
|
||||
if not remote:
|
||||
return None
|
||||
|
||||
@@ -1,139 +0,0 @@
|
||||
import os
|
||||
import math
|
||||
import threading
|
||||
|
||||
import boto3
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.progress import progress
|
||||
from dvc.cloud.base import DataCloudBase, DataCloudError
|
||||
from dvc.config import Config
|
||||
|
||||
|
||||
def sizeof_fmt(num, suffix='B'):
|
||||
""" Convert number of bytes to human-readable string """
|
||||
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
||||
if abs(num) < 1024.0:
|
||||
return "%3.1f%s%s" % (num, unit, suffix)
|
||||
num /= 1024.0
|
||||
return "%.1f%s%s" % (num, 'Y', suffix)
|
||||
|
||||
|
||||
def percent_cb(name, complete, total):
|
||||
""" Callback for updating target progress """
|
||||
Logger.debug('{}: {} transferred out of {}'.format(name,
|
||||
sizeof_fmt(complete),
|
||||
sizeof_fmt(total)))
|
||||
progress.update_target(name, complete, total)
|
||||
|
||||
|
||||
def create_cb(name):
|
||||
""" Create callback function for multipart object """
|
||||
return (lambda cur, tot: percent_cb(name, cur, tot))
|
||||
|
||||
|
||||
class Callback(object):
|
||||
def __init__(self, name, total):
|
||||
self.name = name
|
||||
self.total = total
|
||||
self.current = 0
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def __call__(self, byts):
|
||||
with self.lock:
|
||||
self.current += byts
|
||||
progress.update_target(self.name, self.current, self.total)
|
||||
|
||||
|
||||
class AWSKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
|
||||
class DataCloudAWS(DataCloudBase):
|
||||
""" DataCloud class for Amazon Web Services """
|
||||
REGEX = r'^s3://(?P<path>.*)$'
|
||||
|
||||
@property
|
||||
def profile(self):
|
||||
return self._cloud_settings.cloud_config.get(Config.SECTION_AWS_PROFILE, None)
|
||||
|
||||
def connect(self):
|
||||
session = boto3.Session(profile_name=self.profile)
|
||||
self.s3 = session.resource('s3')
|
||||
bucket = self.s3.Bucket(self.storage_bucket)
|
||||
if bucket is None:
|
||||
raise DataCloudError('Storage path {} is not setup correctly'.format(self.storage_bucket))
|
||||
|
||||
def create_cb_pull(self, name, key):
|
||||
total = self.s3.Object(bucket_name=key.bucket, key=key.name).content_length
|
||||
return Callback(name, total)
|
||||
|
||||
def create_cb_push(self, name, fname):
|
||||
total = os.path.getsize(fname)
|
||||
return Callback(name, total)
|
||||
|
||||
def _pull_key(self, key, fname, no_progress_bar=False):
|
||||
Logger.debug("Pulling key '{}' from bucket '{}' to file '{}'".format(key.name,
|
||||
key.bucket,
|
||||
fname))
|
||||
self._makedirs(fname)
|
||||
|
||||
tmp_file = self.tmp_file(fname)
|
||||
name = os.path.relpath(fname, self._cloud_settings.cache.local.cache_dir)
|
||||
|
||||
if self._cmp_checksum(key, fname):
|
||||
Logger.debug('File "{}" matches with "{}".'.format(fname, key.name))
|
||||
return fname
|
||||
|
||||
Logger.debug('Downloading cache file from S3 "{}/{}" to "{}"'.format(key.bucket,
|
||||
key.name,
|
||||
fname))
|
||||
|
||||
if no_progress_bar:
|
||||
cb = None
|
||||
else:
|
||||
cb = self.create_cb_pull(name, key)
|
||||
|
||||
|
||||
try:
|
||||
self.s3.Object(key.bucket, key.name).download_file(tmp_file, Callback=cb)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to download "{}": {}'.format(key.name, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, fname)
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
Logger.debug('Downloading completed')
|
||||
|
||||
return fname
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
try:
|
||||
self.s3.Object(self.storage_bucket, key_name).get()
|
||||
return AWSKey(self.storage_bucket, key_name)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return AWSKey(self.storage_bucket, key_name)
|
||||
|
||||
def _push_key(self, key, path):
|
||||
""" push, aws version """
|
||||
name = os.path.relpath(path, self._cloud_settings.cache.local.cache_dir)
|
||||
cb = self.create_cb_push(name, path)
|
||||
try:
|
||||
self.s3.Object(key.bucket, key.name).upload_file(path, Callback=cb)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to upload "{}": {}'.format(path, exc))
|
||||
return None
|
||||
|
||||
progress.finish_target(name)
|
||||
|
||||
return path
|
||||
@@ -1,228 +0,0 @@
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
|
||||
from dvc.config import Config
|
||||
from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.remote.local import RemoteLOCAL
|
||||
|
||||
|
||||
STATUS_UNKNOWN = 0
|
||||
STATUS_OK = 1
|
||||
STATUS_MODIFIED = 2
|
||||
STATUS_NEW = 3
|
||||
STATUS_DELETED = 4
|
||||
|
||||
|
||||
STATUS_MAP = {
|
||||
# (local_exists, remote_exists, cmp)
|
||||
(True, True, True) : STATUS_OK,
|
||||
(True, True, False) : STATUS_MODIFIED,
|
||||
(True, False, None) : STATUS_NEW,
|
||||
(False, True, None) : STATUS_DELETED,
|
||||
}
|
||||
|
||||
|
||||
class DataCloudError(DvcException):
|
||||
""" Data Cloud exception """
|
||||
def __init__(self, msg):
|
||||
super(DataCloudError, self).__init__('Data sync error: {}'.format(msg))
|
||||
|
||||
|
||||
class CloudSettings(object):
|
||||
def __init__(self, cache=None, global_storage_path=None, cloud_config=None):
|
||||
self.cache = cache
|
||||
self.cloud_config = cloud_config
|
||||
self.global_storage_path = global_storage_path
|
||||
|
||||
|
||||
class DataCloudBase(object):
|
||||
""" Base class for DataCloud """
|
||||
|
||||
REGEX = r''
|
||||
|
||||
def __init__(self, cloud_settings):
|
||||
self._cloud_settings = cloud_settings
|
||||
|
||||
@property
|
||||
def url(self):
|
||||
config = self._cloud_settings.cloud_config
|
||||
return config.get(Config.SECTION_REMOTE_URL, None)
|
||||
|
||||
@classmethod
|
||||
def match(cls, url):
|
||||
return re.match(cls.REGEX, url)
|
||||
|
||||
def group(self, group):
|
||||
return self.match(self.url).group(group)
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return self.group('path')
|
||||
|
||||
@classmethod
|
||||
def supported(cls, url):
|
||||
return cls.match(url) != None
|
||||
|
||||
# backward compatibility
|
||||
@property
|
||||
def storage_path(self):
|
||||
""" get storage path
|
||||
|
||||
Precedence: Storage, then cloud specific
|
||||
"""
|
||||
if self._cloud_settings.global_storage_path:
|
||||
return self._cloud_settings.global_storage_path
|
||||
|
||||
if not self.url:
|
||||
path = self._cloud_settings.cloud_config.get(Config.SECTION_CORE_STORAGEPATH, None)
|
||||
if path:
|
||||
Logger.warn('Using obsoleted config format. Consider updating.')
|
||||
else:
|
||||
path = self.path
|
||||
|
||||
return path
|
||||
|
||||
def _storage_path_parts(self):
|
||||
"""
|
||||
Split storage path into parts. I.e. 'dvc-test/myrepo' -> ['dvc', 'myrepo']
|
||||
"""
|
||||
return self.storage_path.strip('/').split('/', 1)
|
||||
|
||||
@property
|
||||
def storage_bucket(self):
|
||||
""" Data -> StoragePath takes precedence; if doesn't exist, use cloud-specific """
|
||||
return self._storage_path_parts()[0]
|
||||
|
||||
@property
|
||||
def storage_prefix(self):
|
||||
"""
|
||||
Prefix within the bucket. I.e. 'myrepo' in 'dvc-test/myrepo'.
|
||||
"""
|
||||
parts = self._storage_path_parts()
|
||||
if len(parts) > 1:
|
||||
return parts[1]
|
||||
return ''
|
||||
|
||||
def cache_file_key(self, fname):
|
||||
""" Key of a file within the bucket """
|
||||
relpath = os.path.relpath(fname, self._cloud_settings.cache.local.cache_dir)
|
||||
relpath = relpath.replace('\\', '/')
|
||||
return '{}/{}'.format(self.storage_prefix, relpath).strip('/')
|
||||
|
||||
@staticmethod
|
||||
def tmp_file(fname):
|
||||
""" Temporary name for a partial download """
|
||||
return fname + '.part'
|
||||
|
||||
def sanity_check(self):
|
||||
"""
|
||||
Cloud-specific method to check config for basic requirements.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _push_key(self, key, path):
|
||||
pass
|
||||
|
||||
def collect(self, arg):
|
||||
path, local = arg
|
||||
ret = [path]
|
||||
|
||||
if not RemoteLOCAL.is_dir_cache(path):
|
||||
return ret
|
||||
|
||||
if local:
|
||||
if not os.path.isfile(path):
|
||||
return ret
|
||||
dir_path = path
|
||||
else:
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
Logger.debug("File '{}' does not exist in the cloud".format(path))
|
||||
return ret
|
||||
tmp = os.path.join(tempfile.mkdtemp(), os.path.basename(path))
|
||||
self._pull_key(key, tmp, no_progress_bar=True)
|
||||
dir_path = tmp
|
||||
|
||||
for relpath, md5 in RemoteLOCAL.get_dir_cache(dir_path).items():
|
||||
cache = self._cloud_settings.cache.local.get(md5)
|
||||
ret.append(cache)
|
||||
|
||||
return ret
|
||||
|
||||
def _cmp_checksum(self, blob, fname):
|
||||
md5 = self._cloud_settings.cache.local.path_to_md5(fname)
|
||||
if self._cloud_settings.cache.local.state.changed(fname, md5=md5):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def push(self, path):
|
||||
key = self._get_key(path)
|
||||
if key:
|
||||
Logger.debug("File '{}' already uploaded to the cloud. Validating checksum...".format(path))
|
||||
if self._cmp_checksum(key, path):
|
||||
Logger.debug('File checksum matches. No uploading is needed.')
|
||||
return []
|
||||
Logger.debug('Checksum mismatch. Reuploading is required.')
|
||||
|
||||
key = self._new_key(path)
|
||||
return self._push_key(key, path)
|
||||
|
||||
def _makedirs(self, fname):
|
||||
dname = os.path.dirname(fname)
|
||||
try:
|
||||
os.makedirs(dname)
|
||||
except OSError as e:
|
||||
if e.errno != os.errno.EEXIST:
|
||||
raise
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
""" Cloud-specific method of pulling keys """
|
||||
pass
|
||||
|
||||
def _get_key(self, path):
|
||||
""" Cloud-specific method of getting keys """
|
||||
pass
|
||||
|
||||
def pull(self, path):
|
||||
""" Generic method for pulling data from the cloud """
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
Logger.error("File '{}' does not exist in the cloud".format(path))
|
||||
return None
|
||||
|
||||
return self._pull_key(key, path)
|
||||
|
||||
def _status(self, key, path):
|
||||
remote_exists = key != None
|
||||
local_exists = os.path.exists(path)
|
||||
|
||||
diff = None
|
||||
if remote_exists and local_exists:
|
||||
diff = self._cmp_checksum(key, path)
|
||||
|
||||
return STATUS_MAP.get((local_exists, remote_exists, diff), STATUS_UNKNOWN)
|
||||
|
||||
def status(self, path):
|
||||
"""
|
||||
Generic method for checking data item status.
|
||||
"""
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
return STATUS_NEW
|
||||
|
||||
return self._status(key, path)
|
||||
|
||||
def connect(self):
|
||||
pass
|
||||
|
||||
def disconnect(self):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
self.connect()
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
self.disconnect()
|
||||
+22
-28
@@ -6,26 +6,26 @@ from dvc.exceptions import DvcException
|
||||
from dvc.config import Config, ConfigError
|
||||
from dvc.utils import map_progress
|
||||
|
||||
from dvc.cloud.aws import DataCloudAWS
|
||||
from dvc.cloud.gcp import DataCloudGCP
|
||||
from dvc.cloud.ssh import DataCloudSSH
|
||||
from dvc.cloud.local import DataCloudLOCAL
|
||||
from dvc.cloud.base import DataCloudBase, CloudSettings
|
||||
from dvc.cloud.base import STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED
|
||||
from dvc.remote.s3 import RemoteS3
|
||||
from dvc.remote.gs import RemoteGS
|
||||
from dvc.remote.ssh import RemoteSSH
|
||||
from dvc.remote.local import RemoteLOCAL
|
||||
from dvc.remote.base import RemoteBase
|
||||
from dvc.remote.base import STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED
|
||||
|
||||
|
||||
class DataCloud(object):
|
||||
""" Generic class to do initial config parsing and redirect to proper DataCloud methods """
|
||||
|
||||
CLOUD_MAP = {
|
||||
'aws' : DataCloudAWS,
|
||||
'gcp' : DataCloudGCP,
|
||||
'ssh' : DataCloudSSH,
|
||||
'local' : DataCloudLOCAL,
|
||||
'aws' : RemoteS3,
|
||||
'gcp' : RemoteGS,
|
||||
'ssh' : RemoteSSH,
|
||||
'local' : RemoteLOCAL,
|
||||
}
|
||||
|
||||
def __init__(self, cache=None, config=None):
|
||||
self._cache = cache
|
||||
def __init__(self, project, config=None):
|
||||
self.project = project
|
||||
self._config = config
|
||||
|
||||
remote = self._config[Config.SECTION_CORE].get(Config.SECTION_CORE_REMOTE, '')
|
||||
@@ -41,9 +41,9 @@ class DataCloud(object):
|
||||
self._cloud = self._init_remote(remote)
|
||||
|
||||
@staticmethod
|
||||
def supported(url):
|
||||
def supported(config):
|
||||
for cloud in DataCloud.CLOUD_MAP.values():
|
||||
if cloud.supported(url):
|
||||
if cloud.supported(config):
|
||||
return cloud
|
||||
return None
|
||||
|
||||
@@ -53,10 +53,9 @@ class DataCloud(object):
|
||||
if not cloud_config:
|
||||
raise ConfigError("Can't find remote section '{}' in config".format(section))
|
||||
|
||||
url = cloud_config[Config.SECTION_REMOTE_URL]
|
||||
cloud_type = self.supported(url)
|
||||
cloud_type = self.supported(cloud_config)
|
||||
if not cloud_type:
|
||||
raise ConfigError("Unsupported url '{}'".format(url))
|
||||
raise ConfigError("Unsupported cloud '{}'".format(cloud_config))
|
||||
|
||||
return self._init_cloud(cloud_config, cloud_type)
|
||||
|
||||
@@ -81,12 +80,7 @@ class DataCloud(object):
|
||||
if global_storage_path:
|
||||
Logger.warn('Using obsoleted config format. Consider updating.')
|
||||
|
||||
cloud_settings = CloudSettings(cache=self._cache,
|
||||
global_storage_path=global_storage_path,
|
||||
cloud_config=cloud_config)
|
||||
|
||||
cloud = cloud_type(cloud_settings)
|
||||
cloud.sanity_check()
|
||||
cloud = cloud_type(self.project, cloud_config)
|
||||
return cloud
|
||||
|
||||
def _collect(self, cloud, targets, jobs, local):
|
||||
@@ -104,12 +98,14 @@ class DataCloud(object):
|
||||
if remote:
|
||||
return self._init_remote(remote)
|
||||
|
||||
return self._cloud
|
||||
if self._cloud:
|
||||
return self._cloud
|
||||
|
||||
raise ConfigError("No remote repository specified. Setup default repository " \
|
||||
"with 'dvc config core.remote <name>' or use '-r <name>'.")
|
||||
|
||||
def _filter(self, func, status, targets, jobs, remote):
|
||||
cloud = self._get_cloud(remote)
|
||||
if not cloud:
|
||||
return []
|
||||
|
||||
with cloud:
|
||||
filtered = []
|
||||
@@ -146,8 +142,6 @@ class DataCloud(object):
|
||||
Check status of data items in a cloud-agnostic way.
|
||||
"""
|
||||
cloud = self._get_cloud(remote)
|
||||
if not cloud:
|
||||
return []
|
||||
|
||||
with cloud:
|
||||
return self._status(cloud, targets, jobs)
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
import base64
|
||||
import os
|
||||
|
||||
from google.cloud import storage as gc
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.config import Config, ConfigError
|
||||
from dvc.cloud.base import DataCloudError, DataCloudBase
|
||||
from dvc.progress import progress
|
||||
|
||||
|
||||
class DataCloudGCP(DataCloudBase):
|
||||
""" DataCloud class for Google Cloud Platform """
|
||||
REGEX = r'^gs://(?P<path>.*)$'
|
||||
|
||||
@property
|
||||
def gc_project_name(self):
|
||||
"""
|
||||
Get project name from config.
|
||||
"""
|
||||
return self._cloud_settings.cloud_config.get(Config.SECTION_GCP_PROJECTNAME, None)
|
||||
|
||||
def connect(self):
|
||||
client = gc.Client(project=self.gc_project_name)
|
||||
self.bucket = client.bucket(self.storage_bucket)
|
||||
if not self.bucket.exists():
|
||||
raise DataCloudError('sync up: google cloud bucket {} '
|
||||
'doesn\'t exist'.format(self.storage_bucket))
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
name = os.path.relpath(path, self._cloud_settings.cache.local.cache_dir)
|
||||
tmp_file = self.tmp_file(path)
|
||||
|
||||
if self._cmp_checksum(key, path):
|
||||
Logger.debug('File "{}" matches with "{}".'.format(path, key.name))
|
||||
return path
|
||||
|
||||
Logger.debug('Downloading cache file from gc "{}/{}"'.format(key.bucket.name, key.name))
|
||||
|
||||
if not no_progress_bar:
|
||||
# percent_cb is not available for download_to_filename, so
|
||||
# lets at least update progress at keypoints(start, finish)
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
try:
|
||||
key.download_to_filename(tmp_file)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to download "{}": {}'.format(key.name, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
Logger.debug('Downloading completed')
|
||||
|
||||
return path
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return self.bucket.get_blob(key_name)
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return self.bucket.blob(key_name)
|
||||
|
||||
def _push_key(self, key, path):
|
||||
""" push, gcp version """
|
||||
name = os.path.relpath(path, self._cloud_settings.cache.local.cache_dir)
|
||||
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
key.upload_from_filename(path)
|
||||
|
||||
progress.finish_target(name)
|
||||
Logger.debug('uploading %s completed' % path)
|
||||
|
||||
return path
|
||||
@@ -1,58 +0,0 @@
|
||||
import os
|
||||
import filecmp
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.utils import copyfile, file_md5
|
||||
from dvc.cloud.base import DataCloudError, DataCloudBase
|
||||
|
||||
|
||||
class LocalKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return os.path.join(self.bucket, self.name)
|
||||
|
||||
|
||||
class DataCloudLOCAL(DataCloudBase):
|
||||
"""
|
||||
Driver for local storage.
|
||||
"""
|
||||
REGEX = r'^(?P<path>(/+|.:\\+).*)$'
|
||||
|
||||
def cache_file_key(self, path):
|
||||
return os.path.relpath(os.path.abspath(path), self._cloud_settings.cache.local.cache_dir)
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = LocalKey(self.storage_path, key_name)
|
||||
if os.path.exists(key.path) and os.path.isfile(key.path):
|
||||
return key
|
||||
return None
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = LocalKey(self.storage_path, key_name)
|
||||
self._makedirs(key.path)
|
||||
return key
|
||||
|
||||
def _push_key(self, key, path):
|
||||
self._makedirs(key.path)
|
||||
copyfile(path, key.path)
|
||||
return path
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
tmp_file = self.tmp_file(path)
|
||||
try:
|
||||
copyfile(key.path, tmp_file, no_progress_bar=no_progress_bar)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to copy "{}": {}'.format(key.path, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
|
||||
return path
|
||||
@@ -1,110 +0,0 @@
|
||||
import os
|
||||
import errno
|
||||
import filecmp
|
||||
import paramiko
|
||||
import posixpath
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.utils import copyfile, file_md5
|
||||
from dvc.cloud.base import DataCloudError, DataCloudBase
|
||||
from dvc.cloud.aws import create_cb
|
||||
from dvc.progress import progress
|
||||
|
||||
|
||||
class SSHKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return os.path.join(self.bucket, self.name)
|
||||
|
||||
|
||||
class DataCloudSSH(DataCloudBase):
|
||||
"""
|
||||
Driver for remote storage over ssh.
|
||||
"""
|
||||
#NOTE: ~/ paths are temporarily forbidden
|
||||
REGEX = r'^(?P<user>.*)@(?P<host>.*):(?P<path>/+.*)$'
|
||||
|
||||
@property
|
||||
def hostname(self):
|
||||
return self.group('host')
|
||||
|
||||
@property
|
||||
def username(self):
|
||||
return self.group('user')
|
||||
|
||||
def connect(self):
|
||||
self.ssh = paramiko.SSHClient()
|
||||
|
||||
self.ssh.load_system_host_keys()
|
||||
self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
self.ssh.connect(self.hostname, username=self.username)
|
||||
|
||||
def disconnect(self):
|
||||
self.ssh.close()
|
||||
|
||||
def get_sftp(self):
|
||||
#NOTE: SFTP doesn't seem to be thread safe
|
||||
return self.ssh.open_sftp()
|
||||
|
||||
def cache_file_key(self, path):
|
||||
relpath = os.path.relpath(os.path.abspath(path), self._cloud_settings.cache.local.cache_dir)
|
||||
return relpath.replace('\\', '/')
|
||||
|
||||
def _isfile_remote(self, path):
|
||||
try:
|
||||
self.get_sftp().open(path, 'r')
|
||||
return True
|
||||
except IOError as exc:
|
||||
if exc.errno != errno.ENOENT:
|
||||
raise
|
||||
return False
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = SSHKey(self.path, key_name)
|
||||
if self._isfile_remote(key.path):
|
||||
return key
|
||||
return None
|
||||
|
||||
def _makedirs_remote(self, dname):
|
||||
try:
|
||||
self.get_sftp().chdir(dname)
|
||||
except IOError as exc:
|
||||
if exc.errno != errno.ENOENT:
|
||||
raise
|
||||
|
||||
parent = posixpath.dirname(dname)
|
||||
if len(parent):
|
||||
self._makedirs_remote(parent)
|
||||
self.get_sftp().mkdir(dname)
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = SSHKey(self.path, key_name)
|
||||
self._makedirs_remote(posixpath.dirname(key.path))
|
||||
return key
|
||||
|
||||
def _push_key(self, key, path):
|
||||
self.get_sftp().put(path, key.path, callback=create_cb(key.name))
|
||||
progress.finish_target(key.name)
|
||||
return path
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
tmp_file = self.tmp_file(path)
|
||||
try:
|
||||
self.get_sftp().get(key.path, tmp_file, callback=create_cb(key.name))
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to copy "{}": {}'.format(key.path, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
progress.finish_target(key.name)
|
||||
|
||||
return path
|
||||
@@ -1,7 +1,3 @@
|
||||
import os
|
||||
|
||||
import dvc.cloud.base as cloud
|
||||
|
||||
from dvc.command.common.base import CmdBase
|
||||
|
||||
|
||||
|
||||
@@ -5,14 +5,13 @@ try:
|
||||
except ImportError:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from dvc.cloud.gcp import DataCloudGCP
|
||||
from dvc.dependency.s3 import DependencyS3
|
||||
from dvc.remote.gs import RemoteGS
|
||||
from dvc.config import Config
|
||||
|
||||
|
||||
class DependencyGS(DependencyS3):
|
||||
REGEX = DataCloudGCP.REGEX
|
||||
REGEX = RemoteGS.REGEX
|
||||
|
||||
def __init__(self, stage, path, info=None, remote=None):
|
||||
super(DependencyGS, self).__init__(stage, path)
|
||||
|
||||
@@ -11,7 +11,6 @@ except ImportError:
|
||||
from dvc.dependency.base import *
|
||||
from dvc.logger import Logger
|
||||
from dvc.utils import remove
|
||||
from dvc.cloud.local import DataCloudLOCAL
|
||||
from dvc.config import Config
|
||||
from dvc.remote.local import RemoteLOCAL
|
||||
|
||||
|
||||
@@ -6,13 +6,12 @@ except ImportError:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from dvc.dependency.base import DependencyBase
|
||||
from dvc.cloud.aws import DataCloudAWS
|
||||
from dvc.remote.s3 import RemoteS3
|
||||
from dvc.config import Config
|
||||
|
||||
|
||||
class DependencyS3(DependencyBase):
|
||||
REGEX = DataCloudAWS.REGEX
|
||||
REGEX = RemoteS3.REGEX
|
||||
|
||||
def __init__(self, stage, path, info=None, remote=None):
|
||||
super(DependencyS3, self).__init__(stage, path)
|
||||
|
||||
+3
-3
@@ -5,8 +5,6 @@ import json
|
||||
import networkx as nx
|
||||
from jsonpath_rw import parse
|
||||
|
||||
import dvc.cloud.base as cloud
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.stage import Stage
|
||||
@@ -54,7 +52,7 @@ class Project(object):
|
||||
self.link_state = LinkState(self.root_dir, self.dvc_dir)
|
||||
self.logger = Logger(self.config._config[Config.SECTION_CORE].get(Config.SECTION_CORE_LOGLEVEL, None))
|
||||
self.cache = Cache(self)
|
||||
self.cloud = DataCloud(cache=self.cache, config=self.config._config)
|
||||
self.cloud = DataCloud(self, config=self.config._config)
|
||||
self.updater = Updater(self.dvc_dir)
|
||||
|
||||
self._ignore()
|
||||
@@ -271,6 +269,8 @@ class Project(object):
|
||||
return status
|
||||
|
||||
def _cloud_status(self, target=None, jobs=1, remote=None):
|
||||
import dvc.remote.base as cloud
|
||||
|
||||
status = {}
|
||||
for target, ret in self.cloud.status(self._used_cache(target), jobs, remote=remote):
|
||||
if ret == cloud.STATUS_UNKNOWN or ret == cloud.STATUS_OK:
|
||||
|
||||
+155
-4
@@ -1,6 +1,33 @@
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
|
||||
from dvc.config import Config
|
||||
from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
|
||||
|
||||
STATUS_UNKNOWN = 0
|
||||
STATUS_OK = 1
|
||||
STATUS_MODIFIED = 2
|
||||
STATUS_NEW = 3
|
||||
STATUS_DELETED = 4
|
||||
|
||||
|
||||
STATUS_MAP = {
|
||||
# (local_exists, remote_exists, cmp)
|
||||
(True, True, True) : STATUS_OK,
|
||||
(True, True, False) : STATUS_MODIFIED,
|
||||
(True, False, None) : STATUS_NEW,
|
||||
(False, True, None) : STATUS_DELETED,
|
||||
}
|
||||
|
||||
|
||||
class DataCloudError(DvcException):
|
||||
""" Data Cloud exception """
|
||||
def __init__(self, msg):
|
||||
super(DataCloudError, self).__init__('Data sync error: {}'.format(msg))
|
||||
|
||||
|
||||
|
||||
class RemoteBase(object):
|
||||
@@ -19,16 +46,140 @@ class RemoteBase(object):
|
||||
return re.match(cls.REGEX, url)
|
||||
|
||||
def save_info(self, path_info):
|
||||
pass
|
||||
raise NotImplementedError
|
||||
|
||||
def save(self, path_info):
|
||||
pass
|
||||
raise NotImplementedError
|
||||
|
||||
def checkout(self, path_info, checksum_info):
|
||||
pass
|
||||
raise NotImplementedError
|
||||
|
||||
def download(self, path_info, path):
|
||||
pass
|
||||
raise NotImplementedError
|
||||
|
||||
def upload(self, path, path_info):
|
||||
raise NotImplementedError
|
||||
|
||||
# Old code starting from here
|
||||
|
||||
def cache_file_key(self, fname):
|
||||
""" Key of a file within the bucket """
|
||||
relpath = os.path.relpath(fname, self.project.cache.local.cache_dir)
|
||||
relpath = relpath.replace('\\', '/')
|
||||
return '{}/{}'.format(self.prefix, relpath).strip('/')
|
||||
|
||||
def cache_key_name(self, path):
|
||||
relpath = os.path.relpath(path, self.project.cache.local.cache_dir)
|
||||
return relpath.replace('\\', '').replace('/', '')
|
||||
|
||||
@staticmethod
|
||||
def tmp_file(fname):
|
||||
""" Temporary name for a partial download """
|
||||
return fname + '.part'
|
||||
|
||||
def _push_key(self, key, path):
|
||||
pass
|
||||
|
||||
def collect(self, arg):
|
||||
from dvc.remote.local import RemoteLOCAL
|
||||
|
||||
path, local = arg
|
||||
ret = [path]
|
||||
|
||||
if not RemoteLOCAL.is_dir_cache(path):
|
||||
return ret
|
||||
|
||||
if local:
|
||||
if not os.path.isfile(path):
|
||||
return ret
|
||||
dir_path = path
|
||||
else:
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
Logger.debug("File '{}' does not exist in the cloud".format(path))
|
||||
return ret
|
||||
tmp = os.path.join(tempfile.mkdtemp(), os.path.basename(path))
|
||||
self._pull_key(key, tmp, no_progress_bar=True)
|
||||
dir_path = tmp
|
||||
|
||||
for relpath, md5 in RemoteLOCAL.get_dir_cache(dir_path).items():
|
||||
cache = self.project.cache.local.get(md5)
|
||||
ret.append(cache)
|
||||
|
||||
return ret
|
||||
|
||||
def _cmp_checksum(self, blob, fname):
|
||||
md5 = self.project.cache.local.path_to_md5(fname)
|
||||
if self.project.cache.local.state.changed(fname, md5=md5):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def push(self, path):
|
||||
key = self._get_key(path)
|
||||
if key:
|
||||
Logger.debug("File '{}' already uploaded to the cloud. Validating checksum...".format(path))
|
||||
if self._cmp_checksum(key, path):
|
||||
Logger.debug('File checksum matches. No uploading is needed.')
|
||||
return []
|
||||
Logger.debug('Checksum mismatch. Reuploading is required.')
|
||||
|
||||
key = self._new_key(path)
|
||||
return self._push_key(key, path)
|
||||
|
||||
def _makedirs(self, fname):
|
||||
dname = os.path.dirname(fname)
|
||||
try:
|
||||
os.makedirs(dname)
|
||||
except OSError as e:
|
||||
if e.errno != os.errno.EEXIST:
|
||||
raise
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
""" Cloud-specific method of pulling keys """
|
||||
pass
|
||||
|
||||
def _get_key(self, path):
|
||||
""" Cloud-specific method of getting keys """
|
||||
pass
|
||||
|
||||
def pull(self, path):
|
||||
""" Generic method for pulling data from the cloud """
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
Logger.error("File '{}' does not exist in the cloud".format(path))
|
||||
return None
|
||||
|
||||
return self._pull_key(key, path)
|
||||
|
||||
def _status(self, key, path):
|
||||
remote_exists = key != None
|
||||
local_exists = os.path.exists(path)
|
||||
|
||||
diff = None
|
||||
if remote_exists and local_exists:
|
||||
diff = self._cmp_checksum(key, path)
|
||||
|
||||
return STATUS_MAP.get((local_exists, remote_exists, diff), STATUS_UNKNOWN)
|
||||
|
||||
def status(self, path):
|
||||
"""
|
||||
Generic method for checking data item status.
|
||||
"""
|
||||
key = self._get_key(path)
|
||||
if not key:
|
||||
return STATUS_NEW
|
||||
|
||||
return self._status(key, path)
|
||||
|
||||
def connect(self):
|
||||
pass
|
||||
|
||||
def disconnect(self):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
self.connect()
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
self.disconnect()
|
||||
|
||||
+59
-1
@@ -1,3 +1,4 @@
|
||||
import os
|
||||
import posixpath
|
||||
from google.cloud import storage
|
||||
|
||||
@@ -6,8 +7,10 @@ try:
|
||||
except ImportError:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.remote.base import RemoteBase
|
||||
from dvc.config import Config
|
||||
from dvc.progress import progress
|
||||
|
||||
|
||||
class RemoteGS(RemoteBase):
|
||||
@@ -17,7 +20,8 @@ class RemoteGS(RemoteBase):
|
||||
|
||||
def __init__(self, project, config):
|
||||
self.project = project
|
||||
self.url = config.get(Config.SECTION_REMOTE_URL, '/')
|
||||
storagepath = 'gs://' + config.get(Config.SECTION_AWS_STORAGEPATH, '/').lstrip('/')
|
||||
self.url = config.get(Config.SECTION_REMOTE_URL, storagepath)
|
||||
self.projectname = config.get(Config.SECTION_GCP_PROJECTNAME, None)
|
||||
|
||||
@property
|
||||
@@ -98,3 +102,57 @@ class RemoteGS(RemoteBase):
|
||||
|
||||
blob = self.gs.bucket(path_info['bucket']).get_blob(path_info['key'])
|
||||
blob.download_to_filename(path)
|
||||
|
||||
# old code from here
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
name = self.cache_key_name(path)
|
||||
tmp_file = self.tmp_file(path)
|
||||
|
||||
if self._cmp_checksum(key, path):
|
||||
Logger.debug('File "{}" matches with "{}".'.format(path, key.name))
|
||||
return path
|
||||
|
||||
Logger.debug('Downloading cache file from gc "{}/{}"'.format(key.bucket.name, key.name))
|
||||
|
||||
if not no_progress_bar:
|
||||
# percent_cb is not available for download_to_filename, so
|
||||
# lets at least update progress at keypoints(start, finish)
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
try:
|
||||
key.download_to_filename(tmp_file)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to download "{}": {}'.format(key.name, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
Logger.debug('Downloading completed')
|
||||
|
||||
return path
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return self.gs.bucket(self.bucket).get_blob(key_name)
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return self.gs.bucket(self.bucket).blob(key_name)
|
||||
|
||||
def _push_key(self, key, path):
|
||||
""" push, gcp version """
|
||||
name = self.cache_key_name(path)
|
||||
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
key.upload_from_filename(path)
|
||||
|
||||
progress.finish_target(name)
|
||||
Logger.debug('uploading %s completed' % path)
|
||||
|
||||
return path
|
||||
|
||||
+53
-2
@@ -2,15 +2,26 @@ import os
|
||||
import uuid
|
||||
import json
|
||||
import shutil
|
||||
import filecmp
|
||||
|
||||
from dvc.system import System
|
||||
from dvc.remote.base import RemoteBase
|
||||
from dvc.state import State
|
||||
from dvc.logger import Logger
|
||||
from dvc.utils import remove, move
|
||||
from dvc.utils import remove, move, copyfile, file_md5
|
||||
from dvc.config import Config
|
||||
|
||||
|
||||
class LocalKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return os.path.join(self.bucket, self.name)
|
||||
|
||||
|
||||
class RemoteLOCAL(RemoteBase):
|
||||
scheme = ''
|
||||
REGEX = r'^(?P<path>(/+|.:\\+).*)$'
|
||||
@@ -27,7 +38,8 @@ class RemoteLOCAL(RemoteBase):
|
||||
def __init__(self, project, config):
|
||||
self.project = project
|
||||
self.link_state = project.link_state
|
||||
self.cache_dir = config.get(Config.SECTION_REMOTE_URL, None)
|
||||
storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, None)
|
||||
self.cache_dir = config.get(Config.SECTION_REMOTE_URL, storagepath)
|
||||
self.cache_type = config.get(Config.SECTION_CACHE_TYPE, None)
|
||||
|
||||
if self.cache_dir != None and not os.path.exists(self.cache_dir):
|
||||
@@ -262,3 +274,42 @@ class RemoteLOCAL(RemoteBase):
|
||||
raise NotImplementedError
|
||||
|
||||
self._copy(path_info['path'], path)
|
||||
|
||||
# old code starting from here
|
||||
def cache_file_key(self, path):
|
||||
return os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir)
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = LocalKey(self.prefix, key_name)
|
||||
if os.path.exists(key.path) and os.path.isfile(key.path):
|
||||
return key
|
||||
return None
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = LocalKey(self.prefix, key_name)
|
||||
self._makedirs(key.path)
|
||||
return key
|
||||
|
||||
def _push_key(self, key, path):
|
||||
self._makedirs(key.path)
|
||||
name = self.cache_key_name(path)
|
||||
copyfile(path, key.path, name=name)
|
||||
return path
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
name = self.cache_key_name(path)
|
||||
|
||||
tmp_file = self.tmp_file(path)
|
||||
try:
|
||||
copyfile(key.path, tmp_file, no_progress_bar=no_progress_bar, name=name)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to copy "{}": {}'.format(key.path, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
|
||||
return path
|
||||
|
||||
+99
-1
@@ -1,4 +1,7 @@
|
||||
import os
|
||||
import math
|
||||
import boto3
|
||||
import threading
|
||||
import posixpath
|
||||
|
||||
try:
|
||||
@@ -6,10 +9,31 @@ try:
|
||||
except ImportError:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.progress import progress
|
||||
from dvc.config import Config
|
||||
from dvc.remote.base import RemoteBase
|
||||
|
||||
|
||||
class Callback(object):
|
||||
def __init__(self, name, total):
|
||||
self.name = name
|
||||
self.total = total
|
||||
self.current = 0
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def __call__(self, byts):
|
||||
with self.lock:
|
||||
self.current += byts
|
||||
progress.update_target(self.name, self.current, self.total)
|
||||
|
||||
|
||||
class AWSKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
|
||||
class RemoteS3(RemoteBase):
|
||||
scheme = 's3'
|
||||
REGEX = r'^s3://(?P<path>.*)$'
|
||||
@@ -17,7 +41,8 @@ class RemoteS3(RemoteBase):
|
||||
|
||||
def __init__(self, project, config):
|
||||
self.project = project
|
||||
self.url = config.get(Config.SECTION_REMOTE_URL, '/')
|
||||
storagepath = 's3://' + config.get(Config.SECTION_AWS_STORAGEPATH, '').lstrip('/')
|
||||
self.url = config.get(Config.SECTION_REMOTE_URL, storagepath)
|
||||
self.region = config.get(Config.SECTION_AWS_REGION, None)
|
||||
self.profile = config.get(Config.SECTION_AWS_PROFILE, None)
|
||||
self.credentialpath = config.get(Config.SECTION_AWS_CREDENTIALPATH, None)
|
||||
@@ -97,3 +122,76 @@ class RemoteS3(RemoteBase):
|
||||
raise NotImplementedError
|
||||
|
||||
self.s3.Object(path_info['bucket'], path_info['key']).download_file(path)
|
||||
|
||||
# Old code starting from here
|
||||
def create_cb_pull(self, name, key):
|
||||
total = self.s3.Object(bucket_name=key.bucket, key=key.name).content_length
|
||||
return Callback(name, total)
|
||||
|
||||
def create_cb_push(self, name, fname):
|
||||
total = os.path.getsize(fname)
|
||||
return Callback(name, total)
|
||||
|
||||
def _pull_key(self, key, fname, no_progress_bar=False):
|
||||
Logger.debug("Pulling key '{}' from bucket '{}' to file '{}'".format(key.name,
|
||||
key.bucket,
|
||||
fname))
|
||||
self._makedirs(fname)
|
||||
|
||||
tmp_file = self.tmp_file(fname)
|
||||
name = self.cache_key_name(fname)
|
||||
|
||||
if self._cmp_checksum(key, fname):
|
||||
Logger.debug('File "{}" matches with "{}".'.format(fname, key.name))
|
||||
return fname
|
||||
|
||||
Logger.debug('Downloading cache file from S3 "{}/{}" to "{}"'.format(key.bucket,
|
||||
key.name,
|
||||
fname))
|
||||
|
||||
if no_progress_bar:
|
||||
cb = None
|
||||
else:
|
||||
cb = self.create_cb_pull(name, key)
|
||||
|
||||
|
||||
try:
|
||||
self.s3.Object(key.bucket, key.name).download_file(tmp_file, Callback=cb)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to download "{}": {}'.format(key.name, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, fname)
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
Logger.debug('Downloading completed')
|
||||
|
||||
return fname
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
try:
|
||||
self.s3.Object(self.bucket, key_name).get()
|
||||
return AWSKey(self.bucket, key_name)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
return AWSKey(self.bucket, key_name)
|
||||
|
||||
def _push_key(self, key, path):
|
||||
""" push, aws version """
|
||||
name = self.cache_key_name(path)
|
||||
cb = self.create_cb_push(name, path)
|
||||
try:
|
||||
self.s3.Object(key.bucket, key.name).upload_file(path, Callback=cb)
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to upload "{}": {}'.format(path, exc))
|
||||
return None
|
||||
|
||||
progress.finish_target(name)
|
||||
|
||||
return path
|
||||
|
||||
+143
-1
@@ -1,8 +1,150 @@
|
||||
import os
|
||||
import errno
|
||||
import filecmp
|
||||
import paramiko
|
||||
import posixpath
|
||||
|
||||
from dvc.logger import Logger
|
||||
from dvc.utils import copyfile, file_md5
|
||||
from dvc.progress import progress
|
||||
from dvc.remote.base import RemoteBase
|
||||
from dvc.config import Config
|
||||
|
||||
|
||||
def sizeof_fmt(num, suffix='B'):
|
||||
""" Convert number of bytes to human-readable string """
|
||||
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
||||
if abs(num) < 1024.0:
|
||||
return "%3.1f%s%s" % (num, unit, suffix)
|
||||
num /= 1024.0
|
||||
return "%.1f%s%s" % (num, 'Y', suffix)
|
||||
|
||||
|
||||
def percent_cb(name, complete, total):
|
||||
""" Callback for updating target progress """
|
||||
Logger.debug('{}: {} transferred out of {}'.format(name,
|
||||
sizeof_fmt(complete),
|
||||
sizeof_fmt(total)))
|
||||
progress.update_target(name, complete, total)
|
||||
|
||||
|
||||
def create_cb(name):
|
||||
""" Create callback function for multipart object """
|
||||
return (lambda cur, tot: percent_cb(name, cur, tot))
|
||||
|
||||
|
||||
class SSHKey(object):
|
||||
def __init__(self, bucket, name):
|
||||
self.name = name
|
||||
self.bucket = bucket
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return os.path.join(self.bucket, self.name)
|
||||
|
||||
|
||||
class RemoteSSH(RemoteBase):
|
||||
scheme = 'ssh'
|
||||
|
||||
#NOTE: ~/ paths are temporarily forbidden
|
||||
REGEX = r'^(?P<user>.*)@(?P<host>.*):(?P<path>/+.*)$'
|
||||
pass
|
||||
|
||||
|
||||
def __init__(self, project, config):
|
||||
self.project = project
|
||||
storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, '/')
|
||||
self.url = config.get(Config.SECTION_REMOTE_URL, storagepath)
|
||||
|
||||
def group(self, group):
|
||||
return self.match(self.url).group(group)
|
||||
|
||||
@property
|
||||
def path(self):
|
||||
return self.group('path')
|
||||
|
||||
@property
|
||||
def hostname(self):
|
||||
return self.group('host')
|
||||
|
||||
@property
|
||||
def username(self):
|
||||
return self.group('user')
|
||||
|
||||
@property
|
||||
def ssh(self):
|
||||
ssh = paramiko.SSHClient()
|
||||
|
||||
ssh.load_system_host_keys()
|
||||
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
ssh.connect(self.hostname, username=self.username)
|
||||
|
||||
return ssh
|
||||
|
||||
def disconnect(self):
|
||||
self.ssh.close()
|
||||
|
||||
def get_sftp(self):
|
||||
#NOTE: SFTP doesn't seem to be thread safe
|
||||
return self.ssh.open_sftp()
|
||||
|
||||
def cache_file_key(self, path):
|
||||
relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir)
|
||||
return relpath.replace('\\', '/')
|
||||
|
||||
def _isfile_remote(self, path):
|
||||
try:
|
||||
self.get_sftp().open(path, 'r')
|
||||
return True
|
||||
except IOError as exc:
|
||||
if exc.errno != errno.ENOENT:
|
||||
raise
|
||||
return False
|
||||
|
||||
def _get_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = SSHKey(self.path, key_name)
|
||||
if self._isfile_remote(key.path):
|
||||
return key
|
||||
return None
|
||||
|
||||
def _makedirs_remote(self, dname):
|
||||
try:
|
||||
self.get_sftp().chdir(dname)
|
||||
except IOError as exc:
|
||||
if exc.errno != errno.ENOENT:
|
||||
raise
|
||||
|
||||
parent = posixpath.dirname(dname)
|
||||
if len(parent):
|
||||
self._makedirs_remote(parent)
|
||||
self.get_sftp().mkdir(dname)
|
||||
|
||||
def _new_key(self, path):
|
||||
key_name = self.cache_file_key(path)
|
||||
key = SSHKey(self.path, key_name)
|
||||
self._makedirs_remote(posixpath.dirname(key.path))
|
||||
return key
|
||||
|
||||
def _push_key(self, key, path):
|
||||
name = self.cache_key_name(path)
|
||||
self.get_sftp().put(path, key.path, callback=create_cb(name))
|
||||
progress.finish_target(key.name)
|
||||
return path
|
||||
|
||||
def _pull_key(self, key, path, no_progress_bar=False):
|
||||
self._makedirs(path)
|
||||
|
||||
name = self.cache_key_name(path)
|
||||
|
||||
tmp_file = self.tmp_file(path)
|
||||
try:
|
||||
self.get_sftp().get(key.path, tmp_file, callback=create_cb(name))
|
||||
except Exception as exc:
|
||||
Logger.error('Failed to copy "{}": {}'.format(key.path, exc))
|
||||
return None
|
||||
|
||||
os.rename(tmp_file, path)
|
||||
progress.finish_target(key.name)
|
||||
|
||||
return path
|
||||
|
||||
+3
-3
@@ -81,16 +81,16 @@ def cached_property(func):
|
||||
return property(get)
|
||||
|
||||
|
||||
def copyfile(src, dest, no_progress_bar=False):
|
||||
def copyfile(src, dest, no_progress_bar=False, name=None):
|
||||
'''Copy file with progress bar'''
|
||||
copied = 0
|
||||
name = os.path.basename(src)
|
||||
name = name if name else os.path.basename(dest)
|
||||
total = os.stat(src).st_size
|
||||
|
||||
fsrc = open(src, 'rb')
|
||||
|
||||
if os.path.isdir(dest):
|
||||
fdest = open(dest + '/' + name, 'wb+')
|
||||
fdest = open(dest + '/' + os.path.basename(src), 'wb+')
|
||||
else:
|
||||
fdest = open(dest, 'wb+')
|
||||
|
||||
|
||||
@@ -1,139 +1,9 @@
|
||||
import os
|
||||
import configobj
|
||||
from unittest import TestCase
|
||||
import mock
|
||||
|
||||
try:
|
||||
from StringIO import StringIO
|
||||
except ImportError:
|
||||
# Python 3
|
||||
from io import StringIO
|
||||
|
||||
try:
|
||||
import __builtin__
|
||||
builtin_module_name = '__builtin__'
|
||||
except ImportError:
|
||||
# Python 3
|
||||
builtin_module_name = 'builtins'
|
||||
|
||||
from dvc.cloud.data_cloud import DataCloud
|
||||
from dvc.cloud.aws import DataCloudAWS
|
||||
from dvc.cloud.gcp import DataCloudGCP
|
||||
from dvc.main import main
|
||||
from tests.basic_env import TestDvc
|
||||
|
||||
|
||||
class TestConfigTest(TestCase):
|
||||
def setUp(self):
|
||||
pass
|
||||
|
||||
def _conf(self, s):
|
||||
c = configobj.ConfigObj(s)
|
||||
return c
|
||||
|
||||
def storage_path_hierarchy_test(self):
|
||||
""" StoragePath should be read first from global and then from cloud section """
|
||||
|
||||
c = ("[core]",
|
||||
"loglevel =",
|
||||
"cloud = aws",
|
||||
"storagepath = globalsb/global_storage_path",
|
||||
"[aws]",
|
||||
"storagepath = awssb/aws_storage_path",
|
||||
"credentialpath =",
|
||||
"[gcp]",
|
||||
"storagepath = googlesb/google_storage_path"
|
||||
)
|
||||
s = StringIO('\n'.join(c))
|
||||
conf = self._conf(s)
|
||||
cloud = DataCloud(None, conf)
|
||||
self.assertIsInstance(cloud._cloud, DataCloudAWS)
|
||||
self.assertEqual(cloud._cloud.storage_bucket, 'globalsb')
|
||||
self.assertEqual(cloud._cloud.storage_prefix, 'global_storage_path')
|
||||
|
||||
c = ("[core]",
|
||||
"loglevel =",
|
||||
"cloud = Aws",
|
||||
"[aws]",
|
||||
"storagepath = awssb/aws_storage_path",
|
||||
"credentialpath =",
|
||||
"[gcp]",
|
||||
"storagepath = googlesb/google_storage_path"
|
||||
)
|
||||
s = StringIO('\n'.join(c))
|
||||
conf = self._conf(s)
|
||||
cloud = DataCloud(None, conf)
|
||||
self.assertIsInstance(cloud._cloud, DataCloudAWS)
|
||||
self.assertEqual(cloud._cloud.storage_bucket, 'awssb')
|
||||
self.assertEqual(cloud._cloud.storage_prefix, 'aws_storage_path')
|
||||
|
||||
def mocked_open_aws_default_credentials(klass, file, asdf):
|
||||
""" mocked open for ~/.aws/credentials """
|
||||
if file == os.path.expanduser('~/.aws/credentials'):
|
||||
creds0 = ("[default]",
|
||||
"aws_access_key_id = default_access_id",
|
||||
"aws_secret_access_key = default_sekret")
|
||||
return StringIO('\n'.join(creds0))
|
||||
|
||||
creds1 = ("[default]",
|
||||
"aws_access_key_id = override_access_id",
|
||||
"aws_secret_access_key = override_sekret")
|
||||
|
||||
return StringIO('\n'.join(creds1))
|
||||
|
||||
def aws_credentials_default_tes_t_DISABLED(self):
|
||||
""" in absence of [AWS] -> CredentialPath, aws creds should be read from ~/.aws/credentials """
|
||||
|
||||
default_path = os.path.expanduser('~/.aws/credentials')
|
||||
c = ("[core]",
|
||||
"loglevel =",
|
||||
"cloud = aws",
|
||||
"[aws]",
|
||||
"storagepath = awssb/aws_storage_path",
|
||||
"credentialpath =",
|
||||
)
|
||||
s = StringIO('\n'.join(c))
|
||||
|
||||
patcher = mock.patch(builtin_module_name + '.open', side_effect=self.mocked_open_aws_default_credentials)
|
||||
|
||||
# patcher.start()
|
||||
conf = self._conf(s)
|
||||
cloud = DataCloud(None, conf)
|
||||
aws_creds = cloud._cloud._get_credentials()
|
||||
patcher.stop()
|
||||
|
||||
self.assertEqual(aws_creds[0], 'default_access_id')
|
||||
self.assertEqual(aws_creds[1], 'default_sekret')
|
||||
|
||||
def mocked_isfile(file):
|
||||
""" tell os some_credential_path exists """
|
||||
return file == 'some_credential_path'
|
||||
|
||||
@mock.patch('dvc.config.os.path.isfile', side_effect = mocked_isfile)
|
||||
def aws_credentials_specified_tes_t_DISABLED(self, isfile_function):
|
||||
""" in presence of [AWS] -> CredentialPath, use those credentials """
|
||||
|
||||
c = ("[core]",
|
||||
"loglevel =",
|
||||
"cloud = aws",
|
||||
"[aws]",
|
||||
"storagePath = awssb/aws_storage_path",
|
||||
"credentialpath = some_credential_path",
|
||||
)
|
||||
s = StringIO('\n'.join(c))
|
||||
|
||||
patcher = mock.patch(builtin_module_name + '.open', side_effect=self.mocked_open_aws_default_credentials)
|
||||
|
||||
patcher.start()
|
||||
conf = self._conf(s)
|
||||
cloud = DataCloud(None, conf)
|
||||
aws_creds = cloud._cloud._get_credentials()
|
||||
patcher.stop()
|
||||
|
||||
self.assertEqual(aws_creds[0], 'override_access_id')
|
||||
self.assertEqual(aws_creds[1], 'override_sekret')
|
||||
|
||||
|
||||
class TestConfigCLI(TestDvc):
|
||||
def _contains(self, section, field, value):
|
||||
config = configobj.ConfigObj(self.dvc.config.config_file)
|
||||
|
||||
+18
-20
@@ -8,8 +8,8 @@ import tempfile
|
||||
|
||||
from dvc.main import main
|
||||
from dvc.config import Config, ConfigError
|
||||
from dvc.cloud.data_cloud import DataCloud, DataCloudAWS, DataCloudLOCAL, DataCloudGCP
|
||||
from dvc.cloud.base import STATUS_UNKNOWN, STATUS_OK, STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED, CloudSettings
|
||||
from dvc.cloud.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL
|
||||
from dvc.remote.base import STATUS_UNKNOWN, STATUS_OK, STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED
|
||||
|
||||
from tests.basic_env import TestDvc
|
||||
|
||||
@@ -102,21 +102,21 @@ def sleep():
|
||||
|
||||
class TestDataCloud(TestDvc):
|
||||
def _test_cloud(self, config, cl):
|
||||
cloud = DataCloud(cache=self.dvc.cache, config=config)
|
||||
cloud = DataCloud(self.dvc, config=config)
|
||||
self.assertIsInstance(cloud._cloud, cl)
|
||||
|
||||
def test(self):
|
||||
config = TEST_CONFIG
|
||||
|
||||
for scheme, cl in [('s3://', DataCloudAWS), ('gs://', DataCloudGCP), ('/', DataCloudLOCAL)]:
|
||||
config[TEST_SECTION][Config.SECTION_REMOTE_URL] = scheme + 'a/b'
|
||||
for scheme, cl in [('s3://', RemoteS3), ('gs://', RemoteGS), (tempfile.mkdtemp(), RemoteLOCAL)]:
|
||||
config[TEST_SECTION][Config.SECTION_REMOTE_URL] = scheme + str(uuid.uuid4())
|
||||
self._test_cloud(config, cl)
|
||||
|
||||
def test_unsupported(self):
|
||||
with self.assertRaises(ConfigError) as cx:
|
||||
config = TEST_CONFIG
|
||||
config[TEST_SECTION][Config.SECTION_REMOTE_URL] = 'notsupportedscheme://a/b'
|
||||
DataCloud(cache=self.dvc.cache, config=config)
|
||||
DataCloud(self.dvc, config=config)
|
||||
|
||||
|
||||
class TestDataCloudBase(TestDvc):
|
||||
@@ -138,9 +138,7 @@ class TestDataCloudBase(TestDvc):
|
||||
|
||||
config = TEST_CONFIG
|
||||
config[TEST_SECTION][Config.SECTION_REMOTE_URL] = repo
|
||||
cloud_settings = CloudSettings(cache=self.dvc.cache,
|
||||
cloud_config=config[TEST_SECTION])
|
||||
self.cloud = self._get_cloud_class()(cloud_settings)
|
||||
self.cloud = self._get_cloud_class()(self.dvc, config[TEST_SECTION])
|
||||
|
||||
def _test_cloud(self):
|
||||
self._setup_cloud()
|
||||
@@ -215,7 +213,7 @@ class TestDataCloudBase(TestDvc):
|
||||
self._test_cloud()
|
||||
|
||||
|
||||
class TestDataCloudAWS(TestDataCloudBase):
|
||||
class TestRemoteS3(TestDataCloudBase):
|
||||
def _should_test(self):
|
||||
return _should_test_aws()
|
||||
|
||||
@@ -223,10 +221,10 @@ class TestDataCloudAWS(TestDataCloudBase):
|
||||
return get_aws_url()
|
||||
|
||||
def _get_cloud_class(self):
|
||||
return DataCloudAWS
|
||||
return RemoteS3
|
||||
|
||||
|
||||
class TestDataCloudGCP(TestDataCloudBase):
|
||||
class TestRemoteGS(TestDataCloudBase):
|
||||
def _should_test(self):
|
||||
return _should_test_gcp()
|
||||
|
||||
@@ -234,10 +232,10 @@ class TestDataCloudGCP(TestDataCloudBase):
|
||||
return get_gcp_url()
|
||||
|
||||
def _get_cloud_class(self):
|
||||
return DataCloudGCP
|
||||
return RemoteGS
|
||||
|
||||
|
||||
class TestDataCloudLOCAL(TestDataCloudBase):
|
||||
class TestRemoteLOCAL(TestDataCloudBase):
|
||||
def _should_test(self):
|
||||
return True
|
||||
|
||||
@@ -246,10 +244,10 @@ class TestDataCloudLOCAL(TestDataCloudBase):
|
||||
return self.dname
|
||||
|
||||
def _get_cloud_class(self):
|
||||
return DataCloudLOCAL
|
||||
return RemoteLOCAL
|
||||
|
||||
def test(self):
|
||||
super(TestDataCloudLOCAL, self).test()
|
||||
super(TestRemoteLOCAL, self).test()
|
||||
self.assertTrue(os.path.isdir(self.dname))
|
||||
|
||||
|
||||
@@ -314,7 +312,7 @@ class TestDataCloudCLIBase(TestDvc):
|
||||
self._test_compat()
|
||||
|
||||
|
||||
class TestDataCloudLOCALCLI(TestDataCloudCLIBase):
|
||||
class TestRemoteLOCALCLI(TestDataCloudCLIBase):
|
||||
def _test_compat(self):
|
||||
storagepath = get_local_storagepath()
|
||||
self.main(['config', 'core.cloud', 'local'])
|
||||
@@ -330,7 +328,7 @@ class TestDataCloudLOCALCLI(TestDataCloudCLIBase):
|
||||
self._test_cloud(TEST_REMOTE)
|
||||
|
||||
|
||||
class TestDataCloudSSHCLI(TestDataCloudCLIBase):
|
||||
class TestRemoteSSHCLI(TestDataCloudCLIBase):
|
||||
def _should_test(self):
|
||||
return _should_test_ssh()
|
||||
|
||||
@@ -346,7 +344,7 @@ class TestDataCloudSSHCLI(TestDataCloudCLIBase):
|
||||
self._test_cloud(TEST_REMOTE)
|
||||
|
||||
|
||||
class TestDataCloudAWSCLI(TestDataCloudCLIBase):
|
||||
class TestRemoteS3CLI(TestDataCloudCLIBase):
|
||||
def _should_test(self):
|
||||
return _should_test_aws()
|
||||
|
||||
@@ -365,7 +363,7 @@ class TestDataCloudAWSCLI(TestDataCloudCLIBase):
|
||||
self._test_cloud(TEST_REMOTE)
|
||||
|
||||
|
||||
class TestDataCloudGCPCLI(TestDataCloudCLIBase):
|
||||
class TestRemoteGSCLI(TestDataCloudCLIBase):
|
||||
def _should_test(self):
|
||||
return _should_test_gcp()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user