zenml-nightly 0.58.2.dev20240624__py3-none-any.whl → 0.58.2.dev20240625__py3-none-any.whl

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/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",

zenml/models/v2/core/action.py

@@ -0,0 +1,276 @@
+#  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.
+"""Collection of all models concerning actions."""
+
+import copy
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Optional,
+    TypeVar,
+)
+from uuid import UUID
+
+from pydantic import Field
+
+from zenml.constants import STR_FIELD_MAX_LENGTH
+from zenml.enums import PluginSubType
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
+    WorkspaceScopedFilter,
+    WorkspaceScopedRequest,
+    WorkspaceScopedResponse,
+    WorkspaceScopedResponseBody,
+    WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
+)
+from zenml.models.v2.core.user import UserResponse
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+
+# ------------------ Request Model ------------------
+
+
+class ActionRequest(WorkspaceScopedRequest):
+    """Model for creating a new action."""
+
+    name: str = Field(
+        title="The name of the action.", max_length=STR_FIELD_MAX_LENGTH
+    )
+    description: str = Field(
+        default="",
+        title="The description of the action",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    flavor: str = Field(
+        title="The flavor of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    plugin_subtype: PluginSubType = Field(
+        title="The subtype of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Dict[str, Any] = Field(
+        title="The configuration for the action.",
+    )
+    service_account_id: UUID = Field(
+        title="The service account that is used to execute the action.",
+    )
+    auth_window: Optional[int] = Field(
+        default=None,
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action. Set this to 0 to authorize the "
+        "service account indefinitely (not recommended). If not set, a "
+        "default value defined for each individual action type is used.",
+    )
+
+
+# ------------------ Update Model ------------------
+
+
+class ActionUpdate(BaseUpdate):
+    """Update model for actions."""
+
+    name: Optional[str] = Field(
+        default=None,
+        title="The new name for the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    description: Optional[str] = Field(
+        default=None,
+        title="The new description for the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The configuration for the action.",
+    )
+    service_account_id: Optional[UUID] = Field(
+        default=None,
+        title="The service account that is used to execute the action.",
+    )
+    auth_window: Optional[int] = Field(
+        default=None,
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action. Set this to 0 to authorize the "
+        "service account indefinitely (not recommended). If not set, a "
+        "default value defined for each individual action type is used.",
+    )
+
+    @classmethod
+    def from_response(cls, response: "ActionResponse") -> "ActionUpdate":
+        """Create an update model from a response model.
+
+        Args:
+            response: The response model to create the update model from.
+
+        Returns:
+            The update model.
+        """
+        return ActionUpdate(
+            configuration=copy.deepcopy(response.configuration),
+        )
+
+
+# ------------------ Response Model ------------------
+
+
+class ActionResponseBody(WorkspaceScopedResponseBody):
+    """Response body for actions."""
+
+    flavor: str = Field(
+        title="The flavor of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    plugin_subtype: PluginSubType = Field(
+        title="The subtype of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+
+
+class ActionResponseMetadata(WorkspaceScopedResponseMetadata):
+    """Response metadata for actions."""
+
+    description: str = Field(
+        default="",
+        title="The description of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    configuration: Dict[str, Any] = Field(
+        title="The configuration for the action.",
+    )
+    auth_window: int = Field(
+        title="The time window in minutes for which the service account is "
+        "authorized to execute the action."
+    )
+
+
+class ActionResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the action entity."""
+
+    service_account: UserResponse = Field(
+        title="The service account that is used to execute the action.",
+    )
+
+
+class ActionResponse(
+    WorkspaceScopedResponse[
+        ActionResponseBody, ActionResponseMetadata, ActionResponseResources
+    ]
+):
+    """Response model for actions."""
+
+    name: str = Field(
+        title="The name of the action.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+
+    def get_hydrated_version(self) -> "ActionResponse":
+        """Get the hydrated version of this action.
+
+        Returns:
+            An instance of the same entity with the metadata field attached.
+        """
+        from zenml.client import Client
+
+        return Client().zen_store.get_action(self.id)
+
+    # Body and metadata properties
+    @property
+    def flavor(self) -> str:
+        """The `flavor` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().flavor
+
+    @property
+    def plugin_subtype(self) -> PluginSubType:
+        """The `plugin_subtype` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().plugin_subtype
+
+    @property
+    def description(self) -> str:
+        """The `description` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().description
+
+    @property
+    def auth_window(self) -> int:
+        """The `auth_window` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().auth_window
+
+    @property
+    def configuration(self) -> Dict[str, Any]:
+        """The `configuration` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().configuration
+
+    def set_configuration(self, configuration: Dict[str, Any]) -> None:
+        """Set the `configuration` property.
+
+        Args:
+            configuration: The value to set.
+        """
+        self.get_metadata().configuration = configuration
+
+    # Resource properties
+    @property
+    def service_account(self) -> "UserResponse":
+        """The `service_account` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().service_account
+
+
+# ------------------ Filter Model ------------------
+
+
+class ActionFilter(WorkspaceScopedFilter):
+    """Model to enable advanced filtering of all actions."""
+
+    name: Optional[str] = Field(
+        default=None,
+        description="Name of the action.",
+    )
+    flavor: Optional[str] = Field(
+        default=None,
+        title="The flavor of the action.",
+    )
+    plugin_subtype: Optional[str] = Field(
+        default=None,
+        title="The subtype of the action.",
+    )