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

zenml/integrations/prodigy/flavors/prodigy_annotator_flavor.py

@@ -0,0 +1,101 @@
+#  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.
+"""Prodigy annotator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.annotators.base_annotator import (
+    BaseAnnotatorConfig,
+    BaseAnnotatorFlavor,
+)
+from zenml.integrations.prodigy import PRODIGY_ANNOTATOR_FLAVOR
+from zenml.stack.authentication_mixin import AuthenticationConfigMixin
+
+if TYPE_CHECKING:
+    from zenml.integrations.prodigy.annotators import ProdigyAnnotator
+
+
+class ProdigyAnnotatorConfig(BaseAnnotatorConfig, AuthenticationConfigMixin):
+    """Config for the Prodigy annotator.
+
+    See https://prodi.gy/docs/install#config for more on custom config files,
+    but this allows you to override the default Prodigy config.
+
+    Attributes:
+        custom_config_path: The path to a custom config file for Prodigy.
+    """
+
+    custom_config_path: Optional[str] = None
+
+
+class ProdigyAnnotatorFlavor(BaseAnnotatorFlavor):
+    """Prodigy annotator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return PRODIGY_ANNOTATOR_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/annotator/prodigy.png"
+
+    @property
+    def config_class(self) -> Type[ProdigyAnnotatorConfig]:
+        """Returns `ProdigyAnnotatorConfig` config class.
+
+        Returns:
+                The config class.
+        """
+        return ProdigyAnnotatorConfig
+
+    @property
+    def implementation_class(self) -> Type["ProdigyAnnotator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.prodigy.annotators import (
+            ProdigyAnnotator,
+        )
+
+        return ProdigyAnnotator

zenml/integrations/pycaret/__init__.py

@@ -28,6 +28,12 @@ class PyCaretIntegration(Integration):
         "catboost",
         "lightgbm",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "scikit-learn",
+        "xgboost",
+        "catboost",
+        "lightgbm",
+    ]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/pycaret/materializers/model_materializer.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """PyCaret materializer."""
 
-import tempfile
 from typing import (
     Any,
     Type,
@@ -65,7 +64,6 @@ from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 from xgboost import XGBClassifier, XGBRegressor
 
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils import io_utils
 
@@ -133,19 +131,10 @@ class PyCaretMaterializer(BaseMaterializer):
         Returns:
             A PyCaret model.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-
-        # Copy from artifact store to temporary directory
-        io_utils.copy_dir(self.uri, temp_dir.name)
-
-        # Load the model from the temporary directory
-        model = load_model(temp_dir.name)
-
-        # Cleanup and return
-        fileio.rmtree(temp_dir.name)
-
-        return model
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            io_utils.copy_dir(self.uri, temp_dir)
+            model = load_model(temp_dir)
+            return model
 
     def save(self, model: Any) -> None:
         """Writes a PyCaret model to the artifact store.
@@ -153,10 +142,6 @@ class PyCaretMaterializer(BaseMaterializer):
         Args:
             model: Any of the supported models.
         """
-        # Create a temporary directory to store the model
-        temp_dir = tempfile.TemporaryDirectory()
-        save_model(model, temp_dir.name)
-        io_utils.copy_dir(temp_dir.name, self.uri)
-
-        # Remove the temporary directory
-        fileio.rmtree(temp_dir.name)
+        with self.get_temporary_directory(delete_at_exit=True) as temp_dir:
+            save_model(model, temp_dir)
+            io_utils.copy_dir(temp_dir, self.uri)

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/registry.py

@@ -111,7 +111,7 @@ class IntegrationRegistry(object):
                 )
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )
@@ -124,6 +124,43 @@ class IntegrationRegistry(object):
                 )
             ]
 
+    def select_uninstall_requirements(
+        self,
+        integration_name: Optional[str] = None,
+        target_os: Optional[str] = None,
+    ) -> List[str]:
+        """Select the uninstall requirements for a given integration or all integrations.
+
+        Args:
+            integration_name: Name of the integration to check.
+            target_os: Target OS for the requirements.
+
+        Returns:
+            List of requirements for the integration uninstall.
+
+        Raises:
+            KeyError: If the integration is not found.
+        """
+        if integration_name:
+            if integration_name in self.list_integration_names:
+                return self._integrations[
+                    integration_name
+                ].get_uninstall_requirements(target_os=target_os)
+            else:
+                raise KeyError(
+                    f"Integration {integration_name} does not exist. "
+                    f"Currently the following integrations are implemented. "
+                    f"{self.list_integration_names}"
+                )
+        else:
+            return [
+                requirement
+                for name in self.list_integration_names
+                for requirement in self._integrations[
+                    name
+                ].get_uninstall_requirements(target_os=target_os)
+            ]
+
     def is_installed(self, integration_name: Optional[str] = None) -> bool:
         """Checks if all requirements for an integration are installed.
 

zenml/integrations/s3/__init__.py

@@ -18,7 +18,6 @@ operations on S3 buckets.
 """
 from typing import List, Type
 
-from zenml.enums import StackComponentType
 from zenml.integrations.constants import S3
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
@@ -42,8 +41,8 @@ class S3Integration(Integration):
     # The above command installs boto3==1.26.76, so we use the same version
     # here to avoid the dependency resolution overhead.
     REQUIREMENTS = [
-        "s3fs>2022.3.0,<=2023.4.0",
-        "boto3<=1.26.76",
+        "s3fs>2022.3.0",
+        "boto3",
         # The following dependencies are only required for the AWS connector.
         "aws-profile-manager",
     ]

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,23 +27,125 @@ 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]
 
 
+class ZenMLS3Filesystem(s3fs.S3FileSystem):  # type: ignore[misc]
+    """Modified s3fs.S3FileSystem to disable caching.
+
+    The original s3fs.S3FileSystem caches all class instances based on the
+    constructor input arguments and it never releases them. This is problematic
+    in the context of the ZenML server, because the server is a long-running
+    process that instantiates many S3 filesystems with different credentials,
+    especially when the credentials are generated by service connectors.
+
+    The caching behavior of s3fs causes the server to slowly consume more and
+    more memory over time until it crashes. This class disables the caching
+    behavior of s3fs by setting the `cachable` attribute to `False`.
+
+    In addition to disabling instance caching, this class also provides a
+    correct cleanup implementation by overriding the `close_session` method
+    the S3 aiobotocore client. The original one provided by s3fs was causing
+    memory leaks by creating a new event loop in the destructor instead of
+    using the existing one.
+
+    A `close` method is also provided to allow for synchronous on-demand cleanup
+    of the S3 client.
+    """
+
+    cachable = False
+
+    async def _close(self) -> None:
+        """Close the S3 client."""
+        if self._s3creator is not None:  # type: ignore[has-type]
+            await self._s3creator.__aexit__(None, None, None)  # type: ignore[has-type]
+            self._s3creator = None
+            self._s3 = None
+
+    close = sync_wrapper(_close)
+
+    @staticmethod
+    def close_session(loop: Any, s3: Any) -> None:
+        """Close the S3 client session.
+
+        Args:
+            loop: The event loop to use for closing the session.
+            s3: The S3 client to close.
+        """
+        # IMPORTANT: This method is a copy of the original close_session method
+        # from s3fs.S3FileSystem. The only difference is that it uses the
+        # provided event loop instead of creating a new one.
+        if loop is not None and loop.is_running():
+            try:
+                # NOTE: this is the line in the original method that causes
+                # the memory leak
+                # loop = asyncio.get_event_loop()
+                loop.create_task(s3.__aexit__(None, None, None))
+                return
+            except RuntimeError:
+                pass
+            try:
+                sync(loop, s3.__aexit__, None, None, None, timeout=0.1)
+                return
+            except FSTimeoutError:
+                pass
+        try:
+            # close the actual socket
+            s3._client._endpoint.http_session._connector._close()
+        except AttributeError:
+            # but during shutdown, it may have gone
+            pass
+
+
 class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
     """Artifact Store for S3 based artifacts."""
 
-    _filesystem: Optional[s3fs.S3FileSystem] = None
+    _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:
@@ -54,7 +158,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
     def get_credentials(
         self,
-    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
         """Gets authentication credentials.
 
         If an authentication secret is configured, the secret values are
@@ -62,8 +166,8 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         attributes.
 
         Returns:
-            Tuple (key, secret, token) of credentials used to authenticate with
-            the S3 filesystem.
+            Tuple (key, secret, token, region) of credentials used to
+            authenticate with the S3 filesystem.
 
         Raises:
             RuntimeError: If the AWS connector behaves unexpectedly.
@@ -83,6 +187,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 credentials.access_key,
                 credentials.secret_key,
                 credentials.token,
+                client.meta.region_name,
             )
 
         secret = self.get_typed_authentication_secret(
@@ -93,12 +198,13 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 secret.aws_access_key_id,
                 secret.aws_secret_access_key,
                 secret.aws_session_token,
+                None,
             )
         else:
-            return self.config.key, self.config.secret, self.config.token
+            return self.config.key, self.config.secret, self.config.token, None
 
     @property
-    def filesystem(self) -> s3fs.S3FileSystem:
+    def filesystem(self) -> ZenMLS3Filesystem:
         """The s3 filesystem to access this artifact store.
 
         Returns:
@@ -108,18 +214,32 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         if self._filesystem and not self.connector_has_expired():
             return self._filesystem
 
-        key, secret, token = self.get_credentials()
+        key, secret, token, region = self.get_credentials()
 
-        self._filesystem = s3fs.S3FileSystem(
+        # Use the region from the connector if available, otherwise some
+        # remote workloads (e.g. Sagemaker) might not work correctly because
+        # they look for the bucket in the wrong region
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+        if self.config.client_kwargs:
+            client_kwargs.update(self.config.client_kwargs)
+
+        self._filesystem = ZenMLS3Filesystem(
             key=key,
             secret=secret,
             token=token,
-            client_kwargs=self.config.client_kwargs,
+            client_kwargs=client_kwargs,
             config_kwargs=self.config.config_kwargs,
             s3_additional_kwargs=self.config.s3_additional_kwargs,
         )
         return self._filesystem
 
+    def cleanup(self) -> None:
+        """Close the filesystem."""
+        if self._filesystem:
+            self._filesystem.close()
+
     def open(self, path: PathType, mode: str = "r") -> Any:
         """Open a file at the given path.
 
@@ -338,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

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Amazon S3 artifact store flavor."""
 
-import json
 import re
 from typing import (
     TYPE_CHECKING,
@@ -23,16 +22,16 @@ from typing import (
     Optional,
     Set,
     Type,
-    Union,
 )
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.artifact_stores import (
     BaseArtifactStoreConfig,
     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 (
@@ -65,50 +64,17 @@ class S3ArtifactStoreConfig(
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"s3://"}
 
-    key: Optional[str] = SecretField()
-    secret: Optional[str] = SecretField()
-    token: Optional[str] = SecretField()
+    key: Optional[str] = SecretField(default=None)
+    secret: Optional[str] = SecretField(default=None)
+    token: Optional[str] = SecretField(default=None)
     client_kwargs: Optional[Dict[str, Any]] = None
     config_kwargs: Optional[Dict[str, Any]] = None
     s3_additional_kwargs: Optional[Dict[str, Any]] = None
 
-    @validator(
-        "client_kwargs", "config_kwargs", "s3_additional_kwargs", pre=True
-    )
-    def _convert_json_string(
-        cls, value: Union[None, str, Dict[str, Any]]
-    ) -> Optional[Dict[str, Any]]:
-        """Converts potential JSON strings passed via the CLI to dictionaries.
-
-        Args:
-            value: The value to convert.
-
-        Returns:
-            The converted value.
-
-        Raises:
-            TypeError: If the value is not a `str`, `Dict` or `None`.
-            ValueError: If the value is an invalid json string or a json string
-                that does not decode into a dictionary.
-        """
-        if isinstance(value, str):
-            try:
-                dict_ = json.loads(value)
-            except json.JSONDecodeError as e:
-                raise ValueError(f"Invalid json string '{value}'") from e
-
-            if not isinstance(dict_, Dict):
-                raise ValueError(
-                    f"Json string '{value}' did not decode into a dictionary."
-                )
+    _bucket: Optional[str] = None
 
-            return dict_
-        elif isinstance(value, Dict) or value is None:
-            return value
-        else:
-            raise TypeError(f"{value} is not a json string or a dictionary.")
-
-    @validator("client_kwargs")
+    @field_validator("client_kwargs")
+    @classmethod
     def _validate_client_kwargs(
         cls, value: Optional[Dict[str, Any]]
     ) -> Optional[Dict[str, Any]]:
@@ -143,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."""
@@ -176,7 +153,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def docs_url(self) -> Optional[str]:
-        """A url to point at docs explaining this flavor.
+        """A URL to point at docs explaining this flavor.
 
         Returns:
             A flavor docs url.
@@ -185,7 +162,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def sdk_docs_url(self) -> Optional[str]:
-        """A url to point at SDK docs explaining this flavor.
+        """A URL to point at SDK docs explaining this flavor.
 
         Returns:
             A flavor SDK docs url.
@@ -194,7 +171,7 @@ class S3ArtifactStoreFlavor(BaseArtifactStoreFlavor):
 
     @property
     def logo_url(self) -> str:
-        """A url to represent the flavor in the dashboard.
+        """A URL to represent the flavor in the dashboard.
 
         Returns:
             The flavor logo.