zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__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/{new/pipelines → login}/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  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.
@@ -11,3 +11,6 @@
 #  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.
+"""ZenML login utilities."""
+
+