Validate cloudpickle content hash

src/zenml/artifacts/utils.py

@@ -624,6 +624,7 @@ def load_artifact_from_response(artifact: "ArtifactVersionResponse") -> Any:
         data_type=artifact.data_type,
         uri=artifact.uri,
         artifact_store=artifact_store,
+        expected_content_hash=artifact.content_hash,
     )
 
 
@@ -807,6 +808,7 @@ def _load_artifact_from_uri(
     data_type: Union["Source", str],
     uri: str,
     artifact_store: Optional["BaseArtifactStore"] = None,
+    expected_content_hash: Optional[str] = None,
 ) -> Any:
     """Load an artifact using the given materializer.
 
@@ -815,6 +817,8 @@ def _load_artifact_from_uri(
         data_type: The source of the artifact data type.
         uri: The uri of the artifact.
         artifact_store: The artifact store used to store this artifact.
+        expected_content_hash: The content hash recorded for the artifact
+            version.
 
     Returns:
         The artifact loaded into memory.
@@ -863,6 +867,7 @@ def _load_artifact_from_uri(
     materializer_object: BaseMaterializer = materializer_class(
         uri, artifact_store
     )
+    materializer_object.expected_content_hash = expected_content_hash
     artifact = materializer_object.load(artifact_class)
     logger.debug("Artifact loaded successfully.")
 
@@ -1112,12 +1117,14 @@ def load_model_from_metadata(model_uri: str) -> Any:
     """
     # Load the model from its metadata
     artifact_versions_by_uri = Client().list_artifact_versions(uri=model_uri)
+    expected_content_hash: Optional[str] = None
     if artifact_versions_by_uri.total == 1:
         artifact_store = (
             _get_artifact_store_from_response_or_from_active_stack(
                 artifact_versions_by_uri.items[0]
             )
         )
+        expected_content_hash = artifact_versions_by_uri.items[0].content_hash
     else:
         artifact_store = Client().active_stack.artifact_store
 
@@ -1132,6 +1139,7 @@ def load_model_from_metadata(model_uri: str) -> Any:
         data_type=data_type,
         uri=model_uri,
         artifact_store=artifact_store,
+        expected_content_hash=expected_content_hash,
     )
 
     # Switch to eval mode if the model is a torch model

src/zenml/materializers/base_materializer.py

@@ -122,6 +122,10 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
 
     _DOCS_BUILDING_MODE: ClassVar[bool] = False
 
+    # Content hash recorded for the artifact version, set by the loading
+    # machinery before `load` so materializers can validate the stored data.
+    expected_content_hash: Optional[str] = None
+
     def __init__(
         self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
     ):

src/zenml/materializers/cloudpickle_materializer.py

@@ -13,8 +13,9 @@
 #  permissions and limitations under the License.
 """Implementation of ZenML's cloudpickle materializer."""
 
+import hashlib
 import os
-from typing import Any, ClassVar, Tuple, Type
+from typing import Any, ClassVar, Optional, Tuple, Type
 
 import cloudpickle
 
@@ -23,6 +24,7 @@
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils.io_utils import (
+    CallbackWriter,
     read_file_contents_as_string,
     write_file_contents_as_string,
 )
@@ -48,12 +50,18 @@ class CloudpickleMaterializer(BaseMaterializer):
     ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
     SKIP_REGISTRATION: ClassVar[bool] = True
 
+    _content_hash: Optional[str] = None
+
     def load(self, data_type: Type[Any]) -> Any:
         """Reads an artifact from a cloudpickle file.
 
         Args:
             data_type: The data type of the artifact.
 
+        Raises:
+            RuntimeError: If the stored data does not match the recorded
+                content hash.
+
         Returns:
             The loaded artifact data.
         """
@@ -72,8 +80,19 @@ def load(self, data_type: Type[Any]) -> Any:
         # load data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
         with self.artifact_store.open(filepath, "rb") as fid:
-            data = cloudpickle.load(fid)
-        return data
+            if self.expected_content_hash is None:
+                return cloudpickle.load(fid)
+
+            serialized_data = fid.read()
+
+        content_hash = hashlib.sha256(serialized_data).hexdigest()
+        if content_hash != self.expected_content_hash:
+            raise RuntimeError(
+                f"The artifact at '{self.uri}' does not match its recorded "
+                "content hash."
+            )
+
+        return cloudpickle.loads(serialized_data)
 
     def _load_python_version(self) -> str:
         """Loads the Python version that was used to materialize the artifact.
@@ -110,11 +129,25 @@ def save(self, data: Any) -> None:
 
         # save data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
+        digest = hashlib.sha256()
         with self.artifact_store.open(filepath, "wb") as fid:
-            cloudpickle.dump(data, fid)
+            cloudpickle.dump(data, CallbackWriter(fid, digest.update))
+        self._content_hash = digest.hexdigest()
 
     def _save_python_version(self) -> None:
         """Saves the Python version used to materialize the artifact."""
         filepath = os.path.join(self.uri, DEFAULT_PYTHON_VERSION_FILENAME)
         current_python_version = Environment().python_version()
         write_file_contents_as_string(filepath, current_python_version)
+
+    def compute_content_hash(self, data: Any) -> Optional[str]:
+        """Returns the content hash of the stored data.
+
+        Args:
+            data: The saved data.
+
+        Returns:
+            The content hash computed during `save`, or `None` if the data
+            was not saved by this instance.
+        """
+        return self._content_hash

src/zenml/orchestrators/step_runner.py

@@ -612,6 +612,7 @@ def _load_artifact(artifact_store: "BaseArtifactStore") -> Any:
             materializer: BaseMaterializer = materializer_class(
                 uri=artifact.uri, artifact_store=artifact_store
             )
+            materializer.expected_content_hash = artifact.content_hash
 
             if artifact.chunk_index is not None:
                 # We need to skip the type compatibility check here because

src/zenml/utils/io_utils.py

@@ -16,7 +16,7 @@
 import fnmatch
 import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Iterable
+from typing import IO, TYPE_CHECKING, Any, AnyStr, Callable, Generic, Iterable
 
 import click
 
@@ -38,6 +38,48 @@
     from zenml.io.filesystem import PathType
 
 
+class CallbackWriter(Generic[AnyStr]):
+    """Write proxy that invokes a callback with everything written through it."""
+
+    _fp: IO[AnyStr]
+    _callback: Callable[[AnyStr], Any]
+
+    def __init__(
+        self, fp: IO[AnyStr], callback: Callable[[AnyStr], Any]
+    ) -> None:
+        """Initializes the proxy.
+
+        Args:
+            fp: The underlying write handle.
+            callback: Callback invoked with every chunk passed to `write`.
+        """
+        self._fp = fp
+        self._callback = callback
+
+    def write(self, data: AnyStr) -> int:
+        """Invokes the callback with `data` and writes it.
+
+        Args:
+            data: The data to write.
+
+        Returns:
+            The number of characters or bytes written.
+        """
+        self._callback(data)
+        return self._fp.write(data)
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegates attribute access to the wrapped handle.
+
+        Args:
+            name: The attribute name.
+
+        Returns:
+            The attribute of the wrapped handle.
+        """
+        return getattr(self._fp, name)
+
+
 def is_root(path: str) -> bool:
     """Returns true if path has no parent in local filesystem.
 

tests/unit/materializers/test_cloudpickle_materializer.py

@@ -17,6 +17,9 @@
 import pickle
 from tempfile import TemporaryDirectory
 
+import cloudpickle
+import pytest
+
 from tests.unit.test_general import _test_materializer
 from zenml.environment import Environment
 from zenml.materializers.cloudpickle_materializer import (
@@ -69,3 +72,61 @@ def test_cloudpickle_materializer_can_load_pickle(clean_client):
         materializer = CloudpickleMaterializer(uri=artifact_uri)
         loaded_object = materializer.load(data_type=Unmaterializable)
         assert loaded_object.cat == "aria"
+
+
+def test_cloudpickle_materializer_load_without_content_hash(clean_client):
+    """Test that loading works if no content hash is recorded."""
+    my_object = Unmaterializable()
+    with TemporaryDirectory(
+        dir=clean_client.active_stack.artifact_store.path
+    ) as artifact_uri:
+        with open(os.path.join(artifact_uri, DEFAULT_FILENAME), "wb") as f:
+            pickle.dump(my_object, f)
+        materializer = CloudpickleMaterializer(uri=artifact_uri)
+        assert materializer.expected_content_hash is None
+        loaded_object = materializer.load(data_type=Unmaterializable)
+        assert loaded_object.cat == "aria"
+
+
+def test_cloudpickle_materializer_load_with_matching_content_hash(
+    clean_client,
+):
+    """Test that loading works if the stored data matches the recorded hash."""
+    my_object = Unmaterializable()
+    with TemporaryDirectory(
+        dir=clean_client.active_stack.artifact_store.path
+    ) as artifact_uri:
+        materializer = CloudpickleMaterializer(uri=artifact_uri)
+        assert materializer.compute_content_hash(my_object) is None
+        materializer.save(my_object)
+        content_hash = materializer.compute_content_hash(my_object)
+        assert content_hash is not None
+
+        loader = CloudpickleMaterializer(uri=artifact_uri)
+        loader.expected_content_hash = content_hash
+        loaded_object = loader.load(data_type=Unmaterializable)
+        assert loaded_object.cat == "aria"
+
+
+def test_cloudpickle_materializer_load_with_content_hash_mismatch(
+    clean_client,
+):
+    """Test that loading fails if the stored data does not match the hash."""
+    my_object = Unmaterializable()
+    with TemporaryDirectory(
+        dir=clean_client.active_stack.artifact_store.path
+    ) as artifact_uri:
+        materializer = CloudpickleMaterializer(uri=artifact_uri)
+        materializer.save(my_object)
+        content_hash = materializer.compute_content_hash(my_object)
+        assert content_hash is not None
+
+        with open(os.path.join(artifact_uri, DEFAULT_FILENAME), "wb") as f:
+            cloudpickle.dump({"other": "object"}, f)
+
+        loader = CloudpickleMaterializer(uri=artifact_uri)
+        loader.expected_content_hash = content_hash
+        with pytest.raises(
+            RuntimeError, match="does not match its recorded content hash"
+        ):
+            loader.load(data_type=Unmaterializable)

tests/unit/test_general.py

@@ -108,6 +108,11 @@ def _test_materializer(
             assert isinstance(key, str)
             assert isinstance(value, MetadataTypeTuple)
 
+        # Set the content hash on the materializer like the loading machinery
+        materializer.expected_content_hash = materializer.compute_content_hash(
+            step_output
+        )
+
         # Assert that materializer loads the data with the correct type
         loaded_data = materializer.load(step_output_type)
         if assert_data_type:

tests/unit/utils/test_io_utils.py

@@ -298,6 +298,19 @@ def test_move_moves_a_directory_from_source_to_destination(tmp_path) -> None:
     )
 
 
+def test_callback_writer_invokes_callback_for_every_write(tmp_path) -> None:
+    """Test that the callback writer invokes the callback for every write"""
+    written_chunks = []
+    file_path = os.path.join(tmp_path, "test.txt")
+    with open(file_path, "w") as f:
+        writer = io_utils.CallbackWriter(f, written_chunks.append)
+        writer.write("aria")
+        writer.write("_cat")
+
+    assert written_chunks == ["aria", "_cat"]
+    assert io_utils.read_file_contents_as_string(file_path) == "aria_cat"
+
+
 def test_get_grandparent_gets_the_grandparent_directory(tmp_path) -> None:
     """Test that get_grandparent gets the grandparent directory"""
     io_utils.create_dir_recursive_if_not_exists(