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

zenml/client.py

@@ -19,7 +19,6 @@ import os
 from abc import ABCMeta
 from collections import Counter
 from datetime import datetime
-from functools import partial
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -82,6 +81,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 +204,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 +310,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 +330,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
@@ -1149,6 +1152,7 @@ class Client(metaclass=ClientMetaClass):
         name: str,
         components: Mapping[StackComponentType, Union[str, UUID]],
         stack_spec_file: Optional[str] = None,
+        labels: Optional[Dict[str, Any]] = None,
     ) -> StackResponse:
         """Registers a stack and its components.
 
@@ -1156,6 +1160,7 @@ class Client(metaclass=ClientMetaClass):
             name: The name of the stack to register.
             components: dictionary which maps component types to component names
             stack_spec_file: path to the stack spec file
+            labels: The labels of the stack.
 
         Returns:
             The model of the registered stack.
@@ -1180,6 +1185,7 @@ class Client(metaclass=ClientMetaClass):
             stack_spec_path=stack_spec_file,
             workspace=self.active_workspace.id,
             user=self.active_user.id,
+            labels=labels,
         )
 
         self._validate_stack_configuration(stack=stack)
@@ -1223,8 +1229,8 @@ class Client(metaclass=ClientMetaClass):
         size: int = PAGE_SIZE_DEFAULT,
         logical_operator: LogicalOperators = LogicalOperators.AND,
         id: Optional[Union[UUID, str]] = None,
-        created: Optional[datetime] = None,
-        updated: Optional[datetime] = None,
+        created: Optional[Union[datetime, str]] = None,
+        updated: Optional[Union[datetime, str]] = None,
         name: Optional[str] = None,
         description: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
@@ -1275,6 +1281,7 @@ class Client(metaclass=ClientMetaClass):
         name_id_or_prefix: Optional[Union[UUID, str]] = None,
         name: Optional[str] = None,
         stack_spec_file: Optional[str] = None,
+        labels: Optional[Dict[str, Any]] = None,
         description: Optional[str] = None,
         component_updates: Optional[
             Dict[StackComponentType, List[Union[UUID, str]]]
@@ -1285,7 +1292,8 @@ class Client(metaclass=ClientMetaClass):
         Args:
             name_id_or_prefix: The name, id or prefix of the stack to update.
             name: the new name of the stack.
-            stack_spec_file: path to the stack spec file
+            stack_spec_file: path to the stack spec file.
+            labels: The new labels of the stack component.
             description: the new description of the stack.
             component_updates: dictionary which maps stack component types to
                 lists of new stack component names or ids.
@@ -1339,6 +1347,15 @@ class Client(metaclass=ClientMetaClass):
                 for c_type, c_list in components_dict.items()
             }
 
+        if labels is not None:
+            existing_labels = stack.labels or {}
+            existing_labels.update(labels)
+
+            existing_labels = {
+                k: v for k, v in existing_labels.items() if v is not None
+            }
+            update_model.labels = existing_labels
+
         updated_stack = self.zen_store.update_stack(
             stack_id=stack.id,
             stack_update=update_model,
@@ -2435,9 +2452,7 @@ class Client(metaclass=ClientMetaClass):
                 )
 
             for pipeline in depaginate(
-                functools.partial(
-                    Client().list_pipelines, name=name_id_or_prefix
-                )
+                Client().list_pipelines, name=name_id_or_prefix
             ):
                 Client().delete_pipeline(pipeline.id)
         else:
@@ -2575,11 +2590,9 @@ class Client(metaclass=ClientMetaClass):
                 )
 
             builds = depaginate(
-                partial(
-                    self.list_builds,
-                    pipeline_id=pipeline.id,
-                    stack_id=stack.id if stack else None,
-                )
+                self.list_builds,
+                pipeline_id=pipeline.id,
+                stack_id=stack.id if stack else None,
             )
 
             for build in builds:
@@ -2754,18 +2767,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 +2948,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 +3173,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 +3204,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 +3226,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 +3249,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 +3270,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 +3279,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 +3297,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):
@@ -3771,7 +3951,7 @@ class Client(metaclass=ClientMetaClass):
         """
         if delete_from_artifact_store:
             unused_artifact_versions = depaginate(
-                partial(self.list_artifact_versions, only_unused=True)
+                self.list_artifact_versions, only_unused=True
             )
             for unused_artifact_version in unused_artifact_versions:
                 self._delete_artifact_from_artifact_store(
@@ -3980,7 +4160,7 @@ class Client(metaclass=ClientMetaClass):
             ValueError: If the artifact version is still used in any runs.
         """
         if artifact_version not in depaginate(
-            partial(self.list_artifact_versions, only_unused=True)
+            self.list_artifact_versions, only_unused=True
         ):
             raise ValueError(
                 "The metadata of artifact versions that are used in runs "
@@ -5495,6 +5675,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 +5689,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 +5705,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 +5744,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: