Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-09 23:10:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #757 from efiop/master

Browse Source 
dvc: initial hdfs support
This commit is contained in:
Ruslan Kuprieiev
2018-06-10 05:36:46 +03:00
committed by GitHub
parent 77b814f039 ff0d91737b
commit 35f5945720
9 changed files with 251 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/cache.py
+1
View File
@@ -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)
dvc/config.py
+6 -2
View File
@@ -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<name>.*)"\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 = {
dvc/dependency/__init__.py
+6 -2
View File
@@ -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),
}
dvc/dependency/hdfs.py
+33
View File
@@ -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
dvc/output/__init__.py
+4 -2
View File
@@ -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}
dvc/output/hdfs.py
+43
View File
@@ -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)
dvc/remote/__init__.py
+14 -1
View File
@@ -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))
dvc/remote/hdfs.py
+103
View File
@@ -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<path>.*)$'
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<checksum>.*)'
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'])
tests/test_repro.py
+41
View File
@@ -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':