mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-10 23:10:42 +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
+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)
|
||||
|
||||
Reference in New Issue
Block a user