From 59dcdb02eafedf7b7d8ba56e2cccaaaee5277aa1 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Mon, 17 Sep 2018 08:03:21 +0300 Subject: [PATCH] state: optimize operations Signed-off-by: Ruslan Kuprieiev --- dvc/project.py | 179 ++++++++++++++++++++----------------- dvc/remote/local.py | 6 +- dvc/state.py | 185 ++++++++++++++++++--------------------- tests/test_data_cloud.py | 81 ++++++++--------- tests/test_state.py | 33 +++---- 5 files changed, 239 insertions(+), 245 deletions(-) diff --git a/dvc/project.py b/dvc/project.py index 95f464cbe..141031993 100644 --- a/dvc/project.py +++ b/dvc/project.py @@ -148,16 +148,17 @@ class Project(object): fnames = [fname] stages = [] - for f in fnames: - stage = Stage.loads(project=self, - outs=[f], - add=True) + with self.state: + for f in fnames: + stage = Stage.loads(project=self, + outs=[f], + add=True) - self._check_output_duplication(stage.outs) + self._check_output_duplication(stage.outs) - stage.save() - stage.dump() - stages.append(stage) + stage.save() + stage.dump() + stages.append(stage) return stages @@ -184,34 +185,35 @@ class Project(object): [from_path])[0] found = False - for stage in self.stages(): - for out in stage.outs: - if out.path != from_out.path: - continue + with self.state: + for stage in self.stages(): + for out in stage.outs: + if out.path != from_out.path: + continue - if not stage.is_data_source: - msg = 'Dvcfile \'{}\' is not a data source.' - raise DvcException(msg.format(stage.rel_path)) + if not stage.is_data_source: + msg = 'Dvcfile \'{}\' is not a data source.' + raise DvcException(msg.format(stage.rel_path)) - found = True - to_out = Output.loads_from(out.stage, - [to_path], - out.cache, - out.metric)[0] - out.move(to_out) + found = True + to_out = Output.loads_from(out.stage, + [to_path], + out.cache, + out.metric)[0] + out.move(to_out) - stage_base = os.path.basename(stage.path) - stage_base = stage_base.rstrip(Stage.STAGE_FILE_SUFFIX) + stage_base = os.path.basename(stage.path) + stage_base = stage_base.rstrip(Stage.STAGE_FILE_SUFFIX) - stage_dir = os.path.dirname(stage.path) - from_base = os.path.basename(from_path) - to_base = os.path.basename(to_path) - if stage_base == from_base: - os.unlink(stage.path) - path = to_base + Stage.STAGE_FILE_SUFFIX - stage.path = os.path.join(stage_dir, path) + stage_dir = os.path.dirname(stage.path) + from_base = os.path.basename(from_path) + to_base = os.path.basename(to_path) + if stage_base == from_base: + os.unlink(stage.path) + path = to_base + Stage.STAGE_FILE_SUFFIX + stage.path = os.path.join(stage_dir, path) - stage.dump() + stage.dump() if not found: msg = 'Unable to find dvcfile with output \'{}\'' @@ -239,9 +241,12 @@ class Project(object): self._check_output_duplication(stage.outs) - if not no_exec: - stage.run() + with self.state: + if not no_exec: + stage.run() + stage.dump() + return stage def imp(self, url, out): @@ -252,7 +257,9 @@ class Project(object): self._check_output_duplication(stage.outs) - stage.run() + with self.state: + stage.run() + stage.dump() return stage @@ -291,15 +298,22 @@ class Project(object): core = config._config[config.SECTION_CORE] interactive = core.get(config.SECTION_CORE_INTERACTIVE, False) - if recursive: - return self._reproduce_stages(G, - stages, - node, - force, - dry, - interactive) + with self.state: + if recursive: + ret = self._reproduce_stages(G, + stages, + node, + force, + dry, + interactive) + else: + ret = self._reproduce_stage(stages, + node, + force, + dry, + interactive) - return self._reproduce_stage(stages, node, force, dry, interactive) + return ret def _reproduce_stages(self, G, stages, node, force, dry, interactive): import networkx as nx @@ -331,15 +345,16 @@ class Project(object): else: stages = all_stages - self._cleanup_unused_links(all_stages) + with self.state: + self._cleanup_unused_links(all_stages) - for stage in stages: - if stage.locked: - msg = 'DVC file \'{}\' is locked. Its dependecies are not ' \ - 'going to be checked out.' - self.logger.warn(msg.format(stage.relpath)) + for stage in stages: + if stage.locked: + msg = 'DVC file \'{}\' is locked. Its dependecies are ' \ + 'not going to be checked out.' + self.logger.warn(msg.format(stage.relpath)) - stage.checkout() + stage.checkout() def _used_cache(self, target=None, all_branches=False, active=True): cache = {} @@ -380,29 +395,30 @@ class Project(object): self.logger.info("No unused {} cache to remove.".format(typ)) def gc(self, all_branches=False, cloud=False, remote=None): - clist = self._used_cache(target=None, - all_branches=all_branches, - active=False) - self._do_gc('local', self.cache.local.gc, clist) + with self.state: + clist = self._used_cache(target=None, + all_branches=all_branches, + active=False) + self._do_gc('local', self.cache.local.gc, clist) - if self.cache.s3: - self._do_gc('s3', self.cache.s3.gc, clist) + if self.cache.s3: + self._do_gc('s3', self.cache.s3.gc, clist) - if self.cache.gs: - self._do_gc('gs', self.cache.gs.gc, clist) + if self.cache.gs: + self._do_gc('gs', self.cache.gs.gc, clist) - if self.cache.ssh: - self._do_gc('ssh', self.cache.ssh.gc, clist) + if self.cache.ssh: + self._do_gc('ssh', self.cache.ssh.gc, clist) - if self.cache.hdfs: - self._do_gc('hdfs', self.cache.hdfs.gc, clist) + if self.cache.hdfs: + self._do_gc('hdfs', self.cache.hdfs.gc, clist) - if self.cache.azure: - self._do_gc('azure', self.cache.azure.gc, clist) + if self.cache.azure: + self._do_gc('azure', self.cache.azure.gc, clist) - if cloud: - self._do_gc('remote', self.cloud._get_cloud(remote, - 'gc -c').gc, clist) + if cloud: + self._do_gc('remote', self.cloud._get_cloud(remote, + 'gc -c').gc, clist) def push(self, target=None, @@ -410,10 +426,11 @@ class Project(object): remote=None, all_branches=False, show_checksums=False): - self.cloud.push(self._used_cache(target, all_branches)['local'], - jobs, - remote=remote, - show_checksums=show_checksums) + with self.state: + self.cloud.push(self._used_cache(target, all_branches)['local'], + jobs, + remote=remote, + show_checksums=show_checksums) def fetch(self, target=None, @@ -421,10 +438,11 @@ class Project(object): remote=None, all_branches=False, show_checksums=False): - self.cloud.pull(self._used_cache(target, all_branches)['local'], - jobs, - remote=remote, - show_checksums=show_checksums) + with self.state: + self.cloud.pull(self._used_cache(target, all_branches)['local'], + jobs, + remote=remote, + show_checksums=show_checksums) def pull(self, target=None, @@ -487,12 +505,13 @@ class Project(object): cloud=False, remote=None, show_checksums=False): - if cloud: - return self._cloud_status(target, - jobs, - remote=remote, - show_checksums=show_checksums) - return self._local_status(target) + with self.state: + if cloud: + return self._cloud_status(target, + jobs, + remote=remote, + show_checksums=show_checksums) + return self._local_status(target) def _read_metric_json(self, fd, json_path): import json diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 2c8fe7a6c..a30bad54e 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -137,7 +137,6 @@ class RemoteLOCAL(RemoteBase): def collect_dir_cache(self, dname): dir_info = [] - db = self.state.load() bar = False for root, dirs, files in os.walk(dname): if len(files) > LARGE_DIR_SIZE: @@ -158,13 +157,10 @@ class RemoteLOCAL(RemoteBase): progress.update_target(title, processed, total) processed += 1 - md5 = self.state.update(path, use_db=db) + md5 = self.state.update(path) dir_info.append({self.PARAM_RELPATH: relpath, self.PARAM_MD5: md5}) - db.commit() - db.close() - if bar: progress.finish_target(title) diff --git a/dvc/state.py b/dvc/state.py index 4d5d0a83b..5eda5330f 100644 --- a/dvc/state.py +++ b/dvc/state.py @@ -36,15 +36,14 @@ class State(object): return self.state_file = os.path.join(self.dvc_dir, self.STATE_FILE) - # Try loading once to check that the file is indeed a database - # and reformat it if it is not. - try: - db = self.load() - db.close() - except sqlite3.DatabaseError: - os.unlink(self.state_file) - db = self.load() - db.close() + self.db = None + self.c = None + + def __enter__(self): + self.load() + + def __exit__(self, type, value, tb): + self.dump() @staticmethod def init(project): @@ -65,16 +64,40 @@ class State(object): return actual.split('.')[0] != md5.split('.')[0] def load(self): - db = sqlite3.connect(self.state_file) - c = db.cursor() - # Check that the state file is indeed a database - cmd = "CREATE TABLE IF NOT EXISTS {} ({})" - c.execute(cmd.format(self.STATE_TABLE, self.STATE_TABLE_LAYOUT)) - c.execute(cmd.format(self.LINK_STATE_TABLE, - self.LINK_STATE_TABLE_LAYOUT)) - db.commit() - c.close() - return db + retries = 1 + while True: + assert self.db is None + self.db = sqlite3.connect(self.state_file) + self.c = self.db.cursor() + + # Try loading once to check that the file is indeed a database + # and reformat it if it is not. + try: + # Check that the state file is indeed a database + cmd = "CREATE TABLE IF NOT EXISTS {} ({})" + self.c.execute(cmd.format(self.STATE_TABLE, + self.STATE_TABLE_LAYOUT)) + self.c.execute(cmd.format(self.LINK_STATE_TABLE, + self.LINK_STATE_TABLE_LAYOUT)) + return + except sqlite3.DatabaseError: + self.db.close() + self.c.close() + self.db = None + self.c = None + if retries > 0: + os.unlink(self.state_file) + retries -= 1 + else: + raise + + def dump(self): + assert self.db is not None + self.db.commit() + self.c.close() + self.db.close() + self.db = None + self.c = None @staticmethod def mtime(path): @@ -93,82 +116,61 @@ class State(object): def inode(path): return System.inode(path) - def _do_update(self, path, use_db=None): + def _do_update(self, path): if not os.path.exists(path): return (None, None) mtime = self.mtime(path) inode = self.inode(path) - if use_db is None: - db = self.load() - else: - db = use_db + cmd = 'SELECT * from {} WHERE inode={}'.format(self.STATE_TABLE, + inode) - c = db.cursor() - - md5 = self._get(inode, mtime, use_db=db) - if md5 is None: + self.c.execute(cmd) + ret = self.c.fetchall() + if len(ret) == 0: md5, info = self._collect(path) + cmd = 'INSERT INTO {}(inode, mtime, md5, timestamp) ' \ + 'VALUES ({}, "{}", "{}", "{}")' + self.c.execute(cmd.format(self.STATE_TABLE, + inode, + mtime, + md5, + int(nanotime.timestamp(time.time())))) else: - info = None - - timestamp = int(nanotime.timestamp(time.time())) - cmd = 'REPLACE INTO {}(inode, mtime, md5, timestamp) ' \ - 'VALUES ({}, "{}", "{}", "{}")'.format(self.STATE_TABLE, - inode, - mtime, - md5, - timestamp) - - c.execute(cmd) - c.close() - if use_db is None: - db.commit() - db.close() + assert len(ret) == 1 + assert len(ret[0]) == 4 + i, m, md5, timestamp = ret[0] + assert i == inode + if mtime != m: + md5, info = self._collect(path) + cmd = 'UPDATE {} SET ' \ + 'mtime = "{}", md5 = "{}", timestamp = "{}" ' \ + 'WHERE inode = {}' + self.c.execute(cmd.format(self.STATE_TABLE, + mtime, + md5, + int(nanotime.timestamp(time.time())), + inode)) + else: + info = None + cmd = 'UPDATE {} SET timestamp = "{}" WHERE inode = {}' + self.c.execute(cmd.format(self.STATE_TABLE, + int(nanotime.timestamp(time.time())), + inode)) return (md5, info) - def update(self, path, use_db=None): - return self._do_update(path, use_db=use_db)[0] + def update(self, path): + return self._do_update(path)[0] - def update_info(self, path, use_db=None): - md5, info = self._do_update(path, use_db=use_db) + def update_info(self, path): + md5, info = self._do_update(path) if not info: info = self.project.cache.local.load_dir_cache(md5) return (md5, info) - def _get(self, inode, mtime, use_db=None): - cmd = 'SELECT * from {} WHERE inode={}'.format(self.STATE_TABLE, - inode) - - if use_db is None: - db = self.load() - else: - db = use_db - - c = db.cursor() - c.execute(cmd) - ret = c.fetchall() - c.close() - - if use_db is None: - db.commit() - db.close() - - if len(ret) == 0: - return None - assert len(ret) == 1 - assert len(ret[0]) == 4 - i, m, md5, timestamp = ret[0] - assert i == inode - - if mtime == m: - return md5 - - return None - - def update_link(self, path, use_db=None): + def update_link(self, path): if not os.path.exists(path): return @@ -176,32 +178,18 @@ class State(object): inode = self.inode(path) relpath = os.path.relpath(path, self.root_dir) - if use_db is None: - db = self.load() - else: - db = use_db - - c = db.cursor() - cmd = 'REPLACE INTO {}(path, inode, mtime) ' \ 'VALUES ("{}", {}, "{}")'.format(self.LINK_STATE_TABLE, relpath, inode, mtime) - c.execute(cmd) - c.close() - - if use_db is None: - db.commit() - db.close() + self.c.execute(cmd) def remove_unused_links(self, used): unused = [] - db = self.load() - c = db.cursor() - c.execute('SELECT * FROM {}'.format(self.LINK_STATE_TABLE)) - for row in c: + self.c.execute('SELECT * FROM {}'.format(self.LINK_STATE_TABLE)) + for row in self.c: p, i, m = row path = os.path.join(self.root_dir, p) @@ -219,11 +207,6 @@ class State(object): remove(path) unused.append(p) - db.commit() for p in unused: cmd = 'DELETE FROM {} WHERE path = "{}"' - c.execute(cmd.format(self.LINK_STATE_TABLE, p)) - - c.close() - db.commit() - db.close() + self.c.execute(cmd.format(self.LINK_STATE_TABLE, p)) diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index 134273983..c62bce045 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -221,55 +221,56 @@ class TestDataCloudBase(TestDvc): info_dir = stage_dir.outs[0].dumpd() md5_dir = info_dir['md5'] - # Check status - sleep() - status = self.cloud.status([info], show_checksums=True) - self.assertEqual([(md5, STATUS_NEW)], status) + with self.cloud.project.state: + # Check status + sleep() + status = self.cloud.status([info], show_checksums=True) + self.assertEqual([(md5, STATUS_NEW)], status) - status_dir = self.cloud.status([info_dir], show_checksums=True) - self.assertTrue((md5_dir, STATUS_NEW) in status_dir) + status_dir = self.cloud.status([info_dir], show_checksums=True) + self.assertTrue((md5_dir, STATUS_NEW) in status_dir) - # Push and check status - self.cloud.push([info]) - self.assertTrue(os.path.exists(cache)) - self.assertTrue(os.path.isfile(cache)) + # Push and check status + self.cloud.push([info]) + self.assertTrue(os.path.exists(cache)) + self.assertTrue(os.path.isfile(cache)) - self.cloud.push([info_dir]) - self.assertTrue(os.path.isfile(cache_dir)) + self.cloud.push([info_dir]) + self.assertTrue(os.path.isfile(cache_dir)) - sleep() - status = self.cloud.status([info], show_checksums=True) - self.assertEqual([(md5, STATUS_OK)], status) + sleep() + status = self.cloud.status([info], show_checksums=True) + self.assertEqual([(md5, STATUS_OK)], status) - status_dir = self.cloud.status([info_dir], show_checksums=True) - self.assertTrue((md5_dir, STATUS_OK) in status_dir) + status_dir = self.cloud.status([info_dir], show_checksums=True) + self.assertTrue((md5_dir, STATUS_OK) in status_dir) - # Remove and check status - sleep() - shutil.rmtree(self.dvc.cache.local.cache_dir) + # Remove and check status + sleep() + shutil.rmtree(self.dvc.cache.local.cache_dir) - status = self.cloud.status([info], show_checksums=True) - self.assertEqual([(md5, STATUS_DELETED)], status) + status = self.cloud.status([info], show_checksums=True) + self.assertEqual([(md5, STATUS_DELETED)], status) - status_dir = self.cloud.status([info_dir], show_checksums=True) - self.assertEqual([(md5_dir, STATUS_DELETED)], status_dir) + status_dir = self.cloud.status([info_dir], show_checksums=True) + self.assertEqual([(md5_dir, STATUS_DELETED)], status_dir) - # Pull and check status - self.cloud.pull([info]) - self.assertTrue(os.path.exists(cache)) - self.assertTrue(os.path.isfile(cache)) - with open(cache, 'r') as fd: - self.assertEqual(fd.read(), self.FOO_CONTENTS) + # Pull and check status + self.cloud.pull([info]) + self.assertTrue(os.path.exists(cache)) + self.assertTrue(os.path.isfile(cache)) + with open(cache, 'r') as fd: + self.assertEqual(fd.read(), self.FOO_CONTENTS) - self.cloud.pull([info_dir]) - self.assertTrue(os.path.isfile(cache_dir)) + self.cloud.pull([info_dir]) + self.assertTrue(os.path.isfile(cache_dir)) - sleep() - status = self.cloud.status([info], show_checksums=True) - self.assertEqual([(md5, STATUS_OK)], status) + sleep() + status = self.cloud.status([info], show_checksums=True) + self.assertEqual([(md5, STATUS_OK)], status) - status_dir = self.cloud.status([info_dir], show_checksums=True) - self.assertTrue((md5_dir, STATUS_OK) in status_dir) + status_dir = self.cloud.status([info_dir], show_checksums=True) + self.assertTrue((md5_dir, STATUS_OK) in status_dir) def test(self): self._ensure_should_run() @@ -335,9 +336,6 @@ class TestRemoteSSH(TestDataCloudBase): def _get_cloud_class(self): return RemoteSSH - def test(self): - super(TestRemoteSSH, self).test() - class TestRemoteHDFS(TestDataCloudBase): def _should_test(self): @@ -349,9 +347,6 @@ class TestRemoteHDFS(TestDataCloudBase): def _get_cloud_class(self): return RemoteHDFS - def test(self): - super(TestRemoteHDFS, self).test() - class TestDataCloudCLIBase(TestDvc): def main(self, args): diff --git a/tests/test_state.py b/tests/test_state.py index 292e5e9f5..ce9bcd5be 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -18,23 +18,24 @@ class TestState(TestDvc): state = State(self.dvc) - entry_md5 = state.update(path) - self.assertEqual(entry_md5, md5) + with state: + entry_md5 = state.update(path) + self.assertEqual(entry_md5, md5) - # Sleep some time to simulate realistic behavior. - # Some filesystems have a bad date resolution for - # mtime(i.e. 1sec for HFS) that cause problems with - # our 'state' system not being able to distinguish - # files that were modified within that delta. - time.sleep(1) + # Sleep some time to simulate realistic behavior. + # Some filesystems have a bad date resolution for + # mtime(i.e. 1sec for HFS) that cause problems with + # our 'state' system not being able to distinguish + # files that were modified within that delta. + time.sleep(1) - os.unlink(path) - with open(path, 'w+') as fd: - fd.write('1') + os.unlink(path) + with open(path, 'w+') as fd: + fd.write('1') - md5 = file_md5(path)[0] - mtime = os.path.getmtime(path) - inode = System.inode(path) + md5 = file_md5(path)[0] + mtime = os.path.getmtime(path) + inode = System.inode(path) - entry_md5 = state.update(path) - self.assertEqual(entry_md5, md5) + entry_md5 = state.update(path) + self.assertEqual(entry_md5, md5)