zenml-nightly 0.66.0.dev20240930__py3-none-any.whl → 0.68.0.dev20241026__py3-none-any.whl

zenml/artifacts/external_artifact_config.py

@@ -13,14 +13,13 @@
 #  permissions and limitations under the License.
 """External artifact definition."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 from uuid import UUID
 
 from pydantic import BaseModel, model_validator
 
 from zenml.logger import get_logger
-from zenml.model.model import Model
-from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -32,26 +31,23 @@ class ExternalArtifactConfiguration(BaseModel):
     """
 
     id: Optional[UUID] = None
-    name: Optional[str] = None
-    version: Optional[str] = None
-    model: Optional[Model] = None
 
-    @model_validator(mode="after")
-    def external_artifact_validator(self) -> "ExternalArtifactConfiguration":
-        """Model validator for the external artifact configuration.
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _remove_old_attributes(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Remove old attributes that are not used anymore.
 
-        Raises:
-            ValueError: if both version and model fields are set.
+        Args:
+            data: The model data.
 
         Returns:
-            the validated instance.
+            Model data without the removed attributes.
         """
-        if self.version and self.model:
-            raise ValueError(
-                "Cannot provide both `version` and `model` when "
-                "creating an external artifact."
-            )
-        return self
+        data.pop("name", None)
+        data.pop("version", None)
+        data.pop("model", None)
+        return data
 
     def get_artifact_version_id(self) -> UUID:
         """Get the artifact.
@@ -71,29 +67,16 @@ class ExternalArtifactConfiguration(BaseModel):
 
         if self.id:
             response = client.get_artifact_version(self.id)
-        elif self.name:
-            if self.version:
-                response = client.get_artifact_version(
-                    self.name, version=self.version
-                )
-            elif self.model:
-                response_ = self.model.get_artifact(self.name)
-                if not isinstance(response_, ArtifactVersionResponse):
-                    raise RuntimeError(
-                        f"Failed to pull artifact `{self.name}` from the Model "
-                        f"(name=`{self.model.name}`, version="
-                        f"`{self.model.version}`). Please validate the "
-                        "input and try again."
-                    )
-                response = response_
-            else:
-                response = client.get_artifact_version(self.name)
         else:
             raise RuntimeError(
-                "Either the ID or name of the artifact must be provided. "
-                "If you created this ExternalArtifact from a value, please "
+                "The ID of the artifact must be provided. "
+                "- If you created this ExternalArtifact from a value, please "
                 "ensure that `upload_by_value` was called before trying to "
-                "fetch the artifact ID."
+                "fetch the artifact ID.\n- If you specified an artifact name "
+                "or model name for this external artifact, this functionality "
+                "was removed from the ExternalArtifact class. Use Client "
+                "methods instead to dynamically fetch an artifact via name or "
+                "from a model instead."
             )
 
         artifact_store_id = client.active_stack.artifact_store.id
@@ -107,6 +90,4 @@ class ExternalArtifactConfiguration(BaseModel):
                 "store."
             )
 
-        self.id = response.id
-
         return self.id

zenml/artifacts/load_directory_materializer.py

@@ -0,0 +1,91 @@
+#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Only-load materializer for directories."""
+
+import os
+import tempfile
+from pathlib import Path
+from typing import Any, ClassVar, Tuple, Type
+
+from zenml.enums import ArtifactType
+from zenml.io import fileio
+from zenml.materializers.base_materializer import BaseMaterializer
+
+
+class PreexistingDataMaterializer(BaseMaterializer):
+    """Materializer to load directories from the artifact store.
+
+    This materializer is very special, since it do not implement save
+    logic at all. The save of the data to some URI inside the artifact store
+    shall happen outside and is in user's responsibility.
+
+    This materializer solely supports the `register_artifact` function.
+    """
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (Path,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+    SKIP_REGISTRATION: ClassVar[bool] = True
+
+    def load(self, data_type: Type[Any]) -> Any:
+        """Copy the artifact file(s) to a local temp directory.
+
+        Args:
+            data_type: Unused.
+
+        Returns:
+            Path to the local directory that contains the artifact files.
+        """
+        directory = tempfile.mkdtemp(prefix="zenml-artifact")
+        if fileio.isdir(self.uri):
+            self._copy_directory(src=self.uri, dst=directory)
+            return Path(directory)
+        else:
+            dst = os.path.join(directory, os.path.split(self.uri)[-1])
+            fileio.copy(src=self.uri, dst=dst)
+            return Path(dst)
+
+    def save(self, data: Any) -> None:
+        """Store the directory in the artifact store.
+
+        Args:
+            data: Path to a local directory to store.
+
+        Raises:
+            NotImplementedError: Always
+        """
+        raise NotImplementedError(
+            "`PreexistingDataMaterializer` can only be used in the "
+            "context of `register_artifact` function, "
+            "which expects the data to be already properly saved in "
+            "the Artifact Store, thus `save` logic makes no sense here."
+        )
+
+    @staticmethod
+    def _copy_directory(src: str, dst: str) -> None:
+        """Recursively copy a directory.
+
+        Args:
+            src: The directory to copy.
+            dst: Where to copy the directory to.
+        """
+        for src_dir, _, files in fileio.walk(src):
+            src_dir_ = str(src_dir)
+            dst_dir = str(os.path.join(dst, os.path.relpath(src_dir_, src)))
+            fileio.makedirs(dst_dir)
+
+            for file in files:
+                file_ = str(file)
+                src_file = os.path.join(src_dir_, file_)
+                dst_file = os.path.join(dst_dir, file_)
+                fileio.copy(src_file, dst_file)

zenml/artifacts/utils.py

@@ -19,15 +19,30 @@ import tempfile
 import time
 import zipfile
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Type,
+    Union,
+    cast,
+)
 from uuid import UUID, uuid4
 
+from zenml.artifacts.artifact_config import ArtifactConfig
+from zenml.artifacts.load_directory_materializer import (
+    PreexistingDataMaterializer,
+)
 from zenml.client import Client
 from zenml.constants import (
     MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION,
     MODEL_METADATA_YAML_FILE_NAME,
 )
 from zenml.enums import (
+    ArtifactType,
     ExecutionStatus,
     MetadataResourceTypes,
     StackComponentType,
@@ -50,8 +65,9 @@ from zenml.models import (
     StepRunResponse,
     StepRunUpdate,
 )
-from zenml.new.steps.step_context import get_step_context
+from zenml.models.v2.core.artifact import ArtifactResponse
 from zenml.stack import StackComponent
+from zenml.steps.step_context import get_step_context
 from zenml.utils import source_utils
 from zenml.utils.yaml_utils import read_yaml, write_yaml
 
@@ -111,10 +127,6 @@ def save_artifact(
 
     Returns:
         The saved artifact response.
-
-    Raises:
-        RuntimeError: If artifact URI already exists.
-        EntityExistsError: If artifact version already exists.
     """
     from zenml.materializers.materializer_registry import (
         materializer_registry,
@@ -123,24 +135,11 @@ def save_artifact(
 
     client = Client()
 
-    # Get or create the artifact
-    try:
-        artifact = client.list_artifacts(name=name)[0]
-        if artifact.has_custom_name != has_custom_name:
-            client.update_artifact(
-                name_id_or_prefix=artifact.id, has_custom_name=has_custom_name
-            )
-    except IndexError:
-        try:
-            artifact = client.zen_store.create_artifact(
-                ArtifactRequest(
-                    name=name,
-                    has_custom_name=has_custom_name,
-                    tags=tags,
-                )
-            )
-        except EntityExistsError:
-            artifact = client.list_artifacts(name=name)[0]
+    artifact = _get_or_create_artifact(
+        name=name,
+        has_custom_name=has_custom_name,
+        tags=tags,
+    )
 
     # Get the current artifact store
     artifact_store = client.active_stack.artifact_store
@@ -151,16 +150,14 @@ def save_artifact(
     if not uri.startswith(artifact_store.path):
         uri = os.path.join(artifact_store.path, uri)
 
-    if manual_save and artifact_store.exists(uri):
+    if manual_save:
         # This check is only necessary for manual saves as we already check
         # it when creating the directory for step output artifacts
-        other_artifacts = client.list_artifact_versions(uri=uri, size=1)
-        if other_artifacts and (other_artifact := other_artifacts[0]):
-            raise RuntimeError(
-                f"Cannot save new artifact {name} version to URI "
-                f"{uri} because the URI is already used by artifact "
-                f"{other_artifact.name} (version {other_artifact.version})."
-            )
+        _check_if_artifact_with_given_uri_already_registered(
+            artifact_store=artifact_store,
+            uri=uri,
+            name=name,
+        )
     artifact_store.makedirs(uri)
 
     # Find and initialize the right materializer class
@@ -211,7 +208,9 @@ def save_artifact(
             )
 
     # Create the artifact version
-    def _create_version() -> Optional[ArtifactVersionResponse]:
+    def _create_version(
+        version: Union[int, str],
+    ) -> Optional[ArtifactVersionResponse]:
         artifact_version = ArtifactVersionRequest(
             artifact_id=artifact.id,
             version=version,
@@ -233,37 +232,12 @@ def save_artifact(
         except EntityExistsError:
             return None
 
-    response = None
-    if not version:
-        retries_made = 0
-        for i in range(MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION):
-            # Get new artifact version
-            version = _get_new_artifact_version(name)
-            if response := _create_version():
-                break
-            # smoothed exponential back-off, it will go as 0.2, 0.3,
-            # 0.45, 0.68, 1.01, 1.52, 2.28, 3.42, 5.13, 7.69, ...
-            sleep = 0.2 * 1.5**i
-            logger.debug(
-                f"Failed to create artifact version `{version}` for "
-                f"artifact `{name}`. Retrying in {sleep}..."
-            )
-            time.sleep(sleep)
-            retries_made += 1
-        if not response:
-            raise EntityExistsError(
-                f"Failed to create new artifact version for artifact "
-                f"`{name}`. Retried {retries_made} times. "
-                "This could be driven by exceptionally high concurrency of "
-                "pipeline runs. Please, reach out to us on ZenML Slack for support."
-            )
-    else:
-        response = _create_version()
-        if not response:
-            raise EntityExistsError(
-                f"Failed to create artifact version `{version}` for artifact "
-                f"`{name}`. Given version already exists."
-            )
+    response = _create_artifact_version_with_retries(
+        name=name,
+        version=version,
+        create_version_fn=_create_version,
+    )
+
     if artifact_metadata:
         client.create_run_metadata(
             metadata=artifact_metadata,
@@ -272,29 +246,112 @@ def save_artifact(
         )
 
     if manual_save:
+        _link_artifact_version_to_the_step_and_model(
+            artifact_version=response,
+            is_model_artifact=is_model_artifact,
+            is_deployment_artifact=is_deployment_artifact,
+        )
+
+    return response
+
+
+def register_artifact(
+    folder_or_file_uri: str,
+    name: str,
+    version: Optional[Union[int, str]] = None,
+    tags: Optional[List[str]] = None,
+    has_custom_name: bool = True,
+    is_model_artifact: bool = False,
+    is_deployment_artifact: bool = False,
+    artifact_metadata: Dict[str, "MetadataType"] = {},
+) -> "ArtifactVersionResponse":
+    """Register existing data stored in the artifact store as a ZenML Artifact.
+
+    Args:
+        folder_or_file_uri: The full URI within the artifact store to the folder
+            or to the file.
+        name: The name of the artifact.
+        version: The version of the artifact. If not provided, a new
+            auto-incremented version will be used.
+        tags: Tags to associate with the artifact.
+        has_custom_name: If the artifact name is custom and should be listed in
+            the dashboard "Artifacts" tab.
+        is_model_artifact: If the artifact is a model artifact.
+        is_deployment_artifact: If the artifact is a deployment artifact.
+        artifact_metadata: Metadata dictionary to attach to the artifact version.
+
+    Returns:
+        The saved artifact response.
+
+    Raises:
+        FileNotFoundError: If the folder URI is outside of the artifact store
+            bounds.
+    """
+    client = Client()
+
+    # Get the current artifact store
+    artifact_store = client.active_stack.artifact_store
+
+    if not folder_or_file_uri.startswith(artifact_store.path):
+        raise FileNotFoundError(
+            f"Folder `{folder_or_file_uri}` is outside of "
+            f"artifact store bounds `{artifact_store.path}`"
+        )
+
+    _check_if_artifact_with_given_uri_already_registered(
+        artifact_store=artifact_store,
+        uri=folder_or_file_uri,
+        name=name,
+    )
+
+    artifact = _get_or_create_artifact(
+        name=name,
+        has_custom_name=has_custom_name,
+        tags=tags,
+    )
+
+    # Create the artifact version
+    def _create_version(
+        version: Union[int, str],
+    ) -> Optional[ArtifactVersionResponse]:
+        artifact_version = ArtifactVersionRequest(
+            artifact_id=artifact.id,
+            version=version,
+            tags=tags,
+            type=ArtifactType.DATA,
+            uri=folder_or_file_uri,
+            materializer=source_utils.resolve(PreexistingDataMaterializer),
+            data_type=source_utils.resolve(Path),
+            user=Client().active_user.id,
+            workspace=Client().active_workspace.id,
+            artifact_store_id=artifact_store.id,
+            has_custom_name=has_custom_name,
+        )
         try:
-            error_message = "step run"
-            step_context = get_step_context()
-            step_run = step_context.step_run
-            client.zen_store.update_run_step(
-                step_run_id=step_run.id,
-                step_run_update=StepRunUpdate(
-                    saved_artifact_versions={name: response.id}
-                ),
+            return client.zen_store.create_artifact_version(
+                artifact_version=artifact_version
             )
-            error_message = "model"
-            model = step_context.model
-            if model:
-                from zenml.model.utils import link_artifact_to_model
-
-                link_artifact_to_model(
-                    artifact_version_id=response.id,
-                    model=model,
-                    is_model_artifact=is_model_artifact,
-                    is_deployment_artifact=is_deployment_artifact,
-                )
-        except (RuntimeError, StepContextError):
-            logger.debug(f"Unable to link saved artifact to {error_message}.")
+        except EntityExistsError:
+            return None
+
+    response = _create_artifact_version_with_retries(
+        name=name,
+        version=version,
+        create_version_fn=_create_version,
+    )
+
+    if artifact_metadata:
+        client.create_run_metadata(
+            metadata=artifact_metadata,
+            resource_id=response.id,
+            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+        )
+
+    _link_artifact_version_to_the_step_and_model(
+        artifact_version=response,
+        is_model_artifact=is_model_artifact,
+        is_deployment_artifact=is_deployment_artifact,
+    )
 
     return response
 
@@ -577,6 +634,184 @@ def get_artifacts_versions_of_pipeline_run(
 # -------------------------
 
 
+def _check_if_artifact_with_given_uri_already_registered(
+    artifact_store: "BaseArtifactStore",
+    uri: str,
+    name: str,
+) -> None:
+    """Check if the given artifact store already contains an artifact with the given URI.
+
+    Args:
+        artifact_store: The artifact store to check.
+        uri: The uri of the artifact.
+        name: The name of the artifact.
+
+    Raises:
+        RuntimeError: If the artifact store already contains an artifact with
+            the given URI.
+    """
+    if artifact_store.exists(uri):
+        # This check is only necessary for manual saves as we already check
+        # it when creating the directory for step output artifacts
+        other_artifacts = Client().list_artifact_versions(uri=uri, size=1)
+        if other_artifacts and (other_artifact := other_artifacts[0]):
+            raise RuntimeError(
+                f"Cannot create new artifact {name} version with URI "
+                f"{uri} because the URI is already used by artifact "
+                f"{other_artifact.name} (version {other_artifact.version})."
+            )
+
+
+def _get_or_create_artifact(
+    name: str, has_custom_name: bool, tags: Optional[List[str]] = None
+) -> ArtifactResponse:
+    """Get or create an artifact with the given name.
+
+    Args:
+        name: The name of the artifact.
+        has_custom_name: If the artifact name is custom and should be listed in
+            the dashboard "Artifacts" tab.
+        tags: Tags to associate with the artifact.
+
+    Returns:
+        The artifact.
+    """
+    client = Client()
+    # Get or create the artifact
+    try:
+        artifact = client.list_artifacts(name=name)[0]
+        if artifact.has_custom_name != has_custom_name:
+            client.update_artifact(
+                name_id_or_prefix=artifact.id, has_custom_name=has_custom_name
+            )
+    except IndexError:
+        try:
+            artifact = client.zen_store.create_artifact(
+                ArtifactRequest(
+                    name=name,
+                    has_custom_name=has_custom_name,
+                    tags=tags,
+                )
+            )
+        except EntityExistsError:
+            artifact = client.list_artifacts(name=name)[0]
+    return artifact
+
+
+def _create_artifact_version_with_retries(
+    name: str,
+    version: Optional[Union[int, str]],
+    create_version_fn: Callable[
+        [
+            Union[int, str],
+        ],
+        Optional[ArtifactVersionResponse],
+    ],
+) -> ArtifactVersionResponse:
+    """Create an artifact version with some retries.
+
+    This function will retry the creation of an artifact version up to
+    MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION times if it fails.
+    It can fail in high-concurrency environments.
+
+    Args:
+        name: The name of the artifact.
+        version: The version of the artifact. If not provided, a new
+            auto-incremented version will be used.
+        create_version_fn: The function to create the artifact version.
+
+    Returns:
+        The created artifact version.
+
+    Raises:
+        EntityExistsError: If the artifact version could not be created
+            after MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION attempts due
+            to collisions.
+
+    """
+    response = None
+    if not version:
+        retries_made = 0
+        for i in range(MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION):
+            # Get new artifact version
+            version = _get_new_artifact_version(name)
+            if response := create_version_fn(version):
+                break
+            # smoothed exponential back-off, it will go as 0.2, 0.3,
+            # 0.45, 0.68, 1.01, 1.52, 2.28, 3.42, 5.13, 7.69, ...
+            sleep = 0.2 * 1.5**i
+            logger.debug(
+                f"Failed to create artifact version `{version}` for "
+                f"artifact `{name}`. Retrying in {sleep}..."
+            )
+            time.sleep(sleep)
+            retries_made += 1
+        if not response:
+            raise EntityExistsError(
+                f"Failed to create new artifact version for artifact "
+                f"`{name}`. Retried {retries_made} times. "
+                "This could be driven by exceptionally high concurrency of "
+                "pipeline runs. Please, reach out to us on ZenML Slack for support."
+            )
+    else:
+        response = create_version_fn(version)
+        if not response:
+            raise EntityExistsError(
+                f"Failed to create artifact version `{version}` for artifact "
+                f"`{name}`. Given version already exists."
+            )
+    return response
+
+
+def _link_artifact_version_to_the_step_and_model(
+    artifact_version: ArtifactVersionResponse,
+    is_model_artifact: bool,
+    is_deployment_artifact: bool,
+) -> None:
+    """Link an artifact version to the step and its' context model.
+
+    This function links the AV to:
+        - the step run
+        - the MV from the step context
+
+    Args:
+        artifact_version: The artifact version to link.
+        is_model_artifact: Whether the artifact is a model artifact.
+        is_deployment_artifact: Whether the artifact is a deployment artifact.
+    """
+    client = Client()
+    try:
+        error_message = "step run"
+        step_context = get_step_context()
+        step_run = step_context.step_run
+        client.zen_store.update_run_step(
+            step_run_id=step_run.id,
+            step_run_update=StepRunUpdate(
+                saved_artifact_versions={
+                    artifact_version.artifact.name: artifact_version.id
+                }
+            ),
+        )
+        error_message = "model"
+
+        if step_context.model_version:
+            from zenml.model.utils import (
+                link_artifact_version_to_model_version,
+            )
+
+            artifact_config = ArtifactConfig(
+                is_model_artifact=is_model_artifact,
+                is_deployment_artifact=is_deployment_artifact,
+            )
+            link_artifact_version_to_model_version(
+                artifact_version=artifact_version,
+                model_version=step_context.model_version,
+                artifact_config=artifact_config,
+            )
+    except (RuntimeError, StepContextError):
+        logger.debug(f"Unable to link saved artifact to {error_message}.")
+
+
 def _load_artifact_from_uri(
     materializer: Union["Source", str],
     data_type: Union["Source", str],