diff --git a/dvc/cloud/aws.py b/dvc/cloud/aws.py index 70beddf74..37b0337dd 100644 --- a/dvc/cloud/aws.py +++ b/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) diff --git a/dvc/cloud/base.py b/dvc/cloud/base.py index 7ea3400e5..e280b1e8d 100644 --- a/dvc/cloud/base.py +++ b/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 diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index 5ce068833..f5cc84711 100644 --- a/dvc/cloud/data_cloud.py +++ b/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) diff --git a/dvc/cloud/gcp.py b/dvc/cloud/gcp.py index c6595f82f..2c54f469b 100644 --- a/dvc/cloud/gcp.py +++ b/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) diff --git a/dvc/cloud/local.py b/dvc/cloud/local.py index b100784f2..714611778 100644 --- a/dvc/cloud/local.py +++ b/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) diff --git a/dvc/output.py b/dvc/output.py index dd6f79eb6..0f574baba 100644 --- a/dvc/output.py +++ b/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)) diff --git a/dvc/project.py b/dvc/project.py index 99363beeb..f7522ecc9 100644 --- a/dvc/project.py +++ b/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() diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 5021768e4..510080a62 100644 --- a/tests/test_data_cloud.py +++ b/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'])