From 4382b80d3ba343a70947cc2d97bc0e41f8d867eb Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Fri, 27 Aug 2021 12:45:25 +0300 Subject: [PATCH] dvcfile: support remote per output (#6486) Related to #2095 --- dvc/output.py | 43 +++++++++++++++++++++++---------- dvc/schema.py | 1 + dvc/stage/serialize.py | 3 +++ tests/func/test_data_cloud.py | 44 ++++++++++++++++++++++++++++++++++ tests/func/test_remote.py | 2 +- tests/remotes/__init__.py | 20 +++++++++++++++- tests/remotes/local.py | 18 ++++++++++---- tests/unit/output/test_load.py | 10 ++++++++ 8 files changed, 121 insertions(+), 20 deletions(-) diff --git a/dvc/output.py b/dvc/output.py index 910ebfafe..51b306c5d 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -81,6 +81,7 @@ def loadd_from(stage, d_list): desc = d.pop(Output.PARAM_DESC, False) isexec = d.pop(Output.PARAM_ISEXEC, False) live = d.pop(Output.PARAM_LIVE, False) + remote = d.pop(Output.PARAM_REMOTE, None) ret.append( _get( stage, @@ -94,6 +95,7 @@ def loadd_from(stage, d_list): desc=desc, isexec=isexec, live=live, + remote=remote, ) ) return ret @@ -109,6 +111,7 @@ def loads_from( checkpoint=False, isexec=False, live=False, + remote=None, ): return [ _get( @@ -122,6 +125,7 @@ def loads_from( checkpoint=checkpoint, isexec=isexec, live=live, + remote=remote, ) for s in s_list ] @@ -162,7 +166,6 @@ def load_from_pipeline(stage, data, typ="outs"): metric = typ == stage.PARAM_METRICS plot = typ == stage.PARAM_PLOTS live = typ == stage.PARAM_LIVE - if live: # `live` is single object data = [data] @@ -185,6 +188,7 @@ def load_from_pipeline(stage, data, typ="outs"): Output.PARAM_CACHE, Output.PARAM_PERSIST, Output.PARAM_CHECKPOINT, + Output.PARAM_REMOTE, ], ) @@ -255,6 +259,7 @@ class Output: PARAM_LIVE = "live" PARAM_LIVE_SUMMARY = "summary" PARAM_LIVE_HTML = "html" + PARAM_REMOTE = "remote" METRIC_SCHEMA = Any( None, @@ -283,6 +288,7 @@ class Output: live=False, desc=None, isexec=False, + remote=None, ): self.repo = stage.repo if stage else None @@ -326,7 +332,7 @@ class Output: self.obj = None self.isexec = False if self.IS_DEPENDENCY else isexec - self.def_remote = None + self.remote = remote def _parse_path(self, fs, path_info): if fs.scheme != "local": @@ -843,7 +849,8 @@ class Output: try: objects.check(self.odb, obj) except FileNotFoundError: - self.repo.cloud.pull([obj.hash_info], **kwargs) + remote = self.repo.cloud.get_remote_odb(self.remote) + self.repo.cloud.pull([obj.hash_info], odb=remote, **kwargs) if self.obj: return self.obj @@ -855,9 +862,9 @@ class Output: return self.obj - def collect_used_dir_cache( + def _collect_used_dir_cache( self, remote=None, force=False, jobs=None, filter_info=None - ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: + ) -> Optional["Tree"]: """Fetch dir cache and return used object IDs for this out.""" try: @@ -878,13 +885,13 @@ class Output: "unable to fully collect used cache" " without cache for directory '{}'".format(self) ) - return {} + return None obj = self.get_obj() if filter_info and filter_info != self.path_info: prefix = filter_info.relative_to(self.path_info).parts obj = obj.filter(prefix) - return {None: set(self._named_obj_ids(obj))} + return obj def get_used_objs( self, **kwargs @@ -917,22 +924,31 @@ class Output: return {} if self.is_dir_checksum: - return self.collect_used_dir_cache(**kwargs) + obj = self._collect_used_dir_cache(**kwargs) + else: + obj = self.get_obj(filter_info=kwargs.get("filter_info")) + if not obj: + obj = self.odb.get(self.hash_info) - obj = self.get_obj(filter_info=kwargs.get("filter_info")) if not obj: - obj = self.odb.get(self.hash_info) + return {} - return {None: set(self._named_obj_ids(obj))} + if self.remote: + remote = self.repo.cloud.get_remote_odb(name=self.remote) + else: + remote = None + + return {remote: self._named_obj_ids(obj)} def _named_obj_ids(self, obj): name = str(self) obj.hash_info.obj_name = name - yield obj.hash_info + oids = {obj.hash_info} if isinstance(obj, Tree): for key, entry_obj in obj: entry_obj.hash_info.obj_name = self.fs.sep.join([name, *key]) - yield entry_obj.hash_info + oids.add(entry_obj.hash_info) + return oids def get_used_external( self, **kwargs @@ -1033,4 +1049,5 @@ SCHEMA = { Output.PARAM_CACHE: bool, Output.PARAM_METRIC: Output.METRIC_SCHEMA, Output.PARAM_DESC: str, + Output.PARAM_REMOTE: str, } diff --git a/dvc/schema.py b/dvc/schema.py index 0d6e910f7..68dcc2a35 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -49,6 +49,7 @@ OUT_PSTAGE_DETAILED_SCHEMA = { Output.PARAM_PERSIST: bool, Output.PARAM_CHECKPOINT: bool, Output.PARAM_DESC: str, + Output.PARAM_REMOTE: str, } } diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index bc1c22391..ec145dbe7 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -28,6 +28,7 @@ PARAM_PLOT = Output.PARAM_PLOT PARAM_PERSIST = Output.PARAM_PERSIST PARAM_CHECKPOINT = Output.PARAM_CHECKPOINT PARAM_DESC = Output.PARAM_DESC +PARAM_REMOTE = Output.PARAM_REMOTE DEFAULT_PARAMS_FILE = ParamsDependency.DEFAULT_PARAMS_FILE @@ -52,6 +53,8 @@ def _get_flags(out): yield from out.plot.items() if out.live and isinstance(out.live, dict): yield from out.live.items() + if out.remote: + yield PARAM_REMOTE, out.remote def _serialize_out(out): diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index fb51701ff..91afcb206 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -609,3 +609,47 @@ def test_pull_no_00_prefix(tmp_dir, dvc, remote, monkeypatch): stats = dvc.pull() assert stats["fetched"] == 2 assert set(stats["added"]) == {"foo", "bar"} + + +def test_output_remote(tmp_dir, dvc, make_remote): + from dvc.utils.serialize import modify_yaml + + make_remote("default", default=True) + make_remote("for_foo", default=False) + make_remote("for_data", default=False) + + tmp_dir.dvc_gen("foo", "foo") + tmp_dir.dvc_gen("bar", "bar") + tmp_dir.dvc_gen("data", {"one": "one", "two": "two"}) + + with modify_yaml("foo.dvc") as d: + d["outs"][0]["remote"] = "for_foo" + + with modify_yaml("data.dvc") as d: + d["outs"][0]["remote"] = "for_data" + + dvc.push() + + default = dvc.cloud.get_remote_odb("default") + for_foo = dvc.cloud.get_remote_odb("for_foo") + for_data = dvc.cloud.get_remote_odb("for_data") + + assert set(default.all()) == {"37b51d194a7513e45b56f6524f2d51f2"} + assert set(for_foo.all()) == {"acbd18db4cc2f85cedef654fccc4a4d8"} + assert set(for_data.all()) == { + "f97c5d29941bfb1b2fdab0874906ab82", + "6b18131dc289fd37006705affe961ef8.dir", + "b8a9f715dbb64fd5c56e7783c6820a61", + } + + clean(["foo", "bar", "data"], dvc) + + dvc.pull() + + assert set(dvc.odb.local.all()) == { + "37b51d194a7513e45b56f6524f2d51f2", + "acbd18db4cc2f85cedef654fccc4a4d8", + "f97c5d29941bfb1b2fdab0874906ab82", + "6b18131dc289fd37006705affe961ef8.dir", + "b8a9f715dbb64fd5c56e7783c6820a61", + } diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index edfc66d85..11729c59f 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -200,7 +200,7 @@ def test_partial_push_n_pull(tmp_dir, dvc, tmp_path_factory, local_remote): dvc.push() remove(dvc.odb.local.cache_dir) - baz.collect_used_dir_cache() + baz._collect_used_dir_cache() with patch.object(LocalFileSystem, "upload", side_effect=Exception): with pytest.raises(DownloadError) as download_error_info: dvc.pull() diff --git a/tests/remotes/__init__.py b/tests/remotes/__init__.py index bd98b9da5..836e784bf 100644 --- a/tests/remotes/__init__.py +++ b/tests/remotes/__init__.py @@ -16,7 +16,7 @@ from .gs import ( # noqa: F401; noqa: F401 ) from .hdfs import HDFS, hadoop, hdfs, hdfs_server, real_hdfs # noqa: F401 from .http import HTTP, http, http_server # noqa: F401 -from .local import Local, local_cloud, local_remote # noqa: F401 +from .local import Local, local_cloud, local_remote, make_local # noqa: F401 from .oss import ( # noqa: F401 OSS, TEST_OSS_REPO_BUCKET, @@ -100,6 +100,24 @@ def docker_services( return Services(executor) +@pytest.fixture +def make_cloud(request): + def _make_cloud(typ): + return request.getfixturevalue(f"make_{typ}")() + + return _make_cloud + + +@pytest.fixture +def make_remote(tmp_dir, dvc, make_cloud): + def _make_remote(name, typ="local", **kwargs): + cloud = make_cloud(typ) + tmp_dir.add_remote(name=name, config=cloud.config, **kwargs) + return cloud + + return _make_remote + + @pytest.fixture def remote(tmp_dir, dvc, request): cloud = request.param diff --git a/tests/remotes/local.py b/tests/remotes/local.py index 8ba3580d5..869efb8a3 100644 --- a/tests/remotes/local.py +++ b/tests/remotes/local.py @@ -17,11 +17,19 @@ class Local(Base): @pytest.fixture -def local_cloud(make_tmp_dir): - ret = make_tmp_dir("local-cloud") - ret.url = str(ret) - ret.config = {"url": ret.url} - return ret +def make_local(make_tmp_dir): + def _make_local(): + ret = make_tmp_dir("local-cloud") + ret.url = str(ret) + ret.config = {"url": ret.url} + return ret + + return _make_local + + +@pytest.fixture +def local_cloud(make_local): + return make_local() @pytest.fixture diff --git a/tests/unit/output/test_load.py b/tests/unit/output/test_load.py index 41e07d361..37a0370a4 100644 --- a/tests/unit/output/test_load.py +++ b/tests/unit/output/test_load.py @@ -79,6 +79,16 @@ def test_load_remote_files_from_pipeline(dvc): assert not out.hash_info +def test_load_remote(dvc): + stage = Stage(dvc) + (foo, bar) = output.load_from_pipeline( + stage, + ["foo", {"bar": {"remote": "myremote"}}], + ) + assert foo.remote is None + assert bar.remote == "myremote" + + @pytest.mark.parametrize("typ", [None, "", "illegal"]) def test_load_from_pipeline_error_on_typ(dvc, typ): with pytest.raises(ValueError):