zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.58.2.dev20240626__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/artifact_stores/base_artifact_store.py

@@ -14,6 +14,7 @@
 """The base interface to extend the ZenML artifact store."""
 
 import inspect
+import os
 import textwrap
 from abc import abstractmethod
 from pathlib import Path
@@ -34,6 +35,7 @@ from typing import (
 
 from pydantic import model_validator
 
+from zenml.constants import ENV_ZENML_SERVER
 from zenml.enums import StackComponentType
 from zenml.exceptions import ArtifactStoreInterfaceError
 from zenml.io import fileio
@@ -433,7 +435,11 @@ class BaseArtifactStore(StackComponent):
             **kwargs: The keyword arguments to pass to the Pydantic object.
         """
         super(BaseArtifactStore, self).__init__(*args, **kwargs)
-        self._register()
+
+        # If running in a ZenML server environment, we don't register
+        # the filesystems. We always use the artifact stores directly.
+        if ENV_ZENML_SERVER not in os.environ:
+            self._register()
 
     def _register(self) -> None:
         """Create and register a filesystem within the filesystem registry."""

zenml/artifacts/utils.py

@@ -436,18 +436,21 @@ def load_artifact_visualization(
     artifact_store = _load_artifact_store(
         artifact_store_id=artifact.artifact_store_id, zen_store=zen_store
     )
-    mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
-    value = _load_file_from_artifact_store(
-        uri=visualization.uri,
-        artifact_store=artifact_store,
-        mode=mode,
-    )
+    try:
+        mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
+        value = _load_file_from_artifact_store(
+            uri=visualization.uri,
+            artifact_store=artifact_store,
+            mode=mode,
+        )
 
-    # Encode image visualizations if requested
-    if visualization.type == VisualizationType.IMAGE and encode_image:
-        value = base64.b64encode(bytes(value))
+        # Encode image visualizations if requested
+        if visualization.type == VisualizationType.IMAGE and encode_image:
+            value = base64.b64encode(bytes(value))
 
-    return LoadedVisualization(type=visualization.type, value=value)
+        return LoadedVisualization(type=visualization.type, value=value)
+    finally:
+        artifact_store.cleanup()
 
 
 def load_artifact_from_response(artifact: "ArtifactVersionResponse") -> Any:

zenml/cli/service_connectors.py

@@ -1109,6 +1109,7 @@ def describe_service_connector(
                 name_id_or_prefix=name_id_or_prefix,
                 resource_type=resource_type,
                 resource_id=resource_id,
+                verify=True,
             )
         except (
             KeyError,

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):
@@ -5495,6 +5666,7 @@ class Client(metaclass=ClientMetaClass):
             name_id_or_prefix=name_id_or_prefix,
             resource_type=resource_type,
             resource_id=resource_id,
+            verify=False,
         )
 
         connector_client.configure_local_client(
@@ -5508,6 +5680,7 @@ class Client(metaclass=ClientMetaClass):
         name_id_or_prefix: Union[UUID, str],
         resource_type: Optional[str] = None,
         resource_id: Optional[str] = None,
+        verify: bool = False,
     ) -> "ServiceConnector":
         """Get the client side of a service connector instance to use with a local client.
 
@@ -5523,6 +5696,8 @@ class Client(metaclass=ClientMetaClass):
                 equivalent to the one requested, a `ValueError` exception is
                 raised. May be omitted for connectors and resource types that do
                 not support multiple resource instances.
+            verify: Whether to verify that the service connector configuration
+                and credentials can be used to gain access to the resource.
 
         Returns:
             The client side of the indicated service connector instance that can
@@ -5560,9 +5735,10 @@ class Client(metaclass=ClientMetaClass):
                 )
             )
 
-            # Verify the connector client on the local machine, because the
-            # server-side implementation may not be able to do so
-            connector_client.verify()
+            if verify:
+                # Verify the connector client on the local machine, because the
+                # server-side implementation may not be able to do so
+                connector_client.verify()
         else:
             connector_instance = (
                 service_connector_registry.instantiate_connector(

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: