zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/artifact_stores/base_artifact_store.py

@@ -12,8 +12,12 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """The base interface to extend the ZenML artifact store."""
+
+import inspect
+import os
 import textwrap
 from abc import abstractmethod
+from pathlib import Path
 from typing import (
     Any,
     Callable,
@@ -29,64 +33,109 @@ from typing import (
     cast,
 )
 
-from pydantic import root_validator
+from pydantic import model_validator
 
+from zenml.constants import ENV_ZENML_SERVER
 from zenml.enums import StackComponentType
 from zenml.exceptions import ArtifactStoreInterfaceError
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent, StackComponentConfig
 from zenml.utils import io_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
 PathType = Union[bytes, str]
 
 
-def _sanitize_potential_path(potential_path: Any) -> Any:
-    """Sanitizes the input if it is a path.
+class _sanitize_paths:
+    """Sanitizes path inputs before calling the original function.
 
-    If the input is a **remote** path, this function replaces backslash path
-    separators by forward slashes.
+    Extra decoration layer is needed to pass in fixed artifact store root
+    path for static methods that are called on filesystems directly.
 
     Args:
-        potential_path: Value that potentially refers to a (remote) path.
+        func: The function to decorate.
+        fixed_root_path: The fixed artifact store root path.
+        is_static: Whether the function is static or not.
 
     Returns:
-        The original input or a sanitized version of it in case of a remote
-        path.
+        Function that calls the input function with sanitized path inputs.
     """
-    if isinstance(potential_path, bytes):
-        path = fileio.convert_to_str(potential_path)
-    elif isinstance(potential_path, str):
-        path = potential_path
-    else:
-        # Neither string nor bytes, this is not a path
-        return potential_path
 
-    if io_utils.is_remote(path):
-        # If we have a remote path, replace windows path separators with
-        # slashes
-        import ntpath
-        import posixpath
+    def __init__(self, func: Callable[..., Any], fixed_root_path: str) -> None:
+        """Initializes the decorator.
 
-        path = path.replace(ntpath.sep, posixpath.sep)
+        Args:
+            func: The function to decorate.
+            fixed_root_path: The fixed artifact store root path.
+        """
+        self.func = func
+        self.fixed_root_path = fixed_root_path
 
-    return path
+        self.path_args: List[int] = []
+        self.path_kwargs: List[str] = []
+        for i, param in enumerate(
+            inspect.signature(self.func).parameters.values()
+        ):
+            if param.annotation == PathType:
+                self.path_kwargs.append(param.name)
+                if param.default == inspect.Parameter.empty:
+                    self.path_args.append(i)
 
+    def _validate_path(self, path: str) -> None:
+        """Validates a path.
 
-def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
-    """Sanitizes path inputs before calling the original function.
+        Args:
+            path: The path to validate.
 
-    Args:
-        _func: The function for which to sanitize the inputs.
+        Raises:
+            FileNotFoundError: If the path is outside of the artifact store
+                bounds.
+        """
+        if not path.startswith(self.fixed_root_path):
+            raise FileNotFoundError(
+                f"File `{path}` is outside of "
+                f"artifact store bounds `{self.fixed_root_path}`"
+            )
 
-    Returns:
-        Function that calls the input function with sanitized path inputs.
-    """
+    def _sanitize_potential_path(self, potential_path: Any) -> Any:
+        """Sanitizes the input if it is a path.
 
-    def inner_function(*args: Any, **kwargs: Any) -> Any:
-        """Inner function.
+        If the input is a **remote** path, this function replaces backslash path
+        separators by forward slashes.
+
+        Args:
+            potential_path: Value that potentially refers to a (remote) path.
+
+        Returns:
+            The original input or a sanitized version of it in case of a remote
+            path.
+        """
+        if isinstance(potential_path, bytes):
+            path = fileio.convert_to_str(potential_path)
+        elif isinstance(potential_path, str):
+            path = potential_path
+        else:
+            # Neither string nor bytes, this is not a path
+            return potential_path
+
+        if io_utils.is_remote(path):
+            # If we have a remote path, replace windows path separators with
+            # slashes
+            import ntpath
+            import posixpath
+
+            path = path.replace(ntpath.sep, posixpath.sep)
+            self._validate_path(path)
+        else:
+            self._validate_path(str(Path(path).absolute().resolve()))
+
+        return path
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Decorator function that sanitizes paths before calling the original function.
 
         Args:
             *args: Positional args.
@@ -95,15 +144,28 @@ def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
         Returns:
             Output of the input function called with sanitized paths.
         """
-        args = tuple(_sanitize_potential_path(arg) for arg in args)
+        # verify if `self` is part of the args
+        has_self = bool(args and isinstance(args[0], BaseArtifactStore))
+
+        # sanitize inputs for relevant args and kwargs, keep rest unchanged
+        args = tuple(
+            self._sanitize_potential_path(
+                arg,
+            )
+            if i + has_self in self.path_args
+            else arg
+            for i, arg in enumerate(args)
+        )
         kwargs = {
-            key: _sanitize_potential_path(value)
+            key: self._sanitize_potential_path(
+                value,
+            )
+            if key in self.path_kwargs
+            else value
             for key, value in kwargs.items()
         }
 
-        return _func(*args, **kwargs)
-
-    return inner_function
+        return self.func(*args, **kwargs)
 
 
 class BaseArtifactStoreConfig(StackComponentConfig):
@@ -112,16 +174,19 @@ class BaseArtifactStoreConfig(StackComponentConfig):
     path: str
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]]
+    IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
 
-    @root_validator(skip_on_failure=True)
-    def _ensure_artifact_store(cls, values: Dict[str, Any]) -> Any:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _ensure_artifact_store(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validator function for the Artifact Stores.
 
         Checks whether supported schemes are defined and the given path is
         supported.
 
         Args:
-            values: The values to validate.
+            data: the input data to construct the artifact store.
 
         Returns:
             The validated values.
@@ -151,18 +216,20 @@ class BaseArtifactStoreConfig(StackComponentConfig):
                     """
                 )
             )
-        values["path"] = values["path"].strip("'\"`")
-        if not any(
-            values["path"].startswith(i) for i in cls.SUPPORTED_SCHEMES
-        ):
-            raise ArtifactStoreInterfaceError(
-                f"The path: '{values['path']}' you defined for your "
-                f"artifact store is not supported by the implementation of "
-                f"{cls.schema()['title']}, because it does not start with "
-                f"one of its supported schemes: {cls.SUPPORTED_SCHEMES}."
-            )
 
-        return values
+        if "path" in data:
+            data["path"] = data["path"].strip("'\"`")
+            if not any(
+                data["path"].startswith(i) for i in cls.SUPPORTED_SCHEMES
+            ):
+                raise ArtifactStoreInterfaceError(
+                    f"The path: '{data['path']}' you defined for your "
+                    f"artifact store is not supported by the implementation of "
+                    f"{cls.schema()['title']}, because it does not start with "
+                    f"one of its supported schemes: {cls.SUPPORTED_SCHEMES}."
+                )
+
+        return data
 
 
 class BaseArtifactStore(StackComponent):
@@ -200,11 +267,11 @@ class BaseArtifactStore(StackComponent):
 
     # --- User interface ---
     @abstractmethod
-    def open(self, name: PathType, mode: str = "r") -> Any:
+    def open(self, path: PathType, mode: str = "r") -> Any:
         """Open a file at the given path.
 
         Args:
-            name: The path of the file to open.
+            path: The path of the file to open.
             mode: The mode to open the file.
 
         Returns:
@@ -322,6 +389,7 @@ class BaseArtifactStore(StackComponent):
             The stat descriptor.
         """
 
+    @abstractmethod
     def size(self, path: PathType) -> Optional[int]:
         """Get the size of a file in bytes.
 
@@ -367,7 +435,11 @@ class BaseArtifactStore(StackComponent):
             **kwargs: The keyword arguments to pass to the Pydantic object.
         """
         super(BaseArtifactStore, self).__init__(*args, **kwargs)
-        self._register()
+
+        # If running in a ZenML server environment, we don't register
+        # the filesystems. We always use the artifact stores directly.
+        if ENV_ZENML_SERVER not in os.environ:
+            self._register()
 
     def _register(self) -> None:
         """Create and register a filesystem within the filesystem registry."""
@@ -375,34 +447,45 @@ class BaseArtifactStore(StackComponent):
         from zenml.io.filesystem_registry import default_filesystem_registry
         from zenml.io.local_filesystem import LocalFilesystem
 
+        overloads: Dict[str, Any] = {
+            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
+        }
+        for abc_method in inspect.getmembers(BaseArtifactStore):
+            if getattr(abc_method[1], "__isabstractmethod__", False):
+                sanitized_method = _sanitize_paths(
+                    getattr(self, abc_method[0]), self.path
+                )
+                # prepare overloads for filesystem methods
+                overloads[abc_method[0]] = staticmethod(sanitized_method)
+
+                # decorate artifact store methods
+                setattr(
+                    self,
+                    abc_method[0],
+                    sanitized_method,
+                )
+
         # Local filesystem is always registered, no point in doing it again.
         if isinstance(self, LocalFilesystem):
             return
 
         filesystem_class = type(
-            self.__class__.__name__,
-            (BaseFilesystem,),
-            {
-                "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
-                "open": staticmethod(_sanitize_paths(self.open)),
-                "copyfile": staticmethod(_sanitize_paths(self.copyfile)),
-                "exists": staticmethod(_sanitize_paths(self.exists)),
-                "glob": staticmethod(_sanitize_paths(self.glob)),
-                "isdir": staticmethod(_sanitize_paths(self.isdir)),
-                "listdir": staticmethod(_sanitize_paths(self.listdir)),
-                "makedirs": staticmethod(_sanitize_paths(self.makedirs)),
-                "mkdir": staticmethod(_sanitize_paths(self.mkdir)),
-                "remove": staticmethod(_sanitize_paths(self.remove)),
-                "rename": staticmethod(_sanitize_paths(self.rename)),
-                "rmtree": staticmethod(_sanitize_paths(self.rmtree)),
-                "size": staticmethod(_sanitize_paths(self.size)),
-                "stat": staticmethod(_sanitize_paths(self.stat)),
-                "walk": staticmethod(_sanitize_paths(self.walk)),
-            },
+            self.__class__.__name__, (BaseFilesystem,), overloads
         )
 
         default_filesystem_registry.register(filesystem_class)
 
+    def _remove_previous_file_versions(self, path: PathType) -> None:
+        """Remove all file versions but the latest in the given path.
+
+        Method is useful for logs stored in versioned file systems
+        like AWS S3.
+
+        Args:
+            path: The path to the file.
+        """
+        return
+
 
 class BaseArtifactStoreFlavor(Flavor):
     """Base class for artifact store flavors."""

zenml/artifact_stores/local_artifact_store.py

@@ -21,7 +21,7 @@ artifacts get stored.
 import os
 from typing import TYPE_CHECKING, ClassVar, Optional, Set, Type, Union
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.artifact_stores import (
     BaseArtifactStore,
@@ -50,7 +50,8 @@ class LocalArtifactStoreConfig(BaseArtifactStoreConfig):
 
     path: str = ""
 
-    @validator("path")
+    @field_validator("path")
+    @classmethod
     def ensure_path_local(cls, path: str) -> str:
         """Pydantic validator which ensures that the given path is a local path.
 

zenml/artifacts/artifact_config.py

@@ -12,18 +12,15 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Artifact Config classes to support Model Control Plane feature."""
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, root_validator
+from typing import Any, Dict, List, Optional, Union
 
-from zenml.enums import ModelStages
-from zenml.exceptions import StepContextError
-from zenml.logger import get_logger
-from zenml.new.steps.step_context import get_step_context
-
-if TYPE_CHECKING:
-    from zenml.model.model import Model
+from pydantic import BaseModel, Field, model_validator
 
+from zenml.enums import ArtifactType
+from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -40,74 +37,86 @@ class ArtifactConfig(BaseModel):
         int, ArtifactConfig(
             name="my_artifact",  # override the default artifact name
             version=42,  # set a custom version
+            artifact_type=ArtifactType.MODEL,  # Specify the artifact type
             tags=["tag1", "tag2"],  # set custom tags
-            model_name="my_model",  # link the artifact to a model
         )
     ]:
         return ...
     ```
 
     Attributes:
-        name: The name of the artifact.
+        name: The name of the artifact:
+            - static string e.g. "name"
+            - dynamic string e.g. "name_{date}_{time}_{custom_placeholder}"
+            If you use any placeholders besides `date` and `time`,
+            you need to provide the values for them in the `substitutions`
+            argument of the step decorator or the `substitutions` argument
+            of `with_options` of the step.
         version: The version of the artifact.
         tags: The tags of the artifact.
-        model_name: The name of the model to link artifact to.
-        model_version: The identifier of a version of the model to link the artifact
-            to. It can be an exact version ("my_version"), exact version number
-            (42), stage (ModelStages.PRODUCTION or "production"), or
-            (ModelStages.LATEST or None) for the latest version (default).
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
+        run_metadata: Metadata to add to the artifact.
+        artifact_type: Optional type of the artifact. If not given, the type
+            specified by the materializer that is used to save this artifact
+            is used.
     """
 
     name: Optional[str] = None
-    version: Optional[Union[str, int]] = None
+    version: Optional[Union[str, int]] = Field(
+        default=None, union_mode="smart"
+    )
     tags: Optional[List[str]] = None
+    run_metadata: Optional[Dict[str, MetadataType]] = None
 
-    model_name: Optional[str] = None
-    model_version: Optional[Union[ModelStages, str, int]] = None
-    is_model_artifact: bool = False
-    is_deployment_artifact: bool = False
+    artifact_type: Optional[ArtifactType] = None
 
-    @root_validator
-    def _root_validator(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        model_name = values.get("model_name", None)
-        if model_name and values.get("model_version", None) is None:
-            raise ValueError(
-                f"Creation of new model version from `{cls}` is not allowed. "
-                "Please either keep `model_name` and `model_version` both "
-                "`None` to get the model version from the step context or "
-                "specify both at the same time. You can use `ModelStages.LATEST` "
-                "as `model_version` when latest model version is desired."
-            )
-        return values
-
-    class Config:
-        """Config class for ArtifactConfig."""
+    @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.
 
-        smart_union = True
+        Args:
+            data: The model data.
 
-    @property
-    def _model(self) -> Optional["Model"]:
-        """The model linked to this artifact.
+        Raises:
+            ValueError: If the artifact is configured to be
+                both a model and a deployment artifact.
 
         Returns:
-            The model or None if the model version cannot be determined.
+            Model data without the removed attributes.
         """
-        try:
-            model_ = get_step_context().model
-        except (StepContextError, RuntimeError):
-            model_ = None
-        # Check if another model name was specified
-        if (self.model_name is not None) and (
-            model_ is None or model_.name != self.model_name
-        ):
-            # Create a new Model instance with the provided model name and version
-            from zenml.model.model import Model
-
-            on_the_fly_config = Model(
-                name=self.model_name, version=self.model_version
+        model_name = data.pop("model_name", None)
+        model_version = data.pop("model_version", None)
+
+        if model_name or model_version:
+            logger.warning(
+                "Specifying a model name or version for a step output "
+                "artifact is not supported anymore."
+            )
+
+        is_model_artifact = data.pop("is_model_artifact", None)
+        is_deployment_artifact = data.pop("is_deployment_artifact", None)
+
+        if is_model_artifact and is_deployment_artifact:
+            raise ValueError(
+                "An artifact can only be a model artifact or deployment "
+                "artifact."
+            )
+        elif is_model_artifact:
+            logger.warning(
+                "`ArtifactConfig(..., is_model_artifact=True)` is deprecated "
+                "and will be removed soon. Use `ArtifactConfig(..., "
+                "artifact_type=ArtifactType.MODEL)` instead. For more info: "
+                "https://docs.zenml.io/user-guide/starter-guide/manage-artifacts"
+            )
+            data.setdefault("artifact_type", ArtifactType.MODEL)
+        elif is_deployment_artifact:
+            logger.warning(
+                "`ArtifactConfig(..., is_deployment_artifact=True)` is "
+                "deprecated and will be removed soon. Use `ArtifactConfig(..., "
+                "artifact_type=ArtifactType.SERVICE)` instead. For more info: "
+                "https://docs.zenml.io/user-guide/starter-guide/manage-artifacts"
             )
-            return on_the_fly_config
+            data.setdefault("artifact_type", ArtifactType.SERVICE)
 
-        return model_
+        return data

zenml/artifacts/external_artifact.py

@@ -12,16 +12,18 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """External artifact definition."""
+
 import os
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Optional, Type, Union
 from uuid import UUID, uuid4
 
-from pydantic import root_validator
+from pydantic import Field, model_validator
 
 from zenml.artifacts.external_artifact_config import (
     ExternalArtifactConfiguration,
 )
 from zenml.config.source import Source
+from zenml.enums import ArtifactSaveType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 
@@ -39,26 +41,11 @@ class ExternalArtifact(ExternalArtifactConfiguration):
     can be used to provide any value as input to a step without needing to
     write an additional step that returns this value.
 
-    The external artifact needs to have either a value associated with it
-    that will be uploaded to the artifact store, or reference an artifact
-    that is already registered in ZenML.
-
-    There are several ways to reference an existing artifact:
-    - By providing an artifact ID.
-    - By providing an artifact name and version. If no version is provided,
-        the latest version of that artifact will be used.
+    The external artifact needs to have a value associated with it
+    that will be uploaded to the artifact store.
 
     Args:
         value: The artifact value.
-        id: The ID of an artifact that should be referenced by this external
-            artifact.
-        name: Name of an artifact to search. If none of
-            `version`, `pipeline_run_name`, or `pipeline_name` are set, the
-            latest version of the artifact will be used.
-        version: Version of the artifact to search. Only used when `name` is
-            provided. Cannot be used together with `model`.
-        model: The model to search in. Only used when `name`
-            is provided. Cannot be used together with `version`.
         materializer: The materializer to use for saving the artifact value
             to the artifact store. Only used when `value` is provided.
         store_artifact_metadata: Whether metadata for the artifact should
@@ -85,27 +72,28 @@ class ExternalArtifact(ExternalArtifactConfiguration):
     """
 
     value: Optional[Any] = None
-    materializer: Optional[MaterializerClassOrSource] = None
+    materializer: Optional[MaterializerClassOrSource] = Field(
+        default=None, union_mode="left_to_right"
+    )
     store_artifact_metadata: bool = True
     store_artifact_visualizations: bool = True
 
-    @root_validator
-    def _validate_all(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        options = [
-            values.get(field, None) is not None
-            for field in ["value", "id", "name"]
-        ]
-        if sum(options) > 1:
-            raise ValueError(
-                "Only one of `value`, `id`, or `name` can be provided when "
-                "creating an external artifact."
-            )
-        elif sum(options) == 0:
+    @model_validator(mode="after")
+    def external_artifact_validator(self) -> "ExternalArtifact":
+        """Model validator for the external artifact.
+
+        Raises:
+            ValueError: If an ID was set.
+
+        Returns:
+            The validated instance.
+        """
+        if self.id:
             raise ValueError(
-                "Either `value`, `id`, or `name` must be provided when "
-                "creating an external artifact."
+                "External artifacts can only be initialized with a value."
             )
-        return values
+
+        return self
 
     def upload_by_value(self) -> UUID:
         """Uploads the artifact by value.
@@ -127,7 +115,7 @@ class ExternalArtifact(ExternalArtifactConfiguration):
             materializer=self.materializer,
             uri=uri,
             has_custom_name=False,
-            manual_save=False,
+            save_type=ArtifactSaveType.EXTERNAL,
         )
 
         # To avoid duplicate uploads, switch to referencing the uploaded
@@ -147,7 +135,4 @@ class ExternalArtifact(ExternalArtifactConfiguration):
         """
         return ExternalArtifactConfiguration(
             id=self.id,
-            name=self.name,
-            version=self.version,
-            model=self.model,
         )