Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

checkout: use index checkout

Browse Source 
Needed for https://github.com/iterative/dvc/pull/9424
This commit is contained in:
Ruslan Kuprieiev
2023-06-04 23:27:01 +03:00
parent f9104275f3
commit 78637713cf
8 changed files with 152 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/repo/checkout.py
+99 -44
View File
@@ -1,16 +1,12 @@
import logging
import os
from typing import TYPE_CHECKING, Dict, List, Set
from typing import Dict, List
from dvc.exceptions import CheckoutError, CheckoutErrorSuggestGit, NoOutputOrStageError
from dvc.utils import relpath
from . import locked
if TYPE_CHECKING:
from . import Repo
from .stage import StageInfo
logger = logging.getLogger(__name__)
@@ -30,37 +26,41 @@ def _remove_unused_links(repo):
return ret
def get_all_files_numbers(pairs):
return sum(stage.get_all_files_number(filter_info) for stage, filter_info in pairs)
def _build_out_changes(index, changes):
from dvc_data.index.checkout import MODIFY
out_keys = []
for out in index.outs:
if not out.use_cache:
continue
def _collect_pairs(
self: "Repo", targets, with_deps: bool, recursive: bool
) -> Set["StageInfo"]:
from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError
ws, key = out.index_key
if ws != "repo":
continue
pairs: Set["StageInfo"] = set()
for target in targets:
try:
pairs.update(
self.stage.collect_granular(
target, with_deps=with_deps, recursive=recursive
)
)
except (
StageFileDoesNotExistError,
StageFileBadNameError,
NoOutputOrStageError,
) as exc:
if not target:
raise
raise CheckoutErrorSuggestGit(target) from exc
out_keys.append(key)
return pairs
out_changes = {}
for objects in changes.values():
for change in objects:
for out_key in out_keys:
if (
len(out_key) > len(change.key)
or change.key[: len(out_key)] != out_key
):
continue
if change.key == out_key:
out_changes[out_key] = change.typ
elif not out_changes.get(out_key):
out_changes[out_key] = MODIFY
break
return out_changes
@locked
def checkout(
def checkout( # noqa: C901
self,
targets=None,
with_deps=False,
@@ -70,13 +70,18 @@ def checkout(
allow_missing=False,
**kwargs,
):
from dvc import prompt
from dvc.fs.callbacks import Callback
from dvc.repo.index import build_data_index
from dvc.stage.exceptions import StageFileBadNameError, StageFileDoesNotExistError
from dvc_data.hashfile.checkout import CheckoutError as IndexCheckoutError
from dvc_data.index.checkout import ADD, DELETE, MODIFY
from dvc_data.index.checkout import checkout as icheckout
stats: Dict[str, List[str]] = {
"added": [],
"deleted": [],
"modified": [],
"failed": [],
}
if not targets:
targets = [None]
@@ -85,28 +90,78 @@ def checkout(
if isinstance(targets, str):
targets = [targets]
pairs = _collect_pairs(self, targets, with_deps, recursive)
total = get_all_files_numbers(pairs)
def onerror(target, exc):
if target and isinstance(
exc,
(
StageFileDoesNotExistError,
StageFileBadNameError,
NoOutputOrStageError,
),
):
raise CheckoutErrorSuggestGit(target) from exc
raise # pylint: disable=misplaced-bare-raise
view = self.index.targets_view(
targets,
recursive=recursive,
with_deps=with_deps,
onerror=onerror,
)
with Callback.as_tqdm_callback(
unit="entry",
desc="Building data index",
) as cb:
old = build_data_index(
view, self.root_dir, self.fs, compute_hash=True, callback=cb
)
new = view.data["repo"]
for out in view.outs:
out.changed_cache()
with Callback.as_tqdm_callback(
unit="file",
desc="Checkout",
disable=total == 0,
) as cb:
cb.set_size(total)
for stage, filter_info in pairs:
result = stage.checkout(
force=force,
progress_callback=cb,
try:
changes = icheckout(
new,
self.root_dir,
self.fs,
old=old,
callback=cb,
delete=True,
prompt=prompt.confirm,
update_meta=False,
relink=relink,
filter_info=filter_info,
force=force,
allow_missing=allow_missing,
state=self.state,
**kwargs,
)
for key, items in result.items():
stats[key].extend(_fspath_dir(path) for path in items)
except IndexCheckoutError as exc:
raise CheckoutError(exc.paths, {}) from exc
if stats.get("failed"):
raise CheckoutError(stats["failed"], stats)
out_changes = _build_out_changes(view, changes)
typ_map = {ADD: "added", DELETE: "deleted", MODIFY: "modified"}
for key, typ in out_changes.items():
out_path = self.fs.path.join(self.root_dir, *key)
self.state.save_link(out_path, self.fs)
stats[typ_map[typ]].append(_fspath_dir(out_path))
failed = []
for out in view.outs:
if not out.use_cache:
continue
if not out.hash_info:
failed.append(_fspath_dir(out.fs_path))
if not allow_missing and failed:
raise CheckoutError(failed, stats)
del stats["failed"]
return stats
dvc/repo/data.py
+2 -4
View File
@@ -73,10 +73,8 @@ def _diff(
):
if (
change.typ == UNCHANGED
and change.old
and change.new
and not change.old.hash_info
and not change.new.hash_info
and (not change.old or not change.old.hash_info)
and (not change.new or not change.new.hash_info)
):
# NOTE: emulating previous behaviour
continue
dvc/repo/imports.py
+1 -1
View File
@@ -110,7 +110,7 @@ def save_imports(
desc="Downloading imports from source",
unit="files",
) as cb:
checkout(data_view, tmpdir, cache.fs, callback=cb, storage="data")
checkout(data_view, tmpdir, cache.fs, callback=cb, storage="remote")
md5(data_view)
save(data_view, odb=cache, hardlink=True)
dvc/repo/index.py
+41 -6
View File
@@ -19,6 +19,7 @@ from typing import (
from funcy.debug import format_time
from dvc.fs import LocalFileSystem
from dvc.fs.callbacks import DEFAULT_CALLBACK
from dvc.utils.objects import cached_property
if TYPE_CHECKING:
@@ -26,6 +27,7 @@ if TYPE_CHECKING:
from pygtrie import Trie
from dvc.dependency import Dependency, ParamsDependency
from dvc.fs.callbacks import Callback
from dvc.output import Output
from dvc.repo import Repo
from dvc.repo.stage import StageInfo
@@ -113,14 +115,18 @@ def collect_files(
def _load_data_from_outs(index, prefix, outs):
from dvc_data.index import DataIndexEntry
from dvc_data.index import DataIndexEntry, Meta
parents = set()
for out in outs:
if not out.use_cache:
continue
ws, key = out.index_key
for key_len in range(1, len(key)):
parents.add((ws, key[:key_len]))
entry = DataIndexEntry(
key=key,
meta=out.meta,
@@ -137,6 +143,11 @@ def _load_data_from_outs(index, prefix, outs):
# index.view() work correctly.
index[(*prefix, ws, *key)] = entry
for ws, key in parents:
index[(*prefix, ws, *key)] = DataIndexEntry(
key=key, meta=Meta(isdir=True), loaded=True
)
def _load_storage_from_out(storage_map, key, out):
from dvc.config import NoRemoteError
@@ -163,7 +174,7 @@ def _load_storage_from_out(storage_map, key, out):
if out.stage.is_import:
dep = out.stage.deps[0]
storage_map.add_data(FileStorage(key, dep.fs, dep.fs_path))
storage_map.add_remote(FileStorage(key, dep.fs, dep.fs_path))
class Index:
@@ -534,10 +545,12 @@ class IndexView:
lambda: _DataPrefixes(set(), set())
)
for out, filter_info in self._filtered_outs:
if not out.use_cache:
continue
workspace, key = out.index_key
if filter_info and out.fs.path.isin(filter_info, out.fs_path):
key = key + out.fs.path.relparts(filter_info, out.fs_path)
entry = self._index.data[workspace][key]
entry = self._index.data[workspace].get(key)
if entry and entry.meta and entry.meta.isdir:
prefixes[workspace].recursive.add(key)
prefixes[workspace].explicit.update(key[:i] for i in range(len(key), 0, -1))
@@ -550,6 +563,9 @@ class IndexView:
ret: Dict[str, Set["DataIndexKey"]] = defaultdict(set)
for out, filter_info in self._filtered_outs:
if not out.use_cache:
continue
workspace, key = out.index_key
if filter_info and out.fs.path.isin(filter_info, out.fs_path):
key = key + out.fs.path.relparts(filter_info, out.fs_path)
@@ -579,14 +595,15 @@ class IndexView:
return data
def build_data_index(
def build_data_index( # noqa: C901
index: Union["Index", "IndexView"],
path: str,
fs: "FileSystem",
workspace: str = "repo",
compute_hash: Optional[bool] = False,
callback: "Callback" = DEFAULT_CALLBACK,
) -> "DataIndex":
from dvc_data.index import DataIndex, DataIndexEntry
from dvc_data.index import DataIndex, DataIndexEntry, Meta
from dvc_data.index.build import build_entries, build_entry
from dvc_data.index.save import build_tree
@@ -595,9 +612,16 @@ def build_data_index(
ignore = index.repo.dvcignore
data = DataIndex()
parents = set()
for key in index.data_keys.get(workspace, set()):
out_path = fs.path.join(path, *key)
for key_len in range(1, len(key)):
parents.add(key[:key_len])
if not fs.exists(out_path):
continue
try:
out_entry = build_entry(
out_path,
@@ -606,10 +630,11 @@ def build_data_index(
state=index.repo.state,
)
except FileNotFoundError:
out_entry = DataIndexEntry()
continue
out_entry.key = key
data.add(out_entry)
callback.relative_update(1)
if not out_entry.meta or not out_entry.meta.isdir:
continue
@@ -628,6 +653,7 @@ def build_data_index(
entry.key = key + entry.key
data.add(entry)
callback.relative_update(1)
if compute_hash:
tree_meta, tree = build_tree(data, key)
@@ -635,5 +661,14 @@ def build_data_index(
out_entry.hash_info = tree.hash_info
out_entry.loaded = True
data.add(out_entry)
callback.relative_update(1)
for key in parents:
parent_path = fs.path.join(path, *key)
if not fs.exists(parent_path):
continue
direntry = DataIndexEntry(key=key, meta=Meta(isdir=True), loaded=True)
data.add(direntry)
callback.relative_update(1)
return data
dvc/stage/__init__.py
-5
View File
@@ -717,11 +717,6 @@ class Stage(params.StageParams):
for out in self.outs
)
def get_all_files_number(self, filter_info=None) -> int:
return sum(
out.get_files_number(filter_info) for out in self.filter_outs(filter_info)
)
def get_used_objs(
self, *args, **kwargs
) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]:
pyproject.toml
+1 -1
View File
@@ -34,7 +34,7 @@ dependencies = [
"configobj>=5.0.6",
"distro>=1.3",
"dpath<3,>=2.1.0",
"dvc-data>=0.53.0,<0.54",
"dvc-data>=0.54.3,<0.55",
"dvc-http>=2.29.0",
"dvc-render>=0.3.1,<1",
"dvc-studio-client>=0.9.2,<1",
tests/func/test_add.py
+4 -4
View File
@@ -336,11 +336,11 @@ def test_should_update_state_entry_for_file_after_add(mocker, dvc, tmp_dir):
os.rename("foo", "foo.back")
ret = main(["checkout"])
assert ret == 0
assert file_md5_counter.mock.call_count == 1
assert file_md5_counter.mock.call_count == 2
ret = main(["status"])
assert ret == 0
assert file_md5_counter.mock.call_count == 1
assert file_md5_counter.mock.call_count == 2
def test_should_update_state_entry_for_directory_after_add(mocker, dvc, tmp_dir):
@@ -362,11 +362,11 @@ def test_should_update_state_entry_for_directory_after_add(mocker, dvc, tmp_dir)
os.rename("data", "data.back")
ret = main(["checkout"])
assert ret == 0
assert file_md5_counter.mock.call_count == 5
assert file_md5_counter.mock.call_count == 6
ret = main(["status"])
assert ret == 0
assert file_md5_counter.mock.call_count == 6
assert file_md5_counter.mock.call_count == 7
def test_add_commit(tmp_dir, dvc):
tests/func/test_checkout.py
+4 -59
View File
@@ -8,12 +8,7 @@ import pytest
from dvc.cli import main
from dvc.dvcfile import PROJECT_FILE, load_file
from dvc.exceptions import (
CheckoutError,
CheckoutErrorSuggestGit,
ConfirmRemoveError,
NoOutputOrStageError,
)
from dvc.exceptions import CheckoutError, CheckoutErrorSuggestGit, NoOutputOrStageError
from dvc.fs import LocalFileSystem, system
from dvc.stage import Stage
from dvc.stage.exceptions import StageFileDoesNotExistError
@@ -82,7 +77,7 @@ def test_checkout_corrupted_cache_file(tmp_dir, dvc):
fd.write("1")
with pytest.raises(CheckoutError):
dvc.checkout(force=True)
dvc.checkout(force=True, relink=True)
assert not os.path.isfile("foo")
assert not os.path.isfile(cache)
@@ -114,7 +109,7 @@ def test_checkout_corrupted_cache_dir(tmp_dir, dvc):
fobj.write("1")
with pytest.raises(CheckoutError):
dvc.checkout(force=True)
dvc.checkout(force=True, relink=True)
assert not os.path.exists(cache)
@@ -294,19 +289,6 @@ def test_checkout_directory(tmp_dir, dvc):
assert os.path.exists("data")
def test_checkout_hook(mocker, tmp_dir, dvc):
"""Test that dvc checkout handles EOFError gracefully, which is what
it will experience when running in a git hook.
"""
tmp_dir.dvc_gen({"data": {"foo": "foo"}})
mocker.patch("sys.stdout.isatty", return_value=True)
mocker.patch("dvc.prompt.input", side_effect=EOFError)
(tmp_dir / "data").gen("test", "test")
with pytest.raises(ConfirmRemoveError):
dvc.checkout()
def test_checkout_suggest_git(tmp_dir, dvc, scm):
with pytest.raises(CheckoutErrorSuggestGit) as e:
dvc.checkout(targets="gitbranch")
@@ -636,27 +618,6 @@ def test_checkout_recursive(tmp_dir, dvc):
}
def test_checkout_for_external_outputs(tmp_dir, dvc, workspace):
workspace.gen("foo", "foo")
file_path = workspace / "foo"
dvc.add("remote://workspace/foo")
odb = dvc.cloud.get_remote_odb("workspace")
odb.fs.remove(str(file_path))
assert not file_path.exists()
stats = dvc.checkout(force=True)
assert stats == {**empty_checkout, "added": ["remote://workspace/foo"]}
assert file_path.exists()
workspace.gen("foo", "foo\nfoo")
stats = dvc.checkout(force=True)
assert stats == {**empty_checkout, "modified": ["remote://workspace/foo"]}
assert file_path.read_text() == "foo"
def test_checkouts_with_different_addressing(tmp_dir, dvc, run_copy):
tmp_dir.gen({"foo": "foo", "lorem": "lorem"})
run_copy("foo", "bar", name="copy-foo-bar")
@@ -725,22 +686,6 @@ def test_checkouts_for_pipeline_tracked_outs(tmp_dir, dvc, scm, run_copy):
assert set(dvc.checkout()["added"]) == {"bar", "ipsum"}
def test_checkout_external_modified_file(tmp_dir, dvc, scm, mocker, workspace):
# regression: happened when file in external output changed and checkout
# was attempted without force, dvc checks if it's present in its cache
# before asking user to remove it.
workspace.gen("foo", "foo")
dvc.add("remote://workspace/foo", external=True)
scm.add(["foo.dvc"])
scm.commit("add foo")
workspace.gen("foo", "foobar") # messing up the external outputs
mocker.patch("dvc.prompt.confirm", return_value=True)
dvc.checkout()
assert (workspace / "foo").read_text() == "foo"
def test_checkout_executable(tmp_dir, dvc):
tmp_dir.dvc_gen("foo", "foo")
@@ -817,7 +762,7 @@ def test_checkout_partial_unchanged(tmp_dir, dvc):
assert len(stats["modified"]) == 1
stats = dvc.checkout(str(data_dir / "empty_sub_dir"))
assert not any(stats.values())
assert stats == {**empty_checkout, "modified": ["data" + os.sep]}
dvc.checkout(str(data_dir))