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

zenml/models/v2/core/trigger.py

@@ -13,15 +13,14 @@
 #  permissions and limitations under the License.
 """Collection of all models concerning triggers."""
 
-import copy
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import Field, model_validator
 
+from zenml.config.schedule import Schedule
 from zenml.constants import STR_FIELD_MAX_LENGTH
-from zenml.enums import PluginSubType
-from zenml.models.v2.base.base import BaseZenModel
+from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.page import Page
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
@@ -32,138 +31,129 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseResources,
 )
 from zenml.models.v2.core.trigger_execution import TriggerExecutionResponse
-from zenml.models.v2.core.user import UserResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
+    from zenml.models.v2.core.action import (
+        ActionResponse,
+    )
     from zenml.models.v2.core.event_source import EventSourceResponse
 
-# ------------------ Base Model ------------------
+
+# ------------------ Request Model ------------------
 
 
-class TriggerBase(BaseModel):
-    """Base model for triggers."""
+class TriggerRequest(WorkspaceScopedRequest):
+    """Model for creating a new trigger."""
 
     name: str = Field(
-        title="The name of the Trigger.", max_length=STR_FIELD_MAX_LENGTH
+        title="The name of the trigger.", max_length=STR_FIELD_MAX_LENGTH
     )
     description: str = Field(
         default="",
         title="The description of the trigger",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    event_source_id: UUID = Field(
-        title="The event source that activates this trigger.",
-    )
-    event_filter: Dict[str, Any] = Field(
-        title="Filter applied to events that activate this trigger.",
-    )
-
-    action: Dict[str, Any] = Field(
-        title="The configuration for the action that is executed by this "
-        "trigger.",
-    )
-    action_flavor: str = Field(
-        title="The flavor of the action that is executed by this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
+    action_id: UUID = Field(
+        title="The action that is executed by this trigger.",
     )
-    action_subtype: PluginSubType = Field(
-        title="The subtype of the action that is executed by this trigger.",
+    schedule: Optional[Schedule] = Field(
+        default=None,
+        title="The schedule for the trigger. Either a schedule or an event "
+        "source is required.",
     )
-    service_account_id: UUID = Field(
-        title="The service account that is used to execute the action.",
+    event_source_id: Optional[UUID] = Field(
+        default=None,
+        title="The event source that activates this trigger. Either a schedule "
+        "or an event source is required.",
     )
-    auth_window: Optional[int] = Field(
+    event_filter: Optional[Dict[str, Any]] = 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.",
+        title="Filter applied to events that activate this trigger. Only "
+        "set if the trigger is activated by an event source.",
     )
 
+    @model_validator(mode="after")
+    def _validate_schedule_or_event_source(self) -> "TriggerRequest":
+        """Validate that either a schedule or an event source is provided.
 
-# ------------------ Request Model ------------------
-class TriggerRequest(TriggerBase, WorkspaceScopedRequest):
-    """Model for creating a new Trigger."""
+        Returns:
+            The validated request.
+
+        Raises:
+            ValueError: If neither a schedule nor an event source is provided,
+                or if both are provided.
+        """
+        if not self.schedule and not self.event_source_id:
+            raise ValueError(
+                "Either a schedule or an event source is required."
+            )
+
+        if self.schedule and self.event_source_id:
+            raise ValueError("Only a schedule or an event source is allowed.")
+
+        return self
 
 
 # ------------------ Update Model ------------------
 
 
-class TriggerUpdate(BaseZenModel):
+class TriggerUpdate(BaseUpdate):
     """Update model for triggers."""
 
     name: Optional[str] = Field(
         default=None,
-        title="The new name for the Trigger.",
+        title="The new name for the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     description: Optional[str] = Field(
         default=None,
-        title="The new description for the trigger",
+        title="The new description for the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
     event_filter: Optional[Dict[str, Any]] = Field(
         default=None,
-        title="New filter applied to events that activate this trigger.",
+        title="New filter applied to events that activate this trigger. Only "
+        "valid if the trigger is already configured to be activated by an "
+        "event source.",
     )
-    action: Optional[Dict[str, Any]] = Field(
+    schedule: Optional[Schedule] = Field(
         default=None,
-        title="The new configuration for the action that is executed by this "
-        "trigger.",
+        title="The updated schedule for the trigger. Only valid if the trigger "
+        "is already configured to be activated by a schedule.",
     )
     is_active: Optional[bool] = Field(
         default=None,
         title="The new status of the trigger.",
     )
-    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: "TriggerResponse") -> "TriggerUpdate":
-        """Create an update model from a response model.
-
-        Args:
-            response: The response model to create the update model from.
-
-        Returns:
-            The update model.
-        """
-        return TriggerUpdate(
-            name=response.name,
-            description=response.description,
-            action=copy.deepcopy(response.action),
-            event_filter=copy.deepcopy(response.event_filter),
-            is_active=response.is_active,
-            service_account_id=response.get_resources().service_account.id,
-        )
 
 
 # ------------------ Response Model ------------------
 
 
 class TriggerResponseBody(WorkspaceScopedResponseBody):
-    """ResponseBody for triggers."""
+    """Response body for triggers."""
 
-    event_source_flavor: str = Field(
-        title="The flavor of the event source that activates this trigger.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     action_flavor: str = Field(
         title="The flavor of the action that is executed by this trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    action_subtype: PluginSubType = Field(
+    action_subtype: str = Field(
         title="The subtype of the action that is executed by this trigger.",
     )
+    event_source_flavor: Optional[str] = Field(
+        default=None,
+        title="The flavor of the event source that activates this trigger. Not "
+        "set if the trigger is activated by a schedule.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    event_source_subtype: Optional[str] = Field(
+        default=None,
+        title="The subtype of the event source that activates this trigger. "
+        "Not set if the trigger is activated by a schedule.",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
     is_active: bool = Field(
         title="Whether the trigger is active.",
     )
@@ -172,33 +162,33 @@ class TriggerResponseBody(WorkspaceScopedResponseBody):
 class TriggerResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for triggers."""
 
-    event_filter: Dict[str, Any] = Field(
-        title="The event that activates this trigger.",
-    )
-    action: Dict[str, Any] = Field(
-        title="The action that is executed by this trigger.",
-    )
     description: str = Field(
         default="",
-        title="The description of the trigger",
+        title="The description of the trigger.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    auth_window: int = Field(
-        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.",
+    event_filter: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The event that activates this trigger. Not set if the trigger "
+        "is activated by a schedule.",
+    )
+    schedule: Optional[Schedule] = Field(
+        default=None,
+        title="The schedule that activates this trigger. Not set if the "
+        "trigger is activated by an event source.",
     )
 
 
 class TriggerResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the trigger entity."""
 
-    event_source: "EventSourceResponse" = Field(
-        title="The event source that activates this trigger.",
+    action: "ActionResponse" = Field(
+        title="The action that is executed by this trigger.",
     )
-    service_account: UserResponse = Field(
-        title="The service account that is used to execute the action.",
+    event_source: Optional["EventSourceResponse"] = Field(
+        default=None,
+        title="The event source that activates this trigger. Not set if the "
+        "trigger is activated by a schedule.",
     )
     executions: Page[TriggerExecutionResponse] = Field(
         title="The executions of this trigger.",
@@ -213,7 +203,7 @@ class TriggerResponse(
     """Response model for models."""
 
     name: str = Field(
-        title="The name of the model",
+        title="The name of the trigger",
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
@@ -227,15 +217,6 @@ class TriggerResponse(
 
         return Client().zen_store.get_trigger(self.id)
 
-    @property
-    def event_source_flavor(self) -> str:
-        """The `event_source_flavor` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().event_source_flavor
-
     @property
     def action_flavor(self) -> str:
         """The `action_flavor` property.
@@ -246,7 +227,7 @@ class TriggerResponse(
         return self.get_body().action_flavor
 
     @property
-    def action_subtype(self) -> PluginSubType:
+    def action_subtype(self) -> str:
         """The `action_subtype` property.
 
         Returns:
@@ -255,48 +236,40 @@ class TriggerResponse(
         return self.get_body().action_subtype
 
     @property
-    def is_active(self) -> bool:
-        """The `is_active` property.
+    def event_source_flavor(self) -> Optional[str]:
+        """The `event_source_flavor` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_body().is_active
+        return self.get_body().event_source_flavor
 
     @property
-    def event_filter(self) -> Dict[str, Any]:
-        """The `event_filter` property.
+    def event_source_subtype(self) -> Optional[str]:
+        """The `event_source_subtype` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().event_filter
+        return self.get_body().event_source_subtype
 
     @property
-    def action(self) -> Dict[str, Any]:
-        """The `action` property.
+    def is_active(self) -> bool:
+        """The `is_active` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().action
-
-    def set_action(self, action: Dict[str, Any]) -> None:
-        """Set the `action` property.
-
-        Args:
-            action: The value to set.
-        """
-        self.get_metadata().action = action
+        return self.get_body().is_active
 
     @property
-    def event_source(self) -> "EventSourceResponse":
-        """The `event_source` property.
+    def event_filter(self) -> Optional[Dict[str, Any]]:
+        """The `event_filter` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_resources().event_source
+        return self.get_metadata().event_filter
 
     @property
     def description(self) -> str:
@@ -308,37 +281,59 @@ class TriggerResponse(
         return self.get_metadata().description
 
     @property
-    def service_account(self) -> UserResponse:
-        """The `service_account` property.
+    def action(self) -> "ActionResponse":
+        """The `action` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_resources().service_account
+        return self.get_resources().action
 
     @property
-    def auth_window(self) -> int:
-        """The `auth_window` property.
+    def event_source(self) -> Optional["EventSourceResponse"]:
+        """The `event_source` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().auth_window
+        return self.get_resources().event_source
+
+    @property
+    def executions(self) -> Page[TriggerExecutionResponse]:
+        """The `event_source` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().executions
 
 
 # ------------------ Filter Model ------------------
 
 
 class TriggerFilter(WorkspaceScopedFilter):
-    """Model to enable advanced filtering of all TriggerModels."""
+    """Model to enable advanced filtering of all triggers."""
+
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "action_flavor",
+        "action_subtype",
+        "event_source_flavor",
+        "event_source_subtype",
+    ]
 
     name: Optional[str] = Field(
         default=None,
-        description="Name of the trigger",
+        description="Name of the trigger.",
     )
     event_source_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="By the event source this trigger is attached to.",
+        description="The event source this trigger is attached to.",
+        union_mode="left_to_right",
+    )
+    action_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="The action this trigger is attached to.",
         union_mode="left_to_right",
     )
     is_active: Optional[bool] = Field(
@@ -353,13 +348,59 @@ class TriggerFilter(WorkspaceScopedFilter):
         default=None,
         title="The subtype of the action that is executed by this trigger.",
     )
-    # TODO: Ignore these in normal filter and handle in sqlzenstore
-    resource_id: Optional[Union[UUID, str]] = Field(
+    event_source_flavor: Optional[str] = Field(
         default=None,
-        description="By the resource this trigger references.",
-        union_mode="left_to_right",
+        title="The flavor of the event source that activates this trigger.",
     )
-    resource_type: Optional[str] = Field(
+    event_source_subtype: Optional[str] = Field(
         default=None,
-        description="By the resource type this trigger references.",
+        title="The subtype of the event source that activates this trigger.",
     )
+
+    def get_custom_filters(
+        self,
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            ActionSchema,
+            EventSourceSchema,
+            TriggerSchema,
+        )
+
+        custom_filters = super().get_custom_filters()
+
+        if self.event_source_flavor:
+            event_source_flavor_filter = and_(
+                EventSourceSchema.id == TriggerSchema.event_source_id,
+                EventSourceSchema.flavor == self.event_source_flavor,
+            )
+            custom_filters.append(event_source_flavor_filter)
+
+        if self.event_source_subtype:
+            event_source_subtype_filter = and_(
+                EventSourceSchema.id == TriggerSchema.event_source_id,
+                EventSourceSchema.plugin_subtype == self.event_source_subtype,
+            )
+            custom_filters.append(event_source_subtype_filter)
+
+        if self.action_flavor:
+            action_flavor_filter = and_(
+                ActionSchema.id == TriggerSchema.action_id,
+                ActionSchema.flavor == self.action_flavor,
+            )
+            custom_filters.append(action_flavor_filter)
+
+        if self.action_subtype:
+            action_subtype_filter = and_(
+                ActionSchema.id == TriggerSchema.action_id,
+                ActionSchema.plugin_subtype == self.action_subtype,
+            )
+            custom_filters.append(action_subtype_filter)
+
+        return custom_filters

zenml/new/pipelines/pipeline.py

@@ -19,7 +19,6 @@ import inspect
 from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -98,11 +97,11 @@ if TYPE_CHECKING:
     from zenml.config.source import Source
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
     from zenml.model.model import Model
+    from zenml.types import HookSpecification
 
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, "Source", FunctionType]
 
 logger = get_logger(__name__)
 
@@ -1307,6 +1306,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         step_configurations: Optional[
             Mapping[str, "StepConfigurationUpdateOrDict"]
         ] = None,
+        steps: Optional[Mapping[str, "StepConfigurationUpdateOrDict"]] = None,
         config_path: Optional[str] = None,
         unlisted: bool = False,
         prevent_build_reuse: bool = False,
@@ -1319,6 +1319,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
             schedule: Optional schedule to use for the run.
             build: Optional build to use for the run.
             step_configurations: Configurations for steps of the pipeline.
+            steps: Configurations for steps of the pipeline. This is equivalent
+                to `step_configurations`, and will be ignored if
+                `step_configurations` is set as well.
             config_path: Path to a yaml configuration file. This file will
                 be parsed as a
                 `zenml.config.pipeline_configurations.PipelineRunConfiguration`
@@ -1334,6 +1337,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
         Returns:
             The copied pipeline instance.
         """
+        if steps and step_configurations:
+            logger.warning(
+                "Step configurations were passed using both the "
+                "`step_configurations` and `steps` keywords, ignoring the "
+                "values passed using the `steps` keyword."
+            )
+
         pipeline_copy = self.copy()
 
         pipeline_copy._parse_config_file(
@@ -1353,7 +1363,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 "run_name": run_name,
                 "schedule": schedule,
                 "build": build,
-                "step_configurations": step_configurations,
+                "step_configurations": step_configurations or steps,
                 "config_path": config_path,
                 "unlisted": unlisted,
                 "prevent_build_reuse": prevent_build_reuse,

zenml/new/pipelines/pipeline_decorator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """ZenML pipeline decorator definition."""
 
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -31,8 +30,8 @@ if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
     from zenml.model.model import Model
     from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.types import HookSpecification
 
-    HookSpecification = Union[str, FunctionType]
     F = TypeVar("F", bound=Callable[..., None])
 
 logger = get_logger(__name__)

zenml/new/pipelines/run_utils.py

@@ -3,16 +3,7 @@
 import time
 from collections import defaultdict
 from datetime import datetime
-from types import FunctionType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Union
 from uuid import UUID
 
 from zenml import constants
@@ -36,13 +27,11 @@ from zenml.utils import cloud_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
-    from zenml.config.source import Source
     from zenml.model.model import Model
 
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, "Source", FunctionType]
 
 logger = get_logger(__name__)
 

zenml/new/steps/step_decorator.py

@@ -30,17 +30,16 @@ from typing import (
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from types import FunctionType
-
     from zenml.config.base_settings import SettingsOrDict
     from zenml.config.retry_config import StepRetryConfig
     from zenml.config.source import Source
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.model.model import Model
     from zenml.steps import BaseStep
+    from zenml.types import HookSpecification
 
     MaterializerClassOrSource = Union[str, Source, Type[BaseMaterializer]]
-    HookSpecification = Union[str, Source, FunctionType]
+
     OutputMaterializersSpecification = Union[
         MaterializerClassOrSource,
         Sequence[MaterializerClassOrSource],

zenml/pipelines/base_pipeline.py

@@ -15,7 +15,6 @@
 
 import inspect
 from abc import ABC, abstractmethod
-from types import FunctionType
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Mapping, Optional, Union
 from uuid import UUID
 
@@ -35,7 +34,6 @@ if TYPE_CHECKING:
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, FunctionType]
 
 logger = get_logger(__name__)
 

zenml/pipelines/pipeline_decorator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Legacy ZenML pipeline decorator definition."""
 
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -46,8 +45,8 @@ from zenml.pipelines.base_pipeline import (
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.types import HookSpecification
 
-    HookSpecification = Union[str, FunctionType]
 
 logger = get_logger(__name__)