zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.61.0.dev20240712__py3-none-any.whl

zenml/zen_server/deploy/helm/Chart.yaml

@@ -1,6 +1,6 @@
 apiVersion: v2
 name: zenml
-version: "0.58.2"
+version: "0.61.0"
 description: Open source MLOps framework for portable production ready ML pipelines
 keywords:
 - mlops

zenml/zen_server/deploy/helm/README.md

@@ -20,8 +20,8 @@ ZenML is an open-source MLOps framework designed to help you create robust, main
 To install the ZenML chart directly from Amazon ECR, use the following command:
 
 ```bash
-# example command for version 0.58.2
-helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.58.2
+# example command for version 0.61.0
+helm install my-zenml oci://public.ecr.aws/zenml/zenml --version 0.61.0
 ```
 
 Note: Ensure you have OCI support enabled in your Helm client and that you are authenticated with Amazon ECR.

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -21,7 +21,7 @@ from pydantic import BaseModel, Field
 from zenml.config.server_config import ServerConfiguration
 from zenml.exceptions import SubscriptionUpgradeRequiredError
 from zenml.logger import get_logger
-from zenml.zen_server.cloud_utils import ZenMLCloudSession
+from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.feature_gate.feature_gate_interface import (
     FeatureGateInterface,
 )
@@ -59,8 +59,12 @@ class RawUsageEvent(BaseModel):
     )
 
 
-class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
-    """Feature Gate interface definition."""
+class ZenMLCloudFeatureGateInterface(FeatureGateInterface):
+    """ZenML Cloud Feature Gate implementation."""
+
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
 
     def check_entitlement(self, resource: ResourceType) -> None:
         """Checks if a user is entitled to create a resource.
@@ -72,7 +76,7 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
             SubscriptionUpgradeRequiredError: in case a subscription limit is reached
         """
         try:
-            response = self._get(
+            response = self._connection.get(
                 endpoint=ENTITLEMENT_ENDPOINT + "/" + resource, params=None
             )
         except SubscriptionUpgradeRequiredError:
@@ -110,7 +114,9 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
                 "resource_id": str(resource_id),
             },
         ).model_dump()
-        response = self._post(endpoint=USAGE_EVENT_ENDPOINT, data=data)
+        response = self._connection.post(
+            endpoint=USAGE_EVENT_ENDPOINT, data=data
+        )
         if response.status_code != 200:
             logger.error(
                 "Usage report not accepted by upstream backend. "

zenml/zen_server/pipeline_deployment/utils.py

@@ -17,7 +17,7 @@ from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
 )
-from zenml.enums import StackComponentType, StoreType
+from zenml.enums import ExecutionStatus, StackComponentType, StoreType
 from zenml.integrations.utils import get_integration_for_module
 from zenml.models import (
     CodeReferenceRequest,
@@ -26,6 +26,7 @@ from zenml.models import (
     PipelineDeploymentRequest,
     PipelineDeploymentResponse,
     PipelineRunResponse,
+    PipelineRunUpdate,
     StackResponse,
 )
 from zenml.new.pipelines.run_utils import (
@@ -88,8 +89,6 @@ def run_pipeline(
     ensure_async_orchestrator(deployment=deployment_request, stack=stack)
 
     new_deployment = zen_store().create_deployment(deployment_request)
-    placeholder_run = create_placeholder_run(deployment=new_deployment)
-    assert placeholder_run
 
     if auth_context.access_token:
         token = auth_context.access_token
@@ -125,52 +124,62 @@ def run_pipeline(
         deployment_id=new_deployment.id
     )
 
-    def _task() -> None:
-        pypi_requirements, apt_packages = get_requirements_for_stack(
-            stack=stack
-        )
+    placeholder_run = create_placeholder_run(deployment=new_deployment)
+    assert placeholder_run
 
-        if build.python_version:
-            version_info = version.parse(build.python_version)
-            python_version = f"{version_info.major}.{version_info.minor}"
-        else:
-            python_version = (
-                f"{sys.version_info.major}.{sys.version_info.minor}"
+    def _task() -> None:
+        try:
+            pypi_requirements, apt_packages = get_requirements_for_stack(
+                stack=stack
             )
 
-        dockerfile = generate_dockerfile(
-            pypi_requirements=pypi_requirements,
-            apt_packages=apt_packages,
-            zenml_version=zenml_version,
-            python_version=python_version,
-        )
+            if build.python_version:
+                version_info = version.parse(build.python_version)
+                python_version = f"{version_info.major}.{version_info.minor}"
+            else:
+                python_version = (
+                    f"{sys.version_info.major}.{sys.version_info.minor}"
+                )
+
+            dockerfile = generate_dockerfile(
+                pypi_requirements=pypi_requirements,
+                apt_packages=apt_packages,
+                zenml_version=zenml_version,
+                python_version=python_version,
+            )
 
-        image_hash = generate_image_hash(dockerfile=dockerfile)
+            image_hash = generate_image_hash(dockerfile=dockerfile)
 
-        runner_image = workload_manager().build_and_push_image(
-            workload_id=new_deployment.id,
-            dockerfile=dockerfile,
-            image_name=f"{RUNNER_IMAGE_REPOSITORY}:{image_hash}",
-            sync=True,
-        )
+            runner_image = workload_manager().build_and_push_image(
+                workload_id=new_deployment.id,
+                dockerfile=dockerfile,
+                image_name=f"{RUNNER_IMAGE_REPOSITORY}:{image_hash}",
+                sync=True,
+            )
 
-        workload_manager().log(
-            workload_id=new_deployment.id,
-            message="Starting pipeline deployment.",
-        )
-        workload_manager().run(
-            workload_id=new_deployment.id,
-            image=runner_image,
-            command=command,
-            arguments=args,
-            environment=environment,
-            timeout_in_seconds=30,
-            sync=True,
-        )
-        workload_manager().log(
-            workload_id=new_deployment.id,
-            message="Pipeline deployed successfully.",
-        )
+            workload_manager().log(
+                workload_id=new_deployment.id,
+                message="Starting pipeline deployment.",
+            )
+            workload_manager().run(
+                workload_id=new_deployment.id,
+                image=runner_image,
+                command=command,
+                arguments=args,
+                environment=environment,
+                timeout_in_seconds=30,
+                sync=True,
+            )
+            workload_manager().log(
+                workload_id=new_deployment.id,
+                message="Pipeline deployed successfully.",
+            )
+        except Exception:
+            zen_store().update_run(
+                run_id=placeholder_run.id,
+                run_update=PipelineRunUpdate(status=ExecutionStatus.FAILED),
+            )
+            raise
 
     if background_tasks:
         background_tasks.add_task(_task)
@@ -357,8 +366,12 @@ def apply_run_config(
         step_config = StepConfiguration.model_validate(step_config_dict)
 
         if update := run_config.steps.get(invocation_id):
+            update_dict = update.model_dump()
+            # Get rid of deprecated name to prevent overriding the step name
+            # with `None`.
+            update_dict.pop("name", None)
             step_config = pydantic_utils.update_model(
-                step_config, update=update
+                step_config, update=update_dict
             )
         steps[invocation_id] = Step(spec=step.spec, config=step_config)
 

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/rbac/zenml_cloud_rbac.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Dict, List, Set, Tuple
 
-from zenml.zen_server.cloud_utils import ZenMLCloudSession
+from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.rbac.models import Action, Resource
 from zenml.zen_server.rbac.rbac_interface import RBACInterface
 from zenml.zen_server.utils import server_config
@@ -74,9 +74,13 @@ def _convert_from_cloud_resource(cloud_resource: str) -> Resource:
         return Resource(type=resource_type_and_id)
 
 
-class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
+class ZenMLCloudRBAC(RBACInterface):
     """RBAC implementation that uses the ZenML Pro Management Plane as a backend."""
 
+    def __init__(self) -> None:
+        """Initialize the object."""
+        self._connection = cloud_connection()
+
     def check_permissions(
         self, user: "UserResponse", resources: Set[Resource], action: Action
     ) -> Dict[Resource, bool]:
@@ -110,7 +114,9 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             ],
             "action": str(action),
         }
-        response = self._get(endpoint=PERMISSIONS_ENDPOINT, params=params)
+        response = self._connection.get(
+            endpoint=PERMISSIONS_ENDPOINT, params=params
+        )
         value = response.json()
 
         assert isinstance(value, dict)
@@ -147,7 +153,7 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             "resource": _convert_to_cloud_resource(resource),
             "action": str(action),
         }
-        response = self._get(
+        response = self._connection.get(
             endpoint=ALLOWED_RESOURCE_IDS_ENDPOINT, params=params
         )
         response_json = response.json()
@@ -177,4 +183,4 @@ class ZenMLCloudRBAC(RBACInterface, ZenMLCloudSession):
             "resource": _convert_to_cloud_resource(resource),
             "actions": [str(action) for action in actions],
         }
-        self._post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)
+        self._connection.post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)

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