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

zenml/integrations/neptune/experiment_trackers/run_state.py

@@ -16,7 +16,7 @@
 from hashlib import md5
 from typing import TYPE_CHECKING, Any, List, Optional
 
-import neptune  # type: ignore[import-untyped]
+import neptune
 
 import zenml
 from zenml.client import Client
@@ -38,7 +38,7 @@ class RunProvider(metaclass=SingletonMetaClass):
 
     def __init__(self) -> None:
         """Initialize RunProvider. Called with no arguments."""
-        self._active_run = None
+        self._active_run: Optional["Run"] = None
         self._project: Optional[str]
         self._run_name: Optional[str]
         self._token: Optional[str]

zenml/integrations/pytorch/materializers/base_pytorch_materializer.py

@@ -41,7 +41,10 @@ class BasePyTorchMaterializer(BaseMaterializer):
             The loaded PyTorch object.
         """
         with fileio.open(os.path.join(self.uri, self.FILENAME), "rb") as f:
-            return torch.load(f)
+            # NOTE (security): The `torch.load` function uses `pickle` as
+            # the default unpickler, which is NOT secure. This materializer
+            # is intended for use with trusted data sources.
+            return torch.load(f)  # nosec
 
     def save(self, obj: Any) -> None:
         """Uses `torch.save` to save a PyTorch object.
@@ -50,7 +53,10 @@ class BasePyTorchMaterializer(BaseMaterializer):
             obj: The PyTorch object to save.
         """
         with fileio.open(os.path.join(self.uri, self.FILENAME), "wb") as f:
-            torch.save(obj, f, pickle_module=cloudpickle)
+            # NOTE (security): The `torch.save` function uses `cloudpickle` as
+            # the default unpickler, which is NOT secure. This materializer
+            # is intended for use with trusted data sources.
+            torch.save(obj, f, pickle_module=cloudpickle)  # nosec
 
 
 # Alias for the BasePyTorchMaterializer class, allowing users that have already used

zenml/integrations/pytorch/materializers/pytorch_module_materializer.py

@@ -61,7 +61,10 @@ class PyTorchModuleMaterializer(BasePyTorchMaterializer):
             with fileio.open(
                 os.path.join(self.uri, CHECKPOINT_FILENAME), "wb"
             ) as f:
-                torch.save(model.state_dict(), f, pickle_module=cloudpickle)
+                # NOTE (security): The `torch.save` function uses `cloudpickle` as
+                # the default unpickler, which is NOT secure. This materializer
+                # is intended for use with trusted data sources.
+                torch.save(model.state_dict(), f, pickle_module=cloudpickle)  # nosec
 
     def extract_metadata(self, model: Module) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given `Model` object.

zenml/integrations/s3/artifact_stores/s3_artifact_store.py

@@ -13,10 +13,12 @@
 #  permissions and limitations under the License.
 """Implementation of the S3 Artifact Store."""
 
+from contextlib import contextmanager
 from typing import (
     Any,
     Callable,
     Dict,
+    Generator,
     Iterable,
     List,
     Optional,
@@ -25,17 +27,23 @@ from typing import (
     cast,
 )
 
+import boto3
 import s3fs
+from botocore.exceptions import ClientError
 from fsspec.asyn import FSTimeoutError, sync, sync_wrapper
 
 from zenml.artifact_stores import BaseArtifactStore
 from zenml.integrations.s3.flavors.s3_artifact_store_flavor import (
     S3ArtifactStoreConfig,
 )
+from zenml.integrations.s3.utils import split_s3_path
 from zenml.io.fileio import convert_to_str
+from zenml.logger import get_logger
 from zenml.secret.schemas import AWSSecretSchema
 from zenml.stack.authentication_mixin import AuthenticationMixin
 
+logger = get_logger(__name__)
+
 PathType = Union[bytes, str]
 
 
@@ -111,6 +119,34 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
     _filesystem: Optional[ZenMLS3Filesystem] = None
 
+    is_versioned: bool = False
+
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes the artifact store.
+
+        Args:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(*args, **kwargs)
+        self._boto3_bucket_holder = None
+
+        # determine bucket versioning status
+        versioning = self._boto3_bucket.Versioning()
+        with self._shield_lack_of_versioning_permissions(
+            "s3:GetBucketVersioning"
+        ):
+            if versioning.status == "Enabled":
+                self.is_versioned = True
+                logger.warning(
+                    f"The artifact store bucket `{self.config.bucket}` is versioned, "
+                    "this may slow down logging process significantly."
+                )
+
     @property
     def config(self) -> S3ArtifactStoreConfig:
         """Get the config of this artifact store.
@@ -422,3 +458,65 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         # TODO [ENG-153]: Additional params
         for directory, subdirectories, files in self.filesystem.walk(path=top):
             yield f"s3://{directory}", subdirectories, files
+
+    def _remove_previous_file_versions(self, path: PathType) -> None:
+        """Keep only the latest file version in the given path.
+
+        Method is useful for logs stored in versioned file systems
+        like AWS S3.
+
+        Args:
+            path: The path to the file.
+        """
+        if self.is_versioned:
+            if isinstance(path, bytes):
+                path = path.decode()
+            _, prefix = split_s3_path(path)
+            with self._shield_lack_of_versioning_permissions(
+                "s3:ListBucketVersions"
+            ):
+                for version in self._boto3_bucket.object_versions.filter(
+                    Prefix=prefix
+                ):
+                    if not version.is_latest:
+                        with self._shield_lack_of_versioning_permissions(
+                            "s3:DeleteObjectVersion"
+                        ):
+                            version.delete()
+
+    @property
+    def _boto3_bucket(self) -> Any:
+        """Get the boto3 bucket object.
+
+        Returns:
+            The boto3 bucket object.
+        """
+        if self._boto3_bucket_holder and not self.connector_has_expired():
+            return self._boto3_bucket_holder
+
+        key, secret, token, region = self.get_credentials()
+        s3 = boto3.resource(
+            "s3",
+            aws_access_key_id=key,
+            aws_secret_access_key=secret,
+            aws_session_token=token,
+            region_name=region,
+        )
+        self._boto3_bucket_holder = s3.Bucket(self.config.bucket)
+        return self._boto3_bucket_holder
+
+    @contextmanager
+    def _shield_lack_of_versioning_permissions(
+        self, auth_missing: str
+    ) -> Generator[Any, None, None]:
+        try:
+            yield
+        except ClientError as e:
+            if "not authorized" in e.args[0] and auth_missing in e.args[0]:
+                logger.warning(
+                    "Your AWS Connector is lacking critical Versioning permissions. "
+                    f"Please check that `{auth_missing}` is granted."
+                    "This is needed to remove previous versions of log files from your "
+                    "Artifact Store bucket."
+                )
+                self.is_versioned = False

zenml/integrations/s3/flavors/s3_artifact_store_flavor.py

@@ -31,6 +31,7 @@ from zenml.artifact_stores import (
     BaseArtifactStoreFlavor,
 )
 from zenml.integrations.s3 import S3_ARTIFACT_STORE_FLAVOR
+from zenml.integrations.s3.utils import split_s3_path
 from zenml.models import ServiceConnectorRequirements
 from zenml.stack.authentication_mixin import AuthenticationConfigMixin
 from zenml.utils.networking_utils import (
@@ -70,6 +71,8 @@ class S3ArtifactStoreConfig(
     config_kwargs: Optional[Dict[str, Any]] = None
     s3_additional_kwargs: Optional[Dict[str, Any]] = None
 
+    _bucket: Optional[str] = None
+
     @field_validator("client_kwargs")
     @classmethod
     def _validate_client_kwargs(
@@ -106,6 +109,17 @@ class S3ArtifactStoreConfig(
             )
         return value
 
+    @property
+    def bucket(self) -> str:
+        """The bucket name of the artifact store.
+
+        Returns:
+            The bucket name of the artifact store.
+        """
+        if self._bucket is None:
+            self._bucket, _ = split_s3_path(self.path)
+        return self._bucket
+
 
 class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
     """Flavor of the S3 artifact store."""

zenml/integrations/s3/utils.py

@@ -0,0 +1,40 @@
+#  Copyright (c) ZenML GmbH 2022. 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.
+"""Utility methods for S3."""
+
+from typing import Tuple
+
+
+def split_s3_path(s3_path: str) -> Tuple[str, str]:
+    """Split S3 URI into bucket and key.
+
+    Args:
+        s3_path: S3 URI (e.g. "s3://bucket/path")
+
+    Returns:
+        A tuple of bucket and key, for "s3://bucket/path/path2"
+            it will return ("bucket","path/path2")
+
+    Raises:
+        ValueError: if the S3 URI is invalid
+    """
+    if not s3_path.startswith("s3://"):
+        raise ValueError(
+            f"Invalid S3 URI given: {s3_path}. "
+            "It should start with `s3://`."
+        )
+    path_parts = s3_path.replace("s3://", "").split("/")
+    bucket = path_parts.pop(0)
+    key = "/".join(path_parts)
+    return bucket, key

zenml/integrations/tensorflow/__init__.py

@@ -40,15 +40,6 @@ class TensorflowIntegration(Integration):
 
         from zenml.integrations.tensorflow import materializers  # noqa
 
-        if sys.version_info.minor == 8:
-            logger.warning(
-                "Python 3.8 with TensorFlow is not fully "
-                "compatible with Pydantic 2 requirements. "
-                "Consider upgrading to a higher Python "
-                "version if you would like to use the "
-                "Tensorflow integration."
-            )
-
     @classmethod
     def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         """Defines platform specific requirements for the integration.

zenml/logging/step_logging.py

@@ -28,7 +28,6 @@ from zenml.artifacts.utils import (
     _load_artifact_store,
     _load_file_from_artifact_store,
 )
-from zenml.client import Client
 from zenml.exceptions import DoesNotExistException
 from zenml.logger import get_logger
 from zenml.logging import (
@@ -212,6 +211,7 @@ class StepLogsStorage:
     def __init__(
         self,
         logs_uri: str,
+        artifact_store: "BaseArtifactStore",
         max_messages: int = STEP_LOGS_STORAGE_MAX_MESSAGES,
         time_interval: int = STEP_LOGS_STORAGE_INTERVAL_SECONDS,
         merge_files_interval: int = STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS,
@@ -220,6 +220,7 @@ class StepLogsStorage:
 
         Args:
             logs_uri: the URI of the log file or folder.
+            artifact_store: Artifact Store from the current step context
             max_messages: the maximum number of messages to save in the buffer.
             time_interval: the amount of seconds before the buffer gets saved
                 automatically.
@@ -237,22 +238,11 @@ class StepLogsStorage:
         self.disabled_buffer: List[str] = []
         self.last_save_time = time.time()
         self.disabled = False
-        self._artifact_store: Optional["BaseArtifactStore"] = None
+        self.artifact_store = artifact_store
 
         # Immutable filesystems state
         self.last_merge_time = time.time()
 
-    @property
-    def artifact_store(self) -> "BaseArtifactStore":
-        """Returns the active artifact store.
-
-        Returns:
-            The active artifact store.
-        """
-        if self._artifact_store is None:
-            self._artifact_store = Client().active_stack.artifact_store
-        return self._artifact_store
-
     def write(self, text: str) -> None:
         """Main write method.
 
@@ -330,6 +320,9 @@ class StepLogsStorage:
                                 file.write(
                                     f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n"
                                 )
+                        self.artifact_store._remove_previous_file_versions(
+                            self.logs_uri
+                        )
 
             except (OSError, IOError) as e:
                 # This exception can be raised if there are issues with the
@@ -406,13 +399,18 @@ class StepLogsStorage:
 class StepLogsStorageContext:
     """Context manager which patches stdout and stderr during step execution."""
 
-    def __init__(self, logs_uri: str) -> None:
+    def __init__(
+        self, logs_uri: str, artifact_store: "BaseArtifactStore"
+    ) -> None:
         """Initializes and prepares a storage object.
 
         Args:
             logs_uri: the URI of the logs file.
+            artifact_store: Artifact Store from the current step context.
         """
-        self.storage = StepLogsStorage(logs_uri=logs_uri)
+        self.storage = StepLogsStorage(
+            logs_uri=logs_uri, artifact_store=artifact_store
+        )
 
     def __enter__(self) -> "StepLogsStorageContext":
         """Enter condition of the context manager.

zenml/materializers/__init__.py

@@ -31,6 +31,7 @@ from zenml.materializers.structured_string_materializer import (
 )
 from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
+from zenml.materializers.uuid_materializer import UUIDMaterializer
 
 __all__ = [
     "BuiltInContainerMaterializer",
@@ -40,4 +41,5 @@ __all__ = [
     "StructuredStringMaterializer",
     "PydanticMaterializer",
     "ServiceMaterializer",
+    "UUIDMaterializer",
 ]

zenml/materializers/uuid_materializer.py

@@ -0,0 +1,79 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Implementation of ZenML's UUID materializer."""
+
+import os
+import uuid
+from typing import Any, ClassVar, Dict, Optional, Tuple, Type
+
+from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
+from zenml.enums import ArtifactType
+from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.metadata_types import MetadataType
+
+DEFAULT_FILENAME = "uuid.txt"
+
+
+class UUIDMaterializer(BaseMaterializer):
+    """Materializer to handle UUID objects."""
+
+    ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (uuid.UUID,)
+    ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
+
+    def __init__(
+        self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
+    ):
+        """Define `self.data_path`.
+
+        Args:
+            uri: The URI where the artifact data is stored.
+            artifact_store: The artifact store where the artifact data is stored.
+        """
+        super().__init__(uri, artifact_store)
+        self.data_path = os.path.join(self.uri, DEFAULT_FILENAME)
+
+    def load(self, _: Type[uuid.UUID]) -> uuid.UUID:
+        """Read UUID from artifact store.
+
+        Args:
+            _: The type of the data to be loaded.
+
+        Returns:
+            The loaded UUID.
+        """
+        with self.artifact_store.open(self.data_path, "r") as f:
+            uuid_str = f.read().strip()
+        return uuid.UUID(uuid_str)
+
+    def save(self, data: uuid.UUID) -> None:
+        """Write UUID to artifact store.
+
+        Args:
+            data: The UUID to be saved.
+        """
+        with self.artifact_store.open(self.data_path, "w") as f:
+            f.write(str(data))
+
+    def extract_metadata(self, data: uuid.UUID) -> Dict[str, MetadataType]:
+        """Extract metadata from the UUID.
+
+        Args:
+            data: The UUID to extract metadata from.
+
+        Returns:
+            A dictionary of metadata extracted from the UUID.
+        """
+        return {
+            "string_representation": str(data),
+        }

zenml/metadata/lazy_load.py

@@ -16,7 +16,6 @@
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
-    from zenml.model.model import Model
     from zenml.models import RunMetadataResponse
 
 
@@ -30,18 +29,21 @@ class RunMetadataLazyGetter:
 
     def __init__(
         self,
-        _lazy_load_model: "Model",
-        _lazy_load_artifact_name: Optional[str],
-        _lazy_load_artifact_version: Optional[str],
+        _lazy_load_model_name: str,
+        _lazy_load_model_version: Optional[str],
+        _lazy_load_artifact_name: Optional[str] = None,
+        _lazy_load_artifact_version: Optional[str] = None,
     ):
         """Initialize a RunMetadataLazyGetter.
 
         Args:
-            _lazy_load_model: The model version.
+            _lazy_load_model_name: The model name.
+            _lazy_load_model_version: The model version.
             _lazy_load_artifact_name: The artifact name.
             _lazy_load_artifact_version: The artifact version.
         """
-        self._lazy_load_model = _lazy_load_model
+        self._lazy_load_model_name = _lazy_load_model_name
+        self._lazy_load_model_version = _lazy_load_model_version
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
@@ -57,7 +59,8 @@ class RunMetadataLazyGetter:
         from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         return LazyRunMetadataResponse(
-            lazy_load_model=self._lazy_load_model,
+            lazy_load_model_name=self._lazy_load_model_name,
+            lazy_load_model_version=self._lazy_load_model_version,
             lazy_load_artifact_name=self._lazy_load_artifact_name,
             lazy_load_artifact_version=self._lazy_load_artifact_version,
             lazy_load_metadata_name=key,

zenml/model/lazy_load.py

@@ -13,11 +13,15 @@
 #  permissions and limitations under the License.
 """Model Version Data Lazy Loader definition."""
 
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, model_validator
 
-from zenml.model.model import Model
+from zenml.pipelines.pipeline_context import get_pipeline_context
+from zenml.utils.pydantic_utils import before_validator_handler
+
+if TYPE_CHECKING:
+    from zenml.models import ModelVersionResponse, PipelineRunResponse
 
 
 class ModelVersionDataLazyLoader(BaseModel):
@@ -28,7 +32,84 @@ class ModelVersionDataLazyLoader(BaseModel):
     model version during runtime time of the step.
     """
 
-    model: Model
+    model_name: str
+    model_version: Optional[str] = None
     artifact_name: Optional[str] = None
     artifact_version: Optional[str] = None
     metadata_name: Optional[str] = None
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _root_validator(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate all in one.
+
+        Args:
+            data: Dict of values.
+
+        Returns:
+            Dict of validated values.
+
+        Raises:
+            ValueError: If the model version id, but call is not internal.
+        """
+        if data.get("model_version", None) is None:
+            try:
+                context = get_pipeline_context()
+                if (
+                    not context.model
+                    or context.model.name != data["model_name"]
+                ):
+                    raise ValueError(
+                        "`version` must be set if you use the `Model` class "
+                        "directly in the pipeline body, otherwise, you can use "
+                        "`get_pipeline_context().model` to lazy load the current "
+                        "Model Version from the pipeline context."
+                    )
+            except RuntimeError:
+                pass
+        data["suppress_class_validation_warnings"] = True
+        return data
+
+    def _get_model_response(
+        self, pipeline_run: "PipelineRunResponse"
+    ) -> "ModelVersionResponse":
+        # if the version/number is None -> return the model in context
+        if self.model_version is None:
+            if mv := pipeline_run.model_version:
+                if mv.model.name != self.model_name:
+                    raise RuntimeError(
+                        "Lazy loading of the model failed, since given name "
+                        f"`{self.model_name}` does not match the model name "
+                        f"in the pipeline context: `{mv.model.name}`."
+                    )
+                return mv
+            else:
+                raise RuntimeError(
+                    "Lazy loading of the model failed, since the model version "
+                    "is not set in the pipeline context."
+                )
+
+        # else return the model version by version
+        else:
+            from zenml.client import Client
+
+            try:
+                return Client().get_model_version(
+                    model_name_or_id=self.model_name,
+                    model_version_name_or_number_or_id=self.model_version,
+                )
+            except KeyError as e:
+                raise RuntimeError(
+                    "Lazy loading of the model version failed: "
+                    f"no model `{self.model_name}` with version "
+                    f"`{self.model_version}` could be found."
+                ) from e