From 68bebe79441c078641e85a1db7f14c16521fe294 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Fri, 19 Jul 2019 04:04:38 +0300 Subject: [PATCH] remote: hdfs: use pyarrow and connection pool Related to #1629 Speeds up HDFS tests significantly. E.g. ExternalHDFS test goes from 600 sec to 190sec. test_open_external[HDFS] from 160sec to 3sec. Signed-off-by: Ruslan Kuprieiev --- dvc/main.py | 4 +- dvc/path_info.py | 2 +- dvc/remote/hdfs.py | 129 ++++++++++++++++------------- dvc/remote/{ssh => }/pool.py | 23 +++-- dvc/remote/ssh/__init__.py | 7 +- setup.py | 6 ++ tests/func/test_output.py | 12 ++- tests/unit/dependency/test_hdfs.py | 8 ++ tests/unit/output/test_hdfs.py | 8 ++ tests/unit/remote/ssh/test_pool.py | 7 +- 10 files changed, 130 insertions(+), 76 deletions(-) rename dvc/remote/{ssh => }/pool.py (64%) diff --git a/dvc/main.py b/dvc/main.py index d6bc5f5d0..71687c477 100644 --- a/dvc/main.py +++ b/dvc/main.py @@ -8,7 +8,7 @@ from dvc.cli import parse_args from dvc.config import ConfigError from dvc.analytics import Analytics from dvc.exceptions import NotDvcRepoError, DvcParserError -from dvc.remote.ssh.pool import close_ssh_pools +from dvc.remote.pool import close_pools logger = logging.getLogger("dvc") @@ -56,7 +56,7 @@ def main(argv=None): logger.setLevel(outerLogLevel) # Python 2 fails to close these clean occasionally and users see # weird error messages, so we do it manually - close_ssh_pools() + close_pools() Analytics().send_cmd(cmd, args, ret) diff --git a/dvc/path_info.py b/dvc/path_info.py index 27f21bcce..c905f47fe 100644 --- a/dvc/path_info.py +++ b/dvc/path_info.py @@ -129,7 +129,7 @@ class _URLPathParents(object): class URLInfo(object): - DEFAULT_PORTS = {"http": 80, "https": 443, "ssh": 22} + DEFAULT_PORTS = {"http": 80, "https": 443, "ssh": 22, "hdfs": 0} def __init__(self, url): self.parsed = urlparse(url) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 8c51d6604..ea13a176a 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -7,12 +7,19 @@ import posixpath import logging from subprocess import Popen, PIPE +try: + import pyarrow +except ImportError: + pyarrow = None + from dvc.config import Config from dvc.scheme import Schemes -from dvc.remote.base import RemoteBASE, RemoteCmdError +from dvc.utils.compat import urlparse from dvc.utils import fix_env, tmp_fname +from .pool import get_connection +from .base import RemoteBASE, RemoteCmdError logger = logging.getLogger(__name__) @@ -21,17 +28,39 @@ class RemoteHDFS(RemoteBASE): scheme = Schemes.HDFS REGEX = r"^hdfs://((?P.*)@)?.*$" PARAM_CHECKSUM = "checksum" + REQUIRES = {"pyarrow": pyarrow} def __init__(self, repo, config): super(RemoteHDFS, self).__init__(repo, config) - url = config.get(Config.SECTION_REMOTE_URL, "/") - self.path_info = self.path_cls(url) + self.path_info = None + url = config.get(Config.SECTION_REMOTE_URL) + if not url: + return - self.user = self.path_info.user - if not self.user: - self.user = config.get( - Config.SECTION_REMOTE_USER, getpass.getuser() - ) + parsed = urlparse(url) + + user = ( + parsed.username + or config.get(Config.SECTION_REMOTE_USER) + or getpass.getuser() + ) + + self.path_info = self.path_cls.from_parts( + scheme=self.scheme, + host=parsed.hostname, + user=user, + port=parsed.port, + path=parsed.path, + ) + + @staticmethod + def hdfs(path_info): + return get_connection( + pyarrow.hdfs.connect, + path_info.host, + path_info.port, + user=path_info.user, + ) def hadoop_fs(self, cmd, user=None): cmd = "hadoop fs -" + cmd @@ -65,6 +94,7 @@ class RemoteHDFS(RemoteBASE): return match.group(gname) def get_file_checksum(self, path_info): + # NOTE: pyarrow doesn't support checksum, so we need to use hadoop regex = r".*\t.*\t(?P.*)" stdout = self.hadoop_fs( "checksum {}".format(path_info.path), user=path_info.user @@ -73,68 +103,55 @@ class RemoteHDFS(RemoteBASE): def copy(self, from_info, to_info, **_kwargs): dname = posixpath.dirname(to_info.path) - self.hadoop_fs("mkdir -p {}".format(dname), user=to_info.user) - self.hadoop_fs( - "cp -f {} {}".format(from_info.path, to_info.path), - user=to_info.user, - ) - - def rm(self, path_info): - self.hadoop_fs("rm -f {}".format(path_info.path), user=path_info.user) + with self.hdfs(to_info) as hdfs: + hdfs.mkdir(dname) + # NOTE: this is how `hadoop fs -cp` works too: it copies through + # your local machine. + with hdfs.open(from_info.path, "rb") as from_fobj: + with hdfs.open(to_info.path, "wb") as to_fobj: + to_fobj.upload(from_fobj) def remove(self, path_info): if path_info.scheme != "hdfs": raise NotImplementedError - assert path_info.path - - logger.debug("Removing {}".format(path_info.path)) - - self.rm(path_info) + if self.exists(path_info): + logger.debug("Removing {}".format(path_info.path)) + with self.hdfs(path_info) as hdfs: + hdfs.rm(path_info.path) def exists(self, path_info): assert not isinstance(path_info, list) assert path_info.scheme == "hdfs" - - try: - self.hadoop_fs("test -e {}".format(path_info.path)) - return True - except RemoteCmdError: - return False + with self.hdfs(path_info) as hdfs: + return hdfs.exists(path_info.path) def _upload(self, from_file, to_info, **_kwargs): - self.hadoop_fs( - "mkdir -p {}".format(to_info.parent.url), user=to_info.user - ) - - tmp_file = tmp_fname(to_info.url) - - self.hadoop_fs( - "copyFromLocal {} {}".format(from_file, tmp_file), - user=to_info.user, - ) - - self.hadoop_fs( - "mv {} {}".format(tmp_file, to_info.url), user=to_info.user - ) + with self.hdfs(to_info) as hdfs: + hdfs.mkdir(posixpath.dirname(to_info.path)) + tmp_file = tmp_fname(to_info.path) + with open(from_file, "rb") as fobj: + hdfs.upload(tmp_file, fobj) + hdfs.rename(tmp_file, to_info.path) def _download(self, from_info, to_file, **_kwargs): - self.hadoop_fs( - "copyToLocal {} {}".format(from_info.url, to_file), - user=from_info.user, - ) + with self.hdfs(from_info) as hdfs: + with open(to_file, "wb+") as fobj: + hdfs.download(from_info.path, fobj) def list_cache_paths(self): - try: - self.hadoop_fs("test -e {}".format(self.path_info.url)) - except RemoteCmdError: + if not self.exists(self.path_info): return [] - stdout = self.hadoop_fs("ls -R {}".format(self.path_info.url)) - lines = stdout.split("\n") - flist = [] - for line in lines: - if not line.startswith("-"): - continue - flist.append(line.split()[-1]) - return flist + files = [] + dirs = [self.path_info.path] + + with self.hdfs(self.path_info) as hdfs: + while dirs: + for entry in hdfs.ls(dirs.pop(), detail=True): + if entry["kind"] == "directory": + dirs.append(urlparse(entry["name"]).path) + elif entry["kind"] == "file": + files.append(urlparse(entry["name"]).path) + + return files diff --git a/dvc/remote/ssh/pool.py b/dvc/remote/pool.py similarity index 64% rename from dvc/remote/ssh/pool.py rename to dvc/remote/pool.py index a40f00606..e15fdb689 100644 --- a/dvc/remote/ssh/pool.py +++ b/dvc/remote/pool.py @@ -2,12 +2,10 @@ from collections import deque from contextlib import contextmanager from funcy import memoize -from .connection import SSHConnection - @contextmanager -def ssh_connection(*conn_args, **conn_kwargs): - pool = get_ssh_pool(*conn_args, **conn_kwargs) +def get_connection(*args, **kwargs): + pool = get_pool(*args, **kwargs) conn = pool.get_connection() try: yield conn @@ -19,18 +17,19 @@ def ssh_connection(*conn_args, **conn_kwargs): @memoize -def get_ssh_pool(*conn_args, **conn_kwargs): - return SSHPool(conn_args, conn_kwargs) +def get_pool(*args, **kwargs): + return Pool(*args, **kwargs) -def close_ssh_pools(): - for pool in get_ssh_pool.memory.values(): +def close_pools(): + for pool in get_pool.memory.values(): pool.close() - get_ssh_pool.memory.clear() + get_pool.memory.clear() -class SSHPool(object): - def __init__(self, conn_args, conn_kwargs): +class Pool(object): + def __init__(self, conn_func, *conn_args, **conn_kwargs): + self._conn_func = conn_func self._conn_args = conn_args self._conn_kwargs = conn_kwargs self._conns = deque() @@ -48,7 +47,7 @@ class SSHPool(object): try: return self._conns.popleft() except IndexError: - return SSHConnection(*self._conn_args, **self._conn_kwargs) + return self._conn_func(*self._conn_args, **self._conn_kwargs) def release(self, conn): if self._closed: diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index bd25f954c..a50d812e8 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -18,7 +18,9 @@ from dvc.utils import to_chunks from dvc.utils.compat import urlparse, StringIO from dvc.remote.base import RemoteBASE from dvc.scheme import Schemes -from .pool import ssh_connection +from dvc.remote.pool import get_connection + +from .connection import SSHConnection logger = logging.getLogger(__name__) @@ -121,7 +123,8 @@ class RemoteSSH(RemoteBASE): ) ) - return ssh_connection( + return get_connection( + SSHConnection, host, username=user, port=port, diff --git a/setup.py b/setup.py index 28d85847a..ba89aedd7 100644 --- a/setup.py +++ b/setup.py @@ -76,8 +76,13 @@ s3 = ["boto3==1.9.115"] azure = ["azure-storage-blob==2.0.1"] oss = ["oss2==2.6.1"] ssh = ["paramiko>=2.5.0"] +hdfs = ["pyarrow>=0.14.0"] all_remotes = gs + s3 + azure + ssh + oss +if os.name != "nt" or sys.version_info[0] != 2: + # NOTE: there are no pyarrow wheels for python2 on windows + all_remotes += hdfs + # Extra dependecies to run tests tests_requirements = [ "PyInstaller==3.4", @@ -125,6 +130,7 @@ setup( "azure": azure, "oss": oss, "ssh": ssh, + "hdfs": hdfs, # NOTE: https://github.com/inveniosoftware/troubleshooting/issues/1 ":python_version=='2.7'": ["futures", "pathlib2"], "tests": tests_requirements, diff --git a/tests/func/test_output.py b/tests/func/test_output.py index 2be628328..20764ea9f 100644 --- a/tests/func/test_output.py +++ b/tests/func/test_output.py @@ -1,3 +1,6 @@ +import os +import sys + import pytest from dvc.stage import Stage @@ -8,7 +11,14 @@ TESTS = [ ("s3://bucket/path", "s3"), ("gs://bucket/path", "gs"), ("ssh://example.com:/dir/path", "ssh"), - ("hdfs://example.com/dir/path", "hdfs"), + pytest.param( + "hdfs://example.com/dir/path", + "hdfs", + marks=pytest.mark.skipif( + sys.version_info[0] == 2 and os.name == "nt", + reason="Not supported for python 2 on Windows.", + ), + ), ("path/to/file", "local"), ("path\\to\\file", "local"), ("file", "local"), diff --git a/tests/unit/dependency/test_hdfs.py b/tests/unit/dependency/test_hdfs.py index 8362381ba..7a88219a3 100644 --- a/tests/unit/dependency/test_hdfs.py +++ b/tests/unit/dependency/test_hdfs.py @@ -1,8 +1,16 @@ +import os +import sys +import pytest + from dvc.dependency.hdfs import DependencyHDFS from tests.unit.dependency.test_local import TestDependencyLOCAL +@pytest.mark.skipif( + sys.version_info[0] == 2 and os.name == "nt", + reason="Not supported for python 2 on Windows.", +) class TestDependencyHDFS(TestDependencyLOCAL): def _get_cls(self): return DependencyHDFS diff --git a/tests/unit/output/test_hdfs.py b/tests/unit/output/test_hdfs.py index 039711260..b54490ecb 100644 --- a/tests/unit/output/test_hdfs.py +++ b/tests/unit/output/test_hdfs.py @@ -1,8 +1,16 @@ +import os +import sys +import pytest + from dvc.output.hdfs import OutputHDFS from tests.unit.output.test_local import TestOutputLOCAL +@pytest.mark.skipif( + sys.version_info[0] == 2 and os.name == "nt", + reason="Not supported for python 2 on Windows.", +) class TestOutputHDFS(TestOutputLOCAL): def _get_cls(self): return OutputHDFS diff --git a/tests/unit/remote/ssh/test_pool.py b/tests/unit/remote/ssh/test_pool.py index fb9e61a58..c3fb43142 100644 --- a/tests/unit/remote/ssh/test_pool.py +++ b/tests/unit/remote/ssh/test_pool.py @@ -1,11 +1,14 @@ import pytest -from dvc.remote.ssh.pool import ssh_connection +from dvc.remote.pool import get_connection +from dvc.remote.ssh.connection import SSHConnection def test_doesnt_swallow_errors(ssh_server): class MyError(Exception): pass - with pytest.raises(MyError), ssh_connection(**ssh_server.test_creds): + with pytest.raises(MyError), get_connection( + SSHConnection, **ssh_server.test_creds + ): raise MyError