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

commit: batch dump stages, do not re-read on save from the workspace (#10839)

Browse Source
This commit is contained in:
skshetry
2025-08-14 16:45:42 +05:45
committed by GitHub
parent 4b1adfd646
commit a68dda4164
6 changed files with 212 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/dvcfile.py
+47 -12
View File
@@ -162,6 +162,9 @@ class FileMixin:
def dump(self, stage, **kwargs):
raise NotImplementedError
def dump_stages(self, stages, **kwargs):
raise NotImplementedError
def merge(self, ancestor, other, allowed=None):
raise NotImplementedError
@@ -198,6 +201,13 @@ class SingleStageFile(FileMixin):
dump_yaml(self.path, serialize.to_single_stage_file(stage, **kwargs))
self.repo.scm_context.track_file(self.relpath)
def dump_stages(self, stages, **kwargs) -> None:
if not stages:
return
assert len(stages) == 1, "SingleStageFile can only dump one stage."
return self.dump(stages[0], **kwargs)
def remove_stage(self, stage): # noqa: ARG002
self.remove()
@@ -228,17 +238,27 @@ class ProjectFile(FileMixin):
def dump(self, stage, update_pipeline=True, update_lock=True, **kwargs):
"""Dumps given stage appropriately in the dvcfile."""
return self.dump_stages(
[stage], update_pipeline=update_pipeline, update_lock=update_lock, **kwargs
)
def dump_stages(self, stages, update_pipeline=True, update_lock=True, **kwargs):
from dvc.stage import PipelineStage
if not stages:
return
for stage in stages:
assert isinstance(stage, PipelineStage)
if self.verify:
check_dvcfile_path(self.repo, self.path)
if update_pipeline and not stage.is_data_source:
self._dump_pipeline_file(stage)
if update_pipeline:
self._dump_pipeline_file(stages)
if update_lock:
self._dump_lockfile(stage, **kwargs)
self._dump_lockfile(stages, **kwargs)
def dump_dataset(self, dataset):
with modify_yaml(self.path, fs=self.repo.fs) as data:
@@ -260,27 +280,32 @@ class ProjectFile(FileMixin):
raw.append(dataset)
self.repo.scm_context.track_file(self.relpath)
def _dump_lockfile(self, stage, **kwargs):
self._lockfile.dump(stage, **kwargs)
def _dump_lockfile(self, stages, **kwargs):
self._lockfile.dump_stages(stages, **kwargs)
@staticmethod
def _check_if_parametrized(stage, action: str = "dump") -> None:
if stage.raw_data.parametrized:
raise ParametrizedDumpError(f"cannot {action} a parametrized {stage}")
def _dump_pipeline_file(self, stage):
def _dump_pipeline_file(self, stages):
stages = stages if isinstance(stages, list) else [stages]
if not stages:
return
for stage in stages:
self._check_if_parametrized(stage)
stage_data = serialize.to_pipeline_file(stage)
with modify_yaml(self.path, fs=self.repo.fs) as data:
if not data:
logger.info("Creating '%s'", self.relpath)
data["stages"] = data.get("stages", {})
for stage in stages:
stage_data = serialize.to_pipeline_file(stage)
existing_entry = stage.name in data["stages"]
action = "Modifying" if existing_entry else "Adding"
logger.info("%s stage '%s' in '%s'", action, stage.name, self.relpath)
if existing_entry:
orig_stage_data = data["stages"][stage.name]
apply_diff(stage_data[stage.name], orig_stage_data)
@@ -399,9 +424,12 @@ class Lockfile(FileMixin):
data.setdefault("stages", {})
self.repo.scm_context.track_file(self.relpath)
def dump(self, stage, **kwargs):
stage_data = serialize.to_lockfile(stage, **kwargs)
def dump_stages(self, stages, **kwargs):
if not stages:
return
is_modified = False
log_updated = False
with modify_yaml(self.path, fs=self.repo.fs) as data:
if not data:
data.update({"schema": "2.0"})
@@ -409,17 +437,24 @@ class Lockfile(FileMixin):
logger.info("Generating lock file '%s'", self.relpath)
data["stages"] = data.get("stages", {})
for stage in stages:
stage_data = serialize.to_lockfile(stage, **kwargs)
modified = data["stages"].get(stage.name, {}) != stage_data.get(
stage.name, {}
)
if modified:
is_modified = True
if not log_updated:
logger.info("Updating lock file '%s'", self.relpath)
log_updated = True
data["stages"].update(stage_data)
if modified:
if is_modified:
self.repo.scm_context.track_file(self.relpath)
def dump(self, stage, **kwargs):
self.dump_stages([stage], **kwargs)
def remove_stage(self, stage):
if not self.exists():
return
dvc/output.py
+27 -9
View File
@@ -1452,25 +1452,43 @@ class Output:
self.remote = other.remote
self.can_push = other.can_push
def merge_version_meta(self, other: "Output"):
def _get_versioned_meta(
self,
) -> Optional[
tuple["HashInfo", Optional["Meta"], Optional[Union["HashFile", "Tree"]]]
]:
if self.files is not None or (
self.meta is not None and self.meta.version_id is not None
):
old_obj = self.obj if self.obj is not None else self.get_obj()
return self.hash_info, self.meta, old_obj
return None
def merge_version_meta(
self,
old_hi: "HashInfo",
old_meta: Optional["Meta"],
old_obj: Optional[Union["HashFile", "Tree"]],
):
"""Merge version meta for files which are unchanged from other."""
if not self.hash_info:
return
if self.hash_info.isdir:
return self._merge_dir_version_meta(other)
if self.hash_info != other.hash_info:
return self._merge_dir_version_meta(old_hi, old_obj)
if self.hash_info != old_hi:
return
self.meta = other.meta
self.meta = old_meta
def _merge_dir_version_meta(self, other: "Output"):
def _merge_dir_version_meta(
self, old_hi: "HashInfo", old_obj: Optional[Union["HashFile", "Tree"]]
):
from dvc_data.hashfile.tree import update_meta
if not self.obj or not other.hash_info.isdir:
if not self.obj or not old_hi.isdir:
return
other_obj = other.obj if other.obj is not None else other.get_obj()
assert isinstance(self.obj, Tree)
assert isinstance(other_obj, Tree)
updated = update_meta(self.obj, other_obj)
assert isinstance(old_obj, Tree)
updated = update_meta(self.obj, old_obj)
assert updated.hash_info == self.obj.hash_info
self.obj = updated
self.files = updated.as_list(with_meta=True)
dvc/repo/commit.py
+15 -5
View File
@@ -1,3 +1,4 @@
from itertools import groupby
from typing import TYPE_CHECKING
from dvc import prompt
@@ -52,14 +53,21 @@ def commit(
data_only=False,
relink=True,
):
stages_info = [
committed_stages = []
groups = groupby(
[
info
for info in self.stage.collect_granular(
target, with_deps=with_deps, recursive=recursive
)
if not data_only or info.stage.is_data_source
]
for stage_info in stages_info:
],
key=lambda info: info.stage.dvcfile,
)
for dvcfile, stages_info_group in groups:
to_dump = []
for stage_info in stages_info_group:
stage = stage_info.stage
if force:
stage.save(allow_missing=allow_missing)
@@ -73,8 +81,10 @@ def commit(
allow_missing=allow_missing,
relink=relink,
)
stage.dump(update_pipeline=False)
return [s.stage for s in stages_info]
to_dump.append(stage)
dvcfile.dump_stages(to_dump, update_pipeline=False)
committed_stages.extend(to_dump)
return committed_stages
@locked
dvc/stage/__init__.py
+8 -21
View File
@@ -509,34 +509,20 @@ class Stage(params.StageParams):
if not allow_missing:
raise
def get_versioned_outs(self) -> dict[str, "Output"]:
from .exceptions import StageFileDoesNotExistError, StageNotFound
try:
old = self.reload()
except (StageFileDoesNotExistError, StageNotFound):
return {}
return {
out.def_path: out
for out in old.outs
if out.files is not None
or (out.meta is not None and out.meta.version_id is not None)
}
def save_outs(self, allow_missing: bool = False):
from dvc.output import OutputDoesNotExistError
old_versioned_outs = self.get_versioned_outs()
for out in self.outs:
# old state just before saving so to merge them later
old_state = out._get_versioned_meta()
try:
out.save()
except OutputDoesNotExistError:
if not allow_missing:
raise
if old_out := old_versioned_outs.get(out.def_path):
out.merge_version_meta(old_out)
if old_state:
out.merge_version_meta(*old_state)
def ignore_outs(self) -> None:
for out in self.outs:
@@ -579,8 +565,9 @@ class Stage(params.StageParams):
from dvc.output import OutputDoesNotExistError
link_failures = []
old_versioned_outs = self.get_versioned_outs()
for out in self.filter_outs(filter_info):
# old state just before saving so to merge them later
old_state = out._get_versioned_meta()
try:
out.add(filter_info, **kwargs)
except (FileNotFoundError, OutputDoesNotExistError):
@@ -589,8 +576,8 @@ class Stage(params.StageParams):
except CacheLinkError:
link_failures.append(filter_info or out.fs_path)
if old_out := old_versioned_outs.get(out.def_path):
out.merge_version_meta(old_out)
if old_state:
out.merge_version_meta(*old_state)
if link_failures:
raise CacheLinkError(link_failures)
tests/func/test_commit.py
+49 -1
View File
@@ -4,7 +4,7 @@ import textwrap
import pytest
from dvc.dependency.base import DependencyDoesNotExistError
from dvc.dvcfile import PROJECT_FILE
from dvc.dvcfile import PROJECT_FILE, Lockfile, ProjectFile, SingleStageFile
from dvc.fs import localfs
from dvc.output import OutputDoesNotExistError
from dvc.stage.exceptions import StageCommitError
@@ -325,3 +325,51 @@ def test_commit_dos2unix(tmp_dir, dvc):
dvc.commit("foo.dvc", force=True)
content = (tmp_dir / "foo.dvc").read_text()
assert "hash: md5" in content
def test_commit_multiple_files(tmp_dir, dvc, mocker):
tmp_dir.gen({"foo": "foo", "bar": "bar"})
stages = dvc.add(["foo", "bar"], no_commit=True)
test1_stage = dvc.stage.add(name="test", cmd="echo test", deps=["foo"])
test2_stage = dvc.stage.add(name="test2", cmd="echo test2", deps=["foo"])
subdir = tmp_dir / "subdir"
subdir.mkdir()
with subdir.chdir():
bar_relpath = os.path.relpath(tmp_dir / "bar", subdir)
test3_stage = dvc.stage.add(name="test3", cmd="echo test3", deps=[bar_relpath])
pointerfile_spy = mocker.spy(SingleStageFile, "dump_stages")
projectfile_spy = mocker.spy(ProjectFile, "dump_stages")
lockfile_spy = mocker.spy(Lockfile, "dump_stages")
assert set(dvc.commit(force=True)) == {
*stages,
test1_stage,
test2_stage,
test3_stage,
}
pointerfile_spy.assert_has_calls(
[
mocker.call(stages[0].dvcfile, [stages[0]], update_pipeline=False),
mocker.call(stages[1].dvcfile, [stages[1]], update_pipeline=False),
],
any_order=True,
)
projectfile_spy.assert_has_calls(
[
mocker.call(
test1_stage.dvcfile, [test1_stage, test2_stage], update_pipeline=False
),
mocker.call(test3_stage.dvcfile, [test3_stage], update_pipeline=False),
],
any_order=True,
)
lockfile_spy.assert_has_calls(
[
mocker.call(test1_stage.dvcfile._lockfile, [test1_stage, test2_stage]),
mocker.call(test3_stage.dvcfile._lockfile, [test3_stage]),
],
any_order=True,
)
assert dvc.status() == {}
tests/unit/test_dvcfile.py
+33
View File
@@ -90,6 +90,39 @@ def test_dump_stage(tmp_dir, dvc):
assert list(dvcfile.stages.values()) == [stage]
def test_dump_multiple_pipeline_stages(tmp_dir, dvc):
stage1 = PipelineStage(dvc, cmd="cmd1", name="stage1", path="dvc.yaml")
stage2 = PipelineStage(dvc, cmd="cmd2", name="stage2", path="dvc.yaml")
dvcfile = load_file(dvc, "dvc.yaml")
dvcfile.dump_stages([stage1, stage2], update_lock=False, update_pipeline=False)
assert not (tmp_dir / LOCK_FILE).exists()
assert not (tmp_dir / PROJECT_FILE).exists()
dvcfile.dump_stages([stage1, stage2], update_pipeline=False)
assert not (tmp_dir / PROJECT_FILE).exists()
assert (tmp_dir / LOCK_FILE).parse() == {
"schema": "2.0",
"stages": {"stage1": {"cmd": "cmd1"}, "stage2": {"cmd": "cmd2"}},
}
dvcfile.dump_stages([stage1, stage2], update_lock=False)
assert (tmp_dir / PROJECT_FILE).parse() == {
"stages": {"stage1": {"cmd": "cmd1"}, "stage2": {"cmd": "cmd2"}}
}
def test_dump_stages_single_stage(tmp_dir, dvc):
stage = dvc.stage.create(
fname="foo.dvc", outs=["out"], deps=["dep"], single_stage=True
)
stage.dvcfile.dump_stages([stage])
assert (tmp_dir / "foo.dvc").parse() == {
"deps": [{"hash": "md5", "path": "dep"}],
"outs": [{"hash": "md5", "path": "out"}],
}
@pytest.mark.parametrize("file", ["stage.dvc", "dvc.yaml"])
def test_stage_load_file_exists_but_dvcignored(tmp_dir, dvc, scm, file):
(tmp_dir / file).write_text("")