dvc: store cached dir structure in a single json file

Fixes #507

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/cloud/aws.py

@@ -166,13 +166,6 @@ class DataCloudAWS(DataCloudBase):
         key_name = self.cache_file_key(path)
         return self.bucket.get_key(key_name)
 
-    def _get_keys(self, path):
-        key_name = self.cache_file_key(path)
-        keys = self.bucket.list(prefix=key_name)
-        if not keys:
-            return None
-        return list(keys)
-
     def _new_key(self, path):
         key_name = self.cache_file_key(path)
         return self.bucket.new_key(key_name)

dvc/cloud/base.py

@@ -1,5 +1,4 @@
 import os
-import yaml
 import tempfile
 from checksumdir import dirhash
 
@@ -93,25 +92,33 @@ class DataCloudBase(object):
     def _push_key(self, key, path):
         pass
 
+    def collect(self, arg):
+        path, local = arg
+        ret = [path]
+
+        if not Output.is_dir_cache(path):
+            return ret
+
+        if local:
+            if not os.path.isfile(path):
+                return ret
+            dir_path = path
+        else:
+            key = self._get_key(path)
+            if not key:
+                Logger.debug("File '{}' does not exist in the cloud".format(path))
+                return ret
+            tmp = os.path.join(tempfile.mkdtemp(), os.path.basename(path))
+            self._pull_key(key, tmp)
+            dir_path = tmp
+
+        for relpath, md5 in Output.get_dir_cache(dir_path).items():
+            cache = self._cloud_settings.cache.get(md5)
+            ret.append(cache)
+
+        return ret
+
     def push(self, path):
-        """ Generic method for pushing data """
-        if os.path.isfile(path):
-            return self._push(path)
-
-        res = []
-        for root, dirs, files in os.walk(path):
-            for fname in files:
-                p = os.path.join(root, fname)
-                res.append(self._push(p))
-
-                with open(p, 'r') as fd:
-                    d = yaml.safe_load(fd)
-                md5 = d[Output.PARAM_MD5]
-                cache = self._cloud_settings.cache.get(md5)
-                res.append(self._push(cache))
-        return res
-
-    def _push(self, path):
         key = self._get_key(path)
         if key:
             Logger.debug("File '{}' already uploaded to the cloud. Validating checksum...".format(path))
@@ -121,7 +128,7 @@ class DataCloudBase(object):
             Logger.debug('Checksum mismatch. Reuploading is required.')
 
         key = self._new_key(path)
-        return [self._push_key(key, path)]
+        return self._push_key(key, path)
 
     def _makedirs(self, fname):
         dname = os.path.dirname(fname)
@@ -139,38 +146,14 @@ class DataCloudBase(object):
         """ Cloud-specific method of getting keys """
         pass
 
-    def _get_keys(self, path):
-        """ Cloud-specific method of getting keys """
-        pass
-
     def pull(self, path):
         """ Generic method for pulling data from the cloud """
         key = self._get_key(path)
-        if key:
-            return [self._pull_key(key, path)]
-
-        keys = self._get_keys(path)
-        if not keys:
+        if not key:
             Logger.error("File '{}' does not exist in the cloud".format(path))
-            return []
+            return None
 
-        res = []
-        for k in keys:
-            fname = os.path.join(path, os.path.relpath(k.name, self.cache_file_key(path)))
-            res.append(self._pull_key(k, fname))
-            with open(fname, 'r') as fd:
-                d = yaml.safe_load(fd)
-            md5 = d[Output.PARAM_MD5]
-
-            cache = self._cloud_settings.cache.get(md5)
-            cache_key = self._get_key(cache)
-            if not cache_key:
-                Logger.error("File '{}' does not exist in the cloud".format(path))
-                continue
-
-            res.append(self._pull_key(cache_key, cache))
-
-        return res
+        return self._pull_key(key, path)
 
     def _status(self, key, path):
         remote_exists = key != None
@@ -187,29 +170,10 @@ class DataCloudBase(object):
         Generic method for checking data item status.
         """
         key = self._get_key(path)
-        if key:
-            return self._status(key, path)
-
-        keys = self._get_keys(path)
-        if not keys or len(keys) == 0:
+        if not key:
             return STATUS_NEW
 
-        res = []
-        for k in keys:
-            fname = os.path.join(path, os.path.relpath(k.name, self.cache_file_key(path)))
-            res.append(self._status(k, fname))
-
-            tmp = os.path.join(tempfile.mkdtemp(), k.name)
-            self._pull_key(k, tmp)
-            with open(tmp, 'r') as fd:
-                d = yaml.safe_load(fd)
-            md5 = d[Output.PARAM_MD5]
-
-            cache = self._cloud_settings.cache.get(md5)
-            cache_key = self._get_key(cache)
-            res.append(self._status(cache_key, cache))
-
-        return max(res)
+        return self._status(key, path)
 
     def connect(self):
         pass

dvc/cloud/data_cloud.py

@@ -1,3 +1,5 @@
+from multiprocessing.pool import ThreadPool
+
 from dvc.cloud.instance_manager import CloudSettings
 from dvc.logger import Logger
 from dvc.exceptions import DvcException
@@ -79,27 +81,45 @@ class DataCloud(object):
 
         self._cloud.sanity_check()
 
-    def _map_targets(self, func, targets, jobs):
+    def _collect(self, targets, jobs, local):
+        collected = set()
+        pool = ThreadPool(processes=jobs)
+        args = zip(targets, [local]*len(targets))
+        ret = pool.map(self._cloud.collect, args)
+
+        for r in ret:
+            collected |= set(r)
+
+        return collected
+
+    def _map_targets(self, func, targets, jobs, collect_local=False, collect_cloud=False):
         """
         Process targets as data items in parallel.
         """
         self._cloud.connect()
-        return map_progress(func, targets, jobs)
+
+        collected = set()
+        if collect_local:
+            collected |= self._collect(targets, jobs, True)
+        if collect_cloud:
+            collected |= self._collect(targets, jobs, False)
+
+        return map_progress(func, list(collected), jobs)
 
     def push(self, targets, jobs=1):
         """
         Push data items in a cloud-agnostic way.
         """
-        return self._map_targets(self._cloud.push, targets, jobs)
+        return self._map_targets(self._cloud.push, targets, jobs, collect_local=True)
 
     def pull(self, targets, jobs=1):
         """
         Pull data items in a cloud-agnostic way.
         """
-        return self._map_targets(self._cloud.pull, targets, jobs)
+        return self._map_targets(self._cloud.pull, targets, jobs, collect_cloud=True)
 
     def status(self, targets, jobs=1):
         """
         Check status of data items in a cloud-agnostic way.
         """
-        return self._map_targets(self._cloud.status, targets, jobs)
+        return self._map_targets(self._cloud.status, targets, jobs, True, True)

dvc/cloud/gcp.py

@@ -79,13 +79,6 @@ class DataCloudGCP(DataCloudBase):
         key_name = self.cache_file_key(path)
         return self.bucket.get_blob(key_name)
 
-    def _get_keys(self, path):
-        key_name = self.cache_file_key(path)
-        keys = self.bucket.list_blobs(prefix=key_name)
-        if not keys:
-            return None
-        return list(keys)
-
     def _new_key(self, path):
         key_name = self.cache_file_key(path)
         return self.bucket.blob(key_name)

dvc/cloud/local.py

@@ -36,21 +36,6 @@ class DataCloudLOCAL(DataCloudBase):
             return key
         return None
 
-    def _get_keys(self, path):
-        key_name = self.cache_file_key(path)
-        key = LocalKey(self.storage_path, key_name)
-
-        if not os.path.isdir(key.path):
-            return None
-
-        res = []
-        for root, dirs, files in os.walk(key.path):
-            for fname in files:
-                p = os.path.join(root, fname)
-                k_name = os.path.relpath(p, self.storage_path)
-                res.append(LocalKey(self.storage_path, k_name))
-        return res
-
     def _new_key(self, path):
         key_name = self.cache_file_key(path)
         key = LocalKey(self.storage_path, key_name)

dvc/output.py

@@ -6,6 +6,7 @@ from checksumdir import dirhash
 from dvc.system import System
 from dvc.utils import file_md5
 from dvc.exceptions import DvcException
+from dvc.logger import Logger
 
 
 class OutputError(DvcException):
@@ -43,8 +44,10 @@ class CmdOutputAlreadyTrackedError(CmdOutputError):
 
 
 class Dependency(object):
+    PARAM_RELPATH = 'relpath'
     PARAM_PATH = 'path'
     PARAM_MD5 = 'md5'
+    MD5_DIR_SUFFIX = '.dir'
 
     def __init__(self, project, path, md5=None):
         self.project = project
@@ -78,9 +81,13 @@ class Dependency(object):
     def changed(self):
         return self._changed_md5()
 
+    @staticmethod
+    def is_dir_cache(cache):
+        return cache.endswith(Output.MD5_DIR_SUFFIX)
+
     def compute_md5(self):
         if os.path.isdir(self.path):
-            return dirhash(self.path, hashfunc='md5')
+            return dirhash(self.path, hashfunc='md5') + self.MD5_DIR_SUFFIX
         else:
             return file_md5(self.path)[0]
 
@@ -185,24 +192,14 @@ class Output(Dependency):
         return True
 
     def _changed_dir(self):
-        if not os.path.isdir(self.path) or not os.path.isdir(self.cache):
+        if not os.path.isdir(self.path) or not os.path.isfile(self.cache):
             return True
 
-        for root, dirs, files in os.walk(self.path):
-            for fname in files:
-                path = os.path.join(root, fname)
-                mtime = os.path.getmtime(path)
-                inode = os.stat(path).st_ino
+        dir_info = self._collect_dir()
+        dir_info_cached = self.load_dir_cache(self.cache)
 
-                state = self.project.state.get(path)
-                if state and state.mtime == mtime and state.inode == inode:
-                    md5 = state.md5
-                else:
-                    md5 = file_md5(path)[0]
-
-                cache = self.project.cache.get(md5)
-                if self._changed_file(path, cache):
-                    return True
+        if dir_info != dir_info_cached:
+            return True
 
         return False
 
@@ -211,12 +208,10 @@ class Output(Dependency):
 
         if not self.use_cache:
             ret = super(Output, self).changed()
-        elif os.path.isfile(self.path) and \
-             os.path.isfile(self.cache):
-            ret = self._changed_file(self.path, self.cache)
-        elif os.path.isdir(self.path) and \
-             os.path.isdir(self.cache):
+        elif self.is_dir_cache(self.cache):
             ret = self._changed_dir()
+        else:
+            ret = self._changed_file(self.path, self.cache)
 
         msg = u'Data file or dir \'{}\' with cache \'{}\' '
         if ret:
@@ -232,21 +227,35 @@ class Output(Dependency):
         System.hardlink(src, link)
         os.chmod(src, stat.S_IREAD)
 
+    @staticmethod
+    def load_dir_cache(path):
+        with open(path, 'r') as fd:
+            d = yaml.safe_load(fd)
+
+        if not isinstance(d, list):
+            msg = u'Dir cache file format error \'{}\': skipping the file'
+            Logger.error(msg.format(path))
+            return []
+
+        return d
+
+    @staticmethod
+    def get_dir_cache(path):
+        res = {}
+        d = Output.load_dir_cache(path)
+
+        for entry in d:
+            res[entry[Output.PARAM_RELPATH]] = entry[Output.PARAM_MD5]
+
+        return res
+
     def dir_cache(self):
         res = {}
-        for root, dirs, files in os.walk(self.cache):
-            for fname in files:
-                path = os.path.join(root, fname)
-                relpath = os.path.relpath(path, self.cache)
-                with open(path, 'r') as fd:
-                    d = yaml.safe_load(fd)
+        dir_cache = self.get_dir_cache(self.cache)
+
+        for relpath, md5 in dir_cache.items():
+            res[relpath] = self.project.cache.get(md5)
 
-                    if not isinstance(d, dict):
-                        msg = u'Dir cache file format error \'{}\': skipping the file'
-                        self.project.logger.error(msg.format(relpath))
-                    else:
-                        md5 = d[Output.PARAM_MD5]
-                        res[relpath] = self.project.cache.get(md5)
         return res
 
     def checkout(self):
@@ -284,6 +293,42 @@ class Output(Dependency):
 
             self.hardlink(cache, path)
 
+    def _collect_dir(self):
+        dir_info = []
+
+        for root, dirs, files in os.walk(self.path):
+            for fname in files:
+                path = os.path.join(root, fname)
+                relpath = os.path.relpath(path, self.path)
+
+                state = self.project.state.get(path)
+                if state and state.mtime == mtime and state.inode == inode:
+                    md5 = state.md5
+                else:
+                    md5 = file_md5(path)[0]
+
+                dir_info.append({self.PARAM_RELPATH: relpath, self.PARAM_MD5: md5})
+
+        return dir_info
+
+    def _save_dir(self):
+        dir_info = self._collect_dir()
+
+        for entry in dir_info:
+            md5 = entry[self.PARAM_MD5]
+            relpath = entry[self.PARAM_RELPATH]
+            path = os.path.join(self.path, relpath)
+            cache = self.project.cache.get(md5)
+
+            if os.path.exists(cache):
+                self._remove(path, None)
+                self.hardlink(cache, path)
+            else:
+                self.hardlink(path, cache)
+
+        with open(self.cache, 'w+') as fd:
+            yaml.safe_dump(dir_info, fd, default_flow_style=False)
+
     def save(self):
         super(Output, self).save()
 
@@ -307,30 +352,10 @@ class Output(Dependency):
             self.checkout()
             return
 
-        if os.path.isfile(self.path):
+        if os.path.isdir(self.path):
+            self._save_dir()
+        else:
             self.hardlink(self.path, self.cache)
-            return
-
-        for root, dirs, files in os.walk(self.path):
-            for fname in files:
-                path = os.path.join(root, fname)
-                relpath = os.path.relpath(path, self.path)
-                md5 = file_md5(path)[0]
-                cache = self.project.cache.get(md5)
-                cache_info = os.path.join(self.cache, relpath)
-                cache_dir = os.path.dirname(cache_info)
-
-                if os.path.exists(cache):
-                    self._remove(path, None)
-                    self.hardlink(cache, path)
-                else:
-                    self.hardlink(path, cache)
-
-                if not os.path.exists(cache_dir):
-                    os.makedirs(cache_dir)
-
-                with open(cache_info, 'w') as fd:
-                    yaml.safe_dump({self.PARAM_MD5: md5}, fd, default_flow_style=False)
 
     def _remove(self, path, cache):
         self.project.logger.debug(u'Removing \'{}\''.format(path))

dvc/project.py

@@ -159,22 +159,18 @@ class Project(object):
         for stage in self.stages():
             stage.checkout()
 
-    def _used_cache(self, collect_dir=False):
+    def _used_cache(self):
         clist = []
         for stage in self.stages():
             for out in stage.outs:
                 if not out.use_cache:
                     continue
-                if not os.path.isdir(out.cache) and \
-                   out.cache not in clist:
+                if out.cache not in clist:
                     clist.append(out.cache)
-                if os.path.isdir(out.cache):
+                if out.is_dir_cache(out.cache) and os.path.isfile(out.cache):
                     dir_cache = out.dir_cache()
                     clist.extend(dir_cache.values())
-                    if collect_dir:
-                        clist.extend([os.path.join(out.cache, f) for f in dir_cache.keys()])
-                    else:
-                        clist.append(out.cache)
+                    clist.append(out.cache)
         return clist
 
     def _remove_cache_file(self, cache):
@@ -196,7 +192,7 @@ class Project(object):
         os.rmdir(cache)
 
     def gc(self):
-        clist = self._used_cache(collect_dir=False)
+        clist = self._used_cache()
         for cache in self.cache.all():
             if cache in clist:
                 continue
@@ -204,14 +200,15 @@ class Project(object):
             self.logger.info(u'\'{}\' was removed'.format(self.to_dvc_path(cache)))
 
     def push(self, jobs=1):
-        self.cloud.push(self._used_cache(collect_dir=True), jobs)
+        return self.cloud.push(self._used_cache(), jobs)
 
     def pull(self, jobs=1):
-        self.cloud.pull(self._used_cache(collect_dir=False), jobs)
+        ret = self.cloud.pull(self._used_cache(), jobs)
         self.checkout()
+        return ret
 
     def status(self, jobs=1):
-        return self.cloud.status(self._used_cache(collect_dir=True), jobs)
+        return self.cloud.status(self._used_cache(), jobs)
 
     def graph(self):
         G = nx.DiGraph()

tests/test_data_cloud.py

@@ -54,8 +54,7 @@ class TestDataCloudBase(TestDvc):
         self.assertTrue(os.path.isfile(cache))
 
         self.cloud.push(cache_dir)
-        self.assertTrue(os.path.exists(cache_dir))
-        self.assertTrue(os.path.isdir(cache_dir))
+        self.assertTrue(os.path.isfile(cache_dir))
 
         status = self.cloud.status(cache)
         self.assertEqual(status, STATUS_OK)
@@ -88,8 +87,7 @@ class TestDataCloudBase(TestDvc):
             self.assertEqual(fd.read(), self.FOO_CONTENTS)
 
         self.cloud.pull(cache_dir)
-        self.assertTrue(os.path.exists(cache_dir))
-        self.assertTrue(os.path.isdir(cache_dir))
+        self.assertTrue(os.path.isfile(cache_dir))
 
         status = self.cloud.status(cache)
         self.assertEqual(status, STATUS_OK)
@@ -211,8 +209,7 @@ class TestDataCloudLocalCli(TestDvc):
         self.main(['cache', 'push'])
         self.assertTrue(os.path.exists(cache))
         self.assertTrue(os.path.isfile(cache))
-        self.assertTrue(os.path.exists(cache_dir))
-        self.assertTrue(os.path.isdir(cache_dir))
+        self.assertTrue(os.path.isfile(cache_dir))
 
         self.main(['cache', 'status'])
 
@@ -227,7 +224,6 @@ class TestDataCloudLocalCli(TestDvc):
         self.assertTrue(os.path.isfile(cache))
         with open(cache, 'r') as fd:
             self.assertEqual(fd.read(), self.FOO_CONTENTS)
-        self.assertTrue(os.path.exists(cache_dir))
-        self.assertTrue(os.path.isdir(cache_dir))
+        self.assertTrue(os.path.isfile(cache_dir))
 
         self.main(['cache', 'status'])