Merge pull request #868 from CatalystCode/azure-blob-storage-remote

Add Azure Blob Storage remote

README.rst

@@ -94,7 +94,7 @@ Related technologies
 #. `Workflow Management Systems <https://en.wikipedia.org/wiki/Workflow_management_system>`_. 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
 =============

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,

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):

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<container_name>[^;]+);)?'
+             r'(?P<connection_string>.+)?$')
+    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)

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

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',

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)