From 6882a781a4219559cff94a0d1dc6eb6144fb0a05 Mon Sep 17 00:00:00 2001 From: Clemens Wolff Date: Mon, 9 Jul 2018 14:39:24 -0400 Subject: [PATCH] Add Azure Blob Storage remote This change adds a new cloud remote based on Azure Blob Storage [1] and the azure-storage-blob SDK [2]. The implementation closely follows the patterns laid out by the existing cloud remotes (AWS and GCP). Notable differences include: - No use of a lock in the progress-bar callback and no requirement to look up total transfer size ahead of time: the Azure Storage SDK provides the total transfer bytes and current transfer offset in its progress callback. - When `exists` is called, check the existence of every blob instead of listing the entire container. Container list operations can be expensive so unless the number of checked files is very large, making separate requests is likely going to be more efficient. - For increased flexibility, access credentials for the Azure Storage backend can be provided either connection-string style via the remote URL or via environment variables. The Azure Storage backend can be emulated via Azurite [3]. [1] https://azure.microsoft.com/en-us/services/storage/blobs/ [2] https://github.com/Azure/azure-storage-python [3] https://github.com/Azure/azurite This change was created together with @EricSchles. --- README.rst | 2 +- dvc/data_cloud.py | 2 + dvc/remote/__init__.py | 3 +- dvc/remote/azure.py | 241 +++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + setup.py | 6 +- tests/test_data_cloud.py | 38 +++++- 7 files changed, 289 insertions(+), 4 deletions(-) create mode 100644 dvc/remote/azure.py diff --git a/README.rst b/README.rst index 43a68d6a2..92aeb751b 100644 --- a/README.rst +++ b/README.rst @@ -94,7 +94,7 @@ Related technologies #. `Workflow Management Systems `_. DVC is workflow management system designed specificaly to manage machine learning experiments. DVC was built on top of Git. DVC is compatible with Git for storing code and the dependency graph (DAG), but not data files cache. -Data files caches can be transferred separately - now data cache transfer throught AWS S3 and GCP storge are supported. +Data files caches can be transferred separately - now data cache transfer throught AWS S3, Azure Blob Storage and GCP storge are supported. How DVC works ============= diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index e608da938..4a163fb68 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -7,6 +7,7 @@ from dvc.config import Config, ConfigError from dvc.remote.s3 import RemoteS3 from dvc.remote.gs import RemoteGS +from dvc.remote.azure import RemoteAzure from dvc.remote.ssh import RemoteSSH from dvc.remote.hdfs import RemoteHDFS from dvc.remote.local import RemoteLOCAL @@ -19,6 +20,7 @@ class DataCloud(object): CLOUD_MAP = { 'aws' : RemoteS3, 'gcp' : RemoteGS, + 'azure' : RemoteAzure, 'ssh' : RemoteSSH, 'hdfs' : RemoteHDFS, 'local' : RemoteLOCAL, diff --git a/dvc/remote/__init__.py b/dvc/remote/__init__.py index 40a44b5d3..a6537c921 100644 --- a/dvc/remote/__init__.py +++ b/dvc/remote/__init__.py @@ -3,12 +3,13 @@ from dvc.remote.s3 import RemoteS3 from dvc.remote.gs import RemoteGS from dvc.remote.hdfs import RemoteHDFS from dvc.remote.ssh import RemoteSSH +from dvc.remote.azure import RemoteAzure from dvc.config import Config from dvc.exceptions import DvcException -REMOTES = [RemoteHDFS, RemoteSSH, RemoteS3, RemoteGS, RemoteLOCAL] +REMOTES = [RemoteHDFS, RemoteSSH, RemoteS3, RemoteGS, RemoteAzure, RemoteLOCAL] def supported_url(url): diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py new file mode 100644 index 000000000..6f6b09f63 --- /dev/null +++ b/dvc/remote/azure.py @@ -0,0 +1,241 @@ +import os +import re +import time + +try: + from azure.storage.blob import BlockBlobService +except ImportError: + BlockBlobService = None + +try: + from urlparse import urlparse +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): + self.name = name + + def __call__(self, current, total): + progress.update_target(self.name, current, total) + + +class RemoteAzure(RemoteBase): + scheme = 'azure' + REGEX = (r'^azure://' + r'(ContainerName=(?P[^;]+);)?' + r'(?P.+)?$') + REQUIRES = {'azure-storage-blob': BlockBlobService} + PARAM_ETAG = 'etag' + COPY_POLL_SECONDS = 5 + + def __init__(self, project, config): + super().__init__(project, config) + self.project = project + + url = config.get(Config.SECTION_REMOTE_URL) + match = re.match(self.REGEX, url) + + self.bucket = ( + match.group('container_name') + or os.getenv('AZURE_STORAGE_CONTAINER_NAME')) + + self.connection_string = ( + match.group('connection_string') + or os.getenv('AZURE_STORAGE_CONNECTION_STRING')) + + if not self.bucket: + raise ValueError('Azure Storage container name missing') + + if not self.connection_string: + raise ValueError('Azure Storage connection string missing') + + self.__blob_service = None + + @property + def blob_service(self): + if self.__blob_service is None: + self.__blob_service = BlockBlobService( + connection_string=self.connection_string) + self.__blob_service.create_container(self.bucket) + return self.__blob_service + + def _get_etag(self, bucket, key): + try: + blob = self.blob_service.get_blob_properties(bucket, key) + return blob.properties.etag + except Exception: + return None + + def save_info(self, path_info): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + return {self.PARAM_ETAG: self._get_etag( + path_info['bucket'], path_info['key'])} + + def save(self, path_info): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + etag = self._get_etag(path_info['bucket'], path_info['key']) + dest_key = '{}/{}'.format(etag[0:2], etag[2:]) + + self._copy_blob(to_bucket=self.bucket, + to_key=dest_key, + from_bucket=path_info['bucket'], + from_key=path_info['key']) + + return {self.PARAM_ETAG: etag} + + def _copy_blob(self, to_bucket, to_key, from_bucket, from_key): + source = self.blob_service.make_blob_url(from_bucket, from_key) + + copy = self.blob_service.copy_blob(to_bucket, to_key, source) + + if self.COPY_POLL_SECONDS <= 0: + return + + while copy.status != 'success': + time.sleep(self.COPY_POLL_SECONDS) + copy = self.blob_service.get_blob_properties( + to_bucket, to_key).properties.copy + + def checkout(self, path_info, checksum_info): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + etag = checksum_info.get(self.PARAM_ETAG, None) + if not etag: + return + + self._copy_blob(to_bucket=path_info['bucket'], + to_key=path_info['key'], + from_bucket=self.bucket, + from_key='{}/{}'.format(etag[0:2], etag[2:])) + + def remove(self, path_info): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + Logger.debug('Removing azure://{}/{}'.format(path_info['bucket'], + path_info['key'])) + + self.blob_service.delete_blob(path_info['bucket'], path_info['key']) + + def md5s_to_path_infos(self, md5s): + return [{ + 'scheme': self.scheme, + 'bucket': self.bucket, + 'key': '{}/{}'.format(md5[0:2], md5[2:]) + } for md5 in md5s] + + def exists(self, path_infos): + # TODO: why does S3 fetch all blobs first? + ret = [] + + for path_info in path_infos: + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + bucket = path_info['bucket'] + key = path_info['key'] + + exists = self.blob_service.exists(bucket, key) + ret.append(exists) + + return ret + + def upload(self, paths, path_infos, names=None): + assert isinstance(paths, list) + assert isinstance(path_infos, list) + assert len(paths) == len(path_infos) + if not names: + names = len(paths) * [None] + else: + assert isinstance(names, list) + assert len(names) == len(paths) + + for path, path_info, name in zip(paths, path_infos, names): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + bucket = path_info['bucket'] + key = path_info['key'] + + Logger.debug("Uploading '{}' to '{}/{}'".format( + path, bucket, key)) + + if not name: + name = os.path.basename(path) + + cb = Callback(name) + + try: + self.blob_service.create_blob_from_path( + bucket, key, path, progress_callback=cb) + except Exception as ex: + Logger.error("Failed to upload '{}'".format(path), ex) + else: + progress.finish_target(name) + + def download(self, path_infos, fnames, no_progress_bar=False, names=None): + assert isinstance(fnames, list) + assert isinstance(path_infos, list) + assert len(fnames) == len(path_infos) + if not names: + names = len(fnames) * [None] + else: + assert isinstance(names, list) + assert len(names) == len(fnames) + + for fname, path_info, name in zip(fnames, path_infos, names): + if path_info['scheme'] != self.scheme: + raise NotImplementedError + + bucket = path_info['bucket'] + key = path_info['key'] + + Logger.debug("Downloading '{}/{}' to '{}'".format( + bucket, key, fname)) + + tmp_file = self.tmp_file(fname) + if not name: + name = os.path.basename(fname) + + cb = None if no_progress_bar else Callback(name) + + self._makedirs(fname) + + try: + self.blob_service.get_blob_to_path( + bucket, key, tmp_file, progress_callback=cb) + except Exception as exc: + Logger.error("Failed to download '{}/{}'".format( + path_info['bucket'], path_info['key']), exc) + return + + os.rename(tmp_file, fname) + + if not no_progress_bar: + progress.finish_target(name) + + def gc(self, checksum_infos): + used_etags = [info[self.PARAM_ETAG] for info in checksum_infos] + + all_blobs = self.blob_service.list_blobs(self.bucket) + + for blob in all_blobs: + etag = blob.properties.etag + if etag in used_etags: + continue + path_info = {'scheme': self.scheme, + 'key': blob.name, + 'bucket': self.bucket} + self.remove(path_info) diff --git a/requirements.txt b/requirements.txt index 938a5ef76..4fdf9385e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +azure-storage-blob==1.3.0 boto3==1.7.4 ply>=3.9 # See https://github.com/pyinstaller/pyinstaller/issues/1945 configparser>=3.5.0 diff --git a/setup.py b/setup.py index f6b2e6886..3f4a0291f 100644 --- a/setup.py +++ b/setup.py @@ -32,10 +32,13 @@ gs = [ s3 = [ "boto3==1.7.4", ] +azure = [ + "azure-storage-blob==1.3.0" +] ssh = [ "paramiko>=2.4.1", ] -all_remotes = gs + s3 + ssh +all_remotes = gs + s3 + azure + ssh setup( name='dvc', @@ -51,6 +54,7 @@ setup( 'all': all_remotes, 'gs': gs, 's3': s3, + 'azure': azure, 'ssh': ssh, }, keywords='data science, data version control, machine learning', diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index a134c754b..66f582d5e 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -9,7 +9,7 @@ import platform from dvc.main import main from dvc.config import Config, ConfigError -from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH, RemoteHDFS +from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteAzure, RemoteLOCAL, RemoteSSH, RemoteHDFS from dvc.remote.base import STATUS_OK, STATUS_NEW, STATUS_DELETED from tests.basic_env import TestDvc @@ -53,6 +53,15 @@ def _should_test_gcp(): return False +def _should_test_azure(): + if os.getenv("DVC_TEST_AZURE") == "true": + return True + elif os.getenv("DVC_TEST_AZURE") == "false": + return False + + return os.getenv("AZURE_STORAGE_CONTAINER") and os.getenv("AZURE_STORAGE_CONNECTION_STRING") + + def _should_test_ssh(): if os.getenv("DVC_TEST_SSH") == "true": return True @@ -102,6 +111,10 @@ def get_gcp_url(): return 'gs://' + get_gcp_storagepath() +def get_azure_url(): + return 'azure://' + + def sleep(): # Sleep some time to simulate realistic behavior. # Some filesystems have a bad date resolution for @@ -242,6 +255,17 @@ class TestRemoteGS(TestDataCloudBase): return RemoteGS +class TestRemoteAzure(TestDataCloudBase): + def _should_test(self): + return _should_test_azure() + + def _get_url(self): + return get_azure_url() + + def _get_cloud_class(self): + return RemoteAzure + + class TestRemoteLOCAL(TestDataCloudBase): def _should_test(self): return True @@ -433,6 +457,18 @@ class TestRemoteGSCLI(TestDataCloudCLIBase): self._test_cloud(TEST_REMOTE) +class TestRemoteAzureCLI(TestDataCloudCLIBase): + def _should_test(self): + return _should_test_azure() + + def _test(self): + url = get_azure_url() + + self.main(['remote', 'add', TEST_REMOTE, url]) + + self._test_cloud(TEST_REMOTE) + + class TestDataCloudErrorCLI(TestDvc): def main_fail(self, args): ret = main(args)