mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
dvc: misc changes for DataCloud, 'data sync' and 'remove'
This patch introduces DataCloud class that used config to determine proper cloud class(DataCloudAWS and DataCloudGCP are awailable) to use in "data_sync" and "remove" commands. This patch focuses on introducing the common class itself saving original code and definetly needs some RFC later. It also adds simple DataCloudLOCAL class that uses directory as a StoragePath and cp as a sync method. Which makes testing so much easier and covers local scenarios. And it also introduces ability to run data sync in parallel by collecting files first and then running sync in a number of threads. Signed-off-by: Ruslan Kuprieiev <[email protected]>
This commit is contained in:
Ruslan Kuprieiev
@@ -68,9 +68,6 @@ class CmdBase(object):
|
||||
def config(self):
|
||||
return self._settings.config
|
||||
|
||||
def cache_file_key(self, file):
|
||||
return '{}/{}'.format(self.config.storage_prefix, file).strip('/')
|
||||
|
||||
@property
|
||||
def dvc_home(self):
|
||||
return self._settings.dvc_home
|
||||
@@ -147,12 +144,6 @@ class CmdBase(object):
|
||||
def run(self):
|
||||
pass
|
||||
|
||||
def get_cache_file_s3_name(self, cache_file):
|
||||
cache_prefix_file_name = os.path.relpath(System.realpath(cache_file), System.realpath(self.git.git_dir))
|
||||
file_name = os.path.relpath(cache_prefix_file_name, self.config.cache_dir)
|
||||
dvc_file_path_trim = file_name.replace(os.sep, '/').strip('/')
|
||||
return self.config.aws_storage_prefix + '/' + dvc_file_path_trim
|
||||
|
||||
@staticmethod
|
||||
def warning_dvc_is_busy():
|
||||
Logger.warn('Cannot perform the cmd since DVC is busy and locked. Please retry the cmd later.')
|
||||
|
||||
+15
-169
@@ -2,6 +2,8 @@ import base64
|
||||
import hashlib
|
||||
import os
|
||||
|
||||
from multiprocessing.pool import ThreadPool
|
||||
|
||||
from boto.s3.connection import S3Connection
|
||||
from google.cloud import storage as gc
|
||||
|
||||
@@ -10,34 +12,12 @@ from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.runtime import Runtime
|
||||
from dvc.system import System
|
||||
|
||||
from dvc.data_cloud import DataCloud
|
||||
|
||||
class DataSyncError(DvcException):
|
||||
def __init__(self, msg):
|
||||
super(DataSyncError, self).__init__('Data sync error: {}'.format(msg))
|
||||
|
||||
|
||||
def sizeof_fmt(num, suffix='B'):
|
||||
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(complete, total):
|
||||
Logger.debug('{} transferred out of {}'.format(sizeof_fmt(complete), sizeof_fmt(total)))
|
||||
|
||||
|
||||
def file_md5(fname):
|
||||
""" get the (md5 hexdigest, md5 digest) of a file """
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(fname, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024*1000), b""):
|
||||
hash_md5.update(chunk)
|
||||
return (hash_md5.hexdigest(), hash_md5.digest())
|
||||
|
||||
|
||||
class CmdDataSync(CmdBase):
|
||||
def __init__(self, settings):
|
||||
super(CmdDataSync, self).__init__(settings)
|
||||
@@ -50,158 +30,24 @@ class CmdDataSync(CmdBase):
|
||||
|
||||
def run(self):
|
||||
with DvcLock(self.is_locker, self.git):
|
||||
self.sync_sanity_check()
|
||||
cloud = DataCloud(self.settings)
|
||||
#FIXME add cmdline argument and config option for POOL_SIZE
|
||||
POOL_SIZE = 10
|
||||
pool = ThreadPool(processes=POOL_SIZE)
|
||||
targets = []
|
||||
|
||||
for target in self.parsed_args.targets:
|
||||
data_item = self.settings.path_factory.data_item(target)
|
||||
if System.islink(target):
|
||||
self.sync_symlink(data_item)
|
||||
targets.append(target)
|
||||
elif os.path.isdir(target):
|
||||
self.sync_dir(target)
|
||||
for root, dirs, files in os.walk(target):
|
||||
for f in files:
|
||||
targets.append(os.path.join(root, f))
|
||||
else:
|
||||
raise DataSyncError('File "{}" does not exit'.format(target))
|
||||
raise DataSyncError('File "{}" does not exit'.format(target))
|
||||
|
||||
pool.map(cloud.sync, targets)
|
||||
pass
|
||||
|
||||
def sync_sanity_check(self):
|
||||
good = self.config.sanity_check()
|
||||
if not good[0]:
|
||||
error = 'config \'%s\' is not correctly setup. Please fix:\n' % Runtime.CONFIG
|
||||
for e in good[1]:
|
||||
error += ' \n' + e
|
||||
raise DataSyncError(error)
|
||||
|
||||
def sync_dir(self, dir):
|
||||
for file in os.listdir(dir):
|
||||
try:
|
||||
fname = None
|
||||
fname = os.path.join(dir, file)
|
||||
self.sync_object(fname)
|
||||
except DataSyncError as ex:
|
||||
Logger.debug(ex)
|
||||
except Exception as ex:
|
||||
Logger.error('Cannot sync file {}: {}'.format(fname, ex))
|
||||
pass
|
||||
|
||||
def sync_object(self, fname):
|
||||
if os.path.isdir(fname):
|
||||
self.sync_dir(fname)
|
||||
elif System.islink(fname):
|
||||
self.sync_symlink(self.settings.path_factory.data_item(fname))
|
||||
else:
|
||||
raise DvcException('Unsupported file type "{}"'.format(fname))
|
||||
|
||||
def sync_symlink(self, data_item):
|
||||
if os.path.isfile(data_item.resolved_cache.dvc):
|
||||
self.sync_to_cloud(data_item)
|
||||
else:
|
||||
self.sync_from_cloud(data_item)
|
||||
|
||||
def _get_bucket_aws(self):
|
||||
""" get a bucket object, aws """
|
||||
conn = S3Connection(self.config.aws_access_key_id, self.config.aws_secret_access_key, host=self.config.aws_region_host)
|
||||
bucket_name = self.config.storage_bucket
|
||||
bucket = conn.lookup(bucket_name)
|
||||
if bucket is None:
|
||||
raise DataSyncError('Storage path is not setup correctly')
|
||||
return bucket
|
||||
|
||||
def _sync_from_cloud_aws(self, item):
|
||||
""" sync from cloud, aws version """
|
||||
|
||||
bucket = self._get_bucket_aws()
|
||||
|
||||
key_name = self.cache_file_key(item.resolved_cache.dvc)
|
||||
key = bucket.get_key(key_name)
|
||||
if not key:
|
||||
raise DataSyncError('File "{}" does not exist in the cloud'.format(key_name))
|
||||
|
||||
Logger.info('Downloading cache file from S3 "{}/{}"'.format(bucket.name, key_name))
|
||||
key.get_contents_to_filename(item.resolved_cache.relative, cb=percent_cb)
|
||||
Logger.info('Downloading completed')
|
||||
|
||||
def _sync_to_cloud_aws(self, data_item):
|
||||
""" sync_to_cloud, aws version """
|
||||
|
||||
aws_key = self.cache_file_key(data_item.resolved_cache.dvc)
|
||||
bucket = self._get_bucket_aws()
|
||||
key = bucket.get_key(aws_key)
|
||||
if key:
|
||||
Logger.debug('File already uploaded to the cloud. Checksum validation...')
|
||||
|
||||
md5_cloud = key.etag[1:-1]
|
||||
md5_local = file_md5(data_item.resolved_cache.dvc)[0]
|
||||
if md5_cloud == md5_local:
|
||||
Logger.debug('File checksum matches. No uploading is needed.')
|
||||
return
|
||||
|
||||
Logger.debug('Checksum miss-match. Re-uploading is required.')
|
||||
|
||||
Logger.info('Uploading cache file "{}" to S3 "{}"'.format(data_item.resolved_cache.relative, aws_key))
|
||||
key = bucket.new_key(aws_key)
|
||||
key.set_contents_from_filename(data_item.resolved_cache.relative, cb=percent_cb)
|
||||
Logger.info('Uploading completed')
|
||||
|
||||
|
||||
def _get_bucket_gc(self):
|
||||
""" get a bucket object, gc """
|
||||
client = gc.Client(project=self.config.gc_project_name)
|
||||
bucket = client.bucket(self.config.storage_bucket)
|
||||
if not bucket.exists():
|
||||
raise DataSyncError('sync up: google cloud bucket {} doesn\'t exist'.format(self.config.storage_bucket))
|
||||
return bucket
|
||||
|
||||
def _sync_from_cloud_gcp(self, item):
|
||||
""" sync from cloud, gcp version """
|
||||
|
||||
bucket = self._get_bucket_gc()
|
||||
key = self.cache_file_key(item.resolved_cache.dvc)
|
||||
|
||||
blob = bucket.get_blob(key)
|
||||
if not blob:
|
||||
raise DataSyncError('File "{}" does not exist in the cloud'.format(key))
|
||||
|
||||
Logger.info('Downloading cache file from gc "{}/{}"'.format(bucket.name, key))
|
||||
|
||||
blob.download_to_filename(item.resolved_cache.dvc)
|
||||
Logger.info('Downloading completed')
|
||||
|
||||
def _sync_to_cloud_gcp(self, data_item):
|
||||
""" sync_to_cloud, gcp version """
|
||||
|
||||
bucket = self._get_bucket_gc()
|
||||
blob_name = self.cache_file_key(data_item.resolved_cache.dvc)
|
||||
|
||||
blob = bucket.get_blob(blob_name)
|
||||
if blob is not None and blob.exists():
|
||||
b64_encoded_md5 = base64.b64encode(file_md5(data_item.resolved_cache.dvc)[1])
|
||||
|
||||
if blob.md5_hash == b64_encoded_md5:
|
||||
Logger.debug('checksum %s matches. Skipping upload' % data_item.cache.relative)
|
||||
return
|
||||
Logger.debug('checksum %s mismatch. re-uploading' % data_item.cache.relative)
|
||||
|
||||
Logger.info('uploading cache file "{} to gc "{}"'.format(data_item.cache.relative, blob_name))
|
||||
|
||||
blob = bucket.blob(blob_name)
|
||||
blob.upload_from_filename(data_item.resolved_cache.relative)
|
||||
Logger.info('uploading %s completed' % data_item.resolved_cache.relative)
|
||||
|
||||
def sync_from_cloud(self, item):
|
||||
cloud = self.settings.config.cloud
|
||||
assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud
|
||||
if cloud == 'aws':
|
||||
return self._sync_from_cloud_aws(item)
|
||||
elif cloud == 'gcp':
|
||||
return self._sync_from_cloud_gcp(item)
|
||||
|
||||
def sync_to_cloud(self, data_item):
|
||||
cloud = self.settings.config.cloud
|
||||
assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud
|
||||
if cloud == 'aws':
|
||||
return self._sync_to_cloud_aws(data_item)
|
||||
elif cloud == 'gcp':
|
||||
return self._sync_to_cloud_gcp(data_item)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Runtime.run(CmdDataSync)
|
||||
|
||||
@@ -4,7 +4,7 @@ import re
|
||||
import requests
|
||||
|
||||
from dvc.command.base import CmdBase, DvcLock
|
||||
from dvc.command.data_sync import sizeof_fmt
|
||||
from dvc.data_cloud import sizeof_fmt
|
||||
from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.runtime import Runtime
|
||||
|
||||
+4
-21
@@ -6,7 +6,7 @@ from dvc.command.base import CmdBase, DvcLock
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.logger import Logger
|
||||
from dvc.runtime import Runtime
|
||||
|
||||
from dvc.data_cloud import DataCloud
|
||||
|
||||
class DataRemoveError(DvcException):
|
||||
def __init__(self, msg):
|
||||
@@ -30,6 +30,8 @@ class CmdDataRemove(CmdBase):
|
||||
|
||||
def run(self):
|
||||
with DvcLock(self.is_locker, self.git):
|
||||
self.cloud = DataCloud(self.settings)
|
||||
|
||||
if not self.remove_all_targets():
|
||||
return 1
|
||||
|
||||
@@ -94,8 +96,7 @@ class CmdDataRemove(CmdBase):
|
||||
|
||||
def _remove_cloud_cache(self, data_item):
|
||||
if not self.parsed_args.keep_in_cloud:
|
||||
aws_key = self.cache_file_key(data_item.cache.dvc)
|
||||
self.remove_from_cloud(aws_key)
|
||||
self.cloud.remove_from_cloud(data_item)
|
||||
|
||||
def _remove_state_file(self, data_item):
|
||||
if os.path.isfile(data_item.state.relative):
|
||||
@@ -119,24 +120,6 @@ class CmdDataRemove(CmdBase):
|
||||
self.remove_dir_if_empty(dvc_path.relative)
|
||||
Logger.debug(u'[Cmd-Remove] Remove {}. Success.'.format(name))
|
||||
|
||||
def remove_from_cloud(self, aws_file_name):
|
||||
Logger.debug(u'[Cmd-Remove] Remove from cloud {}.'.format(aws_file_name))
|
||||
|
||||
if not self.config.aws_access_key_id or not self.config.aws_secret_access_key:
|
||||
Logger.debug('[Cmd-Remove] Unable to check cache file in the cloud')
|
||||
return
|
||||
conn = S3Connection(self.config.aws_access_key_id, self.config.aws_secret_access_key)
|
||||
bucket_name = self.config.storage_bucket
|
||||
bucket = conn.lookup(bucket_name)
|
||||
if bucket:
|
||||
key = bucket.get_key(aws_file_name)
|
||||
if not key:
|
||||
Logger.warn('[Cmd-Remove] S3 remove warning: file "{}" does not exist in S3'.format(aws_file_name))
|
||||
else:
|
||||
key.delete()
|
||||
Logger.info('[Cmd-Remove] File "{}" was removed from S3'.format(aws_file_name))
|
||||
pass
|
||||
|
||||
def remove_dir_file_by_file(self, target):
|
||||
for f in os.listdir(target):
|
||||
file = os.path.join(target, f)
|
||||
|
||||
-186
@@ -31,20 +31,6 @@ class ConfigI(object):
|
||||
def state_dir(self):
|
||||
return self._state_dir
|
||||
|
||||
@property
|
||||
def storage_prefix(self):
|
||||
return ''
|
||||
|
||||
def sanity_check(self):
|
||||
pass
|
||||
|
||||
def aws_access_key_id(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def aws_secret_access_key(self):
|
||||
return None
|
||||
|
||||
|
||||
class Config(ConfigI):
|
||||
CONFIG = 'dvc.conf'
|
||||
@@ -67,8 +53,6 @@ class Config(ConfigI):
|
||||
level = self._config['Global']['LogLevel']
|
||||
Logger.set_level(level)
|
||||
|
||||
self._aws_creds = None
|
||||
|
||||
super(Config, self).__init__(self._config['Global']['DataDir'],
|
||||
self._config['Global']['CacheDir'],
|
||||
self._config['Global']['StateDir'])
|
||||
@@ -77,173 +61,3 @@ class Config(ConfigI):
|
||||
@property
|
||||
def file(self):
|
||||
return self._conf_file
|
||||
|
||||
@property
|
||||
def aws_access_key_id(self):
|
||||
if not self._aws_creds:
|
||||
self._aws_creds = self.get_aws_credentials()
|
||||
if not self._aws_creds:
|
||||
return None
|
||||
return self._aws_creds[0]
|
||||
|
||||
@property
|
||||
def aws_secret_access_key(self):
|
||||
if not self._aws_creds:
|
||||
self._aws_creds = self.get_aws_credentials()
|
||||
if not self._aws_creds:
|
||||
return None
|
||||
return self._aws_creds[1]
|
||||
|
||||
@property
|
||||
def aws_region_host(self):
|
||||
""" get the region host needed for s3 access
|
||||
|
||||
See notes http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region
|
||||
"""
|
||||
|
||||
region = self._config['AWS']['Region']
|
||||
if region is None or region == '':
|
||||
return 's3.amazonaws.com'
|
||||
if region == 'us-east-1':
|
||||
return 's3.amazonaws.com'
|
||||
return 's3.%s.amazonaws.com' % region
|
||||
|
||||
|
||||
@property
|
||||
def storage_path(self):
|
||||
""" get storage path
|
||||
|
||||
Precedence: Storage, then cloud specific
|
||||
"""
|
||||
|
||||
path = self._config['Global'].get('StoragePath', None)
|
||||
if path:
|
||||
return path
|
||||
|
||||
cloud = self.cloud
|
||||
assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud
|
||||
if cloud == 'aws':
|
||||
path = self._config['AWS'].get('StoragePath', None)
|
||||
elif cloud == 'gcp':
|
||||
path = self._config['GCP'].get('StoragePath', None)
|
||||
if path is None:
|
||||
raise ConfigError('invalid StoragePath: not set for Data or cloud specific')
|
||||
return path
|
||||
|
||||
def _storage_path_parts(self):
|
||||
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):
|
||||
parts = self._storage_path_parts()
|
||||
if len(parts) > 1:
|
||||
return parts[1]
|
||||
return ''
|
||||
|
||||
@property
|
||||
def gc_project_name(self):
|
||||
return self._config['GCP']['ProjectName']
|
||||
|
||||
@property
|
||||
def cloud(self):
|
||||
""" get cloud choice: currently one of ['AWS', 'GCP'] """
|
||||
conf = self._config['Global'].get('Cloud', '')
|
||||
if conf.lower() in ['aws']:
|
||||
return 'aws'
|
||||
if conf.lower() in ['gcp']:
|
||||
return 'gcp'
|
||||
return conf
|
||||
|
||||
|
||||
def get_aws_credentials(self):
|
||||
""" gets aws credentials, looking in various places
|
||||
|
||||
Params:
|
||||
|
||||
Searches:
|
||||
1 any override in dvc.conf [AWS] CredentialPath;
|
||||
2 ~/.aws/credentials
|
||||
|
||||
|
||||
Returns:
|
||||
if successfully found, (access_key_id, secret)
|
||||
None otherwise
|
||||
"""
|
||||
default = os.path.expanduser('~/.aws/credentials')
|
||||
|
||||
paths = []
|
||||
credpath = self._config['AWS'].get('CredentialPath', None)
|
||||
if credpath is not None and len(credpath) > 0:
|
||||
credpath = os.path.expanduser(credpath)
|
||||
if os.path.isfile(credpath):
|
||||
paths.append(credpath)
|
||||
else:
|
||||
Logger.warn('AWS CredentialPath "%s" not found; falling back to default "%s"' % (credpath, default))
|
||||
paths.append(default)
|
||||
else:
|
||||
paths.append(default)
|
||||
|
||||
for path in paths:
|
||||
cc = configparser.SafeConfigParser()
|
||||
threw = False
|
||||
try:
|
||||
# use readfp(open( ... to aid mocking.
|
||||
cc.readfp(open(path, 'r'))
|
||||
except Exception as e:
|
||||
threw = True
|
||||
if not threw and 'default' in cc.keys():
|
||||
access_key = cc['default'].get('aws_access_key_id', None)
|
||||
secret = cc['default'].get('aws_secret_access_key', None)
|
||||
|
||||
if access_key is not None and secret is not None:
|
||||
return (access_key, secret)
|
||||
|
||||
return None
|
||||
|
||||
def sanity_check(self):
|
||||
""" sanity check a config
|
||||
|
||||
check that we have a cloud and storagePath
|
||||
if aws, check can read credentials
|
||||
if google, check ProjectName
|
||||
|
||||
Returns:
|
||||
(T,) if good
|
||||
(F, issues) if bad
|
||||
"""
|
||||
errors = []
|
||||
for key in ['Cloud']:
|
||||
if key.lower() not in self._config['Global'].keys() or len(self._config['Global'][key]) < 1:
|
||||
errors.append('Please set %s in section Global in config file %s' % (key, self.file))
|
||||
|
||||
# now that a cloud is chosen, can check StoragePath
|
||||
sp = self.storage_path
|
||||
if sp is None or len(sp) == 0:
|
||||
errors.append('Please set StoragePath = bucket/{optional path} in conf file "%s" '
|
||||
'either in Global or a cloud specific section' % self.CONFIG)
|
||||
|
||||
cloud = self.cloud
|
||||
if cloud == '':
|
||||
# not set; already complained above
|
||||
pass
|
||||
elif cloud == 'aws':
|
||||
creds = self.get_aws_credentials()
|
||||
if creds is None:
|
||||
errors.append('can\'t find aws credentials.')
|
||||
self._aws_creds = creds
|
||||
elif cloud == 'gcp':
|
||||
project = self._config['GCP'].get('ProjectName', None)
|
||||
if project is None or len(project) < 1:
|
||||
errors.append('can\'t read google cloud project name. Please set ProjectName in section GC.')
|
||||
else:
|
||||
errors.append('unknown Cloud %s' % cloud)
|
||||
|
||||
return (len(errors) == 0, errors)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,379 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import configparser
|
||||
import shutil
|
||||
|
||||
from boto.s3.connection import S3Connection
|
||||
from google.cloud import storage as gc
|
||||
|
||||
from dvc.command.base import CmdBase, DvcLock
|
||||
from dvc.logger import Logger
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.runtime import Runtime
|
||||
from dvc.system import System
|
||||
from dvc.config import ConfigError
|
||||
|
||||
class DataCloudError(DvcException):
|
||||
def __init__(self, msg):
|
||||
super(DataCloudError, self).__init__('Data sync error: {}'.format(msg))
|
||||
|
||||
|
||||
def sizeof_fmt(num, suffix='B'):
|
||||
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(complete, total):
|
||||
Logger.debug('{} transferred out of {}'.format(sizeof_fmt(complete), sizeof_fmt(total)))
|
||||
|
||||
|
||||
def file_md5(fname):
|
||||
""" get the (md5 hexdigest, md5 digest) of a file """
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(fname, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024*1000), b""):
|
||||
hash_md5.update(chunk)
|
||||
return (hash_md5.hexdigest(), hash_md5.digest())
|
||||
|
||||
|
||||
class DataCloudBase(object):
|
||||
""" Base class for DataCloud """
|
||||
def __init__(self, settings, config, cloud_config):
|
||||
self._settings = settings
|
||||
self._config = config
|
||||
self._cloud_config = cloud_config
|
||||
|
||||
@property
|
||||
def storage_path(self):
|
||||
""" get storage path
|
||||
|
||||
Precedence: Storage, then cloud specific
|
||||
"""
|
||||
|
||||
path = self._config['Global'].get('StoragePath', None)
|
||||
if path:
|
||||
return path
|
||||
|
||||
path = self._cloud_config.get('StoragePath', None)
|
||||
if path is None:
|
||||
raise ConfigError('invalid StoragePath: not set for Data or cloud specific')
|
||||
|
||||
return path
|
||||
|
||||
def _storage_path_parts(self):
|
||||
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):
|
||||
parts = self._storage_path_parts()
|
||||
if len(parts) > 1:
|
||||
return parts[1]
|
||||
return ''
|
||||
|
||||
def cache_file_key(self, file):
|
||||
return '{}/{}'.format(self.storage_prefix, file).strip('/')
|
||||
|
||||
def sanity_check(self):
|
||||
pass
|
||||
|
||||
def sync_to_cloud(self, item):
|
||||
pass
|
||||
|
||||
def sync_from_cloud(self, item):
|
||||
pass
|
||||
|
||||
def sync(self, fname):
|
||||
item = self._settings.path_factory.data_item(fname)
|
||||
if os.path.isfile(item.resolved_cache.dvc):
|
||||
self.sync_to_cloud(item)
|
||||
else:
|
||||
self.sync_from_cloud(item)
|
||||
|
||||
def remove_from_cloud(self, item):
|
||||
pass
|
||||
|
||||
|
||||
class DataCloudLOCAL(DataCloudBase):
|
||||
def sync_to_cloud(self, item):
|
||||
Logger.debug('sync to cloud ' + item.resolved_cache.dvc + " " + self.storage_path)
|
||||
shutil.copy(item.resolved_cache.dvc, self.storage_path)
|
||||
|
||||
def sync_from_cloud(self, item):
|
||||
Logger.debug('sync from cloud ' + self.storage_path + " " + item.resolved_cache.dvc)
|
||||
shutil.copy(self.storage_path, item.resolved_cache.dvc)
|
||||
|
||||
def remove_from_cloud(self, item):
|
||||
Logger.debug('rm from cloud ' + item.resolved_cache.dvc)
|
||||
os.remove(item.resolved_cache.dvc)
|
||||
|
||||
class DataCloudAWS(DataCloudBase):
|
||||
""" DataCloud class for Amazon Web Services """
|
||||
@property
|
||||
def aws_access_key_id(self):
|
||||
if not self._aws_creds:
|
||||
self._aws_creds = self.get_aws_credentials()
|
||||
if not self._aws_creds:
|
||||
return None
|
||||
return self._aws_creds[0]
|
||||
|
||||
@property
|
||||
def aws_secret_access_key(self):
|
||||
if not self._aws_creds:
|
||||
self._aws_creds = self.get_aws_credentials()
|
||||
if not self._aws_creds:
|
||||
return None
|
||||
return self._aws_creds[1]
|
||||
|
||||
@property
|
||||
def aws_region_host(self):
|
||||
""" get the region host needed for s3 access
|
||||
|
||||
See notes http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region
|
||||
"""
|
||||
|
||||
region = self._cloud_config['Region']
|
||||
if region is None or region == '':
|
||||
return 's3.amazonaws.com'
|
||||
if region == 'us-east-1':
|
||||
return 's3.amazonaws.com'
|
||||
return 's3.%s.amazonaws.com' % region
|
||||
|
||||
def get_aws_credentials(self):
|
||||
""" gets aws credentials, looking in various places
|
||||
|
||||
Params:
|
||||
|
||||
Searches:
|
||||
1 any override in dvc.conf [AWS] CredentialPath;
|
||||
2 ~/.aws/credentials
|
||||
|
||||
|
||||
Returns:
|
||||
if successfully found, (access_key_id, secret)
|
||||
None otherwise
|
||||
"""
|
||||
default = os.path.expanduser('~/.aws/credentials')
|
||||
|
||||
paths = []
|
||||
credpath = self._cloud_config.get('CredentialPath', None)
|
||||
if credpath is not None and len(credpath) > 0:
|
||||
credpath = os.path.expanduser(credpath)
|
||||
if os.path.isfile(credpath):
|
||||
paths.append(credpath)
|
||||
else:
|
||||
Logger.warn('AWS CredentialPath "%s" not found; falling back to default "%s"' % (credpath, default))
|
||||
paths.append(default)
|
||||
else:
|
||||
paths.append(default)
|
||||
|
||||
for path in paths:
|
||||
cc = configparser.SafeConfigParser()
|
||||
threw = False
|
||||
try:
|
||||
# use readfp(open( ... to aid mocking.
|
||||
cc.readfp(open(path, 'r'))
|
||||
except Exception as e:
|
||||
threw = True
|
||||
if not threw and 'default' in cc.keys():
|
||||
access_key = cc['default'].get('aws_access_key_id', None)
|
||||
secret = cc['default'].get('aws_secret_access_key', None)
|
||||
|
||||
if access_key is not None and secret is not None:
|
||||
return (access_key, secret)
|
||||
|
||||
return None
|
||||
|
||||
def sanity_check(self):
|
||||
creds = self.get_aws_credentials()
|
||||
if creds is None:
|
||||
raise ConfigError('can\'t find aws credentials.')
|
||||
self._aws_creds = creds
|
||||
|
||||
def _get_bucket_aws(self):
|
||||
""" get a bucket object, aws """
|
||||
conn = S3Connection(self.aws_access_key_id, self.aws_secret_access_key, host=self.aws_region_host)
|
||||
bucket_name = self.storage_bucket
|
||||
bucket = conn.lookup(bucket_name)
|
||||
if bucket is None:
|
||||
raise DataCloudError('Storage path is not setup correctly')
|
||||
return bucket
|
||||
|
||||
def sync_from_cloud(self, item):
|
||||
""" sync from cloud, aws version """
|
||||
|
||||
bucket = self._get_bucket_aws()
|
||||
|
||||
key_name = self.cache_file_key(item.resolved_cache.dvc)
|
||||
key = bucket.get_key(key_name)
|
||||
if not key:
|
||||
raise DataCloudError('File "{}" does not exist in the cloud'.format(key_name))
|
||||
|
||||
Logger.info('Downloading cache file from S3 "{}/{}"'.format(bucket.name, key_name))
|
||||
key.get_contents_to_filename(item.resolved_cache.relative, cb=percent_cb)
|
||||
Logger.info('Downloading completed')
|
||||
|
||||
def sync_to_cloud(self, data_item):
|
||||
""" sync_to_cloud, aws version """
|
||||
|
||||
aws_key = self.cache_file_key(data_item.resolved_cache.dvc)
|
||||
bucket = self._get_bucket_aws()
|
||||
key = bucket.get_key(aws_key)
|
||||
if key:
|
||||
Logger.debug('File already uploaded to the cloud. Checksum validation...')
|
||||
|
||||
md5_cloud = key.etag[1:-1]
|
||||
md5_local = file_md5(data_item.resolved_cache.dvc)[0]
|
||||
if md5_cloud == md5_local:
|
||||
Logger.debug('File checksum matches. No uploading is needed.')
|
||||
return
|
||||
|
||||
Logger.debug('Checksum miss-match. Re-uploading is required.')
|
||||
|
||||
Logger.info('Uploading cache file "{}" to S3 "{}"'.format(data_item.resolved_cache.relative, aws_key))
|
||||
key = bucket.new_key(aws_key)
|
||||
key.set_contents_from_filename(data_item.resolved_cache.relative, cb=percent_cb)
|
||||
Logger.info('Uploading completed')
|
||||
|
||||
def remove_from_cloud(self, data_item):
|
||||
aws_file_name = self.cache_file_key(data_item.cache.dvc)
|
||||
|
||||
Logger.debug(u'[Cmd-Remove] Remove from cloud {}.'.format(aws_file_name))
|
||||
|
||||
if not self.aws_access_key_id or not self.aws_secret_access_key:
|
||||
Logger.debug('[Cmd-Remove] Unable to check cache file in the cloud')
|
||||
return
|
||||
conn = S3Connection(self.aws_access_key_id, self.aws_secret_access_key)
|
||||
bucket_name = self.storage_bucket
|
||||
bucket = conn.lookup(bucket_name)
|
||||
if bucket:
|
||||
key = bucket.get_key(aws_file_name)
|
||||
if not key:
|
||||
Logger.warn('[Cmd-Remove] S3 remove warning: file "{}" does not exist in S3'.format(aws_file_name))
|
||||
else:
|
||||
key.delete()
|
||||
Logger.info('[Cmd-Remove] File "{}" was removed from S3'.format(aws_file_name))
|
||||
pass
|
||||
|
||||
|
||||
class DataCloudGCP(DataCloudBase):
|
||||
""" DataCloud class for Google Cloud Platform """
|
||||
@property
|
||||
def gc_project_name(self):
|
||||
return self._cloud_config.get('ProjectName', None)
|
||||
|
||||
def sanity_check(self):
|
||||
project = self.gc_project_name
|
||||
if project is None or len(project) < 1:
|
||||
raise ConfigError('can\'t read google cloud project name. Please set ProjectName in section GC.')
|
||||
|
||||
def _get_bucket_gc(self):
|
||||
""" get a bucket object, gc """
|
||||
client = gc.Client(project=self.gc_project_name)
|
||||
bucket = client.bucket(self.storage_bucket)
|
||||
if not bucket.exists():
|
||||
raise DataCloudError('sync up: google cloud bucket {} doesn\'t exist'.format(self.storage_bucket))
|
||||
return bucket
|
||||
|
||||
def sync_from_cloud(self, item):
|
||||
""" sync from cloud, gcp version """
|
||||
|
||||
bucket = self._get_bucket_gc()
|
||||
key = self.cache_file_key(item.resolved_cache.dvc)
|
||||
|
||||
blob = bucket.get_blob(key)
|
||||
if not blob:
|
||||
raise DataCloudError('File "{}" does not exist in the cloud'.format(key))
|
||||
|
||||
Logger.info('Downloading cache file from gc "{}/{}"'.format(bucket.name, key))
|
||||
|
||||
blob.download_to_filename(item.resolved_cache.dvc)
|
||||
Logger.info('Downloading completed')
|
||||
|
||||
def sync_to_cloud(self, data_item):
|
||||
""" sync_to_cloud, gcp version """
|
||||
|
||||
bucket = self._get_bucket_gc()
|
||||
blob_name = self.cache_file_key(data_item.resolved_cache.dvc)
|
||||
|
||||
blob = bucket.get_blob(blob_name)
|
||||
if blob is not None and blob.exists():
|
||||
b64_encoded_md5 = base64.b64encode(file_md5(data_item.resolved_cache.dvc)[1])
|
||||
|
||||
if blob.md5_hash == b64_encoded_md5:
|
||||
Logger.debug('checksum %s matches. Skipping upload' % data_item.cache.relative)
|
||||
return
|
||||
Logger.debug('checksum %s mismatch. re-uploading' % data_item.cache.relative)
|
||||
|
||||
Logger.info('uploading cache file "{} to gc "{}"'.format(data_item.cache.relative, blob_name))
|
||||
|
||||
blob = bucket.blob(blob_name)
|
||||
blob.upload_from_filename(data_item.resolved_cache.relative)
|
||||
Logger.info('uploading %s completed' % data_item.resolved_cache.relative)
|
||||
|
||||
def remove_from_cloud(self, item):
|
||||
raise Exception('NOT IMPLEMENTED YET')
|
||||
|
||||
|
||||
class DataCloud(object):
|
||||
""" Generic class to do initial config parsing and redirect to proper DataCloud methods """
|
||||
|
||||
CLOUD_MAP = {
|
||||
'AWS': DataCloudAWS,
|
||||
'GCP': DataCloudGCP,
|
||||
'LOCAL': DataCloudLOCAL,
|
||||
}
|
||||
|
||||
def __init__(self, settings):
|
||||
#To handle ConfigI case
|
||||
if not hasattr(settings.config, '_config'):
|
||||
self._cloud = DataCloudBase(None, None, None)
|
||||
return
|
||||
|
||||
self._settings = settings
|
||||
self._config = self._settings.config._config
|
||||
|
||||
self.typ = self._config['Global'].get('Cloud', '').strip().upper()
|
||||
if self.typ not in self.CLOUD_MAP.keys():
|
||||
raise ConfigError('Wrong cloud type %s specified' % self.typ)
|
||||
|
||||
self._cloud = self.CLOUD_MAP[self.typ](self._settings, self._config, self._config[self.typ])
|
||||
|
||||
self.sanity_check()
|
||||
|
||||
def sanity_check(self):
|
||||
""" sanity check a config
|
||||
|
||||
check that we have a cloud and storagePath
|
||||
if aws, check can read credentials
|
||||
if google, check ProjectName
|
||||
|
||||
Returns:
|
||||
(T,) if good
|
||||
(F, issues) if bad
|
||||
"""
|
||||
for key in ['Cloud']:
|
||||
if key.lower() not in self._config['Global'].keys() or len(self._config['Global'][key]) < 1:
|
||||
raise ConfigError('Please set %s in section Global in config file %s' % (key, self.file))
|
||||
|
||||
# now that a cloud is chosen, can check StoragePath
|
||||
sp = self._cloud.storage_path
|
||||
if sp is None or len(sp) == 0:
|
||||
raise ConfigError('Please set StoragePath = bucket/{optional path} in conf file "%s" '
|
||||
'either in Global or a cloud specific section' % self.CONFIG)
|
||||
|
||||
self._cloud.sanity_check()
|
||||
|
||||
def sync(self, fname):
|
||||
return self._cloud.sync(fname)
|
||||
|
||||
def remove_from_cloud(self, item):
|
||||
return self._cloud.remove_from_cloud(item)
|
||||
+18
-16
@@ -10,7 +10,7 @@ import mock
|
||||
from dvc.path.path import Path
|
||||
from dvc.system import System
|
||||
from dvc.config import Config
|
||||
|
||||
from dvc.command.data_cloud import DataCloud
|
||||
|
||||
class TestConfigTest(TestCase):
|
||||
def setUp(self):
|
||||
@@ -24,20 +24,20 @@ class TestConfigTest(TestCase):
|
||||
"DataDir = ",
|
||||
"CacheDir = ",
|
||||
"StateDir = ",
|
||||
"Cloud = amazon",
|
||||
"Cloud = aws",
|
||||
"StoragePath = globalsb/global_storage_path",
|
||||
"[AWS]",
|
||||
"StoragePath = awssb/aws_storage_path",
|
||||
"CredentialPath =",
|
||||
"[GC]",
|
||||
"[GCP]",
|
||||
"StoragePath = googlesb/google_storage_path"
|
||||
)
|
||||
s = StringIO.StringIO('\n'.join(c))
|
||||
conf = Config(s, conf_pseudo_file = s)
|
||||
|
||||
self.assertEqual(conf.cloud, 'amazon')
|
||||
self.assertEqual(conf.storage_bucket, 'globalsb')
|
||||
self.assertEqual(conf.storage_prefix, 'global_storage_path')
|
||||
cloud = DataCloud(conf)
|
||||
self.assertEqual(cloud.typ, 'aws')
|
||||
self.assertEqual(cloud._cloud.storage_bucket, 'globalsb')
|
||||
self.assertEqual(cloud._cloud.storage_prefix, 'global_storage_path')
|
||||
|
||||
c = ("[Global]",
|
||||
"LogLevel =",
|
||||
@@ -48,15 +48,15 @@ class TestConfigTest(TestCase):
|
||||
"[AWS]",
|
||||
"StoragePath = awssb/aws_storage_path",
|
||||
"CredentialPath =",
|
||||
"[GC]",
|
||||
"[GCP]",
|
||||
"StoragePath = googlesb/google_storage_path"
|
||||
)
|
||||
s = StringIO.StringIO('\n'.join(c))
|
||||
conf = Config(s, conf_pseudo_file=s)
|
||||
|
||||
self.assertEqual(conf.cloud, 'aws')
|
||||
self.assertEqual(conf.storage_bucket, 'awssb')
|
||||
self.assertEqual(conf.storage_prefix, 'aws_storage_path')
|
||||
cloud = DataCloud(conf)
|
||||
self.assertEqual(cloud.typ, 'aws')
|
||||
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):
|
||||
@@ -82,7 +82,7 @@ class TestConfigTest(TestCase):
|
||||
"DataDir = ",
|
||||
"CacheDir = ",
|
||||
"StateDir = ",
|
||||
"Cloud = amazon",
|
||||
"Cloud = aws",
|
||||
"[AWS]",
|
||||
"StoragePath = awssb/aws_storage_path",
|
||||
"CredentialPath =",
|
||||
@@ -92,7 +92,8 @@ class TestConfigTest(TestCase):
|
||||
patcher = mock.patch('__builtin__.open', side_effect=self.mocked_open_aws_default_credentials)
|
||||
patcher.start()
|
||||
conf = Config(s, conf_pseudo_file=s)
|
||||
aws_creds = conf.get_aws_credentials()
|
||||
cloud = DataCloud(conf)
|
||||
aws_creds = cloud._cloud.get_aws_credentials()
|
||||
patcher.stop()
|
||||
|
||||
self.assertEqual(aws_creds[0], 'default_access_id')
|
||||
@@ -112,7 +113,7 @@ class TestConfigTest(TestCase):
|
||||
"DataDir = ",
|
||||
"CacheDir = ",
|
||||
"StateDir = ",
|
||||
"Cloud = amazon",
|
||||
"Cloud = aws",
|
||||
"[AWS]",
|
||||
"StoragePath = awssb/aws_storage_path",
|
||||
"CredentialPath = some_credential_path",
|
||||
@@ -122,7 +123,8 @@ class TestConfigTest(TestCase):
|
||||
patcher = mock.patch('__builtin__.open', side_effect=self.mocked_open_aws_default_credentials)
|
||||
patcher.start()
|
||||
conf = Config(s, conf_pseudo_file=s)
|
||||
aws_creds = conf.get_aws_credentials()
|
||||
cloud = DataCloud(conf)
|
||||
aws_creds = cloud._cloud.get_aws_credentials()
|
||||
patcher.stop()
|
||||
|
||||
self.assertEqual(aws_creds[0], 'override_access_id')
|
||||
|
||||
Reference in New Issue
Block a user