zenml-nightly 0.66.0.dev20240928__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/config/pipeline_run_configuration.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline run configuration class."""
 
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import Field, SerializeAsAny
@@ -45,6 +45,7 @@ class PipelineRunConfiguration(
     )
     steps: Dict[str, StepConfigurationUpdate] = {}
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
+    tags: Optional[List[str]] = None
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

zenml/config/server_config.py

@@ -29,6 +29,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
+    DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES,
     DEFAULT_ZENML_SERVER_NAME,
     DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
@@ -40,7 +41,6 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
     DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE,
-    DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD,
     ENV_ZENML_SERVER_PREFIX,
 )
 from zenml.enums import AuthScheme
@@ -217,9 +217,6 @@ class ServerConfiguration(BaseModel):
             one of the reserved values `disabled`, `no`, `none`, `false`, `off`
             or to an empty string, the `Permissions-Policy` header will not be
             included in responses.
-        use_legacy_dashboard: Whether to use the legacy dashboard. If set to
-            `True`, the dashboard will be used with the old UI. If set to
-            `False`, the new dashboard will be used.
         server_name: The name of the ZenML server. Used only during initial
             deployment. Can be changed later as a part of the server settings.
         display_announcements: Whether to display announcements about ZenML in
@@ -231,6 +228,8 @@ class ServerConfiguration(BaseModel):
         auto_activate: Whether to automatically activate the server and create a
             default admin user account with an empty password during the initial
             deployment.
+        max_request_body_size_in_bytes: The maximum size of the request body in
+            bytes. If not specified, the default value of 256 Kb will be used.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -310,7 +309,6 @@ class ServerConfiguration(BaseModel):
         default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS,
         union_mode="left_to_right",
     )
-    use_legacy_dashboard: bool = DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD
 
     server_name: str = DEFAULT_ZENML_SERVER_NAME
     display_announcements: bool = True
@@ -319,6 +317,10 @@ class ServerConfiguration(BaseModel):
 
     thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
 
+    max_request_body_size_in_bytes: int = (
+        DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES
+    )
+
     _deployment_id: Optional[UUID] = None
 
     @model_validator(mode="before")

zenml/config/step_configurations.py

@@ -31,6 +31,7 @@ from pydantic import (
     model_validator,
 )
 
+from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.external_artifact_config import (
     ExternalArtifactConfiguration,
 )
@@ -59,6 +60,7 @@ class PartialArtifactConfiguration(StrictBaseModel):
     # TODO: This could be moved to the `PipelineDeployment` as it's the same
     # for all steps/outputs
     default_materializer_source: Optional[SourceWithValidator] = None
+    artifact_config: Optional[ArtifactConfig] = None
 
     @model_validator(mode="before")
     @classmethod

zenml/constants.py

@@ -134,6 +134,7 @@ APP_NAME = "zenml"
 # Environment variables
 ENV_ZENML_LOGGING_COLORS_DISABLED = "ZENML_LOGGING_COLORS_DISABLED"
 ENV_ZENML_ANALYTICS_OPT_IN = "ZENML_ANALYTICS_OPT_IN"
+ENV_ZENML_USER_ID = "ZENML_USER_ID"
 ENV_ZENML_CONFIG_PATH = "ZENML_CONFIG_PATH"
 ENV_ZENML_DEBUG = "ZENML_DEBUG"
 ENV_ZENML_LOGGING_VERBOSITY = "ZENML_LOGGING_VERBOSITY"
@@ -163,7 +164,6 @@ ENV_ZENML_DISABLE_CLIENT_SERVER_MISMATCH_WARNING = (
 ENV_ZENML_DISABLE_WORKSPACE_WARNINGS = "ZENML_DISABLE_WORKSPACE_WARNINGS"
 ENV_ZENML_SKIP_IMAGE_BUILDER_DEFAULT = "ZENML_SKIP_IMAGE_BUILDER_DEFAULT"
 ENV_ZENML_SERVER = "ZENML_SERVER"
-ENV_ZENML_LOCAL_SERVER = "ZENML_LOCAL_SERVER"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
@@ -181,13 +181,12 @@ ENV_ZENML_SERVER_AUTH_SCHEME = f"{ENV_ZENML_SERVER_PREFIX}AUTH_SCHEME"
 ENV_ZENML_SERVER_REPORTABLE_RESOURCES = (
     f"{ENV_ZENML_SERVER_PREFIX}REPORTABLE_RESOURCES"
 )
-ENV_ZENML_SERVER_USE_LEGACY_DASHBOARD = (
-    f"{ENV_ZENML_SERVER_PREFIX}USE_LEGACY_DASHBOARD"
-)
 ENV_ZENML_SERVER_AUTO_ACTIVATE = f"{ENV_ZENML_SERVER_PREFIX}AUTO_ACTIVATE"
 ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK = (
     "ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK"
 )
+ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING = "ZENML_PREVENT_CLIENT_SIDE_CACHING"
+ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING = "DISABLE_CREDENTIALS_DISK_CACHING"
 
 # Logging variables
 IS_DEBUG_ENV: bool = handle_bool_env_var(ENV_ZENML_DEBUG, default=False)
@@ -253,7 +252,7 @@ DEFAULT_SERVICE_START_STOP_TIMEOUT = 60
 DEFAULT_LOCAL_SERVICE_IP_ADDRESS = "127.0.0.1"
 ZEN_SERVER_ENTRYPOINT = "zenml.zen_server.zen_server_api:app"
 
-STEP_SOURCE_PARAMETER_NAME = "step_source"
+CODE_HASH_PARAMETER_NAME = "step_source"
 
 # Server settings
 DEFAULT_ZENML_SERVER_NAME = "default"
@@ -315,8 +314,8 @@ DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS = (
     "payment=(), sync-xhr=(), usb=()"
 )
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_REPORT_TO = "default"
-DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD = False
 DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS = 30
+DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES = 256 * 1024 * 1024
 
 # Configurations to decide which resources report their usage and check for
 # entitlement in the case of a cloud deployment. Expected Format is this:
@@ -338,6 +337,7 @@ API_TOKEN = "/api_token"
 ARTIFACTS = "/artifacts"
 ARTIFACT_VERSIONS = "/artifact_versions"
 ARTIFACT_VISUALIZATIONS = "/artifact_visualizations"
+AUTH = "/auth"
 CODE_REFERENCES = "/code_references"
 CODE_REPOSITORIES = "/code_repositories"
 COMPONENT_TYPES = "/component-types"
@@ -352,7 +352,6 @@ EVENT_FLAVORS = "/event-flavors"
 EVENT_SOURCES = "/event-sources"
 FLAVORS = "/flavors"
 GET_OR_CREATE = "/get-or-create"
-GRAPH = "/graph"
 HEALTH = "/health"
 INFO = "/info"
 LOGIN = "/login"
@@ -445,33 +444,6 @@ KUBERNETES_CLUSTER_RESOURCE_TYPE = "kubernetes-cluster"
 
 # Stack Recipe constants
 STACK_RECIPES_GITHUB_REPO = "https://github.com/zenml-io/mlops-stacks.git"
-ALPHA_MESSAGE = (
-    "The mlstacks tool/package is in alpha and actively being developed. "
-    "Please avoid running mission-critical workloads on resources deployed "
-    "through these commands. If you encounter any problems, create an issue "
-    f"on the repository {STACK_RECIPES_GITHUB_REPO} and we'll help you out!"
-)
-NOT_INSTALLED_MESSAGE = (
-    "The prerequisites for using `mlstacks` (the `mlstacks` and "
-    "`python-terraform` packages seem to be unavailable on your machine "
-    "and/or in your environment. To install the missing dependencies: \n\n"
-    "`pip install mlstacks`"
-)
-TERRAFORM_NOT_INSTALLED_MESSAGE = (
-    "Terraform appears not to be installed on your machine and/or in your "
-    "environment. Please install Terraform and try again."
-)
-STACK_RECIPE_MODULAR_RECIPES = ["aws", "gcp", "k3d"]
-MLSTACKS_SUPPORTED_STACK_COMPONENTS = [
-    "artifact_store",
-    "container_registry",
-    "experiment_tracker",
-    "orchestrator",
-    "model_deployer",
-    "mlops_platform",
-    "step_operator",
-]
-
 
 # Parameters for internal ZenML Models
 TEXT_FIELD_MAX_LENGTH = 65535
@@ -498,11 +470,3 @@ BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
 
 
 STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
-
-# ZenML Pro
-ZENML_PRO_CONNECTION_ISSUES_SUSPENDED_PAUSED_TENANT_HINT = (
-    "\nHINT: Since you are trying to communicate with the ZenML Pro Tenant, "
-    "please make sure that your tenant is in RUNNING state on your "
-    "Organization page. If the tenant is PAUSED you can `Resume` it via UI "
-    "and try again."
-)

zenml/enums.py

@@ -198,11 +198,8 @@ class SecretValidationLevel(StrEnum):
 class ServerProviderType(StrEnum):
     """ZenML server providers."""
 
-    LOCAL = "local"
+    DAEMON = "daemon"
     DOCKER = "docker"
-    AWS = "aws"
-    GCP = "gcp"
-    AZURE = "azure"
 
 
 class AnalyticsEventSource(StrEnum):

zenml/environment.py

@@ -16,7 +16,7 @@
 import os
 import platform
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, cast
+from typing import Dict
 
 import distro
 
@@ -26,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__)
 
 
@@ -128,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]:
@@ -385,298 +366,3 @@ class Environment(metaclass=SingletonMetaClass):
             "LIGHTNING_CLOUD_URL" in os.environ
             and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
         )
-
-    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/how-to/track-metrics-metadata/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.
-
-        Returns:
-            The newly created class.
-        """
-        cls = cast(
-            Type["BaseEnvironmentComponent"],
-            super().__new__(mcs, name, bases, dct),
-        )
-        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/exceptions.py

@@ -13,10 +13,7 @@
 #  permissions and limitations under the License.
 """ZenML specific exception definitions."""
 
-from typing import TYPE_CHECKING, Dict, List, Optional, Type
-
-if TYPE_CHECKING:
-    from zenml.steps import BaseParameters
+from typing import Dict, Optional
 
 
 class ZenMLBaseException(Exception):
@@ -50,6 +47,17 @@ class AuthorizationException(ZenMLBaseException):
     """Raised when an authorization error occurred while trying to access a ZenML resource ."""
 
 
+class CredentialsNotValid(AuthorizationException):
+    """Raised when the credentials provided are invalid.
+
+    This is a subclass of AuthorizationException and should only be raised when
+    the authentication credentials are invalid (e.g. expired API token, invalid
+    username/password, invalid signature). If caught by the ZenML client, it
+    will trigger an invalidation of the currently cached API token and a
+    re-authentication flow.
+    """
+
+
 class DoesNotExistException(ZenMLBaseException):
     """Raises exception when the entity does not exist in the system but an action is being done that requires it to be present."""
 
@@ -132,43 +140,6 @@ class PipelineConfigurationError(ZenMLBaseException):
     """Raises exceptions when a pipeline configuration contains invalid values."""
 
 
-class MissingStepParameterError(ZenMLBaseException):
-    """Raises exceptions when a step parameter is missing when running a pipeline."""
-
-    def __init__(
-        self,
-        step_name: str,
-        missing_parameters: List[str],
-        parameters_class: Type["BaseParameters"],
-    ):
-        """Initializes a MissingStepParameterError object.
-
-        Args:
-            step_name: Name of the step for which one or more parameters
-                are missing.
-            missing_parameters: Names of all parameters which are missing.
-            parameters_class: Class of the parameters object for which
-                the parameters are missing.
-        """
-        import textwrap
-
-        message = textwrap.fill(
-            textwrap.dedent(
-                f"""
-            Missing parameters {missing_parameters} for '{step_name}' step.
-            There are three ways to solve this issue:
-            (1) Specify a default value in the parameters class
-            `{parameters_class.__name__}`
-            (2) Specify the parameters in code when creating the pipeline:
-            `my_pipeline({step_name}(params={parameters_class.__name__}(...))`
-            (3) Specify the parameters in a yaml configuration file and pass
-            it to the pipeline: `my_pipeline(...).run(config_path='path_to_yaml')`
-            """
-            )
-        )
-        super().__init__(message)
-
-
 class IntegrationError(ZenMLBaseException):
     """Raises exceptions when a requested integration can not be activated."""
 

zenml/hooks/hook_validators.py

@@ -43,9 +43,6 @@ def resolve_and_validate_hook(hook: "HookSpecification") -> Source:
     if not callable(func):
         raise ValueError(f"{func} is not a valid function.")
 
-    from zenml.new.steps.step_context import StepContext
-    from zenml.steps.base_parameters import BaseParameters
-
     sig = inspect.getfullargspec(inspect.unwrap(func))
     sig_annotations = sig.annotations
     if "return" in sig_annotations:
@@ -62,11 +59,7 @@ def resolve_and_validate_hook(hook: "HookSpecification") -> Source:
         seen_annotations = set()
         for annotation in annotations:
             if annotation:
-                if annotation not in (
-                    BaseException,
-                    BaseParameters,
-                    StepContext,
-                ):
+                if annotation not in (BaseException,):
                     raise ValueError(
                         "Hook arguments must be of type `BaseException`, not "
                         f"`{annotation}`."

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -50,8 +50,7 @@ if TYPE_CHECKING:
         DagConfiguration,
         TaskConfiguration,
     )
-    from zenml.models import PipelineDeploymentResponse
-    from zenml.pipelines import Schedule
+    from zenml.models import PipelineDeploymentResponse, ScheduleResponse
     from zenml.stack import Stack
 
 logger = get_logger(__name__)
@@ -391,7 +390,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
 
     @staticmethod
     def _translate_schedule(
-        schedule: Optional["Schedule"] = None,
+        schedule: Optional["ScheduleResponse"] = None,
     ) -> Dict[str, Any]:
         """Convert ZenML schedule into Airflow schedule.
 

zenml/integrations/argilla/__init__.py

@@ -26,7 +26,7 @@ class ArgillaIntegration(Integration):
 
     NAME = ARGILLA
     REQUIREMENTS = [
-        "argilla>=1.20.0,<2",
+        "argilla>=2.0.0",
     ]
 
     @classmethod