zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.58.2.dev20240626__py3-none-any.whl

zenml/config/docker_settings.py

@@ -16,7 +16,7 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import Field, model_validator
+from pydantic import BaseModel, Field, model_validator
 from pydantic_settings import SettingsConfigDict
 
 from zenml.config.base_settings import BaseSettings
@@ -65,6 +65,23 @@ class PythonPackageInstaller(Enum):
     UV = "uv"
 
 
+class DockerBuildConfig(BaseModel):
+    """Configuration for a Docker build.
+
+    Attributes:
+        build_options: Additional options that will be passed unmodified to the
+            Docker build call when building an image. You can use this to for
+            example specify build args or a target stage. See
+            https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
+            for a full list of available options.
+        dockerignore: Path to a dockerignore file to use when building the
+            Docker image.
+    """
+
+    build_options: Dict[str, Any] = {}
+    dockerignore: Optional[str] = None
+
+
 class DockerSettings(BaseSettings):
     """Settings for building Docker images to run ZenML pipelines.
 
@@ -116,12 +133,9 @@ class DockerSettings(BaseSettings):
         build_context_root: Build context root for the Docker build, only used
             when the `dockerfile` attribute is set. If this is left empty, the
             build context will only contain the Dockerfile.
-        build_options: Additional options that will be passed unmodified to the
-            Docker build call when building an image using the specified
-            `dockerfile`. You can use this to for example specify build
-            args or a target stage. See
-            https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
-            for a full list of available options.
+        parent_image_build_config: Configuration for the parent image build.
+        build_options: DEPRECATED, use parent_image_build_config.build_options
+            instead.
         skip_build: If set to `True`, the parent image will be used directly to
             run the steps of your pipeline.
         target_repository: Name of the Docker repository to which the
@@ -158,8 +172,8 @@ class DockerSettings(BaseSettings):
         apt_packages: APT packages to install inside the Docker image.
         environment: Dictionary of environment variables to set inside the
             Docker image.
-        dockerignore: Path to a dockerignore file to use when building the
-            Docker image.
+        build_config: Configuration for the main image build.
+        dockerignore: DEPRECATED, use build_config.dockerignore instead.
         copy_files: DEPRECATED, use the `source_files` attribute instead.
         copy_global_config: DEPRECATED/UNUSED.
         user: If not `None`, will set the user, make it owner of the `/app`
@@ -184,6 +198,7 @@ class DockerSettings(BaseSettings):
     dockerfile: Optional[str] = None
     build_context_root: Optional[str] = None
     build_options: Dict[str, Any] = {}
+    parent_image_build_config: Optional[DockerBuildConfig] = None
     skip_build: bool = False
     target_repository: str = "zenml"
     python_package_installer: PythonPackageInstaller = (
@@ -205,6 +220,7 @@ class DockerSettings(BaseSettings):
     copy_files: bool = True
     copy_global_config: bool = True
     user: Optional[str] = None
+    build_config: Optional[DockerBuildConfig] = None
 
     source_files: SourceFileMode = SourceFileMode.DOWNLOAD_OR_INCLUDE
 

zenml/constants.py

@@ -326,7 +326,7 @@ REQUIRES_CUSTOM_RESOURCE_REPORTING = ["pipeline", "pipeline_run"]
 
 # API Endpoint paths:
 ACTIVATE = "/activate"
-ACTIONS = "/action-flavors"
+ACTIONS = "/actions"
 API = "/api"
 API_KEYS = "/api_keys"
 API_KEY_ROTATE = "/rotate"

zenml/event_hub/base_event_hub.py

@@ -121,7 +121,7 @@ class BaseEventHub(ABC):
             trigger=trigger.id, event_metadata=event.model_dump()
         )
 
-        action_config = trigger.get_metadata().action
+        action_config = trigger.action.configuration
 
         trigger_execution = self.zen_store.create_trigger_execution(request)
 
@@ -130,16 +130,16 @@ class BaseEventHub(ABC):
         # is associated with the service account configured for the trigger
         # and has a validity defined by the trigger's authentication window.
         token = JWTToken(
-            user_id=trigger.service_account.id,
+            user_id=trigger.action.service_account.id,
         )
         expires: Optional[datetime] = None
-        if trigger.auth_window:
+        if trigger.action.auth_window:
             expires = datetime.utcnow() + timedelta(
-                minutes=trigger.auth_window
+                minutes=trigger.action.auth_window
             )
         encoded_token = token.encode(expires=expires)
         auth_context = AuthContext(
-            user=trigger.service_account,
+            user=trigger.action.service_account,
             access_token=token,
             encoded_access_token=encoded_token,
         )

zenml/event_hub/event_hub.py

@@ -14,7 +14,9 @@
 """Base class for all the Event Hub."""
 
 from functools import partial
-from typing import TYPE_CHECKING, List
+from typing import List
+
+from pydantic import ValidationError
 
 from zenml import EventSourceResponse
 from zenml.enums import PluginType
@@ -35,9 +37,6 @@ from zenml.zen_server.utils import plugin_flavor_registry
 
 logger = get_logger(__name__)
 
-if TYPE_CHECKING:
-    pass
-
 
 class InternalEventHub(BaseEventHub):
     """Internal in-server event hub implementation.
@@ -158,9 +157,19 @@ class InternalEventHub(BaseEventHub):
 
             assert issubclass(plugin_flavor, BaseEventSourceFlavor)
 
-            # Get the filter class from the plugin flavor class
             event_filter_config_class = plugin_flavor.EVENT_FILTER_CONFIG_CLASS
-            event_filter = event_filter_config_class(**trigger.event_filter)
+            try:
+                event_filter = event_filter_config_class(
+                    **trigger.event_filter if trigger.event_filter else {}
+                )
+            except ValidationError:
+                logger.exception(
+                    f"Could not instantiate event filter config class for "
+                    f"event source {event_source.id}. Skipping trigger "
+                    f"{trigger.id}."
+                )
+                continue
+
             if event_filter.event_matches_filter(event=event):
                 trigger_list.append(trigger)
 

zenml/event_sources/base_event.py

@@ -13,19 +13,8 @@
 #  permissions and limitations under the License.
 """Base implementation for events."""
 
-from typing import (
-    TYPE_CHECKING,
-)
-
 from pydantic import BaseModel
 
-from zenml.logger import get_logger
-
-if TYPE_CHECKING:
-    pass
-
-logger = get_logger(__name__)
-
 # -------------------- Event Models -----------------------------------
 
 

zenml/event_sources/base_event_source.py

@@ -449,6 +449,13 @@ class BaseEventSourceHandler(BasePlugin, ABC):
             event: The event to dispatch.
             event_source: The event source that produced the event.
         """
+        if not event_source.is_active:
+            logger.debug(
+                f"Event source {event_source.id} is not active. Skipping event "
+                f"dispatch."
+            )
+            return
+
         self.event_hub.publish_event(
             event=event,
             event_source=event_source,

zenml/event_sources/webhooks/base_webhook_event_source.py

@@ -17,7 +17,7 @@ import hashlib
 import hmac
 import json
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Type
+from typing import Any, ClassVar, Dict, Optional, Type
 
 from zenml.enums import PluginSubType
 from zenml.event_sources.base_event import BaseEvent
@@ -33,9 +33,6 @@ from zenml.models import EventSourceResponse
 
 logger = get_logger(__name__)
 
-if TYPE_CHECKING:
-    pass
-
 
 # -------------------- Event Models -----------------------------------
 

zenml/exceptions.py

@@ -197,6 +197,10 @@ class EntityExistsError(ZenMLBaseException):
     """Raised when trying to register an entity that already exists."""
 
 
+class ActionExistsError(EntityExistsError):
+    """Raised when registering an action with a name that already exists."""
+
+
 class TriggerExistsError(EntityExistsError):
     """Raised when registering a trigger with name that already exists."""
 

zenml/hooks/hook_validators.py

@@ -14,14 +14,13 @@
 """Validation functions for hooks."""
 
 import inspect
-from types import FunctionType
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from zenml.config.source import Source
 from zenml.utils import source_utils
 
 if TYPE_CHECKING:
-    HookSpecification = Union[str, Source, FunctionType]
+    from zenml.types import HookSpecification
 
 
 def resolve_and_validate_hook(hook: "HookSpecification") -> Source:

zenml/integrations/bitbucket/plugins/event_sources/bitbucket_webhook_event_source.py

@@ -152,9 +152,9 @@ class BitbucketEvent(BaseEvent):
 class BitbucketWebhookEventFilterConfiguration(WebhookEventFilterConfig):
     """Configuration for Bitbucket event filters."""
 
-    repo: Optional[str]
-    branch: Optional[str]
-    event_type: Optional[BitbucketEventType]
+    repo: Optional[str] = None
+    branch: Optional[str] = None
+    event_type: Optional[BitbucketEventType] = None
 
     def event_matches_filter(self, event: BaseEvent) -> bool:
         """Checks the filter against the inbound event.

zenml/integrations/mlflow/__init__.py

@@ -33,7 +33,7 @@ class MlflowIntegration(Integration):
     NAME = MLFLOW
 
     REQUIREMENTS = [
-        "mlflow>=2.1.1,<=2.12.2",
+        "mlflow>=2.1.1,<=2.14.1",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed

zenml/integrations/s3/artifact_stores/s3_artifact_store.py

@@ -26,6 +26,7 @@ from typing import (
 )
 
 import s3fs
+from fsspec.asyn import FSTimeoutError, sync, sync_wrapper
 
 from zenml.artifact_stores import BaseArtifactStore
 from zenml.integrations.s3.flavors.s3_artifact_store_flavor import (
@@ -38,10 +39,77 @@ from zenml.stack.authentication_mixin import AuthenticationMixin
 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
 
     @property
     def config(self) -> S3ArtifactStoreConfig:
@@ -98,7 +166,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
             return self.config.key, self.config.secret, self.config.token
 
     @property
-    def filesystem(self) -> s3fs.S3FileSystem:
+    def filesystem(self) -> ZenMLS3Filesystem:
         """The s3 filesystem to access this artifact store.
 
         Returns:
@@ -110,7 +178,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
         key, secret, token = self.get_credentials()
 
-        self._filesystem = s3fs.S3FileSystem(
+        self._filesystem = ZenMLS3Filesystem(
             key=key,
             secret=secret,
             token=token,
@@ -120,6 +188,11 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         )
         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.
 

zenml/logging/step_logging.py

@@ -145,61 +145,64 @@ def fetch_logs(
         )
 
     artifact_store = _load_artifact_store(artifact_store_id, zen_store)
-    if not artifact_store.isdir(logs_uri):
-        return _read_file(logs_uri, offset, length)
-    else:
-        files = artifact_store.listdir(logs_uri)
-        if len(files) == 1:
-            return _read_file(
-                os.path.join(logs_uri, str(files[0])), offset, length
-            )
+    try:
+        if not artifact_store.isdir(logs_uri):
+            return _read_file(logs_uri, offset, length)
         else:
-            is_negative_offset = offset < 0
-            files.sort(reverse=is_negative_offset)
-
-            # search for the first file we need to read
-            latest_file_id = 0
-            for i, file in enumerate(files):
-                file_size: int = artifact_store.size(
-                    os.path.join(logs_uri, str(file))
-                )  # type: ignore[assignment]
-
-                if is_negative_offset:
-                    if file_size >= -offset:
-                        latest_file_id = -(i + 1)
-                        break
+            files = artifact_store.listdir(logs_uri)
+            if len(files) == 1:
+                return _read_file(
+                    os.path.join(logs_uri, str(files[0])), offset, length
+                )
+            else:
+                is_negative_offset = offset < 0
+                files.sort(reverse=is_negative_offset)
+
+                # search for the first file we need to read
+                latest_file_id = 0
+                for i, file in enumerate(files):
+                    file_size: int = artifact_store.size(
+                        os.path.join(logs_uri, str(file))
+                    )  # type: ignore[assignment]
+
+                    if is_negative_offset:
+                        if file_size >= -offset:
+                            latest_file_id = -(i + 1)
+                            break
+                        else:
+                            offset += file_size
                     else:
-                        offset += file_size
-                else:
-                    if file_size > offset:
-                        latest_file_id = i
+                        if file_size > offset:
+                            latest_file_id = i
+                            break
+                        else:
+                            offset -= file_size
+
+                # read the files according to pre-filtering
+                files.sort()
+                ret = []
+                for file in files[latest_file_id:]:
+                    ret.append(
+                        _read_file(
+                            os.path.join(logs_uri, str(file)),
+                            offset,
+                            length,
+                        )
+                    )
+                    offset = 0
+                    length -= len(ret[-1])
+                    if length <= 0:
+                        # stop further reading, if the whole length is already read
                         break
-                    else:
-                        offset -= file_size
-
-            # read the files according to pre-filtering
-            files.sort()
-            ret = []
-            for file in files[latest_file_id:]:
-                ret.append(
-                    _read_file(
-                        os.path.join(logs_uri, str(file)),
-                        offset,
-                        length,
+
+                if not ret:
+                    raise DoesNotExistException(
+                        f"Folder '{logs_uri}' is empty in artifact store "
+                        f"'{artifact_store.name}'."
                     )
-                )
-                offset = 0
-                length -= len(ret[-1])
-                if length <= 0:
-                    # stop further reading, if the whole length is already read
-                    break
-
-            if not ret:
-                raise DoesNotExistException(
-                    f"Folder '{logs_uri}' is empty in artifact store "
-                    f"'{artifact_store.name}'."
-                )
-            return "".join(ret)
+                return "".join(ret)
+    finally:
+        artifact_store.cleanup()
 
 
 class StepLogsStorage:

zenml/models/__init__.py

@@ -51,6 +51,15 @@ from zenml.models.v2.base.filter import (
 from zenml.models.v2.base.page import Page
 
 # V2 Core
+from zenml.models.v2.core.action import (
+    ActionFilter,
+    ActionRequest,
+    ActionResponse,
+    ActionResponseBody,
+    ActionResponseMetadata,
+    ActionResponseResources,
+    ActionUpdate,
+)
 from zenml.models.v2.core.action_flavor import (
     ActionFlavorResponse,
     ActionFlavorResponseBody,
@@ -381,6 +390,7 @@ from zenml.models.v2.core.server_settings import (
 # ----------------------------- Forward References -----------------------------
 
 # V2
+ActionResponseResources.model_rebuild()
 APIKeyResponseBody.model_rebuild()
 ArtifactVersionRequest.model_rebuild()
 ArtifactVersionResponseBody.model_rebuild()
@@ -470,6 +480,13 @@ __all__ = [
     "UUIDFilter",
     "Page",
     # V2 Core
+    "ActionFilter",
+    "ActionRequest",
+    "ActionResponse",
+    "ActionResponseBody",
+    "ActionResponseMetadata",
+    "ActionResponseResources",
+    "ActionUpdate",
     "ActionFlavorResponse",
     "ActionFlavorResponseBody",
     "ActionFlavorResponseMetadata",