From ff0d91737bd4df37971e0b00eb645dbbe2adf687 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 10 Jun 2018 05:10:27 +0300 Subject: [PATCH] dvc: initial hdfs support Signed-off-by: Ruslan Kuprieiev --- dvc/cache.py | 1 + dvc/config.py | 8 ++- dvc/dependency/__init__.py | 8 ++- dvc/dependency/hdfs.py | 33 ++++++++++++ dvc/output/__init__.py | 6 ++- dvc/output/hdfs.py | 43 ++++++++++++++++ dvc/remote/__init__.py | 15 +++++- dvc/remote/hdfs.py | 103 +++++++++++++++++++++++++++++++++++++ tests/test_repro.py | 41 +++++++++++++++ 9 files changed, 251 insertions(+), 7 deletions(-) create mode 100644 dvc/dependency/hdfs.py create mode 100644 dvc/output/hdfs.py create mode 100644 dvc/remote/hdfs.py diff --git a/dvc/cache.py b/dvc/cache.py index b2a03b510..13c7b4760 100644 --- a/dvc/cache.py +++ b/dvc/cache.py @@ -28,6 +28,7 @@ class Cache(object): self.s3 = self._get_remote(project, config, Config.SECTION_CACHE_S3) self.gs = self._get_remote(project, config, Config.SECTION_CACHE_GS) self.ssh = self._get_remote(project, config, Config.SECTION_CACHE_SSH) + self.hdfs = self._get_remote(project, config, Config.SECTION_CACHE_HDFS) def _get_remote(self, project, config, name): remote = config.get(name, None) diff --git a/dvc/config.py b/dvc/config.py index 76777d630..a1c27d9f6 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -15,8 +15,8 @@ class ConfigError(DvcException): def supported_url(url): - from dvc.cloud.data_cloud import DataCloud - return DataCloud.supported(url) != None + from dvc.remote import supported_url as supported + return supported(url) class Config(object): @@ -36,10 +36,12 @@ class Config(object): SECTION_CACHE_S3 = 's3' SECTION_CACHE_GS = 'gs' SECTION_CACHE_SSH = 'ssh' + SECTION_CACHE_HDFS = 'hdfs' SECTION_CACHE_SCHEMA = { schema.Optional(SECTION_CACHE_LOCAL): str, schema.Optional(SECTION_CACHE_S3): str, schema.Optional(SECTION_CACHE_GS): str, + schema.Optional(SECTION_CACHE_HDFS): str, # backward compatibility schema.Optional(SECTION_CACHE_DIR, default='cache'): str, @@ -93,6 +95,7 @@ class Config(object): SECTION_REMOTE_REGEX = r'^\s*remote\s*"(?P.*)"\s*$' SECTION_REMOTE_FMT = 'remote "{}"' SECTION_REMOTE_URL = 'url' + SECTION_REMOTE_USER = 'user' SECTION_REMOTE_SCHEMA = { SECTION_REMOTE_URL: supported_url, schema.Optional(SECTION_AWS_REGION): str, @@ -100,6 +103,7 @@ class Config(object): schema.Optional(SECTION_AWS_CREDENTIALPATH, default = ''): str, schema.Optional(SECTION_GCP_PROJECTNAME): str, schema.Optional(SECTION_CACHE_TYPE): SECTION_CACHE_TYPE_SCHEMA, + schema.Optional(SECTION_REMOTE_USER): str, } SCHEMA = { diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index c86f7cada..50222f029 100644 --- a/dvc/dependency/__init__.py +++ b/dvc/dependency/__init__.py @@ -12,23 +12,27 @@ from dvc.dependency.base import DependencyBase from dvc.dependency.s3 import DependencyS3 from dvc.dependency.gs import DependencyGS from dvc.dependency.local import DependencyLOCAL +from dvc.dependency.hdfs import DependencyHDFS from dvc.remote import Remote from dvc.remote.local import RemoteLOCAL from dvc.remote.s3 import RemoteS3 from dvc.remote.gs import RemoteGS from dvc.remote.ssh import RemoteSSH +from dvc.remote.hdfs import RemoteHDFS -DEPS = [DependencyS3, DependencyGS, DependencyLOCAL] +DEPS = [DependencyHDFS, DependencyS3, DependencyGS, DependencyLOCAL] DEP_MAP = {'': DependencyLOCAL, 's3': DependencyS3, - 'gs': DependencyGS} + 'gs': DependencyGS, + 'hdfs': DependencyHDFS,} SCHEMA = { DependencyBase.PARAM_PATH: str, schema.Optional(RemoteLOCAL.PARAM_MD5): schema.Or(str, None), schema.Optional(RemoteS3.PARAM_ETAG): schema.Or(str, None), + schema.Optional(RemoteHDFS.PARAM_CHECKSUM): schema.Or(str, None), } diff --git a/dvc/dependency/hdfs.py b/dvc/dependency/hdfs.py new file mode 100644 index 000000000..9e0e6e5c8 --- /dev/null +++ b/dvc/dependency/hdfs.py @@ -0,0 +1,33 @@ +import posixpath + +try: + from urlparse import urlparse +except ImportError: + from urllib.parse import urlparse + +from dvc.dependency.base import DependencyBase +from dvc.remote.hdfs import RemoteHDFS + + +class DependencyHDFS(DependencyBase): + REGEX = r'^hdfs://.*$' + + def __init__(self, stage, path, info=None, remote=None): + super(DependencyHDFS, self).__init__(stage, path) + self.info = info + self.remote = remote if remote else RemoteHDFS(stage.project, {}) + if remote: + path = posixpath.join(remote.url, urlparse(path).path.lstrip('/')) + self.path_info = {'scheme': 'hdfs', + 'url': path} + + def changed(self): + return self.info != self.remote.save_info(self.path_info) + + def save(self): + self.info = self.remote.save_info(self.path_info) + + def dumpd(self): + ret = self.info + ret[self.PARAM_PATH] = self.path + return ret diff --git a/dvc/output/__init__.py b/dvc/output/__init__.py index 5dae38709..89f6aea91 100644 --- a/dvc/output/__init__.py +++ b/dvc/output/__init__.py @@ -8,13 +8,15 @@ from dvc.dependency.base import DependencyBase from dvc.output.s3 import OutputS3 from dvc.output.gs import OutputGS from dvc.output.local import OutputLOCAL +from dvc.output.hdfs import OutputHDFS from dvc.remote import Remote -OUTS = [OutputS3, OutputGS, OutputLOCAL] +OUTS = [OutputHDFS, OutputS3, OutputGS, OutputLOCAL] -OUTS_MAP = {'s3': OutputS3, +OUTS_MAP = {'hdfs': OutputHDFS, + 's3': OutputS3, 'gs': OutputGS, '': OutputLOCAL} diff --git a/dvc/output/hdfs.py b/dvc/output/hdfs.py new file mode 100644 index 000000000..26620e63e --- /dev/null +++ b/dvc/output/hdfs.py @@ -0,0 +1,43 @@ +from dvc.dependency.hdfs import DependencyHDFS +from dvc.exceptions import DvcException + + +class OutputHDFS(DependencyHDFS): + PARAM_CACHE = 'cache' + + def __init__(self, stage, path, info=None, remote=None, cache=True): + super(OutputHDFS, self).__init__(stage, path, info, remote=remote) + self.use_cache = cache + if cache and self.project.cache.hdfs == None: + raise DvcException("No cache location setup for \'hdfs\' outputs.") + + def dumpd(self): + ret = super(OutputHDFS, self).dumpd() + ret[self.PARAM_CACHE] = self.use_cache + return ret + + def changed(self): + if super(OutputHDFS, self).changed(): + return True + + if self.use_cache and self.info != self.project.cache.hdfs.save_info(self.path_info): + return True + + return False + + def checkout(self): + if not self.use_cache: + return + + self.project.cache.hdfs.checkout(self.path_info, self.info) + + def save(self): + super(OutputHDFS, self).save() + + if not self.use_cache: + return + + self.info = self.project.cache.hdfs.save(self.path_info) + + def remove(self): + self.remote.remove(self.path_info) diff --git a/dvc/remote/__init__.py b/dvc/remote/__init__.py index 10d2c43ed..40a44b5d3 100644 --- a/dvc/remote/__init__.py +++ b/dvc/remote/__init__.py @@ -1,13 +1,26 @@ from dvc.remote.local import RemoteLOCAL 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.config import Config from dvc.exceptions import DvcException +REMOTES = [RemoteHDFS, RemoteSSH, RemoteS3, RemoteGS, RemoteLOCAL] + + +def supported_url(url): + config = {Config.SECTION_REMOTE_URL: url} + for r in REMOTES: + if r.supported(config): + return True + return False + + def Remote(project, config): - for r in [RemoteLOCAL, RemoteS3, RemoteGS]: + for r in REMOTES: if r.supported(config): return r(project, config) raise DvcException('Remote \'{}\' is not supported.'.format(config)) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py new file mode 100644 index 000000000..3f31b41cf --- /dev/null +++ b/dvc/remote/hdfs.py @@ -0,0 +1,103 @@ +import os +import re +import posixpath +from subprocess import Popen, PIPE + +try: + from urlparse import urlparse +except ImportError: + from urllib.parse import urlparse + +from dvc.config import Config +from dvc.remote.base import RemoteBase +from dvc.exceptions import DvcException + + +class RemoteHDFS(RemoteBase): + scheme='hdfs' + REGEX = r'^hdfs://(?P.*)$' + PARAM_CHECKSUM = 'checksum' + + def __init__(self, project, config): + self.project = project + self.url = config.get(Config.SECTION_REMOTE_URL, '/') + self.user = config.get(Config.SECTION_REMOTE_USER, None) + + def hadoop_fs(self, cmd): + cmd = 'hadoop fs -' + cmd + env = os.environ.copy() + if self.user: + env['HADOOP_USER_NAME'] = self.user + p = Popen(cmd, + shell=True, + close_fds=True, + env=env, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + out, err = p.communicate() + if p.returncode != 0: + raise DvcException('HDFS command failed: {}: {}'.format(cmd, err)) + return out + + @staticmethod + def group(regex, s, gname): + match = re.match(regex, s) + assert match != None + return match.group(gname) + + def checksum(self, url): + regex = r'.*\t.*\t(?P.*)' + stdout = self.hadoop_fs('checksum {}'.format(url)) + return self.group(regex, stdout, 'checksum') + + def cp(self, from_url, to_url): + self.hadoop_fs('mkdir -p {}'.format(posixpath.dirname(to_url))) + self.hadoop_fs('cp {} {}'.format(from_url, to_url)) + + def rm(self, url): + self.hadoop_fs('rm {}'.format(url)) + + def save_info(self, path_info): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + assert path_info.get('url') + + return {self.PARAM_CHECKSUM: self.checksum(path_info['url'])} + + def save(self, path_info): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + assert path_info.get('url') + + checksum = self.checksum(path_info['url']) + dest_url = posixpath.join(self.url, checksum[0:2], checksum[2:]) + + self.cp(path_info['url'], dest_url) + + return {self.PARAM_CHECKSUM: checksum} + + def checkout(self, path_info, checksum_info): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + assert path_info.get('url') + + checksum = checksum_info.get(self.PARAM_CHECKSUM, None) + if not checksum: + return + + url = posixpath.join(self.url, checksum[0:2], checksum[2:]) + + self.cp(url, path_info['url']) + + def remove(self, path_info): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + assert path_info.get('url') + + self.rm(path_info['url']) diff --git a/tests/test_repro.py b/tests/test_repro.py index 3ebf1556c..cfc823dca 100644 --- a/tests/test_repro.py +++ b/tests/test_repro.py @@ -3,6 +3,7 @@ import yaml import stat import shutil import filecmp +from subprocess import Popen, PIPE import boto3 import uuid @@ -310,6 +311,46 @@ class TestReproExternalGS(TestReproExternalBase): bucket.blob(key).upload_from_string(body) +class TestReproExternalHDFS(TestReproExternalBase): + def should_test(self): + if os.name == 'nt': + return False + if os.getenv('HADOOP_CONTAINER_IP'): + return True + return False + + @property + def scheme(self): + return 'hdfs' + + @property + def bucket(self): + return os.getenv('HADOOP_CONTAINER_IP') + + @property + def cmd(self): + return 'HADOOP_USER_NAME=root hadoop fs -cp' + + def write(self, bucket, key, body): + url = self.scheme + '://' + bucket + '/' + key + p = Popen('hadoop fs -rm -f {}'.format(url), + shell=True, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + p.communicate() + + p = Popen('echo "{}" | hadoop fs -put - {}'.format(body, url), + shell=True, + executable=os.getenv('SHELL'), + stdin=PIPE, + stdout=PIPE, + stderr=PIPE) + p.communicate() + self.assertEqual(p.returncode, 0) + + class TestReproShell(TestDvc): def test(self): if os.name == 'nt':