From 45181504ea18a11599abb37fb8098ae4297273fa Mon Sep 17 00:00:00 2001 From: Alexander Guschin <1aguschin@gmail.com> Date: Mon, 15 May 2023 10:55:43 +0300 Subject: [PATCH] Add API to add artifacts (#9345) * api call to add/write artifact * return dvcyaml path on add * create dir for dvc.yaml if needed * create dir for dvc.yaml if needed - fix * allow nested artifact names * make dvc get download artifacts * Apply suggestions from code review Co-authored-by: Oded Messer * use feedback * remove 'dvc get' from this PR, keep only adding artifacts * allow first char of artifact name to be a number * add note about constants * fix parsing artifact name * remove 'parse_name' and write tests for 'add' * Apply suggestions from code review Co-authored-by: Oded Messer * updates from comments * fix for studio * fix again * making check_name_format to raise an exception * Update dvc/repo/artifacts.py --------- Co-authored-by: Oded Messer --- dvc/annotations.py | 16 +++++- dvc/repo/artifacts.py | 78 ++++++++++++++++---------- dvc/schema.py | 3 +- tests/func/artifacts/test_artifacts.py | 78 ++++++++++++++++++-------- 4 files changed, 120 insertions(+), 55 deletions(-) diff --git a/dvc/annotations.py b/dvc/annotations.py index 6f0a7e0aa..ed7e63869 100644 --- a/dvc/annotations.py +++ b/dvc/annotations.py @@ -28,9 +28,21 @@ class Annotation: @dataclass -class Artifact(Annotation): +class Artifact: PARAM_PATH: ClassVar[str] = "path" - path: Optional[str] = None + PARAM_DESC: ClassVar[str] = "desc" + PARAM_TYPE: ClassVar[str] = "type" + PARAM_LABELS: ClassVar[str] = "labels" + PARAM_META: ClassVar[str] = "meta" + + path: str + desc: Optional[str] = None + type: Optional[str] = None # noqa: A003 + labels: List[str] = field(default_factory=list) + meta: Dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> Dict[str, str]: + return compact(asdict(self)) ANNOTATION_FIELDS = [field.name for field in fields(Annotation)] diff --git a/dvc/repo/artifacts.py b/dvc/repo/artifacts.py index e39fa67ac..0615486c2 100644 --- a/dvc/repo/artifacts.py +++ b/dvc/repo/artifacts.py @@ -1,19 +1,26 @@ import logging -import os import re -from typing import TYPE_CHECKING, Dict +from pathlib import Path +from typing import Dict, Optional from dvc.annotations import Artifact -from dvc.dvcfile import FileMixin +from dvc.dvcfile import PROJECT_FILE +from dvc.exceptions import InvalidArgumentError +from dvc.repo import Repo from dvc.utils import relpath - -if TYPE_CHECKING: - from dvc.repo import Repo +from dvc.utils.serialize import modify_yaml logger = logging.getLogger(__name__) -NAME_RE = re.compile(r"^[a-z]([a-z0-9-]*[a-z0-9])?$") +# Constants are taken from GTO. +# When we make it a dependency, we can import them instead +SEPARATOR_IN_NAME = ":" +DIRNAME = r"[a-z0-9-_./]+" +NAME = r"[a-z0-9]([a-z0-9-/]*[a-z0-9])?" +NAME_RE = re.compile(f"^{NAME}$") +FULLNAME = f"((?P{DIRNAME}){SEPARATOR_IN_NAME})?(?P{NAME})" +FULLNAME_RE = re.compile(f"^{FULLNAME}$") def name_is_compatible(name: str) -> bool: @@ -22,22 +29,23 @@ def name_is_compatible(name: str) -> bool: def check_name_format(name: str) -> None: if not name_is_compatible(name): - logger.warning( - "Can't use '%s' as artifact name (ID)." + raise InvalidArgumentError( + f"Can't use '{name}' as artifact name (ID)." " You can use letters and numbers, and use '-' as separator" - " (but not at the start or end). The first character must be a letter.", - name, + " (but not at the start or end)." ) -class ArtifactsFile(FileMixin): - from dvc.schema import SINGLE_ARTIFACT_SCHEMA as SCHEMA - - def dump(self, stage, **kwargs): - raise NotImplementedError - - def merge(self, ancestor, other, allowed=None): - raise NotImplementedError +def check_for_nested_dvc_repo(dvcfile: Path): + if dvcfile.is_absolute(): + raise InvalidArgumentError("Use relative path to dvc.yaml.") + path = dvcfile.parent + while path.name: + if (path / Repo.DVC_DIR).is_dir(): + raise InvalidArgumentError( + f"Nested DVC repos like {path} are not supported." + ) + path = path.parent class Artifacts: @@ -50,18 +58,30 @@ class Artifacts: dvcfile, dvcfile_artifacts, ) in self.repo.index._artifacts.items(): # pylint: disable=protected-access - # read the artifacts.yaml file if needed - if isinstance(dvcfile_artifacts, str): - dvcfile_artifacts = ArtifactsFile( - self.repo, - os.path.join(os.path.dirname(dvcfile), dvcfile_artifacts), - verify=False, - ).load() - if not dvcfile_artifacts: - continue dvcyaml = relpath(dvcfile, self.repo.root_dir) artifacts[dvcyaml] = {} for name, value in dvcfile_artifacts.items(): - check_name_format(name) + try: + check_name_format(name) + except InvalidArgumentError as e: + logger.warning(e.msg) artifacts[dvcyaml][name] = Artifact(**value) return artifacts + + def add(self, name: str, artifact: Artifact, dvcfile: Optional[str] = None): + with self.repo.scm_context(quiet=True): + check_name_format(name) + dvcyaml = Path(dvcfile or PROJECT_FILE) + check_for_nested_dvc_repo( + dvcyaml.relative_to(self.repo.root_dir) + if dvcyaml.is_absolute() + else dvcyaml + ) + + with modify_yaml(dvcyaml) as data: + artifacts = data.setdefault("artifacts", {}) + artifacts.update({name: artifact.to_dict()}) + + self.repo.scm_context.track_file(dvcfile) + + return artifacts.get(name) diff --git a/dvc/schema.py b/dvc/schema.py index 677f18292..30fea0d86 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -116,7 +116,6 @@ PLOT_DEFINITION = { SINGLE_PLOT_SCHEMA = {str: Any(PLOT_DEFINITION, None)} ARTIFACTS = "artifacts" SINGLE_ARTIFACT_SCHEMA = Schema({str: ARTIFACT_SCHEMA}) -ARTIFACTS_SCHEMA = Any(str, SINGLE_ARTIFACT_SCHEMA) FOREACH_IN = { Required(FOREACH_KWD): Any(dict, list, str), Required(DO_KWD): STAGE_DEFINITION, @@ -130,7 +129,7 @@ MULTI_STAGE_SCHEMA = { VARS_KWD: VARS_SCHEMA, StageParams.PARAM_PARAMS: [str], StageParams.PARAM_METRICS: [str], - ARTIFACTS: ARTIFACTS_SCHEMA, + ARTIFACTS: SINGLE_ARTIFACT_SCHEMA, } COMPILED_SINGLE_STAGE_SCHEMA = Schema(SINGLE_STAGE_SCHEMA) diff --git a/tests/func/artifacts/test_artifacts.py b/tests/func/artifacts/test_artifacts.py index a632fdc9e..424d7f047 100644 --- a/tests/func/artifacts/test_artifacts.py +++ b/tests/func/artifacts/test_artifacts.py @@ -3,6 +3,7 @@ import os import pytest from dvc.annotations import Artifact +from dvc.exceptions import InvalidArgumentError from dvc.repo.artifacts import name_is_compatible from dvc.utils.strictyaml import YAMLSyntaxError, YAMLValidationError @@ -21,7 +22,7 @@ dvcyaml = { } -def test_reading_artifacts_subdir(tmp_dir, dvc): +def test_artifacts_read_subdir(tmp_dir, dvc): (tmp_dir / "dvc.yaml").dump(dvcyaml) subdir = tmp_dir / "subdir" @@ -38,6 +39,55 @@ def test_reading_artifacts_subdir(tmp_dir, dvc): } +def test_artifacts_add_subdir(tmp_dir, dvc): + subdir = tmp_dir / "subdir" + subdir.mkdir() + + (subdir / "dvc.yaml").dump(dvcyaml) + + new_art = Artifact(path="path") + tmp_dir.dvc.artifacts.add("new", new_art, dvcfile="subdir/dvc.yaml") + + artifacts = { + name: Artifact(**values) for name, values in dvcyaml["artifacts"].items() + } + artifacts["new"] = new_art + assert tmp_dir.dvc.artifacts.read() == { + f"subdir{os.path.sep}dvc.yaml": artifacts, + } + + +def test_artifacts_add_abspath(tmp_dir, dvc): + subdir = tmp_dir / "subdir" + subdir.mkdir() + + new_art = Artifact(path="path") + tmp_dir.dvc.artifacts.add( + "new", new_art, dvcfile=os.path.abspath("subdir/dvc.yaml") + ) + + assert tmp_dir.dvc.artifacts.read() == { + f"subdir{os.path.sep}dvc.yaml": {"new": new_art}, + } + + +def test_artifacts_add_fails_on_dvc_subrepo(tmp_dir, dvc): + # adding artifact to the DVC subrepo from the parent DVC repo + # shouldn't work + subdir = tmp_dir / "subdir" + (subdir / ".dvc").mkdir(parents=True) + + with pytest.raises(InvalidArgumentError): + tmp_dir.dvc.artifacts.add( + "failing", Artifact(path="path"), dvcfile="subdir/dvc.yaml" + ) + + with pytest.raises(InvalidArgumentError): + tmp_dir.dvc.artifacts.add( + "failing", Artifact(path="path"), dvcfile="subdir/dvclive/dvc.yaml" + ) + + bad_dvcyaml_extra_field = { "artifacts": { "lol": {"kek": "cheburek", "path": "lol"}, @@ -71,7 +121,7 @@ artifacts: """ -def test_broken_dvcyaml_id_duplication(tmp_dir, dvc): +def test_artifacts_read_fails_on_id_duplication(tmp_dir, dvc): with open(tmp_dir / "dvc.yaml", "w") as f: f.write(bad_dvcyaml_id_duplication) @@ -79,32 +129,19 @@ def test_broken_dvcyaml_id_duplication(tmp_dir, dvc): tmp_dir.dvc.artifacts.read() -dvcyaml_redirecting = {"artifacts": "artifacts.yaml"} - - -def test_read_artifacts_yaml(tmp_dir, dvc): - (tmp_dir / "dvc.yaml").dump(dvcyaml_redirecting) - (tmp_dir / "artifacts.yaml").dump(dvcyaml["artifacts"]) - - artifacts = { - name: Artifact(**values) for name, values in dvcyaml["artifacts"].items() - } - assert tmp_dir.dvc.artifacts.read() == { - "dvc.yaml": artifacts, - } - - @pytest.mark.parametrize( "name", [ + "1", "m", "nn", "m1", + "1nn", "model-prod", "model-prod-v1", ], ) -def test_check_name_is_valid(name): +def test_name_is_compatible(name): assert name_is_compatible(name) @@ -112,10 +149,8 @@ def test_check_name_is_valid(name): "name", [ "", - "1", "m/", "/m", - "1nn", "###", "@@@", "a model", @@ -125,8 +160,7 @@ def test_check_name_is_valid(name): "model@1", "model#1", "@namespace/model", - "namespace/model", ], ) -def test_check_name_is_invalid(name): +def test_name_is_compatible_fails(name): assert not name_is_compatible(name)