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

zenml/actions/pipeline_run/pipeline_run_action.py

@@ -24,20 +24,21 @@ from zenml.actions.base_action import (
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.enums import PluginSubType
+from zenml.logger import get_logger
 from zenml.models import (
+    ActionRequest,
+    ActionResponse,
+    ActionUpdate,
     TriggerExecutionResponse,
-    TriggerRequest,
-    TriggerResponse,
-    TriggerUpdate,
 )
 from zenml.models.v2.base.base import BaseResponse
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.pipeline_deployment.utils import run_pipeline
-from zenml.zen_server.rbac.models import (  # TODO: Maybe we move these into a common place?
-    ResourceType,
-)
+from zenml.zen_server.rbac.models import ResourceType
 from zenml.zen_server.utils import server_config
 
+logger = get_logger(__name__)
+
 # -------------------- Configuration Models ----------------------------------
 
 
@@ -95,7 +96,7 @@ class PipelineRunActionHandler(BaseActionHandler):
         assert isinstance(config, PipelineRunActionConfiguration)
 
         deployment = zen_store().get_deployment(config.pipeline_deployment_id)
-        print("Running deployment:", deployment)
+        logger.debug("Running deployment:", deployment)
         run_pipeline(
             deployment=deployment,
             run_config=config.run_config,
@@ -121,13 +122,13 @@ class PipelineRunActionHandler(BaseActionHandler):
         except KeyError:
             raise ValueError(f"No deployment found with id {deployment_id}.")
 
-    def _validate_trigger_request(
-        self, trigger: TriggerRequest, config: ActionConfig
+    def _validate_action_request(
+        self, action: ActionRequest, config: ActionConfig
     ) -> None:
-        """Validate a trigger request before it is created in the database.
+        """Validate an action request before it is created in the database.
 
         Args:
-            trigger: Trigger request.
+            action: Action request.
             config: Action configuration instantiated from the request.
         """
         assert isinstance(config, PipelineRunActionConfiguration)
@@ -135,25 +136,25 @@ class PipelineRunActionHandler(BaseActionHandler):
         self._validate_configuration(config)
 
         # If an expiration window is not set, we set it to the default value
-        if trigger.auth_window is None:
-            trigger.auth_window = server_config().pipeline_run_auth_window
+        if action.auth_window is None:
+            action.auth_window = server_config().pipeline_run_auth_window
 
-    def _validate_trigger_update(
+    def _validate_action_update(
         self,
-        trigger: TriggerResponse,
+        action: ActionResponse,
         config: ActionConfig,
-        trigger_update: TriggerUpdate,
+        action_update: ActionUpdate,
         config_update: ActionConfig,
     ) -> None:
-        """Validate a trigger update before it is reflected in the database.
+        """Validate an action update before it is reflected in the database.
 
         Args:
-            trigger: Original trigger before the update.
+            action: Original action before the update.
             config: Action configuration instantiated from the original
-                trigger.
-            trigger_update: Trigger update request.
+                action.
+            action_update: Action update request.
             config_update: Action configuration instantiated from the
-                updated trigger.
+                updated action.
         """
         assert isinstance(config, PipelineRunActionConfiguration)
 
@@ -162,12 +163,14 @@ class PipelineRunActionHandler(BaseActionHandler):
     def extract_resources(
         self,
         action_config: ActionConfig,
+        hydrate: bool = False,
     ) -> Dict[ResourceType, BaseResponse[Any, Any, Any]]:
         """Extract related resources for this action.
 
         Args:
             action_config: Action configuration from which to extract related
                 resources.
+            hydrate: Flag deciding whether to hydrate the resources.
 
         Returns:
             List of resources related to the action.
@@ -181,7 +184,9 @@ class PipelineRunActionHandler(BaseActionHandler):
         zen_store = GlobalConfiguration().zen_store
 
         try:
-            deployment = zen_store.get_deployment(deployment_id=deployment_id)
+            deployment = zen_store.get_deployment(
+                deployment_id=deployment_id, hydrate=hydrate
+            )
         except KeyError:
             raise ValueError(f"No deployment found with id {deployment_id}.")
 
@@ -191,7 +196,7 @@ class PipelineRunActionHandler(BaseActionHandler):
 
         if deployment.pipeline is not None:
             pipeline = zen_store.get_pipeline(
-                pipeline_id=deployment.pipeline.id
+                pipeline_id=deployment.pipeline.id, hydrate=hydrate
             )
             resources[ResourceType.PIPELINE] = pipeline
 

zenml/client.py

@@ -82,6 +82,10 @@ from zenml.exceptions import (
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.models import (
+    ActionFilter,
+    ActionRequest,
+    ActionResponse,
+    ActionUpdate,
     APIKeyFilter,
     APIKeyRequest,
     APIKeyResponse,
@@ -201,7 +205,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
-T = TypeVar("T")
+F = TypeVar("F", bound=Callable[..., Any])
 
 
 class ClientConfiguration(FileSyncModel):
@@ -307,11 +311,11 @@ class ClientMetaClass(ABCMeta):
         return cls._global_client
 
 
-def _fail_for_sql_zen_store(method: Callable[..., T]) -> Callable[..., T]:
-    """Decorator for all methods, that are disallowed when the client is not connected through REST API.
+def _fail_for_sql_zen_store(method: F) -> F:
+    """Decorator for methods that are not allowed with a SQLZenStore.
 
     Args:
-        method: The method
+        method: The method to decorate.
 
     Returns:
         The decorated method.
@@ -327,7 +331,7 @@ def _fail_for_sql_zen_store(method: Callable[..., T]) -> Callable[..., T]:
             )
         return method(self, *args, **kwargs)
 
-    return wrapper
+    return cast(F, wrapper)
 
 
 @evaluate_all_lazy_load_args_in_client_methods
@@ -2754,18 +2758,18 @@ class Client(metaclass=ClientMetaClass):
         self,
         name: str,
         configuration: Dict[str, Any],
-        description: str,
         flavor: str,
         event_source_subtype: PluginSubType,
+        description: str = "",
     ) -> EventSourceResponse:
-        """Registers a event_source.
+        """Registers an event source.
 
         Args:
-            name: The name of the event_source to create.
-            configuration: Configuration for this event source
-            description: The description of the event_source
-            flavor: The flavor of event source
-            event_source_subtype: str
+            name: The name of the event source to create.
+            configuration: Configuration for this event source.
+            flavor: The flavor of event source.
+            event_source_subtype: The event source subtype.
+            description: The description of the event source.
 
         Returns:
             The model of the registered event source.
@@ -2935,52 +2939,218 @@ class Client(metaclass=ClientMetaClass):
         self.zen_store.delete_event_source(event_source_id=event_source.id)
         logger.info("Deleted event_source with name '%s'.", event_source.name)
 
+    # --------------------------------- Actions -------------------------
+
+    @_fail_for_sql_zen_store
+    def create_action(
+        self,
+        name: str,
+        flavor: str,
+        action_type: PluginSubType,
+        configuration: Dict[str, Any],
+        service_account_id: UUID,
+        auth_window: Optional[int] = None,
+        description: str = "",
+    ) -> ActionResponse:
+        """Create an action.
+
+        Args:
+            name: The name of the action.
+            flavor: The flavor of the action,
+            action_type: The action subtype.
+            configuration: The action configuration.
+            service_account_id: The service account that is used to execute the
+                action.
+            auth_window: 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).
+            description: The description of the action.
+
+        Returns:
+            The created action
+        """
+        action = ActionRequest(
+            name=name,
+            description=description,
+            flavor=flavor,
+            plugin_subtype=action_type,
+            configuration=configuration,
+            service_account_id=service_account_id,
+            auth_window=auth_window,
+            user=self.active_user.id,
+            workspace=self.active_workspace.id,
+        )
+
+        return self.zen_store.create_action(action=action)
+
+    @_fail_for_sql_zen_store
+    def get_action(
+        self,
+        name_id_or_prefix: Union[UUID, str],
+        allow_name_prefix_match: bool = True,
+        hydrate: bool = True,
+    ) -> ActionResponse:
+        """Get an action by name, ID or prefix.
+
+        Args:
+            name_id_or_prefix: The name, ID or prefix of the action.
+            allow_name_prefix_match: If True, allow matching by name prefix.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            The action.
+        """
+        return self._get_entity_by_id_or_name_or_prefix(
+            get_method=self.zen_store.get_action,
+            list_method=self.list_actions,
+            name_id_or_prefix=name_id_or_prefix,
+            allow_name_prefix_match=allow_name_prefix_match,
+            hydrate=hydrate,
+        )
+
+    @_fail_for_sql_zen_store
+    def list_actions(
+        self,
+        sort_by: str = "created",
+        page: int = PAGINATION_STARTING_PAGE,
+        size: int = PAGE_SIZE_DEFAULT,
+        logical_operator: LogicalOperators = LogicalOperators.AND,
+        id: Optional[Union[UUID, str]] = None,
+        created: Optional[datetime] = None,
+        updated: Optional[datetime] = None,
+        name: Optional[str] = None,
+        flavor: Optional[str] = None,
+        action_type: Optional[str] = None,
+        workspace_id: Optional[Union[str, UUID]] = None,
+        user_id: Optional[Union[str, UUID]] = None,
+        hydrate: bool = False,
+    ) -> Page[ActionResponse]:
+        """List actions.
+
+        Args:
+            sort_by: The column to sort by
+            page: The page of items
+            size: The maximum size of all pages
+            logical_operator: Which logical operator to use [and, or]
+            id: Use the id of the action to filter by.
+            created: Use to filter by time of creation
+            updated: Use the last updated date for filtering
+            workspace_id: The id of the workspace to filter by.
+            user_id: The id of the user to filter by.
+            name: The name of the action to filter by.
+            flavor: The flavor of the action to filter by.
+            action_type: The type of the action to filter by.
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
+        Returns:
+            A page of actions.
+        """
+        filter_model = ActionFilter(
+            page=page,
+            size=size,
+            sort_by=sort_by,
+            logical_operator=logical_operator,
+            workspace_id=workspace_id,
+            user_id=user_id,
+            name=name,
+            id=id,
+            flavor=flavor,
+            plugin_subtype=action_type,
+            created=created,
+            updated=updated,
+        )
+        filter_model.set_scope_workspace(self.active_workspace.id)
+        return self.zen_store.list_actions(filter_model, hydrate=hydrate)
+
+    @_fail_for_sql_zen_store
+    def update_action(
+        self,
+        name_id_or_prefix: Union[UUID, str],
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        configuration: Optional[Dict[str, Any]] = None,
+        service_account_id: Optional[UUID] = None,
+        auth_window: Optional[int] = None,
+    ) -> ActionResponse:
+        """Update an action.
+
+        Args:
+            name_id_or_prefix: The name, id or prefix of the action to update.
+            name: The new name of the action.
+            description: The new description of the action.
+            configuration: The new configuration of the action.
+            service_account_id: The new service account that is used to execute
+                the action.
+            auth_window: The new 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).
+
+        Returns:
+            The updated action.
+        """
+        action = self.get_action(
+            name_id_or_prefix=name_id_or_prefix, allow_name_prefix_match=False
+        )
+
+        update_model = ActionUpdate(
+            name=name,
+            description=description,
+            configuration=configuration,
+            service_account_id=service_account_id,
+            auth_window=auth_window,
+        )
+
+        return self.zen_store.update_action(
+            action_id=action.id,
+            action_update=update_model,
+        )
+
+    @_fail_for_sql_zen_store
+    def delete_action(self, name_id_or_prefix: Union[str, UUID]) -> None:
+        """Delete an action.
+
+        Args:
+            name_id_or_prefix: The name, id or prefix id of the action
+                to delete.
+        """
+        action = self.get_action(
+            name_id_or_prefix=name_id_or_prefix, allow_name_prefix_match=False
+        )
+
+        self.zen_store.delete_action(action_id=action.id)
+        logger.info("Deleted action with name '%s'.", action.name)
+
     # --------------------------------- Triggers -------------------------
 
     @_fail_for_sql_zen_store
     def create_trigger(
         self,
         name: str,
-        description: str,
         event_source_id: UUID,
         event_filter: Dict[str, Any],
-        action: Dict[str, Any],
-        action_flavor: str,
-        action_subtype: PluginSubType,
-        service_account: Union[str, UUID],
-        auth_window: Optional[int] = None,
+        action_id: UUID,
+        description: str = "",
     ) -> TriggerResponse:
         """Registers a trigger.
 
         Args:
             name: The name of the trigger to create.
-            description: The description of the trigger
             event_source_id: The id of the event source id
             event_filter: The event filter
-            action: The action
-            action_flavor: The action flavor
-            action_subtype: The action subtype
-            service_account: The service account
-            auth_window: The auth window
+            action_id: The ID of the action that should be triggered.
+            description: The description of the trigger
 
         Returns:
-            The model of the registered event source.
+            The created trigger.
         """
-        # Fetch the service account
-        service_account_model = self.get_service_account(
-            name_id_or_prefix=service_account, allow_name_prefix_match=False
-        )
-
         trigger = TriggerRequest(
             name=name,
             description=description,
             event_source_id=event_source_id,
             event_filter=event_filter,
-            action=action,
-            action_flavor=action_flavor,
-            action_subtype=action_subtype,
-            service_account_id=service_account_model.id,
-            auth_window=auth_window,
+            action_id=action_id,
             user=self.active_user.id,
             workspace=self.active_workspace.id,
         )
@@ -2994,10 +3164,10 @@ class Client(metaclass=ClientMetaClass):
         allow_name_prefix_match: bool = True,
         hydrate: bool = True,
     ) -> TriggerResponse:
-        """Get a event source by name, ID or prefix.
+        """Get a trigger by name, ID or prefix.
 
         Args:
-            name_id_or_prefix: The name, ID or prefix of the stack.
+            name_id_or_prefix: The name, ID or prefix of the trigger.
             allow_name_prefix_match: If True, allow matching by name prefix.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
@@ -3025,6 +3195,11 @@ class Client(metaclass=ClientMetaClass):
         updated: Optional[datetime] = None,
         name: Optional[str] = None,
         event_source_id: Optional[UUID] = None,
+        action_id: Optional[UUID] = None,
+        event_source_flavor: Optional[str] = None,
+        event_source_subtype: Optional[str] = None,
+        action_flavor: Optional[str] = None,
+        action_subtype: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
         hydrate: bool = False,
@@ -3042,7 +3217,14 @@ class Client(metaclass=ClientMetaClass):
             workspace_id: The id of the workspace to filter by.
             user_id: The  id of the user to filter by.
             name: The name of the trigger to filter by.
-            event_source_id: The event source associated with the Trigger
+            event_source_id: The event source associated with the trigger.
+            action_id: The action associated with the trigger.
+            event_source_flavor: Flavor of the event source associated with the
+                trigger.
+            event_source_subtype: Type of the event source associated with the
+                trigger.
+            action_flavor: Flavor of the action associated with the trigger.
+            action_subtype: Type of the action associated with the trigger.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -3058,6 +3240,11 @@ class Client(metaclass=ClientMetaClass):
             user_id=user_id,
             name=name,
             event_source_id=event_source_id,
+            action_id=action_id,
+            event_source_flavor=event_source_flavor,
+            event_source_subtype=event_source_subtype,
+            action_flavor=action_flavor,
+            action_subtype=action_subtype,
             id=id,
             created=created,
             updated=updated,
@@ -3074,10 +3261,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         description: Optional[str] = None,
         event_filter: Optional[Dict[str, Any]] = None,
-        action: Optional[Dict[str, Any]] = None,
         is_active: Optional[bool] = None,
-        service_account: Optional[Union[str, UUID]] = None,
-        auth_window: Optional[int] = None,
     ) -> TriggerResponse:
         """Updates a trigger.
 
@@ -3086,11 +3270,7 @@ class Client(metaclass=ClientMetaClass):
             name: the new name of the trigger.
             description: the new description of the trigger.
             event_filter: The event filter configuration.
-            action: The action configuration.
-            is_active: Optional[bool] = Allows for activation/deactivating the
-                event source
-            service_account: The service account
-            auth_window: The auth window
+            is_active: Whether the trigger is active or not.
 
         Returns:
             The model of the updated trigger.
@@ -3108,17 +3288,8 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             description=description,
             event_filter=event_filter,
-            action=action,
             is_active=is_active,
-            auth_window=auth_window,
         )
-        if service_account:
-            # Fetch the service account
-            service_account_model = self.get_service_account(
-                name_id_or_prefix=service_account,
-                allow_name_prefix_match=False,
-            )
-            update_model.service_account_id = service_account_model.id
 
         if name:
             if self.list_triggers(name=name):

zenml/config/compiler.py

@@ -192,10 +192,6 @@ class Compiler:
         Args:
             pipeline: The pipeline to configure.
             config: The run configurations.
-
-        Raises:
-            KeyError: If the run configuration contains options for a
-                non-existent step.
         """
         with pipeline.__suppress_configure_warnings__():
             pipeline.configure(
@@ -209,11 +205,16 @@ class Compiler:
                 parameters=config.parameters,
             )
 
-        for invocation_id in config.steps:
-            if invocation_id not in pipeline.invocations:
-                raise KeyError(
-                    f"Configuration for step {invocation_id} cannot be applied to any pipeline step. Make sure that all configured steps are present in your pipeline."
-                )
+        invalid_step_configs = set(config.steps) - set(pipeline.invocations)
+        if invalid_step_configs:
+            logger.warning(
+                f"Configuration for step invocations {invalid_step_configs} "
+                "cannot be applied to any pipeline step invocations, "
+                "ignoring..."
+            )
+
+        for key in invalid_step_configs:
+            config.steps.pop(key)
 
         # Override `enable_cache` of all steps if set at run level
         if config.enable_cache is not None:

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,
         )