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

zenml/steps/base_step.py

@@ -18,7 +18,6 @@ import hashlib
 import inspect
 from abc import abstractmethod
 from collections import defaultdict
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -76,10 +75,10 @@ if TYPE_CHECKING:
     )
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
     from zenml.model.model import Model
+    from zenml.types import HookSpecification
 
     ParametersOrDict = Union["BaseParameters", Dict[str, Any]]
     MaterializerClassOrSource = Union[str, Source, Type["BaseMaterializer"]]
-    HookSpecification = Union[str, Source, FunctionType]
     OutputMaterializersSpecification = Union[
         "MaterializerClassOrSource",
         Sequence["MaterializerClassOrSource"],

zenml/steps/step_decorator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Step decorator function."""
 
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -37,9 +36,9 @@ if TYPE_CHECKING:
     from zenml.config.source import Source
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.model.model import Model
+    from zenml.types import HookSpecification
 
     MaterializerClassOrSource = Union[str, "Source", Type["BaseMaterializer"]]
-    HookSpecification = Union[str, "Source", FunctionType]
     OutputMaterializersSpecification = Union[
         "MaterializerClassOrSource",
         Sequence["MaterializerClassOrSource"],

zenml/types.py

@@ -11,7 +11,16 @@
 #  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.
-"""Custom types that are used to indicate how to handle data."""
+"""Custom ZenML types."""
+
+from typing import TYPE_CHECKING, Callable, Union
+
+if TYPE_CHECKING:
+    from types import FunctionType
+
+    from zenml.config.source import Source
+
+    HookSpecification = Union[str, Source, FunctionType, Callable[..., None]]
 
 
 class HTMLString(str):

zenml/utils/pipeline_docker_image_builder.py

@@ -32,6 +32,7 @@ from typing import (
 import zenml
 from zenml.config import DockerSettings
 from zenml.config.docker_settings import (
+    DockerBuildConfig,
     PythonEnvironmentExportMethod,
     PythonPackageInstaller,
 )
@@ -212,8 +213,13 @@ class PipelineDockerImageBuilder:
                 # used directly, so we tag it with the requested target name.
                 user_image_name = target_image_name
 
+            build_config = (
+                docker_settings.parent_image_build_config
+                or DockerBuildConfig()
+            )
             build_context = build_context_class(
-                root=docker_settings.build_context_root
+                root=docker_settings.build_context_root,
+                dockerignore_file=build_config.dockerignore,
             )
             build_context.add_file(
                 source=docker_settings.dockerfile, destination="Dockerfile"
@@ -222,7 +228,8 @@ class PipelineDockerImageBuilder:
             image_name_or_digest = image_builder.build(
                 image_name=user_image_name,
                 build_context=build_context,
-                docker_build_options=docker_settings.build_options,
+                docker_build_options=build_config.build_options
+                or docker_settings.build_options,
                 container_registry=container_registry if push else None,
             )
 
@@ -247,13 +254,18 @@ class PipelineDockerImageBuilder:
 
         if requires_zenml_build:
             logger.info("Building Docker image `%s`.", target_image_name)
+            build_config = docker_settings.build_config or DockerBuildConfig()
+
             # Leave the build context empty if we don't want to include any files
             build_context_root = (
                 source_utils.get_source_root() if include_files else None
             )
+            dockerignore = (
+                build_config.dockerignore or docker_settings.dockerignore
+            )
             build_context = build_context_class(
                 root=build_context_root,
-                dockerignore_file=docker_settings.dockerignore,
+                dockerignore_file=dockerignore,
             )
 
             requirements_files = self.gather_requirements_files(
@@ -304,8 +316,11 @@ class PipelineDockerImageBuilder:
                     parent_image
                 )
 
-            build_options = {"pull": pull_parent_image, "rm": False}
-
+            build_options = {
+                "pull": pull_parent_image,
+                "rm": False,
+                **build_config.build_options,
+            }
             dockerfile = self._generate_zenml_pipeline_dockerfile(
                 parent_image=parent_image,
                 docker_settings=docker_settings,

zenml/zen_server/rbac/models.py

@@ -46,6 +46,7 @@ class Action(StrEnum):
 class ResourceType(StrEnum):
     """Resource types of the server API."""
 
+    ACTION = "action"
     ARTIFACT = "artifact"
     ARTIFACT_VERSION = "artifact_version"
     CODE_REPOSITORY = "code_repository"

zenml/zen_server/rbac/utils.py

@@ -189,7 +189,12 @@ def get_permission_denied_model(model: AnyResponse) -> AnyResponse:
         The permission denied model.
     """
     return model.model_copy(
-        update={"body": None, "metadata": None, "permission_denied": True}
+        update={
+            "body": None,
+            "metadata": None,
+            "resources": None,
+            "permission_denied": True,
+        }
     )
 
 
@@ -384,10 +389,12 @@ def get_resource_type_for_model(
         is not associated with any resource type.
     """
     from zenml.models import (
+        ActionResponse,
         ArtifactResponse,
         ArtifactVersionResponse,
         CodeRepositoryResponse,
         ComponentResponse,
+        EventSourceResponse,
         FlavorResponse,
         ModelResponse,
         ModelVersionResponse,
@@ -402,6 +409,8 @@ def get_resource_type_for_model(
         ServiceResponse,
         StackResponse,
         TagResponse,
+        TriggerExecutionResponse,
+        TriggerResponse,
         UserResponse,
         WorkspaceResponse,
     )
@@ -410,6 +419,8 @@ def get_resource_type_for_model(
         Any,
         ResourceType,
     ] = {
+        ActionResponse: ResourceType.ACTION,
+        EventSourceResponse: ResourceType.EVENT_SOURCE,
         FlavorResponse: ResourceType.FLAVOR,
         ServiceConnectorResponse: ResourceType.SERVICE_CONNECTOR,
         ComponentResponse: ResourceType.STACK_COMPONENT,
@@ -428,6 +439,8 @@ def get_resource_type_for_model(
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,
         TagResponse: ResourceType.TAG,
+        TriggerResponse: ResourceType.TRIGGER,
+        TriggerExecutionResponse: ResourceType.TRIGGER_EXECUTION,
         ServiceAccountResponse: ResourceType.SERVICE_ACCOUNT,
         ServiceResponse: ResourceType.SERVICE,
     }
@@ -522,9 +535,11 @@ def get_schema_for_resource_type(
         The database schema.
     """
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         ArtifactSchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
+        EventSourceSchema,
         FlavorSchema,
         ModelSchema,
         ModelVersionSchema,
@@ -539,6 +554,8 @@ def get_schema_for_resource_type(
         StackComponentSchema,
         StackSchema,
         TagSchema,
+        TriggerExecutionSchema,
+        TriggerSchema,
         UserSchema,
         WorkspaceSchema,
     )
@@ -564,6 +581,10 @@ def get_schema_for_resource_type(
         ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
         ResourceType.RUN_METADATA: RunMetadataSchema,
         ResourceType.USER: UserSchema,
+        ResourceType.ACTION: ActionSchema,
+        ResourceType.EVENT_SOURCE: EventSourceSchema,
+        ResourceType.TRIGGER: TriggerSchema,
+        ResourceType.TRIGGER_EXECUTION: TriggerExecutionSchema,
     }
 
     return mapping[resource_type]

zenml/zen_server/routers/actions_endpoints.py

@@ -0,0 +1,324 @@
+#  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.
+"""Endpoint definitions for actions."""
+
+from uuid import UUID
+
+from fastapi import APIRouter, Depends, Security
+
+from zenml import ActionRequest
+from zenml.actions.base_action import BaseActionHandler
+from zenml.constants import (
+    ACTIONS,
+    API,
+    VERSION_1,
+)
+from zenml.enums import PluginType
+from zenml.models import (
+    ActionFilter,
+    ActionResponse,
+    ActionUpdate,
+    Page,
+)
+from zenml.zen_server.auth import AuthContext, authorize
+from zenml.zen_server.exceptions import error_response
+from zenml.zen_server.rbac.endpoint_utils import (
+    verify_permissions_and_create_entity,
+    verify_permissions_and_list_entities,
+)
+from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.utils import (
+    dehydrate_response_model,
+    verify_permission_for_model,
+)
+from zenml.zen_server.utils import (
+    handle_exceptions,
+    make_dependable,
+    plugin_flavor_registry,
+    zen_store,
+)
+
+router = APIRouter(
+    prefix=API + VERSION_1 + ACTIONS,
+    tags=["actions"],
+    responses={401: error_response, 403: error_response},
+)
+
+
+@router.get(
+    "",
+    response_model=Page[ActionResponse],
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def list_actions(
+    action_filter_model: ActionFilter = Depends(make_dependable(ActionFilter)),
+    hydrate: bool = False,
+    _: AuthContext = Security(authorize),
+) -> Page[ActionResponse]:
+    """List actions.
+
+    Args:
+        action_filter_model: Filter model used for pagination, sorting,
+            filtering.
+        hydrate: Flag deciding whether to hydrate the output model(s)
+            by including metadata fields in the response.
+
+    Returns:
+        Page of actions.
+    """
+
+    def list_actions_fn(
+        filter_model: ActionFilter,
+    ) -> Page[ActionResponse]:
+        """List actions through their associated plugins.
+
+        Args:
+            filter_model: Filter model used for pagination, sorting,
+                filtering.
+
+        Raises:
+            ValueError: If the action handler for flavor/type is not valid.
+
+        Returns:
+            All actions.
+        """
+        actions = zen_store().list_actions(
+            action_filter_model=filter_model, hydrate=hydrate
+        )
+
+        # Process the actions through their associated plugins
+        for idx, action in enumerate(actions.items):
+            action_handler = plugin_flavor_registry().get_plugin(
+                name=action.flavor,
+                _type=PluginType.ACTION,
+                subtype=action.plugin_subtype,
+            )
+
+            # Validate that the flavor and plugin_type correspond to an action
+            # handler implementation
+            if not isinstance(action_handler, BaseActionHandler):
+                raise ValueError(
+                    f"Action handler plugin {action.plugin_subtype} "
+                    f"for flavor {action.flavor} is not a valid action "
+                    "handler plugin."
+                )
+
+            actions.items[idx] = action_handler.get_action(
+                action, hydrate=hydrate
+            )
+
+        return actions
+
+    return verify_permissions_and_list_entities(
+        filter_model=action_filter_model,
+        resource_type=ResourceType.ACTION,
+        list_method=list_actions_fn,
+    )
+
+
+@router.get(
+    "/{action_id}",
+    response_model=ActionResponse,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def get_action(
+    action_id: UUID,
+    hydrate: bool = True,
+    _: AuthContext = Security(authorize),
+) -> ActionResponse:
+    """Returns the requested action.
+
+    Args:
+        action_id: ID of the action.
+        hydrate: Flag deciding whether to hydrate the output model(s)
+            by including metadata fields in the response.
+
+    Raises:
+        ValueError: If the action handler for flavor/type is not valid.
+
+    Returns:
+        The requested action.
+    """
+    action = zen_store().get_action(action_id=action_id, hydrate=hydrate)
+
+    verify_permission_for_model(action, action=Action.READ)
+
+    action_handler = plugin_flavor_registry().get_plugin(
+        name=action.flavor,
+        _type=PluginType.ACTION,
+        subtype=action.plugin_subtype,
+    )
+
+    # Validate that the flavor and plugin_type correspond to an action
+    # handler implementation
+    if not isinstance(action_handler, BaseActionHandler):
+        raise ValueError(
+            f"Action handler plugin {action.plugin_subtype} "
+            f"for flavor {action.flavor} is not a valid action "
+            "handler plugin."
+        )
+
+    action = action_handler.get_action(action, hydrate=hydrate)
+
+    return dehydrate_response_model(action)
+
+
+@router.post(
+    "",
+    response_model=ActionResponse,
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_action(
+    action: ActionRequest,
+    _: AuthContext = Security(authorize),
+) -> ActionResponse:
+    """Creates an action.
+
+    Args:
+        action: Action to create.
+
+    Raises:
+        ValueError: If the action handler for flavor/type is not valid.
+
+    Returns:
+        The created action.
+    """
+    service_account = zen_store().get_service_account(
+        service_account_name_or_id=action.service_account_id
+    )
+    verify_permission_for_model(service_account, action=Action.READ)
+
+    action_handler = plugin_flavor_registry().get_plugin(
+        name=action.flavor,
+        _type=PluginType.ACTION,
+        subtype=action.plugin_subtype,
+    )
+
+    # Validate that the flavor and plugin_type correspond to an action
+    # handler implementation
+    if not isinstance(action_handler, BaseActionHandler):
+        raise ValueError(
+            f"Action handler plugin {action.plugin_subtype} "
+            f"for flavor {action.flavor} is not a valid action "
+            "handler plugin."
+        )
+
+    return verify_permissions_and_create_entity(
+        request_model=action,
+        resource_type=ResourceType.ACTION,
+        create_method=action_handler.create_action,
+    )
+
+
+@router.put(
+    "/{action_id}",
+    response_model=ActionResponse,
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def update_action(
+    action_id: UUID,
+    action_update: ActionUpdate,
+    _: AuthContext = Security(authorize),
+) -> ActionResponse:
+    """Update an action.
+
+    Args:
+        action_id: ID of the action to update.
+        action_update: The action update.
+
+    Raises:
+        ValueError: If the action handler for flavor/type is not valid.
+
+    Returns:
+        The updated action.
+    """
+    action = zen_store().get_action(action_id=action_id)
+
+    verify_permission_for_model(action, action=Action.UPDATE)
+
+    if action_update.service_account_id:
+        service_account = zen_store().get_service_account(
+            service_account_name_or_id=action_update.service_account_id
+        )
+        verify_permission_for_model(service_account, action=Action.READ)
+
+    action_handler = plugin_flavor_registry().get_plugin(
+        name=action.flavor,
+        _type=PluginType.ACTION,
+        subtype=action.plugin_subtype,
+    )
+
+    # Validate that the flavor and plugin_type correspond to an action
+    # handler implementation
+    if not isinstance(action_handler, BaseActionHandler):
+        raise ValueError(
+            f"Action handler plugin {action.plugin_subtype} "
+            f"for flavor {action.flavor} is not a valid action "
+            "handler plugin."
+        )
+
+    updated_action = action_handler.update_action(
+        action=action,
+        action_update=action_update,
+    )
+
+    return dehydrate_response_model(updated_action)
+
+
+@router.delete(
+    "/{action_id}",
+    responses={401: error_response, 404: error_response, 422: error_response},
+)
+@handle_exceptions
+def delete_action(
+    action_id: UUID,
+    force: bool = False,
+    _: AuthContext = Security(authorize),
+) -> None:
+    """Delete an action.
+
+    Args:
+        action_id: ID of the action.
+        force: Flag deciding whether to force delete the action.
+
+    Raises:
+        ValueError: If the action handler for flavor/type is not valid.
+    """
+    action = zen_store().get_action(action_id=action_id)
+
+    verify_permission_for_model(action, action=Action.DELETE)
+
+    action_handler = plugin_flavor_registry().get_plugin(
+        name=action.flavor,
+        _type=PluginType.ACTION,
+        subtype=action.plugin_subtype,
+    )
+
+    # Validate that the flavor and plugin_type correspond to an action
+    # handler implementation
+    if not isinstance(action_handler, BaseActionHandler):
+        raise ValueError(
+            f"Action handler plugin {action.plugin_subtype} "
+            f"for flavor {action.flavor} is not a valid action "
+            "handler plugin."
+        )
+
+    action_handler.delete_action(
+        action=action,
+        force=force,
+    )