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

Author: skshetry

dvc/dvcfile.py

@@ -162,6 +162,9 @@ def remove(self, force=False):  # noqa: ARG002
     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 @@ def dump(self, stage, **kwargs) -> None:
         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 @@ def _reset(self):
 
     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
 
-        assert isinstance(stage, 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,32 +280,37 @@ def dump_dataset(self, dataset):
                 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):
-        self._check_if_parametrized(stage)
-        stage_data = serialize.to_pipeline_file(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)
 
         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", {})
-            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)
-            else:
-                data["stages"].update(stage_data)
+            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)
+                else:
+                    data["stages"].update(stage_data)
 
         self.repo.scm_context.track_file(self.relpath)
 
@@ -399,27 +424,37 @@ def dump_dataset(self, dataset: dict):
             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"})
                 # order is important, meta should always be at the top
                 logger.info("Generating lock file '%s'", self.relpath)
 
             data["stages"] = data.get("stages", {})
-            modified = data["stages"].get(stage.name, {}) != stage_data.get(
-                stage.name, {}
-            )
-            if modified:
-                logger.info("Updating lock file '%s'", self.relpath)
-
-            data["stages"].update(stage_data)
+            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

@@ -1452,25 +1452,43 @@ def restore_fields(self, other: "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

@@ -1,3 +1,4 @@
+from itertools import groupby
 from typing import TYPE_CHECKING
 
 from dvc import prompt
@@ -52,29 +53,38 @@ def commit(
     data_only=False,
     relink=True,
 ):
-    stages_info = [
-        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:
-        stage = stage_info.stage
-        if force:
-            stage.save(allow_missing=allow_missing)
-        else:
-            changes = stage.changed_entries()
-            if any(changes):
-                prompt_to_commit(stage, changes, force=force)
+    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
+        ],
+        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)
-        stage.commit(
-            filter_info=stage_info.filter_info,
-            allow_missing=allow_missing,
-            relink=relink,
-        )
-        stage.dump(update_pipeline=False)
-    return [s.stage for s in stages_info]
+            else:
+                changes = stage.changed_entries()
+                if any(changes):
+                    prompt_to_commit(stage, changes, force=force)
+                    stage.save(allow_missing=allow_missing)
+            stage.commit(
+                filter_info=stage_info.filter_info,
+                allow_missing=allow_missing,
+                relink=relink,
+            )
+            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

@@ -509,34 +509,20 @@ def save_deps(self, allow_missing=False):
                 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 @@ def add_outs(self, filter_info=None, allow_missing: bool = False, **kwargs):
         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 @@ def add_outs(self, filter_info=None, allow_missing: bool = False, **kwargs):
             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

@@ -4,7 +4,7 @@
 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() == {}