checkout/fetch: finer control over steps

See https://github.com/iterative/dvc-data/pull/390 and https://github.com/iterative/dvc-data/pull/391

dvc/repo/checkout.py

@@ -74,9 +74,7 @@ def checkout(  # noqa: C901
     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
+    from dvc_data.index.checkout import ADD, DELETE, MODIFY, apply, compare
 
     stats: Dict[str, List[str]] = {
         "added": [],
@@ -119,28 +117,36 @@ def checkout(  # noqa: C901
 
     new = view.data["repo"]
 
+    with Callback.as_tqdm_callback(
+        desc="Comparing indexes",
+        unit="entry",
+    ) as cb:
+        diff = compare(old, new, relink=relink, delete=True, callback=cb)
+
+    failed = []
+
+    def checkout_onerror(src_path, dest_path, _exc):
+        logger.debug(
+            "failed to create {%s} from {%s}", dest_path, src_path, exc_info=True
+        )
+        failed.append(dest_path)
+
     with Callback.as_tqdm_callback(
         unit="file",
         desc="Checkout",
     ) as cb:
-        try:
-            changes = icheckout(
-                new,
-                self.root_dir,
-                self.fs,
-                old=old,
-                callback=cb,
-                delete=True,
-                prompt=prompt.confirm,
-                update_meta=False,
-                relink=relink,
-                force=force,
-                allow_missing=allow_missing,
-                state=self.state,
-                **kwargs,
-            )
-        except IndexCheckoutError as exc:
-            raise CheckoutError(exc.paths, {}) from exc
+        changes = apply(
+            diff,
+            self.root_dir,
+            self.fs,
+            callback=cb,
+            prompt=prompt.confirm,
+            update_meta=False,
+            force=force,
+            onerror=checkout_onerror,
+            state=self.state,
+            **kwargs,
+        )
 
     out_changes = _build_out_changes(view, changes)
 
@@ -150,4 +156,7 @@ def checkout(  # noqa: C901
         self.state.save_link(out_path, self.fs)
         stats[typ_map[typ]].append(_fspath_dir(out_path))
 
+    if failed and not allow_missing:
+        raise CheckoutError(failed, stats)
+
     return stats

dvc/repo/experiments/executor/base.py

@@ -515,7 +515,6 @@ class BaseExecutor(ABC):
                     targets=targets,
                     with_deps=targets is not None,
                     force=True,
-                    quiet=True,
                     allow_missing=True,
                     recursive=kwargs.get("recursive", False),
                 )

dvc/repo/fetch.py

@@ -36,6 +36,7 @@ def fetch(  # noqa: C901, PLR0913
     from fsspec.utils import tokenize
 
     from dvc.fs.callbacks import Callback
+    from dvc_data.index.fetch import collect
     from dvc_data.index.fetch import fetch as ifetch
 
     if isinstance(targets, str):
@@ -76,17 +77,27 @@ def fetch(  # noqa: C901, PLR0913
 
     cache_key = ("fetch", tokenize(sorted(index_keys)))
 
+    with Callback.as_tqdm_callback(
+        desc="Collecting",
+        unit="entry",
+    ) as cb:
+        data = collect(
+            indexes, cache_index=self.data_index, cache_key=cache_key, callback=cb
+        )
+
     with Callback.as_tqdm_callback(
         desc="Fetching",
         unit="file",
     ) as cb:
-        fetch_transferred, fetch_failed = ifetch(
-            indexes,
-            jobs=jobs,
-            callback=cb,
-            cache_index=self.data_index,
-            cache_key=cache_key,
-        )  # pylint: disable=assignment-from-no-return
+        try:
+            fetch_transferred, fetch_failed = ifetch(
+                data,
+                jobs=jobs,
+                callback=cb,
+            )  # pylint: disable=assignment-from-no-return
+        finally:
+            for fs_index in data.values():
+                fs_index.close()
 
     if fetch_transferred:
         # NOTE: dropping cached index to force reloading from newly saved cache

dvc/repo/imports.py

@@ -100,7 +100,8 @@ def save_imports(
         Objects which were downloaded from source location.
     """
     from dvc.stage.exceptions import DataSourceChanged
-    from dvc_data.index import checkout, md5, save
+    from dvc_data.index import md5, save
+    from dvc_data.index.checkout import apply, compare
     from dvc_objects.fs.callbacks import Callback
 
     downloaded: Set["HashInfo"] = set()
@@ -124,7 +125,8 @@ def save_imports(
                     desc="Downloading imports from source",
                     unit="files",
                 ) as cb:
-                    checkout(data_view, tmpdir, cache.fs, callback=cb, storage="remote")
+                    diff = compare(None, data_view)
+                    apply(diff, tmpdir, cache.fs, callback=cb, storage="remote")
                 md5(data_view, name=cache.hash_name)
                 save(data_view, odb=cache, hardlink=True)
 

dvc/repo/worktree.py

@@ -112,8 +112,7 @@ def push_worktree(
     **kwargs: Any,
 ) -> int:
     from dvc.repo.index import build_data_index
-    from dvc_data.index import checkout
-    from dvc_data.index.checkout import VersioningNotSupported
+    from dvc_data.index.checkout import VersioningNotSupported, apply, compare
 
     pushed = 0
     stages: Set["Stage"] = set()
@@ -139,6 +138,18 @@ def push_worktree(
         else:
             diff_kwargs = {}
 
+        with Callback.as_tqdm_callback(
+            unit="entry",
+            desc=f"Comparing indexes for remote {remote_obj.name!r}",
+        ) as cb:
+            diff = compare(
+                old_index,
+                new_index,
+                callback=cb,
+                delete=remote_obj.worktree,
+                **diff_kwargs,
+            )
+
         total = len(new_index)
         with Callback.as_tqdm_callback(
             unit="file",
@@ -147,16 +158,13 @@ def push_worktree(
         ) as cb:
             cb.set_size(total)
             try:
-                stats = checkout(
-                    new_index,
+                stats = apply(
+                    diff,
                     remote_obj.path,
                     remote_obj.fs,
-                    old=old_index,
-                    delete=remote_obj.worktree,
                     callback=cb,
                     latest_only=remote_obj.worktree,
                     jobs=jobs,
-                    **diff_kwargs,
                 )
                 pushed += sum(len(changes) for changes in stats.values())
             except VersioningNotSupported:
@@ -316,23 +324,33 @@ def _fetch_out_changes(
     remote_index: Union["DataIndex", "DataIndexView"],
     remote: "Remote",
 ):
-    from dvc_data.index import checkout
+    from dvc_data.index.checkout import apply, compare
 
     old, new = _get_diff_indexes(out, local_index, remote_index)
+
+    with Callback.as_tqdm_callback(
+        unit="entry",
+        desc="Comparing indexes",
+    ) as cb:
+        diff = compare(
+            old,
+            new,
+            delete=True,
+            meta_only=True,
+            meta_cmp_key=partial(_meta_checksum, remote.fs),
+            callback=cb,
+        )
+
     total = len(new)
     with Callback.as_tqdm_callback(
         unit="file", desc=f"Updating '{out}'", disable=total == 0
     ) as cb:
         cb.set_size(total)
-        checkout(
-            new,
+        apply(
+            diff,
             out.repo.root_dir,
             out.fs,
-            old=old,
-            delete=True,
             update_meta=False,
-            meta_only=True,
-            meta_cmp_key=partial(_meta_checksum, remote.fs),
             storage="data",
             callback=cb,
         )

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=1.11.0,<1.12.0",
+    "dvc-data>=2.0.2,<2.1.0",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<1",
     "dvc-studio-client>=0.9.2,<1",