Merge pull request #2836 from Suor/lazy-cache-paths

perf: optimize cache listing for local, ssh and hdfs

dvc/remote/hdfs.py

@@ -1,24 +1,20 @@
 from __future__ import unicode_literals
 
+from dvc.utils.compat import FileNotFoundError, urlparse
 import io
 import logging
 import os
 import posixpath
 import re
-from contextlib import closing
-from contextlib import contextmanager
-from subprocess import PIPE
-from subprocess import Popen
+from collections import deque
+from contextlib import closing, contextmanager
+import subprocess
 
-from .base import RemoteBASE
-from .base import RemoteCmdError
+from .base import RemoteBASE, RemoteCmdError
 from .pool import get_connection
 from dvc.config import Config
 from dvc.scheme import Schemes
-from dvc.utils import fix_env
-from dvc.utils import tmp_fname
-from dvc.utils.compat import FileNotFoundError
-from dvc.utils.compat import urlparse
+from dvc.utils import fix_env, tmp_fname
 
 logger = logging.getLogger(__name__)
 
@@ -68,15 +64,15 @@ def hadoop_fs(self, cmd, user=None):
         close_fds = os.name != "nt"
 
         executable = os.getenv("SHELL") if os.name != "nt" else None
-        p = Popen(
+        p = subprocess.Popen(
             cmd,
             shell=True,
             close_fds=close_fds,
             executable=executable,
             env=fix_env(os.environ),
-            stdin=PIPE,
-            stdout=PIPE,
-            stderr=PIPE,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
         )
         out, err = p.communicate()
         if p.returncode != 0:
@@ -156,17 +152,14 @@ def open(self, path_info, mode="r", encoding=None):
 
     def list_cache_paths(self):
         if not self.exists(self.path_info):
-            return []
+            return
 
-        files = []
-        dirs = [self.path_info.path]
+        dirs = deque([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
+                        yield urlparse(entry["name"]).path

dvc/remote/local.py

@@ -83,17 +83,7 @@ def supported(cls, config):
 
     def list_cache_paths(self):
         assert self.path_info is not None
-
-        clist = []
-        for entry in os.listdir(fspath_py35(self.path_info)):
-            subdir = self.path_info / entry
-            if not os.path.isdir(fspath_py35(subdir)):
-                continue
-            clist.extend(
-                subdir / cache for cache in os.listdir(fspath_py35(subdir))
-            )
-
-        return clist
+        return walk_files(self.path_info, None)
 
     def get(self, md5):
         if not md5:

dvc/remote/ssh/__init__.py

@@ -264,7 +264,9 @@ def open(self, path_info, mode="r", encoding=None):
 
     def list_cache_paths(self):
         with self.ssh(self.path_info) as ssh:
-            return list(ssh.walk_files(self.path_info.path))
+            # If we simply return an iterator then with above closes instantly
+            for path in ssh.walk_files(self.path_info.path):
+                yield path
 
     def walk_files(self, path_info):
         with self.ssh(path_info) as ssh: