From ff1809b0602b2828967c58c03fce649e277b84af Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 26 Jun 2019 12:15:56 +0200 Subject: [PATCH 1/3] state: bulk sql for unused links deletion --- dvc/state.py | 7 ++++--- tests/func/test_state.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/dvc/state.py b/dvc/state.py index 3c978355b..15a4b0678 100644 --- a/dvc/state.py +++ b/dvc/state.py @@ -449,6 +449,7 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes remove(path) unused.append(relpath) - for relpath in unused: - cmd = "DELETE FROM {} WHERE path = ?".format(self.LINK_STATE_TABLE) - self._execute(cmd, (relpath,)) + cmd = "DELETE FROM {} WHERE path IN ({})".format( + self.LINK_STATE_TABLE, ",".join(["?"] * len(unused)) + ) + self._execute(cmd, tuple(unused)) diff --git a/tests/func/test_state.py b/tests/func/test_state.py index edc2012bf..e3d914eeb 100644 --- a/tests/func/test_state.py +++ b/tests/func/test_state.py @@ -77,3 +77,21 @@ class TestGetStateRecordForInode(TestDvc): state.save(PathInfo(path), md5) ret = state.get_state_record_for_inode(inode) self.assertIsNotNone(ret) + + +def test_remove_unused_links(repo_dir, dvc_repo): + stages = dvc_repo.add(repo_dir.FOO) + assert len(stages) == 1 + + stages = dvc_repo.add(repo_dir.BAR) + assert len(stages) == 1 + + cmd_count_links = "SELECT count(*) FROM {}".format(State.LINK_STATE_TABLE) + with dvc_repo.state: + result = dvc_repo.state._execute(cmd_count_links).fetchone()[0] + assert result == 2 + + dvc_repo.state.remove_unused_links([]) + + result = dvc_repo.state._execute(cmd_count_links).fetchone()[0] + assert result == 0 From 3d5c42e2fe2b6b2e80396b50d200993ce6268036 Mon Sep 17 00:00:00 2001 From: pared Date: Wed, 26 Jun 2019 12:41:06 +0200 Subject: [PATCH 2/3] state: test: refactor existing state test to pytest --- dvc/state.py | 12 ++-- dvc/utils/__init__.py | 5 ++ tests/func/test_state.py | 105 ++++++++++++++++----------------- tests/unit/utils/test_utils.py | 16 +++++ 4 files changed, 78 insertions(+), 60 deletions(-) create mode 100644 tests/unit/utils/test_utils.py diff --git a/dvc/state.py b/dvc/state.py index 15a4b0678..bc55d7194 100644 --- a/dvc/state.py +++ b/dvc/state.py @@ -7,11 +7,12 @@ import sqlite3 import logging from dvc.config import Config -from dvc.utils import remove, current_timestamp, relpath +from dvc.utils import remove, current_timestamp, relpath, chunk from dvc.exceptions import DvcException from dvc.utils.fs import get_mtime_and_size, get_inode from dvc.utils.compat import fspath_py35 +SQLITE_MAX_VARIABLES_NUMBER = 999 logger = logging.getLogger(__name__) @@ -449,7 +450,8 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes remove(path) unused.append(relpath) - cmd = "DELETE FROM {} WHERE path IN ({})".format( - self.LINK_STATE_TABLE, ",".join(["?"] * len(unused)) - ) - self._execute(cmd, tuple(unused)) + for chunk_unused in chunk(unused, SQLITE_MAX_VARIABLES_NUMBER): + cmd = "DELETE FROM {} WHERE path IN ({})".format( + self.LINK_STATE_TABLE, ",".join(["?"] * len(chunk_unused)) + ) + self._execute(cmd, tuple(chunk_unused)) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index 7e83d2485..8c76c36dd 100644 --- a/dvc/utils/__init__.py +++ b/dvc/utils/__init__.py @@ -382,3 +382,8 @@ def relpath(path, start=os.curdir): ): return path return os.path.relpath(path, start) + + +def chunk(iterable, chunk_size): + for i in range(0, len(iterable), chunk_size): + yield iterable[i : i + chunk_size] diff --git a/tests/func/test_state.py b/tests/func/test_state.py index e3d914eeb..5ea87fc1e 100644 --- a/tests/func/test_state.py +++ b/tests/func/test_state.py @@ -8,75 +8,70 @@ from dvc.utils import file_md5 from dvc.main import main from dvc.path_info import PathInfo -from tests.basic_env import TestDvc +def test_state(dvc_repo, repo_dir): + path = os.path.join(dvc_repo.root_dir, repo_dir.FOO) + path_info = PathInfo(path) + md5 = file_md5(path)[0] + + state = State(dvc_repo, dvc_repo.config.config) + + with state: + state.save(path_info, md5) + entry_md5 = state.get(path_info) + assert entry_md5 == md5 + + os.unlink(path) + with open(path, "a") as fd: + fd.write("1") + + entry_md5 = state.get(path_info) + assert entry_md5 is None -class TestState(TestDvc): - def test_update(self): - path = os.path.join(self.dvc.root_dir, self.FOO) - path_info = PathInfo(path) md5 = file_md5(path)[0] + state.save(path_info, md5) - state = State(self.dvc, self.dvc.config.config) - - with state: - state.save(path_info, md5) - entry_md5 = state.get(path_info) - self.assertEqual(entry_md5, md5) - - os.unlink(path) - with open(path, "a") as fd: - fd.write("1") - - entry_md5 = state.get(path_info) - self.assertTrue(entry_md5 is None) - - md5 = file_md5(path)[0] - state.save(path_info, md5) - - entry_md5 = state.get(path_info) - self.assertEqual(entry_md5, md5) + entry_md5 = state.get(path_info) + assert entry_md5 == md5 -class TestStateOverflow(TestDvc): - def test(self): - # NOTE: trying to add more entries than state can handle, - # to see if it will clean up and vacuum successfully - ret = main(["config", "state.row_limit", "10"]) - self.assertEqual(ret, 0) +def test_state_overflow(dvc_repo): + # NOTE: trying to add more entries than state can handle, + # to see if it will clean up and vacuum successfully + ret = main(["config", "state.row_limit", "10"]) + assert ret == 0 - dname = "dir" - os.mkdir(dname) - for i in range(20): - with open(os.path.join(dname, str(i)), "w+") as fobj: - fobj.write(str(i)) + dname = "dir" + os.mkdir(dname) + for i in range(20): + with open(os.path.join(dname, str(i)), "w+") as fobj: + fobj.write(str(i)) - ret = main(["add", "dir"]) - self.assertEqual(ret, 0) + ret = main(["add", "dir"]) + assert ret == 0 -class TestGetStateRecordForInode(TestDvc): - @staticmethod - def mock_get_inode(inode): - def get_inode_mocked(path): - return inode +def mock_get_inode(inode): + def get_inode_mocked(path): + return inode - return get_inode_mocked + return get_inode_mocked - @mock.patch("dvc.state.get_inode", autospec=True) - def test_transforms_inode(self, get_inode_mock): - state = State(self.dvc, self.dvc.config.config) - inode = state.MAX_INT + 2 - self.assertNotEqual(inode, state._to_sqlite(inode)) - path = os.path.join(self.dvc.root_dir, self.FOO) - md5 = file_md5(path)[0] - get_inode_mock.side_effect = self.mock_get_inode(inode) +@mock.patch("dvc.state.get_inode", autospec=True) +def test_get_state_record_for_inode(get_inode_mock, dvc_repo, repo_dir): + state = State(dvc_repo, dvc_repo.config.config) + inode = state.MAX_INT + 2 + assert inode != state._to_sqlite(inode) - with state: - state.save(PathInfo(path), md5) - ret = state.get_state_record_for_inode(inode) - self.assertIsNotNone(ret) + path = os.path.join(dvc_repo.root_dir, repo_dir.FOO) + md5 = file_md5(path)[0] + get_inode_mock.side_effect = mock_get_inode(inode) + + with state: + state.save(PathInfo(path), md5) + ret = state.get_state_record_for_inode(inode) + assert ret is not None def test_remove_unused_links(repo_dir, dvc_repo): diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py new file mode 100644 index 000000000..8e3b3de83 --- /dev/null +++ b/tests/unit/utils/test_utils.py @@ -0,0 +1,16 @@ +import pytest + +from dvc.utils import chunk + + +@pytest.mark.parametrize( + "list_to_chunk, chunk_size, expected_chunks", + [ + ([1, 2, 3, 4], 1, [[1], [2], [3], [4]]), + ([1, 2, 3, 4], 2, [[1, 2], [3, 4]]), + ([1, 2, 3, 4], 3, [[1, 2, 3], [4]]), + ], +) +def test_chunk(list_to_chunk, chunk_size, expected_chunks): + result = list(chunk(list_to_chunk, chunk_size)) + assert result == expected_chunks From 7f8da34cf22fbf5af425453b60085d3db71fed2d Mon Sep 17 00:00:00 2001 From: pared Date: Thu, 27 Jun 2019 13:58:49 +0200 Subject: [PATCH 3/3] utils: extend to_chunks functionality, can specify number of chunks or length of chunk --- dvc/remote/base.py | 2 +- dvc/remote/local/__init__.py | 6 +++--- dvc/remote/ssh/__init__.py | 2 +- dvc/state.py | 6 ++++-- dvc/utils/__init__.py | 36 ++++++++++++++++++++++------------ tests/unit/utils/test_utils.py | 32 +++++++++++++++++++++--------- 6 files changed, 56 insertions(+), 28 deletions(-) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 5f2659fc5..fa8eb2998 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -636,7 +636,7 @@ class RemoteBASE(object): if self.no_traverse and hasattr(self, "batch_exists"): with ThreadPoolExecutor(max_workers=self.JOBS) as executor: path_infos = [self.checksum_to_path_info(x) for x in checksums] - chunks = to_chunks(path_infos, self.JOBS) + chunks = to_chunks(path_infos, num_chunks=self.JOBS) results = executor.map(exists_with_progress, chunks) in_remote = itertools.chain.from_iterable(results) ret = list(itertools.compress(checksums, in_remote)) diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index 7bfd52067..1f8269015 100644 --- a/dvc/remote/local/__init__.py +++ b/dvc/remote/local/__init__.py @@ -360,9 +360,9 @@ class RemoteLOCAL(RemoteBASE): from_infos = cache return ( - to_chunks(from_infos, jobs), - to_chunks(to_infos, jobs), - to_chunks(names, jobs), + to_chunks(from_infos, num_chunks=jobs), + to_chunks(to_infos, num_chunks=jobs), + to_chunks(names, num_chunks=jobs), ) def _process( diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index 04a04a811..07fdb910f 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -157,7 +157,7 @@ class RemoteSSH(RemoteBASE): with ThreadPoolExecutor(max_workers=max_workers) as executor: paths = [path_info.path for path_info in path_infos] - chunks = to_chunks(paths, max_workers) + chunks = to_chunks(paths, num_chunks=max_workers) chunks_and_channels = zip(chunks, channels) outcome = executor.map(_exists, chunks_and_channels) results = list(itertools.chain.from_iterable(outcome)) diff --git a/dvc/state.py b/dvc/state.py index bc55d7194..3f03d084e 100644 --- a/dvc/state.py +++ b/dvc/state.py @@ -7,7 +7,7 @@ import sqlite3 import logging from dvc.config import Config -from dvc.utils import remove, current_timestamp, relpath, chunk +from dvc.utils import remove, current_timestamp, relpath, to_chunks from dvc.exceptions import DvcException from dvc.utils.fs import get_mtime_and_size, get_inode from dvc.utils.compat import fspath_py35 @@ -450,7 +450,9 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes remove(path) unused.append(relpath) - for chunk_unused in chunk(unused, SQLITE_MAX_VARIABLES_NUMBER): + for chunk_unused in to_chunks( + unused, chunk_size=SQLITE_MAX_VARIABLES_NUMBER + ): cmd = "DELETE FROM {} WHERE path IN ({})".format( self.LINK_STATE_TABLE, ",".join(["?"] * len(chunk_unused)) ) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index 8c76c36dd..a372d90d4 100644 --- a/dvc/utils/__init__.py +++ b/dvc/utils/__init__.py @@ -183,16 +183,33 @@ def remove(path): raise -def to_chunks(l, jobs): - n = int(math.ceil(len(l) / jobs)) +def _split(list_to_split, chunk_size): + return [ + list_to_split[i : i + chunk_size] + for i in range(0, len(list_to_split), chunk_size) + ] - if len(l) == 1: - return [l] - if n == 0: - n = 1 +def _to_chunks_by_chunks_number(list_to_split, num_chunks): + chunk_size = int(math.ceil(float(len(list_to_split)) / num_chunks)) - return [l[x : x + n] for x in range(0, len(l), n)] + if len(list_to_split) == 1: + return [list_to_split] + + if chunk_size == 0: + chunk_size = 1 + + return _split(list_to_split, chunk_size) + + +def to_chunks(list_to_split, num_chunks=None, chunk_size=None): + if (num_chunks and chunk_size) or (not num_chunks and not chunk_size): + raise ValueError( + "One and only one of 'num_chunks', 'chunk_size' must be defined" + ) + if chunk_size: + return _split(list_to_split, chunk_size) + return _to_chunks_by_chunks_number(list_to_split, num_chunks) # NOTE: Check if we are in a bundle @@ -382,8 +399,3 @@ def relpath(path, start=os.curdir): ): return path return os.path.relpath(path, start) - - -def chunk(iterable, chunk_size): - for i in range(0, len(iterable), chunk_size): - yield iterable[i : i + chunk_size] diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py index 8e3b3de83..3f27a9f06 100644 --- a/tests/unit/utils/test_utils.py +++ b/tests/unit/utils/test_utils.py @@ -1,16 +1,30 @@ import pytest -from dvc.utils import chunk +from dvc.utils import to_chunks @pytest.mark.parametrize( - "list_to_chunk, chunk_size, expected_chunks", - [ - ([1, 2, 3, 4], 1, [[1], [2], [3], [4]]), - ([1, 2, 3, 4], 2, [[1, 2], [3, 4]]), - ([1, 2, 3, 4], 3, [[1, 2, 3], [4]]), - ], + "chunk_size, expected_chunks", + [(1, [[1], [2], [3], [4]]), (2, [[1, 2], [3, 4]]), (3, [[1, 2, 3], [4]])], ) -def test_chunk(list_to_chunk, chunk_size, expected_chunks): - result = list(chunk(list_to_chunk, chunk_size)) +def test_to_chunks_chunk_size(chunk_size, expected_chunks): + list_to_chunk = [1, 2, 3, 4] + result = list(to_chunks(list_to_chunk, chunk_size=chunk_size)) + assert result == expected_chunks + + +@pytest.mark.parametrize("num_chunks, chunk_size", [(1, 2), (None, None)]) +def test_to_chunks_should_raise(num_chunks, chunk_size): + list_to_chunk = [1, 2, 3] + with pytest.raises(ValueError): + to_chunks(list_to_chunk, num_chunks, chunk_size) + + +@pytest.mark.parametrize( + "num_chunks, expected_chunks", + [(4, [[1], [2], [3], [4]]), (3, [[1, 2], [3, 4]]), (2, [[1, 2], [3, 4]])], +) +def test_to_chunks_num_chunks(num_chunks, expected_chunks): + list_to_chunk = [1, 2, 3, 4] + result = to_chunks(list_to_chunk, num_chunks=num_chunks) assert result == expected_chunks