Merge pull request #834 from efiop/master

RFC/perf optimization for push/pull/status/fetch and support for HDFS for push/pull/status/fetch

dvc/data_cloud.py

@@ -8,6 +8,7 @@ from dvc.config import Config, ConfigError
 from dvc.remote.s3 import RemoteS3
 from dvc.remote.gs import RemoteGS
 from dvc.remote.ssh import RemoteSSH
+from dvc.remote.hdfs import RemoteHDFS
 from dvc.remote.local import RemoteLOCAL
 from dvc.remote.base import RemoteBase
 
@@ -19,6 +20,7 @@ class DataCloud(object):
         'aws'   : RemoteS3,
         'gcp'   : RemoteGS,
         'ssh'   : RemoteSSH,
+        'hdfs'  : RemoteHDFS,
         'local' : RemoteLOCAL,
     }
 

dvc/remote/gs.py

@@ -102,14 +102,17 @@ class RemoteGS(RemoteBase):
                  'bucket': self.bucket,
                  'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s]
 
-    def _exists(self, gs, path_info):
-        return gs.bucket(path_info['bucket']).get_blob(path_info['key']) != None
-
     def exists(self, path_infos):
         ret = []
         gs = self.gs
+
+        keys = [blob.name for blob in gs.bucket(self.bucket).list_blobs(prefix=self.prefix)]
+
         for path_info in path_infos:
-            ret.append(self._exists(gs, path_info))
+            exists = False
+            if path_info['key'] in keys:
+                exists = True
+            ret.append(exists)
         return ret
 
     def upload(self, paths, path_infos, names=None):
@@ -128,9 +131,6 @@ class RemoteGS(RemoteBase):
             if path_info['scheme'] != 'gs':
                 raise NotImplementedError
 
-            if not os.path.exists(path) or self._exists(gs, path_info):
-                continue
-
             Logger.debug("Uploading '{}' to '{}/{}'".format(path,
                                                             path_info['bucket'],
                                                             path_info['key']))
@@ -166,9 +166,6 @@ class RemoteGS(RemoteBase):
             if path_info['scheme'] != 'gs':
                 raise NotImplementedError
 
-            if os.path.exists(path) or not self._exists(gs, path_info):
-                continue
-
             Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'],
                                                               path_info['key'],
                                                               path))

dvc/remote/hdfs.py

@@ -61,13 +61,6 @@ class RemoteHDFS(RemoteBase):
     def rm(self, path_info):
         self.hadoop_fs('rm {}'.format(path_info['url']), user=path_info['user'])
 
-    def exists(self, path_info):
-        try:
-            self.hadoop_fs('test -e {}'.format(path_info['url']), user=path_info['user'])
-            return True
-        except DvcException:
-            return False
-
     def save_info(self, path_info):
         if path_info['scheme'] != 'hdfs':
             raise NotImplementedError
@@ -117,8 +110,31 @@ class RemoteHDFS(RemoteBase):
 
     def md5s_to_path_infos(self, md5s):
         return [{'scheme': 'hdfs',
+                 'user': self.user,
                  'url': posixpath.join(self.url, md5[0:2], md5[2:])} for md5 in md5s]
 
+    def exists(self, path_infos):
+        try:
+            stdout = self.hadoop_fs('ls -R {}'.format(self.url))
+        except DvcException:
+            return len(path_infos) * [False]
+
+        lines = stdout.split('\n')
+        lurl = []
+        for line in lines:
+            if not line.startswith('-'):
+                continue
+            lurl.append(line.split()[-1])
+
+        ret = []
+        for path_info in path_infos:
+            exists = False
+            if path_info['url'] in lurl:
+                exists = True
+            ret.append(exists)
+
+        return ret
+
     def upload(self, paths, path_infos, names=None):
         assert isinstance(paths, list)
         assert isinstance(path_infos, list)
@@ -133,9 +149,6 @@ class RemoteHDFS(RemoteBase):
             if path_info['scheme'] != 'hdfs':
                 raise NotImplementedError
 
-            if not os.path.exists(path) or self.exists(path_info):
-                continue
-
             self.hadoop_fs('mkdir -p {}'.format(posixpath.dirname(path_info['url'])), user=path_info['user'])
             self.hadoop_fs('copyFromLocal {} {}'.format(path, path_info['url']), user=path_info['user'])
 
@@ -153,9 +166,6 @@ class RemoteHDFS(RemoteBase):
             if path_info['scheme'] != 'hdfs':
                 raise NotImplementedError
 
-            if os.path.exists(path) or not self.exists(path_info):
-                continue
-
             dname = os.path.dirname(path)
             if not os.path.exists(dname):
                 os.makedirs(dname)

dvc/remote/local.py

@@ -267,9 +267,6 @@ class RemoteLOCAL(RemoteBase):
             if path_info['scheme'] != 'local':
                 raise NotImplementedError
 
-            if not os.path.exists(path) or os.path.exists(path_info['path']):
-                continue
-
             Logger.debug("Uploading '{}' to '{}'".format(path, path_info['path']))
 
             if not name:
@@ -296,9 +293,6 @@ class RemoteLOCAL(RemoteBase):
             if path_info['scheme'] != 'local':
                 raise NotImplementedError
 
-            if os.path.exists(path) or not os.path.exists(path_info['path']):
-                continue
-
             Logger.debug("Downloading '{}' to '{}'".format(path_info['path'], path))
 
             if not name:
@@ -339,28 +333,18 @@ class RemoteLOCAL(RemoteBase):
     def status(self, checksum_infos, remote, jobs=1):
         checksum_infos = self._collect(checksum_infos)[0]
         md5s = [info[self.PARAM_MD5] for info in checksum_infos]
-        chunks = to_chunks(md5s, jobs)
+        path_infos = remote.md5s_to_path_infos(md5s)
+        remote_exists = remote.exists(path_infos)
+        local_exists = [not self.changed(md5) for md5 in md5s]
 
-        def worker(chunk):
-            local_exists = [os.path.exists(self.get(md5)) for md5 in chunk]
-            path_infos = remote.md5s_to_path_infos(chunk)
-            remote_exists = remote.exists(path_infos)
-            return [(md5, STATUS_MAP[(l,r)]) for md5, l,r in zip(chunk, local_exists, remote_exists)]
-
-        futures = []
-        with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
-            for chunk in chunks:
-                resp = executor.submit(worker, chunk)
-                futures.append(resp)
-
-        status = []
-        for future in futures:
-            status.extend(future.result())
-
-        return status
+        return [(md5, STATUS_MAP[l,r]) for md5, l, r in zip(md5s, local_exists, remote_exists)]
 
     def _do_pull(self, checksum_infos, remote, jobs=1, no_progress_bar=False):
         md5s = [info[self.PARAM_MD5] for info in checksum_infos]
+
+        # NOTE: filter files that are not corrupted
+        md5s = list(filter(lambda md5: self.changed(md5), md5s))
+
         cache = [self.get(md5) for md5 in md5s]
         path_infos = remote.md5s_to_path_infos(md5s)
 
@@ -393,8 +377,15 @@ class RemoteLOCAL(RemoteBase):
 
     def push(self, checksum_infos, remote, jobs=1):
         md5s = [info[self.PARAM_MD5] for info in self._collect(checksum_infos)[0]]
+
         # NOTE: verifying that our cache is not corrupted
         md5s = list(filter(lambda md5: not self.changed(md5), md5s))
+
+        # NOTE: filter files that are already uploaded
+        path_infos = remote.md5s_to_path_infos(md5s)
+        md5s_exist = filter(lambda x: not x[1], list(zip(md5s, remote.exists(path_infos))))
+        md5s = [md5 for md5, exists in md5s_exist]
+
         cache = [self.get(md5) for md5 in md5s]
         path_infos = remote.md5s_to_path_infos(md5s)
 

dvc/remote/s3.py

@@ -117,19 +117,35 @@ class RemoteS3(RemoteBase):
                  'bucket': self.bucket,
                  'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s]
 
-    def _exists(self, s3, path_info):
-        try:
-            s3.Object(path_info['bucket'], path_info['key']).get()
-            return True
-        except Exception:
-            return False
-
     def exists(self, path_infos):
         ret = []
         session = boto3.session.Session()
-        s3 = session.resource('s3')
+        s3 = session.client('s3')
+
+        keys = []
+        kwargs = {'Bucket': self.bucket,
+                  'Prefix': self.prefix}
+        while True:
+            resp = s3.list_objects_v2(**kwargs)
+            contents = resp.get('Contents', None)
+            if not contents:
+                break
+
+            for obj in contents:
+                keys.append(obj['Key'])
+
+            token = resp.get('NextContinuationToken', None)
+            if not token:
+                break
+
+            kwargs['ContinuationToken'] = token
+
         for path_info in path_infos:
-            ret.append(self._exists(s3, path_info))
+            exists = False
+            if path_info['key'] in keys:
+                exists = True
+            ret.append(exists)
+
         return ret
 
     def upload(self, paths, path_infos, names=None):
@@ -143,15 +159,12 @@ class RemoteS3(RemoteBase):
             assert len(names) == len(paths)
 
         session = boto3.session.Session()
-        s3 = session.resource('s3')
+        s3 = session.client('s3')
 
         for path, path_info, name in zip(paths, path_infos, names):
             if path_info['scheme'] != 's3':
                 raise NotImplementedError
 
-            if not os.path.exists(path) or self._exists(s3, path_info):
-                continue
-
             Logger.debug("Uploading '{}' to '{}/{}'".format(path,
                                                             path_info['bucket'],
                                                             path_info['key']))
@@ -163,7 +176,7 @@ class RemoteS3(RemoteBase):
             cb = Callback(name, total)
 
             try:
-                s3.Object(path_info['bucket'], path_info['key']).upload_file(path, Callback=cb)
+                s3.upload_file(path, path_info['bucket'], path_info['key'], Callback=cb)
             except Exception as exc:
                 Logger.error("Failed to upload '{}'".format(path), exc)
                 continue
@@ -181,15 +194,12 @@ class RemoteS3(RemoteBase):
             assert len(names) == len(fnames)
 
         session = boto3.session.Session()
-        s3 = session.resource('s3')
+        s3 = session.client('s3')
 
         for fname, path_info, name in zip(fnames, path_infos, names):
             if path_info['scheme'] != 's3':
                 raise NotImplementedError
 
-            if os.path.exists(fname) or not self._exists(s3, path_info):
-                continue
-
             Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'],
                                                               path_info['key'],
                                                               fname))
@@ -201,14 +211,14 @@ class RemoteS3(RemoteBase):
             if no_progress_bar:
                 cb = None
             else:
-                total = s3.Object(bucket_name=path_info['bucket'],
-                                  key=path_info['key']).content_length
+                total = s3.head_object(Bucket=path_info['bucket'],
+                                       Key=path_info['key'])['ContentLength']
                 cb = Callback(name, total)
 
             self._makedirs(fname)
 
             try:
-                s3.Object(path_info['bucket'], path_info['key']).download_file(tmp_file, Callback=cb)
+                s3.download_file(path_info['bucket'], path_info['key'], tmp_file, Callback=cb)
             except Exception as exc:
                 Logger.error("Failed to download '{}/{}'".format(path_info['bucket'],
                                                                  path_info['key']), exc)

dvc/remote/ssh.py

@@ -86,10 +86,17 @@ class RemoteSSH(RemoteBase):
 
     def exists(self, path_infos):
         ret = []
-        ssh = self.ssh(path_infos[0]['host'], path_infos[0]['user'])
-        for path_info in path_infos:
-            ret.append(self._exists(ssh, path_info))
+        ssh = self.ssh(host=self.host, user=self.user)
+        stdout = self._exec(ssh, 'find {} -type f -follow -print'.format(self.prefix))
+        plist = stdout.split()
         ssh.close()
+
+        for path_info in path_infos:
+            exists = False
+            if path_info['path'] in plist:
+                exists = True
+            ret.append(exists)
+
         return ret
 
     def _exec(self, ssh, cmd):
@@ -197,9 +204,6 @@ class RemoteSSH(RemoteBase):
             if path_info['scheme'] != 'ssh':
                 raise NotImplementedError
 
-            if os.path.exists(path) or not self._exists(ssh, path_info):
-                continue
-
             Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['host'],
                                                               path_info['path'],
                                                               path))
@@ -238,9 +242,6 @@ class RemoteSSH(RemoteBase):
             if path_info['scheme'] != 'ssh':
                 raise NotImplementedError
 
-            if not os.path.exists(path) or self._exists(ssh, path_info):
-                continue
-
             Logger.debug("Uploading '{}' to '{}/{}'".format(path,
                                                             path_info['host'],
                                                             path_info['path']))

tests/test_data_cloud.py

@@ -8,7 +8,7 @@ import tempfile
 
 from dvc.main import main
 from dvc.config import Config, ConfigError
-from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH
+from dvc.data_cloud import DataCloud, RemoteS3, RemoteGS, RemoteLOCAL, RemoteSSH, RemoteHDFS
 from dvc.remote.base import STATUS_OK, STATUS_NEW, STATUS_DELETED
 
 from tests.basic_env import TestDvc
@@ -63,6 +63,14 @@ def _should_test_ssh():
     return True
 
 
+def _should_test_hdfs():
+    if os.name == 'nt':
+        return False
+    if os.getenv('HADOOP_CONTAINER_IP'):
+        return True
+    return False
+
+
 def get_local_storagepath():
     return tempfile.mkdtemp()
 
@@ -75,6 +83,10 @@ def get_ssh_url():
     return 'ssh://{}@127.0.0.1:{}'.format(getpass.getuser(), get_local_storagepath())
 
 
+def get_hdfs_url():
+    return 'hdfs://root@{}{}'.format(os.getenv('HADOOP_CONTAINER_IP'), get_local_storagepath())
+
+
 def get_aws_storagepath():
     return TEST_AWS_REPO_BUCKET + '/' + str(uuid.uuid4())
 
@@ -261,6 +273,20 @@ class TestRemoteSSH(TestDataCloudBase):
         super(TestRemoteSSH, self).test()
 
 
+class TestRemoteHDFS(TestDataCloudBase):
+    def _should_test(self):
+        return _should_test_hdfs()
+
+    def _get_url(self):
+        return get_hdfs_url()
+
+    def _get_cloud_class(self):
+        return RemoteHDFS
+
+    def test(self):
+        super(TestRemoteHDFS, self).test()
+
+
 class TestDataCloudCLIBase(TestDvc):
     def main(self, args):
         ret = main(args)
@@ -348,6 +374,18 @@ class TestRemoteSSHCLI(TestDataCloudCLIBase):
         self._test_cloud(TEST_REMOTE)
 
 
+class TestRemoteHDFSCLI(TestDataCloudCLIBase):
+    def _should_test(self):
+        return _should_test_hdfs()
+
+    def _test(self):
+        url = get_hdfs_url()
+
+        self.main(['remote', 'add', TEST_REMOTE, url])
+
+        self._test_cloud(TEST_REMOTE)
+
+
 class TestCompatRemoteS3CLI(TestDataCloudCLIBase):
     def _should_test(self):
         return _should_test_aws()

tests/test_repro.py

@@ -22,7 +22,7 @@ from dvc.stage import Stage
 from tests.basic_env import TestDvc
 from tests.test_data_cloud import _should_test_aws, TEST_AWS_REPO_BUCKET
 from tests.test_data_cloud import _should_test_gcp, TEST_GCP_REPO_BUCKET
-from tests.test_data_cloud import _should_test_ssh
+from tests.test_data_cloud import _should_test_ssh, _should_test_hdfs
 
 
 class TestRepro(TestDvc):
@@ -393,11 +393,7 @@ class TestReproExternalGS(TestReproExternalBase):
 
 class TestReproExternalHDFS(TestReproExternalBase):
     def should_test(self):
-        if os.name == 'nt':
-            return False
-        if os.getenv('HADOOP_CONTAINER_IP'):
-            return True
-        return False
+        return _should_test_hdfs()
 
     @property
     def scheme(self):