From 5a3543b36fe6b61244772c84612e8df0401b4910 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sat, 30 Jun 2018 22:10:03 +0300 Subject: [PATCH 1/2] RFC/perf: optimise push/pull/status Signed-off-by: Ruslan Kuprieiev --- dvc/remote/gs.py | 17 ++++++--------- dvc/remote/hdfs.py | 13 ------------ dvc/remote/local.py | 39 +++++++++++++--------------------- dvc/remote/s3.py | 52 +++++++++++++++++++++++++++------------------ dvc/remote/ssh.py | 19 +++++++++-------- 5 files changed, 63 insertions(+), 77 deletions(-) diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index c6e03e2e7..5285f4b5d 100644 --- a/dvc/remote/gs.py +++ b/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)) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index e25d69910..ea19db0f1 100644 --- a/dvc/remote/hdfs.py +++ b/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 @@ -133,9 +126,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 +143,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) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 042c5d60e..7a6550881 100644 --- a/dvc/remote/local.py +++ b/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) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 88c516bcc..b56586e86 100644 --- a/dvc/remote/s3.py +++ b/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) diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 749f0b750..e23de8e14 100644 --- a/dvc/remote/ssh.py +++ b/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'])) From 12edcbaa04c2203f6cb01337f08c993c1421f542 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sat, 30 Jun 2018 22:54:49 +0300 Subject: [PATCH 2/2] remote: support HDFS for push/pull/status from local repo Signed-off-by: Ruslan Kuprieiev --- dvc/data_cloud.py | 2 ++ dvc/remote/hdfs.py | 23 +++++++++++++++++++++++ tests/test_data_cloud.py | 40 +++++++++++++++++++++++++++++++++++++++- tests/test_repro.py | 8 ++------ 4 files changed, 66 insertions(+), 7 deletions(-) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index 4f80bf2e4..e608da938 100644 --- a/dvc/data_cloud.py +++ b/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, } diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index ea19db0f1..773a8dbe6 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -110,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) diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 00165c759..066f6c6a6 100644 --- a/tests/test_data_cloud.py +++ b/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() diff --git a/tests/test_repro.py b/tests/test_repro.py index a6688a929..b178c1268 100644 --- a/tests/test_repro.py +++ b/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):