diff --git a/dvc/dependency/__init__.py b/dvc/dependency/__init__.py index 9e7c474cd..47fd63062 100644 --- a/dvc/dependency/__init__.py +++ b/dvc/dependency/__init__.py @@ -5,6 +5,7 @@ from typing import Any from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output from .base import Dependency +from .dataset import DatasetDependency from .db import DbDependency from .param import ParamsDependency from .repo import RepoDependency @@ -33,6 +34,8 @@ def _get(stage, p, info, **kwargs): return DbDependency(stage, d) assert p + if DatasetDependency.is_dataset(p): + return DatasetDependency(stage, p, info) if repo: return RepoDependency(repo, stage, p, info) return Dependency(stage, p, info, **kwargs) diff --git a/dvc/dependency/dataset.py b/dvc/dependency/dataset.py new file mode 100644 index 000000000..37dcea30e --- /dev/null +++ b/dvc/dependency/dataset.py @@ -0,0 +1,69 @@ +from typing import TYPE_CHECKING, Any, ClassVar +from urllib.parse import urlparse + +from funcy import compact, merge + +from dvc_data.hashfile.hash_info import HashInfo + +from .db import AbstractDependency + +if TYPE_CHECKING: + from dvc.stage import Stage + + +class DatasetDependency(AbstractDependency): + PARAM_DATASET = "dataset" + DATASET_SCHEMA: ClassVar[dict] = {PARAM_DATASET: dict} + + def __init__(self, stage: "Stage", p, info, *args, **kwargs): + super().__init__(stage, info, *args, **kwargs) + self.def_path = p + self.name = urlparse(p).netloc + dataset_info = info.get(self.PARAM_DATASET) or {} + self.hash_info = HashInfo(self.PARAM_DATASET, dataset_info) # type: ignore[arg-type] + self.hash_name = self.PARAM_DATASET + + def __repr__(self): + return f"{self.__class__.__name__}({self.def_path!r})" + + def __str__(self): + return self.def_path + + @classmethod + def is_dataset(cls, p: str): + return urlparse(p).scheme == "ds" + + @property + def protocol(self): + return None + + def dumpd(self, **kwargs): + return compact({self.PARAM_PATH: self.def_path, **self.hash_info.to_dict()}) + + def fill_values(self, values=None): + """Load params values dynamically.""" + self.hash_info = HashInfo( + self.PARAM_DATASET, merge(self.hash_info.value, values or {}) + ) + + def workspace_status(self): + registered = self.repo.index.datasets.get(self.name, {}) + info: dict[str, Any] = self.hash_info.value if self.hash_info else {} # type: ignore[assignment] + if info != registered: + return {str(self): "modified"} + return {} + + def status(self): + return self.workspace_status() + + def get_hash(self): + return HashInfo(self.PARAM_DATASET, self.repo.index.datasets.get(self.name, {})) + + def save(self): + self.hash_info = self.get_hash() + + def download(self, to, jobs=None): + raise NotImplementedError + + def update(self, rev=None): + raise NotImplementedError diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index e568b843d..311b3bfb7 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -179,6 +179,7 @@ class SingleStageFile(FileMixin): from dvc.schema import COMPILED_SINGLE_STAGE_SCHEMA as SCHEMA from dvc.stage.loader import SingleStageLoader as LOADER # noqa: N814 + datasets: ClassVar[list[dict[str, Any]]] = [] metrics: ClassVar[list[str]] = [] plots: ClassVar[Any] = {} params: ClassVar[list[str]] = [] @@ -311,6 +312,10 @@ class ProjectFile(FileMixin): def params(self) -> list[str]: return self.contents.get("params", []) + @property + def datasets(self) -> list[dict[str, Any]]: + return self.contents.get("datasets", []) + @property def artifacts(self) -> dict[str, Optional[dict[str, Any]]]: return self.contents.get("artifacts", {}) diff --git a/dvc/repo/graph.py b/dvc/repo/graph.py index 6534d2300..c0570a348 100644 --- a/dvc/repo/graph.py +++ b/dvc/repo/graph.py @@ -117,6 +117,7 @@ def build_graph(stages, outs_trie=None): """ import networkx as nx + from dvc.dependency import DatasetDependency from dvc.exceptions import StagePathAsOutputError from .trie import build_outs_trie @@ -140,6 +141,8 @@ def build_graph(stages, outs_trie=None): continue for dep in stage.deps: + if isinstance(dep, DatasetDependency): + continue dep_key = dep.fs.parts(dep.fs_path) overlapping = [n.value for n in outs_trie.prefixes(dep_key)] if outs_trie.has_subtrie(dep_key): @@ -157,6 +160,8 @@ def build_graph(stages, outs_trie=None): def build_outs_graph(graph, outs_trie): import networkx as nx + from dvc.dependency import DatasetDependency + outs_graph = nx.DiGraph() outs_graph.add_nodes_from(outs_trie.values()) @@ -166,6 +171,8 @@ def build_outs_graph(graph, outs_trie): if stage.is_db_import: continue for dep in stage.deps: + if isinstance(dep, DatasetDependency): + continue dep_key = dep.fs.parts(dep.fs_path) overlapping = [n.value for n in outs_trie.prefixes(dep_key)] if outs_trie.has_subtrie(dep_key): diff --git a/dvc/repo/index.py b/dvc/repo/index.py index c32e537c4..582f40c0b 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -288,6 +288,7 @@ class Index: plots: Optional[dict[str, list[str]]] = None, params: Optional[dict[str, Any]] = None, artifacts: Optional[dict[str, Any]] = None, + datasets: Optional[dict[str, list[dict[str, Any]]]] = None, ) -> None: self.repo = repo self.stages = stages or [] @@ -295,6 +296,7 @@ class Index: self._plots = plots or {} self._params = params or {} self._artifacts = artifacts or {} + self._datasets: dict[str, list[dict[str, Any]]] = datasets or {} self._collected_targets: dict[int, list["StageInfo"]] = {} @cached_property @@ -318,6 +320,7 @@ class Index: plots = {} params = {} artifacts = {} + datasets = {} onerror = onerror or repo.stage_collection_error_handler for _, idx in collect_files(repo, onerror=onerror): @@ -326,6 +329,7 @@ class Index: plots.update(idx._plots) params.update(idx._params) artifacts.update(idx._artifacts) + datasets.update(idx._datasets) return cls( repo, stages=stages, @@ -333,6 +337,7 @@ class Index: plots=plots, params=params, artifacts=artifacts, + datasets=datasets, ) @classmethod @@ -347,6 +352,7 @@ class Index: plots={path: dvcfile.plots} if dvcfile.plots else {}, params={path: dvcfile.params} if dvcfile.params else {}, artifacts={path: dvcfile.artifacts} if dvcfile.artifacts else {}, + datasets={path: dvcfile.datasets} if dvcfile.datasets else {}, ) def update(self, stages: Iterable["Stage"]) -> "Index": @@ -361,6 +367,7 @@ class Index: plots=self._plots, params=self._params, artifacts=self._artifacts, + datasets=self._datasets, ) @cached_property @@ -398,6 +405,11 @@ class Index: for stage in self.stages: yield from stage.outs + @cached_property + def datasets(self) -> dict[str, dict[str, Any]]: + datasets = chain.from_iterable(self._datasets.values()) + return {dataset["name"]: dataset for dataset in datasets} + @cached_property def out_data_keys(self) -> dict[str, set["DataIndexKey"]]: by_workspace: dict[str, set["DataIndexKey"]] = defaultdict(set) diff --git a/dvc/schema.py b/dvc/schema.py index ca110566c..5695a1b41 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -35,6 +35,7 @@ DATA_SCHEMA: dict[Any, Any] = { Output.PARAM_CLOUD: CLOUD_SCHEMA, Output.PARAM_FILES: [DIR_FILES_SCHEMA], Output.PARAM_HASH: str, + **dependency.DatasetDependency.DATASET_SCHEMA, } LOCK_FILE_STAGE_SCHEMA = { vol.Required(StageParams.PARAM_CMD): vol.Any(str, list), @@ -122,7 +123,12 @@ FOREACH_IN = { SINGLE_PIPELINE_STAGE_SCHEMA = { str: either_or(STAGE_DEFINITION, FOREACH_IN, [FOREACH_KWD, DO_KWD]) } + +DATASET_SCHEMA = vol.Schema( + {vol.Required("type"): str, vol.Required("name"): str}, extra=vol.ALLOW_EXTRA +) MULTI_STAGE_SCHEMA = { + "datasets": [DATASET_SCHEMA], PLOTS: [vol.Any(str, SINGLE_PLOT_SCHEMA)], STAGES: SINGLE_PIPELINE_STAGE_SCHEMA, VARS_KWD: VARS_SCHEMA, diff --git a/dvc/stage/loader.py b/dvc/stage/loader.py index a63f9152b..3b6a64423 100644 --- a/dvc/stage/loader.py +++ b/dvc/stage/loader.py @@ -73,6 +73,10 @@ class StageLoader(Mapping): info = info.copy() info.pop("path", None) + if isinstance(item, dependency.DatasetDependency): + item.fill_values(info.get(dependency.DatasetDependency.PARAM_DATASET)) + continue + hash_name = info.pop(Output.PARAM_HASH, None) item.meta = Meta.from_dict(merge_file_meta_from_cloud(info)) item.hash_name, item.hash_info = item._compute_hash_info_from_meta( diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index 4b7c4e89b..a984d480e 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -148,6 +148,7 @@ def to_pipeline_file(stage: "PipelineStage"): def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict: from dvc.cachemgr import LEGACY_HASH_NAMES + from dvc.dependency import DatasetDependency from dvc.output import ( _serialize_hi_to_dict, _serialize_tree_obj_to_files, @@ -158,6 +159,9 @@ def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict: assert stage.cmd def _dumpd(item: "Output"): + if isinstance(item, DatasetDependency): + return item.dumpd() + ret: dict[str, Any] = {item.PARAM_PATH: item.def_path} if item.hash_name not in LEGACY_HASH_NAMES: ret[item.PARAM_HASH] = "md5"