diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index 1e6b8743c..d2f70dd4d 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -8,6 +8,7 @@ from dvc.utils.objects import cached_property from dvc_data.hashfile.db import get_index if TYPE_CHECKING: + from dvc.fs import FileSystem from dvc_data.hashfile.db import HashFileDB from dvc_data.hashfile.hash_info import HashInfo from dvc_data.hashfile.status import CompareStatusResult @@ -17,11 +18,12 @@ logger = logging.getLogger(__name__) class Remote: - def __init__(self, path, fs, **config): + def __init__(self, name: str, path: str, fs: "FileSystem", **config): self.path = path self.fs = fs + self.name = name - self.worktree = config.pop("worktree", False) + self.worktree: bool = config.pop("worktree", False) self.config = config if self.worktree: version_aware = self.config.get("version_aware") @@ -71,7 +73,7 @@ class DataCloud: cls, config, fs_path = get_cloud_fs(self.repo, name=name) fs = cls(**config) config["tmp_dir"] = self.repo.index_db_dir - return Remote(fs_path, fs, **config) + return Remote(name, fs_path, fs, **config) if bool(self.repo.config["remote"]): error_msg = ( diff --git a/dvc/output.py b/dvc/output.py index 97507b79b..e202627e5 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -4,7 +4,7 @@ from collections import defaultdict from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Type, cast from urllib.parse import urlparse -from funcy import collecting, project +from funcy import collecting, first, project from voluptuous import And, Any, Coerce, Length, Lower, Required, SetTo from dvc import prompt @@ -204,6 +204,30 @@ def load_from_pipeline(stage, data, typ="outs"): ) +def split_file_meta_from_cloud(entry: Dict) -> Dict: + if remote_name := entry.pop(Meta.PARAM_REMOTE, None): + remote_meta = {} + for key in ( + S3_PARAM_CHECKSUM, + HDFS_PARAM_CHECKSUM, + Meta.PARAM_VERSION_ID, + ): + if value := entry.pop(key, None): + remote_meta[key] = value + + if remote_meta: + entry[Output.PARAM_CLOUD] = {remote_name: remote_meta} + return entry + + +def merge_file_meta_from_cloud(entry: Dict) -> Dict: + cloud_meta = entry.pop(Output.PARAM_CLOUD, {}) + if remote_name := first(cloud_meta): + entry.update(cloud_meta[remote_name]) + entry[Meta.PARAM_REMOTE] = remote_name + return entry + + class OutputDoesNotExistError(DvcException): def __init__(self, path): msg = f"output '{path}' does not exist" @@ -258,6 +282,7 @@ class Output: PARAM_PERSIST = "persist" PARAM_REMOTE = "remote" PARAM_PUSH = "push" + PARAM_CLOUD = "cloud" METRIC_SCHEMA = Any( None, @@ -273,7 +298,7 @@ class Output: IsStageFileError: Type[DvcException] = OutputIsStageFileError IsIgnoredError: Type[DvcException] = OutputIsIgnoredError - def __init__( + def __init__( # noqa: C901 self, stage, path, @@ -340,6 +365,9 @@ class Output: # should be absolute and don't contain remote:// refs. self.stage = stage self.meta = meta + + if files is not None: + files = [merge_file_meta_from_cloud(f) for f in files] self.files = files self.use_cache = False if self.IS_DEPENDENCY else cache self.metric = False if self.IS_DEPENDENCY else metric @@ -752,24 +780,24 @@ class Output: return checkout_obj def dumpd(self, **kwargs): # noqa: C901 - meta = self.meta.to_dict() - meta.pop("isdir", None) ret: Dict[str, Any] = {} if ( (not self.IS_DEPENDENCY or self.stage.is_import) and self.hash_info.isdir and (kwargs.get("with_files") or self.files is not None) ): - obj: Optional["HashFile"] - if self.obj: - obj = self.obj - else: - obj = self.get_obj() + obj = self.obj or self.get_obj() if obj: - obj = cast(Tree, obj) - ret[self.PARAM_FILES] = obj.as_list(with_meta=True) + obj = cast("Tree", obj) + ret[self.PARAM_FILES] = [ + split_file_meta_from_cloud(f) + for f in obj.as_list(with_meta=True) + ] else: - ret = {**self.hash_info.to_dict(), **meta} + meta_d = self.meta.to_dict() + meta_d.pop("isdir", None) + ret.update(self.hash_info.to_dict()) + ret.update(split_file_meta_from_cloud(meta_d)) if self.is_in_repo: path = self.fs.path.as_posix( @@ -1261,6 +1289,7 @@ DIR_FILES_SCHEMA: Dict[str, Any] = { **CHECKSUMS_SCHEMA, **META_SCHEMA, Required(Tree.PARAM_RELPATH): str, + Output.PARAM_CLOUD: {str: {**META_SCHEMA, **CHECKSUMS_SCHEMA}}, } SCHEMA = { diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 6020970a9..ffebc06b4 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -212,7 +212,7 @@ class Repo: # used by DVCFileSystem to determine if it should traverse subrepos self.subrepos = subrepos - self.cloud = DataCloud(self) + self.cloud: "DataCloud" = DataCloud(self) self.stage: "StageLoad" = StageLoad(self) self.lock: "LockBase" diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py index 46470d87d..da5ba5f3f 100644 --- a/dvc/repo/worktree.py +++ b/dvc/repo/worktree.py @@ -1,6 +1,18 @@ import logging from functools import partial -from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + Optional, + Set, + Tuple, + Union, + cast, +) + +from funcy import first from dvc.fs.callbacks import Callback @@ -28,11 +40,44 @@ def _meta_checksum(fs: "FileSystem", meta: "Meta") -> Any: return getattr(meta, cast(str, fs.PARAM_CHECKSUM)) +def worktree_view_by_remotes( + index: "Index", + targets: Optional["TargetType"] = None, + push: bool = False, + **kwargs: Any, +) -> Iterable[Tuple[Optional[str], "IndexView"]]: + # pylint: disable=protected-access + + from dvc.repo.index import IndexView + + def outs_filter(view: "IndexView", remote: Optional[str]): + def _filter(out: "Output") -> bool: + if out.remote != remote: + return False + if view._outs_filter: + return view._outs_filter(out) + return True + + return _filter + + view = worktree_view(index, targets=targets, push=push, **kwargs) + remotes = {out.remote for out in view.outs} + + if len(remotes) <= 1: + yield first(remotes), view + return + + for remote in remotes: + yield remote, IndexView( + index, view._stage_infos, outs_filter(view, remote) + ) + + def worktree_view( index: "Index", targets: Optional["TargetType"] = None, push: bool = False, - **kwargs, + **kwargs: Any, ) -> "IndexView": """Return view of data that can be stored in worktree remotes. @@ -65,29 +110,44 @@ def worktree_view( ) +def _get_remote( + repo: "Repo", name: Optional[str], default: "Remote", command: str +) -> "Remote": + if name in (None, default.name): + return default + return repo.cloud.get_remote(name, command) + + def fetch_worktree( repo: "Repo", remote: "Remote", targets: Optional["TargetType"] = None, jobs: Optional[int] = None, - **kwargs, + **kwargs: Any, ) -> int: from dvc_data.index import save - view = worktree_view(repo.index, targets=targets, **kwargs) - index = view.data["repo"] - for key, entry in index.iteritems(): - entry.fs = remote.fs - entry.path = remote.fs.path.join( - remote.path, - *key, - ) - total = len(index) - with Callback.as_tqdm_callback( - unit="file", desc="Fetch", disable=total == 0 - ) as cb: - cb.set_size(total) - return save(index, callback=cb, jobs=jobs) + transferred = 0 + for remote_name, view in worktree_view_by_remotes( + repo.index, push=True, targets=targets, **kwargs + ): + remote_obj = _get_remote(repo, remote_name, remote, "fetch") + index = view.data["repo"] + for key, entry in index.iteritems(): + entry.fs = remote_obj.fs + entry.path = remote_obj.fs.path.join( + remote_obj.path, + *key, + ) + total = len(index) + with Callback.as_tqdm_callback( + unit="file", + desc=f"Fetching from remote {remote_obj.name!r}", + disable=total == 0, + ) as cb: + cb.set_size(total) + transferred += save(index, callback=cb, jobs=jobs) + return transferred def push_worktree( @@ -95,61 +155,70 @@ def push_worktree( remote: "Remote", targets: Optional["TargetType"] = None, jobs: Optional[int] = None, - **kwargs, + **kwargs: Any, ) -> int: from dvc.repo.index import build_data_index from dvc_data.index import checkout - view = worktree_view( - repo.index, - push=True, - targets=targets, - **kwargs, - ) - new_index = view.data["repo"] - if remote.worktree: - logger.debug("indexing latest worktree for '%s'", remote.path) - old_index = build_data_index(view, remote.path, remote.fs) - logger.debug("Pushing worktree changes to '%s'", remote.path) - else: - old_index = None - logger.debug("Pushing version-aware files to '%s'", remote.path) + pushed = 0 + stages: Set["Stage"] = set() - if remote.worktree: - diff_kwargs: Dict[str, Any] = { - "meta_only": True, - "meta_cmp_key": partial(_meta_checksum, remote.fs), - } - else: - diff_kwargs = {} + for remote_name, view in worktree_view_by_remotes( + repo.index, push=True, targets=targets, **kwargs + ): + remote_obj = _get_remote(repo, remote_name, remote, "push") + new_index = view.data["repo"] + if remote_obj.worktree: + logger.debug("indexing latest worktree for '%s'", remote_obj.path) + old_index = build_data_index(view, remote_obj.path, remote_obj.fs) + logger.debug("Pushing worktree changes to '%s'", remote_obj.path) + else: + old_index = None + logger.debug( + "Pushing version-aware files to '%s'", remote_obj.path + ) - total = len(new_index) - with Callback.as_tqdm_callback( - unit="file", desc="Push", disable=total == 0 - ) as cb: - cb.set_size(total) - pushed = checkout( - new_index, - remote.path, - remote.fs, - old=old_index, - delete=remote.worktree, - callback=cb, - latest_only=remote.worktree, - jobs=jobs, - **diff_kwargs, - ) - if pushed: - for stage in view.stages: - for out in stage.outs: - workspace, _key = out.index_key - _update_out_meta(out, repo.index.data[workspace]) - stage.dump(with_files=True, update_pipeline=False) + if remote_obj.worktree: + diff_kwargs: Dict[str, Any] = { + "meta_only": True, + "meta_cmp_key": partial(_meta_checksum, remote_obj.fs), + } + else: + diff_kwargs = {} + + total = len(new_index) + with Callback.as_tqdm_callback( + unit="file", + desc=f"Pushing to remote {remote_obj.name!r}", + disable=total == 0, + ) as cb: + cb.set_size(total) + pushed += checkout( + new_index, + remote_obj.path, + remote_obj.fs, + old=old_index, + delete=remote_obj.worktree, + callback=cb, + latest_only=remote_obj.worktree, + jobs=jobs, + **diff_kwargs, + ) + + for out in view.outs: + workspace, _key = out.index_key + _update_out_meta(out, repo.index.data[workspace], remote_obj.name) + stages.add(out.stage) + + for stage in stages: + stage.dump(with_files=True, update_pipeline=False) return pushed def _update_out_meta( - out: "Output", index: Union["DataIndex", "DataIndexView"] + out: "Output", + index: Union["DataIndex", "DataIndexView"], + remote: Optional[str] = None, ): from dvc_data.index.save import build_tree @@ -168,6 +237,8 @@ def _update_out_meta( repo.fs.path.relparts(fs_path, out.fs_path) ) or (None, None) entry.hash_info = hash_info + if entry.meta: + entry.meta.remote = remote if meta is not None and meta.version_id is not None: # preserve existing version IDs for unchanged files in # this dir (entry will have the latest remote version @@ -182,6 +253,8 @@ def _update_out_meta( out.hash_info = entry.hash_info if out.meta.version_id is None: out.meta = entry.meta + if out.meta: + out.meta.remote = remote def update_worktree_stages( diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index 0ab5ad988..515af5b03 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -2,11 +2,13 @@ from collections import OrderedDict from operator import attrgetter from typing import ( TYPE_CHECKING, + Any, Dict, Iterable, List, Optional, Union, + cast, no_type_check, ) @@ -22,6 +24,7 @@ from .utils import resolve_wdir, split_params_deps if TYPE_CHECKING: from dvc.stage import PipelineStage, Stage + from dvc_data.hashfile.tree import Tree PARAM_PARAMS = ParamsDependency.PARAM_PARAMS PARAM_PATH = ParamsDependency.PARAM_PATH @@ -153,23 +156,26 @@ def to_pipeline_file(stage: "PipelineStage"): def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict: + from dvc.output import split_file_meta_from_cloud + assert stage.cmd - def _dumpd(item): - meta_d = item.meta.to_dict() - meta_d.pop("isdir", None) - + def _dumpd(item: "Output"): + ret: Dict[str, Any] = {item.PARAM_PATH: item.def_path} if item.hash_info.isdir and kwargs.get("with_files"): - if item.obj: - obj = item.obj - else: - obj = item.get_obj() - ret = [((item.PARAM_FILES, obj.as_list(with_meta=True)))] + obj = item.obj or item.get_obj() + if obj: + obj = cast("Tree", obj) + ret[item.PARAM_FILES] = [ + split_file_meta_from_cloud(f) + for f in obj.as_list(with_meta=True) + ] else: - ret = [*item.hash_info.to_dict().items(), *meta_d.items()] - ret.insert(0, (item.PARAM_PATH, item.def_path)) - - return OrderedDict(ret) + meta_d = item.meta.to_dict() + meta_d.pop("isdir", None) + ret.update(item.hash_info.to_dict()) + ret.update(split_file_meta_from_cloud(meta_d)) + return ret res = OrderedDict([("cmd", stage.cmd)]) params, deps = split_params_deps(stage) diff --git a/pyproject.toml b/pyproject.toml index 15f451784..c53b70e55 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ dependencies = [ "dvc-render==0.0.17", "dvc-task==0.1.11", "dvclive>=1.2.2", - "dvc-data==0.35.3", + "dvc-data==0.35.4", "dvc-http", "hydra-core>=1.1.0", "iterative-telemetry==0.0.6",