datachain 0.33.1__tar.gz → 0.34.0__tar.gz

{datachain-0.33.1 → datachain-0.34.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.13.1'
+    rev: 'v0.13.2'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.33.1 → datachain-0.34.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.33.1
+Version: 0.34.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.33.1 → datachain-0.34.0}/src/datachain/data_storage/metastore.py

@@ -21,6 +21,7 @@ from sqlalchemy import (
     Table,
     Text,
     UniqueConstraint,
+    desc,
     select,
 )
 from sqlalchemy.sql import func as f
@@ -399,6 +400,7 @@ class AbstractMetastore(ABC, Serializable):
         workers: int = 1,
         python_version: Optional[str] = None,
         params: Optional[dict[str, str]] = None,
+        parent_job_id: Optional[str] = None,
     ) -> str:
         """
         Creates a new job.
@@ -443,6 +445,10 @@ class AbstractMetastore(ABC, Serializable):
     def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
         """Returns all checkpoints related to some job"""
 
+    @abstractmethod
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Optional[Checkpoint]:
+        """Get last created checkpoint for some job."""
+
     @abstractmethod
     def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
         """Gets single checkpoint by id"""
@@ -1548,6 +1554,7 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("error_stack", Text, nullable=False, default=""),
             Column("params", JSON, nullable=False),
             Column("metrics", JSON, nullable=False),
+            Column("parent_job_id", Text, nullable=True),
         ]
 
     @cached_property
@@ -1595,6 +1602,7 @@ class AbstractDBMetastore(AbstractMetastore):
         workers: int = 1,
         python_version: Optional[str] = None,
         params: Optional[dict[str, str]] = None,
+        parent_job_id: Optional[str] = None,
         conn: Optional[Any] = None,
     ) -> str:
         """
@@ -1616,6 +1624,7 @@ class AbstractDBMetastore(AbstractMetastore):
                 error_stack="",
                 params=json.dumps(params or {}),
                 metrics=json.dumps({}),
+                parent_job_id=parent_job_id,
             ),
             conn=conn,
         )
@@ -1770,7 +1779,7 @@ class AbstractDBMetastore(AbstractMetastore):
         )
         return self.get_checkpoint_by_id(checkpoint_id)
 
-    def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
         """List checkpoints by job id."""
         query = self._checkpoints_query().where(self._checkpoints.c.job_id == job_id)
         rows = list(self.db.execute(query, conn=conn))
@@ -1800,3 +1809,15 @@ class AbstractDBMetastore(AbstractMetastore):
         if not rows:
             return None
         return self.checkpoint_class.parse(*rows[0])
+
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Optional[Checkpoint]:
+        query = (
+            self._checkpoints_query()
+            .where(self._checkpoints.c.job_id == job_id)
+            .order_by(desc(self._checkpoints.c.created_at))
+            .limit(1)
+        )
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            return None
+        return self.checkpoint_class.parse(*rows[0])

{datachain-0.33.1 → datachain-0.34.0}/src/datachain/error.py

@@ -101,3 +101,7 @@ class OutdatedDatabaseSchemaError(DataChainError):
 
 class CheckpointNotFoundError(NotFoundError):
     pass
+
+
+class JobNotFoundError(NotFoundError):
+    pass

{datachain-0.33.1 → datachain-0.34.0}/src/datachain/job.py

@@ -22,6 +22,7 @@ class Job:
     python_version: Optional[str] = None
     error_message: str = ""
     error_stack: str = ""
+    parent_job_id: Optional[str] = None
 
     @classmethod
     def parse(
@@ -39,6 +40,7 @@ class Job:
         error_stack: str,
         params: str,
         metrics: str,
+        parent_job_id: Optional[str],
     ) -> "Job":
         return cls(
             str(id),
@@ -54,4 +56,5 @@ class Job:
             python_version,
             error_message,
             error_stack,
+            parent_job_id,
         )

{datachain-0.33.1 → datachain-0.34.0}/src/datachain/lib/dc/datachain.py

@@ -19,7 +19,6 @@ from typing import (
     cast,
     overload,
 )
-from uuid import uuid4
 
 import sqlalchemy
 import ujson as json
@@ -30,10 +29,15 @@ from tqdm import tqdm
 from datachain import semver
 from datachain.dataset import DatasetRecord
 from datachain.delta import delta_disabled
-from datachain.error import ProjectCreateNotAllowedError, ProjectNotFoundError
+from datachain.error import (
+    JobNotFoundError,
+    ProjectCreateNotAllowedError,
+    ProjectNotFoundError,
+)
 from datachain.func import literal
 from datachain.func.base import Function
 from datachain.func.func import Func
+from datachain.job import Job
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.data_model import (
     DataModel,
@@ -50,11 +54,12 @@ from datachain.lib.signal_schema import SignalResolvingError, SignalSchema
 from datachain.lib.udf import Aggregator, BatchMapper, Generator, Mapper, UDFBase
 from datachain.lib.udf_signature import UdfSignature
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
+from datachain.project import Project
 from datachain.query import Session
 from datachain.query.dataset import DatasetQuery, PartitionByType
 from datachain.query.schema import DEFAULT_DELIMITER, Column
 from datachain.sql.functions import path as pathfunc
-from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
+from datachain.utils import batched_it, env2bool, inside_notebook, row_to_nested_dict
 
 from .database import DEFAULT_DATABASE_BATCH_SIZE
 from .utils import (
@@ -578,6 +583,19 @@ class DataChain:
             query=self._query.save(project=project, feature_schema=schema)
         )
 
+    def _calculate_job_hash(self, job_id: str) -> str:
+        """
+        Calculates hash of the job at the place of this chain's save method.
+        Hash is calculated using previous job checkpoint hash (if exists) and
+        adding hash of this chain to produce new hash.
+        """
+        last_checkpoint = self.session.catalog.metastore.get_last_checkpoint(job_id)
+
+        return hashlib.sha256(
+            (bytes.fromhex(last_checkpoint.hash) if last_checkpoint else b"")
+            + bytes.fromhex(self.hash())
+        ).hexdigest()
+
     def save(  # type: ignore[override]
         self,
         name: str,
@@ -602,101 +620,171 @@ class DataChain:
             update_version: which part of the dataset version to automatically increase.
                 Available values: `major`, `minor` or `patch`. Default is `patch`.
         """
+
         catalog = self.session.catalog
-        if version is not None:
-            semver.validate(version)
 
-        if update_version is not None and update_version not in [
-            "patch",
-            "major",
-            "minor",
-        ]:
-            raise ValueError(
-                "update_version can have one of the following values: major, minor or"
-                " patch"
-            )
+        result = None  # result chain that will be returned at the end
+
+        # Version validation
+        self._validate_version(version)
+        self._validate_update_version(update_version)
 
         namespace_name, project_name, name = catalog.get_full_dataset_name(
             name,
             namespace_name=self._settings.namespace,
             project_name=self._settings.project,
         )
+        project = self._get_or_create_project(namespace_name, project_name)
+
+        # Checkpoint handling
+        job, _hash, result = self._resolve_checkpoint(name, project, kwargs)
+
+        # Schema preparation
+        schema = self.signals_schema.clone_without_sys_signals().serialize()
+
+        # Handle retry and delta functionality
+        if not result:
+            result = self._handle_delta(name, version, project, schema, kwargs)
+
+        if not result:
+            # calculate chain if we already don't have result from checkpoint or delta
+            result = self._evolve(
+                query=self._query.save(
+                    name=name,
+                    version=version,
+                    project=project,
+                    description=description,
+                    attrs=attrs,
+                    feature_schema=schema,
+                    update_version=update_version,
+                    **kwargs,
+                )
+            )
+
+        if job:
+            catalog.metastore.create_checkpoint(job.id, _hash)  # type: ignore[arg-type]
+
+        return result
+
+    def _validate_version(self, version: Optional[str]) -> None:
+        """Validate dataset version if provided."""
+        if version is not None:
+            semver.validate(version)
 
+    def _validate_update_version(self, update_version: Optional[str]) -> None:
+        """Ensure update_version is one of: major, minor, patch."""
+        allowed = ["major", "minor", "patch"]
+        if update_version not in allowed:
+            raise ValueError(f"update_version must be one of {allowed}")
+
+    def _get_or_create_project(self, namespace: str, project_name: str) -> Project:
+        """Get project or raise if creation not allowed."""
         try:
-            project = self.session.catalog.metastore.get_project(
+            return self.session.catalog.metastore.get_project(
                 project_name,
-                namespace_name,
+                namespace,
                 create=is_studio(),
             )
         except ProjectNotFoundError as e:
-            # not being able to create it as creation is not allowed
             raise ProjectCreateNotAllowedError("Creating project is not allowed") from e
 
-        schema = self.signals_schema.clone_without_sys_signals().serialize()
+    def _resolve_checkpoint(
+        self,
+        name: str,
+        project: Project,
+        kwargs: dict,
+    ) -> tuple[Optional[Job], Optional[str], Optional["DataChain"]]:
+        """Check if checkpoint exists and return cached dataset if possible."""
+        from .datasets import read_dataset
 
-        # Handle retry and delta functionality
-        if self.delta and name:
-            from datachain.delta import delta_retry_update
+        metastore = self.session.catalog.metastore
 
-            # Delta chains must have delta_on defined (ensured by _as_delta method)
-            assert self._delta_on is not None, "Delta chain must have delta_on defined"
+        job_id = os.getenv("DATACHAIN_JOB_ID")
+        checkpoints_reset = env2bool("DATACHAIN_CHECKPOINTS_RESET", undefined=True)
 
-            result_ds, dependencies, has_changes = delta_retry_update(
-                self,
-                namespace_name,
-                project_name,
-                name,
-                on=self._delta_on,
-                right_on=self._delta_result_on,
-                compare=self._delta_compare,
-                delta_retry=self._delta_retry,
+        if not job_id:
+            return None, None, None
+
+        job = metastore.get_job(job_id)
+        if not job:
+            raise JobNotFoundError(f"Job with id {job_id} not found")
+
+        _hash = self._calculate_job_hash(job.id)
+
+        if (
+            job.parent_job_id
+            and not checkpoints_reset
+            and metastore.find_checkpoint(job.parent_job_id, _hash)
+        ):
+            # checkpoint found → reuse dataset
+            chain = read_dataset(
+                name, namespace=project.namespace.name, project=project.name, **kwargs
             )
+            return job, _hash, chain
 
-            if result_ds:
-                return self._evolve(
-                    query=result_ds._query.save(
-                        name=name,
-                        version=version,
-                        project=project,
-                        feature_schema=schema,
-                        dependencies=dependencies,
-                        **kwargs,
-                    )
-                )
+        return job, _hash, None
 
-            if not has_changes:
-                # sources have not been changed so new version of resulting dataset
-                # would be the same as previous one. To avoid duplicating exact
-                # datasets, we won't create new version of it and we will return
-                # current latest version instead.
-                from .datasets import read_dataset
+    def _handle_delta(
+        self,
+        name: str,
+        version: Optional[str],
+        project: Project,
+        schema: dict,
+        kwargs: dict,
+    ) -> Optional["DataChain"]:
+        """Try to save as a delta dataset.
+        Returns:
+            A DataChain if delta logic could handle it, otherwise None to fall back
+            to the regular save path (e.g., on first dataset creation).
+        """
+        from datachain.delta import delta_retry_update
 
-                return read_dataset(
-                    name, namespace=namespace_name, project=project_name, **kwargs
-                )
+        from .datasets import read_dataset
 
-        result = self._evolve(
-            query=self._query.save(
-                name=name,
-                version=version,
-                project=project,
-                description=description,
-                attrs=attrs,
-                feature_schema=schema,
-                update_version=update_version,
-                **kwargs,
-            )
+        if not self.delta or not name:
+            return None
+
+        assert self._delta_on is not None, "Delta chain must have delta_on defined"
+
+        result_ds, dependencies, has_changes = delta_retry_update(
+            self,
+            project.namespace.name,
+            project.name,
+            name,
+            on=self._delta_on,
+            right_on=self._delta_result_on,
+            compare=self._delta_compare,
+            delta_retry=self._delta_retry,
         )
 
-        if job_id := os.getenv("DATACHAIN_JOB_ID"):
-            catalog.metastore.create_checkpoint(
-                job_id,
-                _hash=hashlib.sha256(  # TODO this will be replaced with self.hash()
-                    str(uuid4()).encode()
-                ).hexdigest(),
+        # Case 1: delta produced a new dataset
+        if result_ds:
+            return self._evolve(
+                query=result_ds._query.save(
+                    name=name,
+                    version=version,
+                    project=project,
+                    feature_schema=schema,
+                    dependencies=dependencies,
+                    **kwargs,
+                )
             )
 
-        return result
+        # Case 2: no changes → reuse last version
+        if not has_changes:
+            # sources have not been changed so new version of resulting dataset
+            # would be the same as previous one. To avoid duplicating exact
+            # datasets, we won't create new version of it and we will return
+            # current latest version instead.
+            return read_dataset(
+                name,
+                namespace=project.namespace.name,
+                project=project.name,
+                **kwargs,
+            )
+
+        # Case 3: first creation of dataset
+        return None
 
     def apply(self, func, *args, **kwargs):
         """Apply any function to the chain.

{datachain-0.33.1 → datachain-0.34.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.33.1
+Version: 0.34.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.33.1 → datachain-0.34.0}/src/datachain.egg-info/SOURCES.txt

@@ -394,6 +394,7 @@ tests/unit/lib/__init__.py
 tests/unit/lib/conftest.py
 tests/unit/lib/test_arrow.py
 tests/unit/lib/test_audio.py
+tests/unit/lib/test_checkpoints.py
 tests/unit/lib/test_clip.py
 tests/unit/lib/test_datachain.py
 tests/unit/lib/test_datachain_bootstrap.py

{datachain-0.33.1 → datachain-0.34.0}/tests/conftest.py

@@ -456,9 +456,12 @@ def cloud_server(request, tmp_upath_factory, cloud_type, version_aware, tree):
     return make_cloud_server(src_path, cloud_type, tree)
 
 
-@pytest.fixture()
-def datachain_job_id(monkeypatch):
-    job_id = str(uuid.uuid4())
+@pytest.fixture
+def datachain_job_id(test_session, monkeypatch):
+    job_id = test_session.catalog.metastore.create_job(
+        "my-job",
+        'import datachain as dc; dc.read_values(num=[1, 2, 3].save("nums")',
+    )
     monkeypatch.setenv("DATACHAIN_JOB_ID", job_id)
     return job_id
 

datachain-0.34.0/tests/unit/lib/test_checkpoints.py

@@ -0,0 +1,200 @@
+import pytest
+
+import datachain as dc
+from datachain.error import DatasetNotFoundError, JobNotFoundError
+from datachain.lib.utils import DataChainError
+
+
+def mapper_fail(num) -> int:
+    raise Exception("Error")
+
+
+@pytest.fixture
+def nums_dataset(test_session):
+    return dc.read_values(num=[1, 2, 3], session=test_session).save("nums")
+
+
+@pytest.mark.parametrize("reset_checkpoints", [True, False])
+@pytest.mark.parametrize("with_delta", [True, False])
+def test_checkpoints(
+    test_session, monkeypatch, nums_dataset, reset_checkpoints, with_delta
+):
+    catalog = test_session.catalog
+
+    monkeypatch.setenv("DATACHAIN_CHECKPOINTS_RESET", reset_checkpoints)
+
+    if with_delta:
+        chain = dc.read_dataset(
+            "nums", delta=True, delta_on=["num"], session=test_session
+        )
+    else:
+        chain = dc.read_dataset("nums", session=test_session)
+
+    # -------------- FIRST RUN -------------------
+    first_job_id = catalog.metastore.create_job("my-job", "echo 1;")
+    monkeypatch.setenv("DATACHAIN_JOB_ID", first_job_id)
+    chain.save("nums1")
+    chain.save("nums2")
+    with pytest.raises(DataChainError):
+        chain.map(new=mapper_fail).save("nums3")
+
+    catalog.get_dataset("nums1")
+    catalog.get_dataset("nums2")
+    with pytest.raises(DatasetNotFoundError):
+        catalog.get_dataset("nums3")
+
+    # -------------- SECOND RUN -------------------
+    second_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=first_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", second_job_id)
+    chain.save("nums1")
+    chain.save("nums2")
+    chain.save("nums3")
+
+    assert len(catalog.get_dataset("nums1").versions) == 2 if reset_checkpoints else 1
+    assert len(catalog.get_dataset("nums2").versions) == 2 if reset_checkpoints else 1
+    assert len(catalog.get_dataset("nums3").versions) == 1
+
+    assert len(list(catalog.metastore.list_checkpoints(first_job_id))) == 2
+    assert len(list(catalog.metastore.list_checkpoints(second_job_id))) == 3
+
+
+@pytest.mark.parametrize("reset_checkpoints", [True, False])
+def test_checkpoints_modified_chains(
+    test_session, monkeypatch, nums_dataset, reset_checkpoints
+):
+    catalog = test_session.catalog
+    monkeypatch.setenv("DATACHAIN_CHECKPOINTS_RESET", reset_checkpoints)
+
+    chain = dc.read_dataset("nums", session=test_session)
+
+    # -------------- FIRST RUN -------------------
+    first_job_id = catalog.metastore.create_job("my-job", "echo 1;")
+    monkeypatch.setenv("DATACHAIN_JOB_ID", first_job_id)
+    chain.save("nums1")
+    chain.save("nums2")
+    chain.save("nums3")
+
+    # -------------- SECOND RUN -------------------
+    second_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=first_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", second_job_id)
+    chain.save("nums1")
+    chain.filter(dc.C("num") > 1).save("nums2")  # added change from first run
+    chain.save("nums3")
+
+    assert len(catalog.get_dataset("nums1").versions) == 2 if reset_checkpoints else 1
+    assert len(catalog.get_dataset("nums2").versions) == 2
+    assert len(catalog.get_dataset("nums3").versions) == 2
+
+    assert len(list(catalog.metastore.list_checkpoints(first_job_id))) == 3
+    assert len(list(catalog.metastore.list_checkpoints(second_job_id))) == 3
+
+
+@pytest.mark.parametrize("reset_checkpoints", [True, False])
+def test_checkpoints_multiple_runs(
+    test_session, monkeypatch, nums_dataset, reset_checkpoints
+):
+    catalog = test_session.catalog
+
+    monkeypatch.setenv("DATACHAIN_CHECKPOINTS_RESET", reset_checkpoints)
+
+    chain = dc.read_dataset("nums", session=test_session)
+
+    # -------------- FIRST RUN -------------------
+    first_job_id = catalog.metastore.create_job("my-job", "echo 1;")
+    monkeypatch.setenv("DATACHAIN_JOB_ID", first_job_id)
+    chain.save("nums1")
+    chain.save("nums2")
+    with pytest.raises(DataChainError):
+        chain.map(new=mapper_fail).save("nums3")
+
+    catalog.get_dataset("nums1")
+    catalog.get_dataset("nums2")
+    with pytest.raises(DatasetNotFoundError):
+        catalog.get_dataset("nums3")
+
+    # -------------- SECOND RUN -------------------
+    second_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=first_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", second_job_id)
+    chain.save("nums1")
+    chain.save("nums2")
+    chain.save("nums3")
+
+    # -------------- THIRD RUN -------------------
+    third_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=second_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", third_job_id)
+    chain.save("nums1")
+    chain.filter(dc.C("num") > 1).save("nums2")
+    with pytest.raises(DataChainError):
+        chain.map(new=mapper_fail).save("nums3")
+
+    # -------------- FOURTH RUN -------------------
+    fourth_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=third_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", fourth_job_id)
+    chain.save("nums1")
+    chain.filter(dc.C("num") > 1).save("nums2")
+    chain.save("nums3")
+
+    num1_versions = len(catalog.get_dataset("nums1").versions)
+    num2_versions = len(catalog.get_dataset("nums2").versions)
+    num3_versions = len(catalog.get_dataset("nums3").versions)
+
+    if reset_checkpoints:
+        assert num1_versions == 4
+        assert num2_versions == 4
+        assert num3_versions == 2
+
+    else:
+        assert num1_versions == 1
+        assert num2_versions == 2
+        assert num3_versions == 2
+
+    assert len(list(catalog.metastore.list_checkpoints(first_job_id))) == 2
+    assert len(list(catalog.metastore.list_checkpoints(second_job_id))) == 3
+    assert len(list(catalog.metastore.list_checkpoints(third_job_id))) == 2
+    assert len(list(catalog.metastore.list_checkpoints(fourth_job_id))) == 3
+
+
+def test_checkpoints_check_valid_chain_is_returned(
+    test_session,
+    monkeypatch,
+    nums_dataset,
+):
+    catalog = test_session.catalog
+
+    monkeypatch.setenv("DATACHAIN_CHECKPOINTS_RESET", False)
+    chain = dc.read_dataset("nums", session=test_session)
+
+    # -------------- FIRST RUN -------------------
+    first_job_id = catalog.metastore.create_job("my-job", "echo 1;")
+    monkeypatch.setenv("DATACHAIN_JOB_ID", first_job_id)
+    chain.save("nums1")
+
+    # -------------- SECOND RUN -------------------
+    second_job_id = catalog.metastore.create_job(
+        "my-job", "echo 1;", parent_job_id=first_job_id
+    )
+    monkeypatch.setenv("DATACHAIN_JOB_ID", second_job_id)
+    ds = chain.save("nums1")
+
+    # checking that we return expected DataChain even though we skipped chain creation
+    # because of the checkpoints
+    assert ds.dataset.name == "nums1"
+    assert len(ds.dataset.versions) == 1
+    assert ds.order_by("num").to_list("num") == [(1,), (2,), (3,)]
+
+
+def test_checkpoints_invalid_parent_job_id(test_session, monkeypatch, nums_dataset):
+    # setting wrong job id
+    monkeypatch.setenv("DATACHAIN_JOB_ID", "caee6c54-6328-4bcd-8ca6-2b31cb4fff94")
+    with pytest.raises(JobNotFoundError):
+        dc.read_dataset("nums", session=test_session).save("nums1")

{datachain-0.33.1 → datachain-0.34.0}/tests/unit/lib/test_datachain.py

@@ -4019,7 +4019,7 @@ def test_update_versions_wrong_value(test_session):
         chain.save(ds_name, update_version="wrong")
 
     assert str(excinfo.value) == (
-        "update_version can have one of the following values: major, minor or patch"
+        "update_version must be one of ['major', 'minor', 'patch']"
     )