zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/actions/pipeline_run/pipeline_run_action.py

@@ -0,0 +1,223 @@
+#  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.
+"""Example file of what an action Plugin could look like."""
+
+from typing import Any, ClassVar, Dict, Optional, Type
+from uuid import UUID
+
+from zenml.actions.base_action import (
+    ActionConfig,
+    BaseActionFlavor,
+    BaseActionHandler,
+)
+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,
+)
+from zenml.models.v2.base.base import BaseResponse
+from zenml.zen_server.auth import AuthContext
+from zenml.zen_server.rbac.models import ResourceType
+from zenml.zen_server.template_execution.utils import run_template
+from zenml.zen_server.utils import server_config
+
+logger = get_logger(__name__)
+
+# -------------------- Configuration Models ----------------------------------
+
+
+class PipelineRunActionConfiguration(ActionConfig):
+    """Configuration class to configure a pipeline run action."""
+
+    template_id: UUID
+    run_config: Optional[PipelineRunConfiguration] = None
+
+
+# -------------------- Pipeline Run Plugin -----------------------------------
+
+
+class PipelineRunActionHandler(BaseActionHandler):
+    """Action handler for running pipelines."""
+
+    @property
+    def config_class(self) -> Type[PipelineRunActionConfiguration]:
+        """Returns the `BasePluginConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return PipelineRunActionConfiguration
+
+    @property
+    def flavor_class(self) -> Type[BaseActionFlavor]:
+        """Returns the flavor class of the plugin.
+
+        Returns:
+            The flavor class of the plugin.
+        """
+        return PipelineRunActionFlavor
+
+    def run(
+        self,
+        config: ActionConfig,
+        trigger_execution: TriggerExecutionResponse,
+        auth_context: AuthContext,
+    ) -> None:
+        """Execute an action.
+
+        Args:
+            config: The action configuration
+            trigger_execution: The trigger execution
+            auth_context: Authentication context with an API token that can be
+                used by external workloads to authenticate with the server
+                during the execution of the action. This API token is associated
+                with the service account that was configured for the trigger
+                that activated the action and has a validity defined by the
+                trigger's authentication window.
+        """
+        from zenml.zen_server.utils import zen_store
+
+        assert isinstance(config, PipelineRunActionConfiguration)
+
+        template = zen_store().get_run_template(config.template_id)
+        logger.debug("Running template:", template)
+        run_template(
+            template=template,
+            run_config=config.run_config,
+            auth_context=auth_context,
+        )
+
+    def _validate_configuration(
+        self, config: PipelineRunActionConfiguration
+    ) -> None:
+        """Validate a pipeline run action configuration.
+
+        Args:
+            config: Pipeline run action configuration.
+
+        Raises:
+            ValueError: In case no deployment can be found with the deployment_id
+        """
+        zen_store = GlobalConfiguration().zen_store
+
+        try:
+            zen_store.get_run_template(template_id=config.template_id)
+        except KeyError:
+            raise ValueError(
+                f"No template found with id {config.template_id}."
+            )
+
+    def _validate_action_request(
+        self, action: ActionRequest, config: ActionConfig
+    ) -> None:
+        """Validate an action request before it is created in the database.
+
+        Args:
+            action: Action request.
+            config: Action configuration instantiated from the request.
+        """
+        assert isinstance(config, PipelineRunActionConfiguration)
+
+        self._validate_configuration(config)
+
+        # If an expiration window is not set, we set it to the default value
+        if action.auth_window is None:
+            action.auth_window = server_config().pipeline_run_auth_window
+
+    def _validate_action_update(
+        self,
+        action: ActionResponse,
+        config: ActionConfig,
+        action_update: ActionUpdate,
+        config_update: ActionConfig,
+    ) -> None:
+        """Validate an action update before it is reflected in the database.
+
+        Args:
+            action: Original action before the update.
+            config: Action configuration instantiated from the original
+                action.
+            action_update: Action update request.
+            config_update: Action configuration instantiated from the
+                updated action.
+        """
+        assert isinstance(config, PipelineRunActionConfiguration)
+
+        self._validate_configuration(config)
+
+    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.
+
+        Raises:
+            ValueError: In case the specified template does not exist.
+        """
+        assert isinstance(action_config, PipelineRunActionConfiguration)
+
+        zen_store = GlobalConfiguration().zen_store
+
+        try:
+            template = zen_store.get_run_template(
+                template_id=action_config.template_id, hydrate=hydrate
+            )
+        except KeyError:
+            raise ValueError(
+                f"No template found with id {action_config.template_id}."
+            )
+
+        resources: Dict[ResourceType, BaseResponse[Any, Any, Any]] = {
+            ResourceType.RUN_TEMPLATE: template
+        }
+
+        if template.pipeline is not None:
+            pipeline = zen_store.get_pipeline(
+                pipeline_id=template.pipeline.id, hydrate=hydrate
+            )
+            resources[ResourceType.PIPELINE] = pipeline
+
+        return resources
+
+
+# -------------------- Pipeline Run Flavor -----------------------------------
+
+
+class PipelineRunActionFlavor(BaseActionFlavor):
+    """Enables users to configure pipeline run action."""
+
+    FLAVOR: ClassVar[str] = "builtin"
+    SUBTYPE: ClassVar[PluginSubType] = PluginSubType.PIPELINE_RUN
+    PLUGIN_CLASS: ClassVar[Type[PipelineRunActionHandler]] = (
+        PipelineRunActionHandler
+    )
+
+    # EventPlugin specific
+    ACTION_CONFIG_CLASS: ClassVar[Type[PipelineRunActionConfiguration]] = (
+        PipelineRunActionConfiguration
+    )

zenml/analytics/context.py

@@ -16,13 +16,20 @@
 This module is based on the 'analytics-python' package created by Segment.
 The base functionalities are adapted to work with the ZenML analytics server.
 """
+
+import datetime
+import locale
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 from uuid import UUID
 
 from zenml import __version__
 from zenml.analytics.client import default_client
-from zenml.constants import ENV_ZENML_SERVER, handle_bool_env_var
+from zenml.analytics.utils import is_analytics_disabled_internally
+from zenml.constants import (
+    ENV_ZENML_SERVER,
+    handle_bool_env_var,
+)
 from zenml.environment import Environment, get_environment
 from zenml.logger import get_logger
 
@@ -52,8 +59,11 @@ class AnalyticsContext:
 
         self.user_id: Optional[UUID] = None
         self.external_user_id: Optional[UUID] = None
+        self.executed_by_service_account: Optional[bool] = None
         self.client_id: Optional[UUID] = None
         self.server_id: Optional[UUID] = None
+        self.external_server_id: Optional[UUID] = None
+        self.server_metadata: Optional[Dict[str, str]] = None
 
         self.database_type: Optional["ServerDatabaseType"] = None
         self.deployment_type: Optional["ServerDeploymentType"] = None
@@ -64,29 +74,57 @@ class AnalyticsContext:
         Returns:
             The analytics context.
         """
+        if is_analytics_disabled_internally():
+            self.analytics_opt_in = False
+            return self
+
         # Fetch the analytics opt-in setting
         from zenml.config.global_config import GlobalConfiguration
 
-        gc = GlobalConfiguration()
-        self.analytics_opt_in = gc.analytics_opt_in
+        try:
+            gc = GlobalConfiguration()
 
-        if not self.analytics_opt_in:
-            return self
+            if not gc.is_initialized:
+                # If the global configuration is not initialized, using the
+                # zen store can lead to multiple initialization issues, because
+                # the analytics are triggered during the initialization of the
+                # zen store.
+                return self
+
+            store_info = gc.zen_store.get_store_info()
+
+            # For local ZenML servers, we always use the client's analytics
+            # opt-in configuration.
+            if self.in_server and not store_info.is_local():
+                self.analytics_opt_in = store_info.analytics_enabled
+            else:
+                self.analytics_opt_in = gc.analytics_opt_in
+
+            if not self.analytics_opt_in:
+                return self
 
-        try:
             # Fetch the `user_id`
             if self.in_server:
                 from zenml.zen_server.auth import get_auth_context
+                from zenml.zen_server.utils import server_config
 
                 # If the code is running on the server, use the auth context.
                 auth_context = get_auth_context()
                 if auth_context is not None:
                     self.user_id = auth_context.user.id
+                    self.executed_by_service_account = (
+                        auth_context.user.is_service_account
+                    )
                     self.external_user_id = auth_context.user.external_user_id
+
+                self.external_server_id = server_config().external_server_id
             else:
                 # If the code is running on the client, use the default user.
                 active_user = gc.zen_store.get_user()
                 self.user_id = active_user.id
+                self.executed_by_service_account = (
+                    active_user.is_service_account
+                )
                 self.external_user_id = active_user.external_user_id
 
             # Fetch the `client_id`
@@ -97,12 +135,10 @@ class AnalyticsContext:
                 # If the code is running on the client, attach the client id.
                 self.client_id = gc.user_id
 
-            # Fetch the store information including the `server_id`
-            store_info = gc.zen_store.get_store_info()
-
             self.server_id = store_info.id
             self.deployment_type = store_info.deployment_type
             self.database_type = store_info.database_type
+            self.server_metadata = store_info.metadata
         except Exception as e:
             self.analytics_opt_in = False
             logger.debug(f"Analytics initialization failed: {e}")
@@ -192,15 +228,10 @@ class AnalyticsContext:
         """
         success = False
         if self.analytics_opt_in and self.user_id is not None:
-            if traits is None:
-                traits = {}
-
-            traits.update({"group_id": group_id})
-
             success, _ = default_client.group(
                 user_id=self.user_id,
                 group_id=group_id,
-                traits=traits,
+                traits=traits or {},
             )
 
         return success
@@ -247,12 +278,36 @@ class AnalyticsContext:
                 "server_id": str(self.server_id),
                 "deployment_type": str(self.deployment_type),
                 "database_type": str(self.database_type),
+                "executed_by_service_account": self.executed_by_service_account,
             }
         )
 
+        try:
+            # Timezone as tzdata
+            tz = (
+                datetime.datetime.now(datetime.timezone.utc)
+                .astimezone()
+                .tzname()
+            )
+            if tz is not None:
+                properties.update({"timezone": tz})
+
+            # Language code such as "en_DE"
+            language_code, encoding = locale.getlocale()
+            if language_code is not None:
+                properties.update({"locale": language_code})
+        except Exception:
+            pass
+
         if self.external_user_id:
             properties["external_user_id"] = self.external_user_id
 
+        if self.external_server_id:
+            properties["external_server_id"] = self.external_server_id
+
+        if self.server_metadata:
+            properties.update(self.server_metadata)
+
         for k, v in properties.items():
             if isinstance(v, UUID):
                 properties[k] = str(v)

zenml/analytics/enums.py

@@ -12,14 +12,22 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Collection of analytics events for ZenML."""
+
 from enum import Enum
 
 
 class AnalyticsEvent(str, Enum):
     """Enum of events to track in segment."""
 
+    # Login
+    DEVICE_VERIFIED = "Device verified"
+
+    # Onboarding
+    USER_ENRICHED = "User Enriched"
+
     # Pipelines
     RUN_PIPELINE = "Pipeline run"
+    RUN_PIPELINE_ENDED = "Pipeline run ended"
     CREATE_PIPELINE = "Pipeline created"
     BUILD_PIPELINE = "Pipeline built"
 
@@ -36,9 +44,18 @@ class AnalyticsEvent(str, Enum):
     REGISTERED_STACK = "Stack registered"
     UPDATED_STACK = "Stack updated"
 
+    # Trigger
+    CREATED_TRIGGER = "Trigger created"
+    UPDATED_TRIGGER = "Trigger updated"
+
+    # Templates
+    CREATED_RUN_TEMPLATE = "Run template created"
+    EXECUTED_RUN_TEMPLATE = "Run templated executed"
+
     # Model Control Plane
     MODEL_DEPLOYED = "Model deployed"
     CREATED_MODEL = "Model created"
+    CREATED_MODEL_VERSION = "Model Version created"
 
     # Analytics opt in and out
     OPT_IN_ANALYTICS = "Analytics opt-in"
@@ -63,24 +80,11 @@ class AnalyticsEvent(str, Enum):
     # Service account and API keys
     CREATED_SERVICE_ACCOUNT = "Service account created"
 
-    # Stack recipes
-    RUN_STACK_RECIPE = "Stack recipe ran"
-    DEPLOY_STACK = "Stack deployed"
-    DESTROY_STACK = "Stack destroyed"
-
-    # Stack component deploy
-    DEPLOY_STACK_COMPONENT = "Stack component deployed"
-    DESTROY_STACK_COMPONENT = "Stack component destroyed"
+    # Full stack infrastructure deployment
+    DEPLOY_FULL_STACK = "Full stack deployed"
 
     # Tag created
     CREATED_TAG = "Tag created"
 
-    # ZenML server events
-    ZENML_SERVER_DEPLOYED = "ZenML server deployed"
-    ZENML_SERVER_DESTROYED = "ZenML server destroyed"
-
-    # ZenML Hub events
-    ZENML_HUB_PLUGIN_INSTALL = "ZenML Hub plugin installed"
-    ZENML_HUB_PLUGIN_UNINSTALL = "ZenML Hub plugin uninstalled"
-    ZENML_HUB_PLUGIN_CLONE = "ZenML Hub plugin pulled"
-    ZENML_HUB_PLUGIN_SUBMIT = "ZenML Hub plugin pushed"
+    # Server Settings
+    SERVER_SETTINGS_UPDATED = "Server Settings Updated"

zenml/analytics/models.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Helper models for ZenML analytics."""
+
 from typing import Any, ClassVar, Dict, List
 
 from pydantic import BaseModel

zenml/analytics/utils.py

@@ -12,8 +12,8 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utility functions and classes for ZenML analytics."""
+
 import json
-import os
 from functools import wraps
 from typing import Any, Callable, Dict, Optional, TypeVar, cast
 from uuid import UUID
@@ -21,13 +21,15 @@ from uuid import UUID
 from zenml.analytics import identify, track
 from zenml.analytics.enums import AnalyticsEvent
 from zenml.analytics.models import AnalyticsTrackedModelMixin
-from zenml.constants import ENV_ZENML_ANALYTICS_OPT_IN
 from zenml.logger import get_logger
 
 logger = get_logger(__name__)
 F = TypeVar("F", bound=Callable[..., Any])
 
 
+_INTERNAL_DISABLE_ANALYTICS = False
+
+
 class AnalyticsEncoder(json.JSONEncoder):
     """Helper encoder class for JSON serialization."""
 
@@ -100,11 +102,13 @@ class analytics_disabler:
 
     def __init__(self) -> None:
         """Initialization of the context manager."""
-        self.original_value = os.environ.get(ENV_ZENML_ANALYTICS_OPT_IN)
+        self.original_value: Optional[bool] = None
 
     def __enter__(self) -> None:
         """Disable the analytics."""
-        os.environ[ENV_ZENML_ANALYTICS_OPT_IN] = str(False)
+        global _INTERNAL_DISABLE_ANALYTICS
+        self.original_value = _INTERNAL_DISABLE_ANALYTICS
+        _INTERNAL_DISABLE_ANALYTICS = True
 
     def __exit__(
         self,
@@ -119,10 +123,18 @@ class analytics_disabler:
             exc_value: The instance of the exception
             traceback: The traceback of the exception
         """
+        global _INTERNAL_DISABLE_ANALYTICS
         if self.original_value is not None:
-            os.environ[ENV_ZENML_ANALYTICS_OPT_IN] = self.original_value
-        else:
-            del os.environ[ENV_ZENML_ANALYTICS_OPT_IN]
+            _INTERNAL_DISABLE_ANALYTICS = self.original_value
+
+
+def is_analytics_disabled_internally() -> bool:
+    """Whether analytics are disabled by an internal helper function.
+
+    Returns:
+        Whether analytics are disabled by an internal helper function.
+    """
+    return _INTERNAL_DISABLE_ANALYTICS
 
 
 def track_decorator(event: AnalyticsEvent) -> Callable[[F], F]:

zenml/annotators/base_annotator.py

@@ -14,7 +14,7 @@
 """Base class for ZenML annotator stack components."""
 
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar, List, Optional, Tuple, Type, cast
+from typing import Any, ClassVar, List, Tuple, Type, cast
 
 from zenml.enums import StackComponentType
 from zenml.stack import Flavor, StackComponent
@@ -91,11 +91,12 @@ class BaseAnnotator(StackComponent, ABC):
         """
 
     @abstractmethod
-    def launch(self, url: Optional[str]) -> None:
+    def launch(self, **kwargs: Any) -> None:
         """Launches the annotation interface.
 
         Args:
-            url: The URL of the annotation interface.
+            **kwargs: Additional keyword arguments to pass to the
+                annotation client.
         """
 
     @abstractmethod