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

zenml/environment.py

@@ -15,9 +15,8 @@
 
 import os
 import platform
-from importlib.util import find_spec
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, cast
+from typing import Dict
 
 import distro
 
@@ -27,9 +26,6 @@ from zenml.enums import EnvironmentType
 from zenml.logger import get_logger
 from zenml.utils.singleton import SingletonMetaClass
 
-if TYPE_CHECKING:
-    from zenml.steps import StepEnvironment
-
 logger = get_logger(__name__)
 
 
@@ -70,6 +66,12 @@ def get_environment() -> str:
         return EnvironmentType.BITBUCKET_CI
     elif Environment.in_ci():
         return EnvironmentType.GENERIC_CI
+    elif Environment.in_github_codespaces():
+        return EnvironmentType.GITHUB_CODESPACES
+    elif Environment.in_vscode_remote_container():
+        return EnvironmentType.VSCODE_REMOTE_CONTAINER
+    elif Environment.in_lightning_ai_studio():
+        return EnvironmentType.LIGHTNING_AI_STUDIO
     elif Environment.in_docker():
         return EnvironmentType.DOCKER
     elif Environment.in_container():
@@ -123,22 +125,6 @@ class Environment(metaclass=SingletonMetaClass):
         only get called once. All following `Environment()` calls will return
         the previously initialized instance.
         """
-        self._components: Dict[str, "BaseEnvironmentComponent"] = {}
-
-    @property
-    def step_is_running(self) -> bool:
-        """Returns if a step is currently running.
-
-        Returns:
-            `True` if a step is currently running, `False` otherwise.
-        """
-        from zenml.steps import STEP_ENVIRONMENT_NAME
-
-        logger.warning(
-            "`Environment().step_is_running` is deprecated and will be "
-            "removed in a future release."
-        )
-        return self.has_component(STEP_ENVIRONMENT_NAME)
 
     @staticmethod
     def get_system_info() -> Dict[str, str]:
@@ -257,17 +243,46 @@ class Environment(metaclass=SingletonMetaClass):
         if Environment.in_google_colab():
             return True
 
-        if find_spec("IPython") is not None:
-            from IPython import get_ipython
+        try:
+            ipython = get_ipython()  # type: ignore[name-defined]
+        except NameError:
+            return False
 
-            if get_ipython().__class__.__name__ in [
-                "TerminalInteractiveShell",
-                "ZMQInteractiveShell",
-                "DatabricksShell",
-            ]:
-                return True
+        if ipython.__class__.__name__ in [
+            "TerminalInteractiveShell",
+            "ZMQInteractiveShell",
+            "DatabricksShell",
+        ]:
+            return True
         return False
 
+    @staticmethod
+    def in_github_codespaces() -> bool:
+        """If the current Python process is running in GitHub Codespaces.
+
+        Returns:
+            `True` if the current Python process is running in GitHub Codespaces,
+            `False` otherwise.
+        """
+        return (
+            "CODESPACES" in os.environ
+            or "GITHUB_CODESPACE_TOKEN" in os.environ
+            or "GITHUB_CODESPACES_PORT_FORWARDING_DOMAIN" in os.environ
+        )
+
+    @staticmethod
+    def in_vscode_remote_container() -> bool:
+        """If the current Python process is running in a VS Code Remote Container.
+
+        Returns:
+            `True` if the current Python process is running in a VS Code Remote Container,
+            `False` otherwise.
+        """
+        return (
+            "REMOTE_CONTAINERS" in os.environ
+            or "VSCODE_REMOTE_CONTAINERS_SESSION" in os.environ
+        )
+
     @staticmethod
     def in_paperspace_gradient() -> bool:
         """If the current Python process is running in Paperspace Gradient.
@@ -339,297 +354,15 @@ class Environment(metaclass=SingletonMetaClass):
         """
         return "microsoft-standard" in platform.uname().release
 
-    def register_component(
-        self, component: "BaseEnvironmentComponent"
-    ) -> "BaseEnvironmentComponent":
-        """Registers an environment component.
-
-        Args:
-            component: a BaseEnvironmentComponent instance.
-
-        Returns:
-            The newly registered environment component, or the environment
-            component that was already registered under the given name.
-        """
-        if component.NAME not in self._components:
-            self._components[component.NAME] = component
-            logger.debug(f"Registered environment component {component.NAME}")
-            return component
-        else:
-            logger.warning(
-                f"Ignoring attempt to overwrite an existing Environment "
-                f"component registered under the name {component.NAME}."
-            )
-            return self._components[component.NAME]
-
-    def deregister_component(
-        self, component: "BaseEnvironmentComponent"
-    ) -> None:
-        """Deregisters an environment component.
-
-        Args:
-            component: a BaseEnvironmentComponent instance.
-        """
-        if self._components.get(component.NAME) is component:
-            del self._components[component.NAME]
-            logger.debug(
-                f"Deregistered environment component {component.NAME}"
-            )
-
-        else:
-            logger.warning(
-                f"Ignoring attempt to deregister an inexistent Environment "
-                f"component with the name {component.NAME}."
-            )
-
-    def get_component(self, name: str) -> Optional["BaseEnvironmentComponent"]:
-        """Get the environment component with a known name.
-
-        Args:
-            name: the environment component name.
-
-        Returns:
-            The environment component that is registered under the given name,
-            or None if no such component is registered.
-        """
-        return self._components.get(name)
-
-    def get_components(
-        self,
-    ) -> Dict[str, "BaseEnvironmentComponent"]:
-        """Get all registered environment components.
-
-        Returns:
-            A dictionary containing all registered environment components.
-        """
-        return self._components.copy()
-
-    def has_component(self, name: str) -> bool:
-        """Check if the environment component with a known name is available.
-
-        Args:
-            name: the environment component name.
-
-        Returns:
-            `True` if an environment component with the given name is
-            currently registered for the given name, `False` otherwise.
-        """
-        return name in self._components
-
-    def __getitem__(self, name: str) -> "BaseEnvironmentComponent":
-        """Get the environment component with the given name.
-
-        Args:
-            name: the environment component name.
-
-        Returns:
-            `BaseEnvironmentComponent` instance that was registered for the
-            given name.
-
-        Raises:
-            KeyError: if no environment component is registered for the given
-                name.
-        """
-        from zenml.steps import STEP_ENVIRONMENT_NAME
-
-        if name == STEP_ENVIRONMENT_NAME:
-            logger.warning(
-                "The `StepEnvironment` class and corresponding "
-                "`Environment.step_environment` property are deprecated and "
-                "will be removed in a future release. Please use the "
-                " `StepContext` to access information about the current run "
-                "instead, as shown here: "
-                "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/fetch-metadata-within-steps"
-            )
-        if name in self._components:
-            return self._components[name]
-        else:
-            raise KeyError(
-                f"No environment component with name {name} is currently "
-                f"registered. This could happen for example if you're trying "
-                f"to access an environment component that is only available "
-                f"in the context of a step function, or, in the case of "
-                f"globally available environment components, if a relevant "
-                f"integration has not been activated yet."
-            )
-
-    @property
-    def step_environment(self) -> "StepEnvironment":
-        """Get the current step environment component, if one is available.
-
-        This should only be called in the context of a step function.
-
-        Returns:
-            The `StepEnvironment` that describes the current step.
-        """
-        from zenml.steps import STEP_ENVIRONMENT_NAME, StepEnvironment
-
-        return cast(StepEnvironment, self[STEP_ENVIRONMENT_NAME])
-
-
-_BASE_ENVIRONMENT_COMPONENT_NAME = "base_environment_component"
-
-
-class EnvironmentComponentMeta(type):
-    """Metaclass registering environment components in the global Environment."""
-
-    def __new__(
-        mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
-    ) -> "EnvironmentComponentMeta":
-        """Hook into creation of an BaseEnvironmentComponent class.
-
-        Args:
-            name: the name of the class being created.
-            bases: the base classes of the class being created.
-            dct: the dictionary of attributes of the class being created.
+    @staticmethod
+    def in_lightning_ai_studio() -> bool:
+        """If the current Python process is running in Lightning.ai studios.
 
         Returns:
-            The newly created class.
+            `True` if the current Python process is running in Lightning.ai studios,
+            `False` otherwise.
         """
-        cls = cast(
-            Type["BaseEnvironmentComponent"],
-            super().__new__(mcs, name, bases, dct),
+        return (
+            "LIGHTNING_CLOUD_URL" in os.environ
+            and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
         )
-        if name != "BaseEnvironmentComponent":
-            assert cls.NAME and cls.NAME != _BASE_ENVIRONMENT_COMPONENT_NAME, (
-                "You should specify a unique NAME when creating an "
-                "EnvironmentComponent!"
-            )
-        return cls
-
-
-class BaseEnvironmentComponent(metaclass=EnvironmentComponentMeta):
-    """Base Environment component class.
-
-    All Environment components must inherit from this class and provide a unique
-    value for the `NAME` attribute.
-
-
-    Different code components can independently contribute with information to
-    the global Environment by extending and instantiating this class:
-
-    ```python
-    from zenml.environment import BaseEnvironmentComponent
-
-    MY_ENV_NAME = "my_env"
-
-    class MyEnvironmentComponent(BaseEnvironmentComponent):
-
-        NAME = MY_ENV_NAME
-
-        def __init__(self, my_env_attr: str) -> None:
-            super().__init__()
-            self._my_env_attr = my_env_attr
-
-        @property
-        def my_env_attr(self) -> str:
-            return self._my_env_attr
-
-    my_env = MyEnvironmentComponent()
-    ```
-
-
-    There are two ways to register and deregister a `BaseEnvironmentComponent`
-    instance with the global Environment:
-
-    1. by explicitly calling its `activate` and `deactivate` methods:
-
-    ```python
-    my_env.activate()
-
-    # ... environment component is active
-    # and registered in the global Environment
-
-    my_env.deactivate()
-
-    # ... environment component is not active
-    ```
-
-    2. by using the instance as a context:
-
-    ```python
-    with my_env:
-        # ... environment component is active
-        # and registered in the global Environment
-
-    # ... environment component is not active
-    ```
-
-    While active, environment components can be discovered and accessed from
-    the global environment:
-
-    ```python
-    from foo.bar.my_env import MY_ENV_NAME
-    from zenml.environment import Environment
-
-    my_env = Environment.get_component(MY_ENV_NAME)
-
-    # this works too, but throws an error if the component is not active:
-
-    my_env = Environment[MY_ENV_NAME]
-    ```
-
-    Attributes:
-        NAME: a unique name for this component. This name will be used to
-            register this component in the global Environment and to
-            subsequently retrieve it by calling `Environment().get_component`.
-    """
-
-    NAME: str = _BASE_ENVIRONMENT_COMPONENT_NAME
-
-    def __init__(self) -> None:
-        """Initialize an environment component."""
-        self._active = False
-
-    def activate(self) -> None:
-        """Activate the environment component and register it in the global Environment.
-
-        Raises:
-            RuntimeError: if the component is already active.
-        """
-        if self._active:
-            raise RuntimeError(
-                f"Environment component {self.NAME} is already active."
-            )
-        Environment().register_component(self)
-        self._active = True
-
-    def deactivate(self) -> None:
-        """Deactivate the environment component and deregister it from the global Environment.
-
-        Raises:
-            RuntimeError: if the component is not active.
-        """
-        if not self._active:
-            raise RuntimeError(
-                f"Environment component {self.NAME} is not active."
-            )
-        Environment().deregister_component(self)
-        self._active = False
-
-    @property
-    def active(self) -> bool:
-        """Check if the environment component is currently active.
-
-        Returns:
-            `True` if the environment component is currently active, `False`
-            otherwise.
-        """
-        return self._active
-
-    def __enter__(self) -> "BaseEnvironmentComponent":
-        """Environment component context entry point.
-
-        Returns:
-            The BaseEnvironmentComponent instance.
-        """
-        self.activate()
-        return self
-
-    def __exit__(self, *args: Any) -> None:
-        """Environment component context exit point.
-
-        Args:
-            *args: the arguments passed to the context exit point.
-        """
-        self.deactivate()

zenml/{steps/base_parameters.py → event_hub/__init__.py}

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  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.
@@ -11,10 +11,8 @@
 #  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.
-"""Step parameters."""
+"""ZenML Event Hub module.
 
-from pydantic import BaseModel
-
-
-class BaseParameters(BaseModel):
-    """Base class to pass parameters into a step."""
+The Event Hub is responsible for receiving all Events and dispatching them to
+the relevant Subscribers/Triggers.
+"""

zenml/event_hub/base_event_hub.py

@@ -0,0 +1,196 @@
+#  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.
+"""Base class for event hub implementations."""
+
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
+
+from zenml import EventSourceResponse
+from zenml.config.global_config import GlobalConfiguration
+from zenml.event_sources.base_event import (
+    BaseEvent,
+)
+from zenml.logger import get_logger
+from zenml.models import (
+    TriggerExecutionRequest,
+    TriggerExecutionResponse,
+    TriggerResponse,
+)
+from zenml.zen_server.auth import AuthContext
+from zenml.zen_server.jwt import JWTToken
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+logger = get_logger(__name__)
+
+ActionHandlerCallback = Callable[
+    [Dict[str, Any], TriggerExecutionResponse, AuthContext], None
+]
+
+
+class BaseEventHub(ABC):
+    """Base class for event hub implementations.
+
+    The event hub is responsible for relaying events from event sources to
+    action handlers. It functions similarly to a pub/sub system where
+    event source handlers publish events and action handlers subscribe to them.
+
+    The event hub also serves to decouple event sources from action handlers,
+    allowing them to be configured independently and their implementations to be
+    unaware of each other.
+    """
+
+    action_handlers: Dict[Tuple[str, str], ActionHandlerCallback] = {}
+
+    @property
+    def zen_store(self) -> "SqlZenStore":
+        """Returns the active zen store.
+
+        Returns:
+            The active zen store.
+
+        Raises:
+            ValueError: If the active zen store is not a SQL zen store.
+        """
+        from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+        zen_store = GlobalConfiguration().zen_store
+
+        if not isinstance(zen_store, SqlZenStore):
+            raise ValueError("The event hub requires a SQL zen store.")
+
+        return zen_store
+
+    def subscribe_action_handler(
+        self,
+        action_flavor: str,
+        action_subtype: str,
+        callback: ActionHandlerCallback,
+    ) -> None:
+        """Subscribe an action handler to the event hub.
+
+        Args:
+            action_flavor: the flavor of the action to trigger.
+            action_subtype: the subtype of the action to trigger.
+            callback: the action to trigger when the trigger is activated.
+        """
+        self.action_handlers[(action_flavor, action_subtype)] = callback
+
+    def unsubscribe_action_handler(
+        self,
+        action_flavor: str,
+        action_subtype: str,
+    ) -> None:
+        """Unsubscribe an action handler from the event hub.
+
+        Args:
+            action_flavor: the flavor of the action to trigger.
+            action_subtype: the subtype of the action to trigger.
+        """
+        self.action_handlers.pop((action_flavor, action_subtype), None)
+
+    def trigger_action(
+        self,
+        event: BaseEvent,
+        event_source: EventSourceResponse,
+        trigger: TriggerResponse,
+        action_callback: ActionHandlerCallback,
+    ) -> None:
+        """Trigger an action.
+
+        Args:
+            event: The event.
+            event_source: The event source that produced the event.
+            trigger: The trigger that was activated.
+            action_callback: The action to trigger.
+        """
+        request = TriggerExecutionRequest(
+            trigger=trigger.id, event_metadata=event.model_dump()
+        )
+
+        action_config = trigger.action.configuration
+
+        trigger_execution = self.zen_store.create_trigger_execution(request)
+
+        # Generate an API token that can be used by external workloads
+        # implementing the action to authenticate with the server. This token
+        # is associated with the service account configured for the trigger
+        # and has a validity defined by the trigger's authentication window.
+        token = JWTToken(
+            user_id=trigger.action.service_account.id,
+        )
+        expires: Optional[datetime] = None
+        if trigger.action.auth_window:
+            expires = datetime.utcnow() + timedelta(
+                minutes=trigger.action.auth_window
+            )
+        encoded_token = token.encode(expires=expires)
+        auth_context = AuthContext(
+            user=trigger.action.service_account,
+            access_token=token,
+            encoded_access_token=encoded_token,
+        )
+
+        try:
+            # TODO: We need to make this async, as this might take quite some
+            # time per trigger. We can either use threads starting here, or
+            # use fastapi background tasks that get passed here instead of
+            # running the event hub as a background tasks in the webhook
+            # endpoints
+            action_callback(
+                action_config,
+                trigger_execution,
+                auth_context,
+            )
+        except Exception:
+            # Don't let action errors stop the event hub from working
+            logger.exception(
+                f"An error occurred while executing trigger {trigger}."
+            )
+
+    @abstractmethod
+    def activate_trigger(self, trigger: TriggerResponse) -> None:
+        """Add a trigger to the event hub.
+
+        Configure the event hub to trigger an action when an event is received.
+
+        Args:
+            trigger: the trigger to activate.
+        """
+
+    @abstractmethod
+    def deactivate_trigger(self, trigger: TriggerResponse) -> None:
+        """Remove a trigger from the event hub.
+
+        Configure the event hub to stop triggering an action when an event is
+        received.
+
+        Args:
+            trigger: the trigger to deactivate.
+        """
+
+    @abstractmethod
+    def publish_event(
+        self,
+        event: BaseEvent,
+        event_source: EventSourceResponse,
+    ) -> None:
+        """Publish an event to the event hub.
+
+        Args:
+            event: The event.
+            event_source: The event source that produced the event.
+        """