zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/cli/stack_components.py

@@ -136,9 +136,18 @@ def generate_stack_component_describe_command(
             if active_components:
                 active_component_id = active_components[0].id
 
+            if component_.connector:
+                # We also need the flavor to get the connector requirements
+                connector_requirements = client.get_flavor_by_name_and_type(
+                    name=component_.flavor, component_type=component_type
+                ).connector_requirements
+            else:
+                connector_requirements = None
+
             cli_utils.print_stack_component_configuration(
                 component=component_,
                 active_status=component_.id == active_component_id,
+                connector_requirements=connector_requirements,
             )
 
             print_model_url(get_component_url(component_))

zenml/cli/text_utils.py

@@ -15,7 +15,9 @@
 
 from typing import List
 
-from rich.markdown import Markdown
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.markdown import Heading, Markdown
+from rich.text import Text
 
 zenml_cli_welcome_message = Markdown(
     """
@@ -82,3 +84,35 @@ links below.\n
 
 """
     )
+
+
+class OldSchoolMarkdownHeading(Heading):
+    """A traditional markdown heading."""
+
+    def __rich_console__(
+        self, console: Console, options: ConsoleOptions
+    ) -> RenderResult:
+        """Render the heading.
+
+        Args:
+            console: The console rendering the content.
+            options: The console options.
+
+        Yields:
+            RenderResult: The rendered content.
+        """
+        text = self.text
+        text.justify = "left"
+        if self.tag == "h1":
+            # Underline and bold h1s
+            yield Text("", style="bold")
+            yield text
+            yield "=" * len(text)
+        else:
+            if self.tag == "h2":
+                # Just underline h2s
+                yield Text("", style="underline")
+            else:
+                # Just bold everything else
+                yield Text("", style="bold")
+            yield text

zenml/cli/utils.py

@@ -17,6 +17,7 @@ import contextlib
 import datetime
 import json
 import os
+import platform
 import re
 import subprocess
 import sys
@@ -43,11 +44,13 @@ import pkg_resources
 import yaml
 from pydantic import BaseModel, SecretStr
 from rich import box, table
+from rich.console import Console
 from rich.emoji import Emoji, NoEmoji
 from rich.markdown import Markdown
 from rich.markup import escape
-from rich.prompt import Confirm
+from rich.prompt import Confirm, Prompt
 from rich.style import Style
+from rich.table import Table
 
 from zenml.client import Client
 from zenml.console import console, zenml_style_defaults
@@ -69,6 +72,7 @@ from zenml.models import (
     BoolFilter,
     NumericFilter,
     Page,
+    ServiceConnectorRequirements,
     StrFilter,
     UUIDFilter,
 )
@@ -582,6 +586,23 @@ def print_stack_configuration(stack: "StackResponse", active: bool) -> None:
         for component in rich_table.columns[0]._cells
     ]
     console.print(rich_table)
+
+    if not stack.labels:
+        declare("No labels are set for this stack.")
+    else:
+        rich_table = table.Table(
+            box=box.HEAVY_EDGE,
+            title="Labels",
+            show_lines=True,
+        )
+        rich_table.add_column("LABEL")
+        rich_table.add_column("VALUE", overflow="fold")
+
+        for label, value in stack.labels.items():
+            rich_table.add_row(label, str(value))
+
+        console.print(rich_table)
+
     declare(
         f"Stack '{stack.name}' with id '{stack.id}' is "
         f"{f'owned by user {stack.user.name}.' if stack.user else 'unowned.'}"
@@ -613,13 +634,18 @@ def print_flavor_list(flavors: Page["FlavorResponse"]) -> None:
 
 
 def print_stack_component_configuration(
-    component: "ComponentResponse", active_status: bool
+    component: "ComponentResponse",
+    active_status: bool,
+    connector_requirements: Optional[ServiceConnectorRequirements] = None,
 ) -> None:
     """Prints the configuration options of a stack component.
 
     Args:
         component: The stack component to print.
         active_status: Whether the stack component is active.
+        connector_requirements: Connector requirements for the component, taken
+            from the component flavor. Only needed if the component has a
+            connector.
     """
     if component.user:
         user_name = component.user.name
@@ -675,7 +701,7 @@ def print_stack_component_configuration(
         rich_table.add_column("VALUE", overflow="fold")
 
         for label, value in component.labels.items():
-            rich_table.add_row(label, value)
+            rich_table.add_row(label, str(value))
 
         console.print(rich_table)
 
@@ -690,11 +716,17 @@ def print_stack_component_configuration(
         rich_table.add_column("PROPERTY")
         rich_table.add_column("VALUE", overflow="fold")
 
+        resource_type = (
+            connector_requirements.resource_type
+            if connector_requirements
+            else component.connector.resource_types[0]
+        )
+
         connector_dict = {
             "ID": str(component.connector.id),
             "NAME": component.connector.name,
             "TYPE": component.connector.type,
-            "RESOURCE TYPE": component.connector.resource_types[0],
+            "RESOURCE TYPE": resource_type,
             "RESOURCE NAME": component.connector_resource_id
             or component.connector.resource_id
             or "N/A",
@@ -1276,9 +1308,11 @@ def pretty_print_model_version_table(
             "NAME": model_version.registered_model.name,
             "MODEL_VERSION": model_version.version,
             "VERSION_DESCRIPTION": model_version.description,
-            "METADATA": model_version.metadata.model_dump()
-            if model_version.metadata
-            else {},
+            "METADATA": (
+                model_version.metadata.model_dump()
+                if model_version.metadata
+                else {}
+            ),
         }
         for model_version in model_versions
     ]
@@ -1318,9 +1352,11 @@ def pretty_print_model_version_details(
             if model_version.last_updated_at
             else "N/A"
         ),
-        "METADATA": model_version.metadata.model_dump()
-        if model_version.metadata
-        else {},
+        "METADATA": (
+            model_version.metadata.model_dump()
+            if model_version.metadata
+            else {}
+        ),
         "MODEL_SOURCE_URI": model_version.model_source_uri,
         "STAGE": model_version.stage.value,
     }
@@ -2748,3 +2784,84 @@ def is_jupyter_installed() -> bool:
         return True
     except pkg_resources.DistributionNotFound:
         return False
+
+
+def multi_choice_prompt(
+    object_type: str,
+    choices: List[List[Any]],
+    headers: List[str],
+    prompt_text: str,
+    allow_zero_be_a_new_object: bool = False,
+    default_choice: Optional[str] = None,
+) -> Optional[int]:
+    """Prompts the user to select a choice from a list of choices.
+
+    Args:
+        object_type: The type of the object
+        choices: The list of choices
+        prompt_text: The prompt text
+        headers: The list of headers.
+        allow_zero_be_a_new_object: Whether to allow zero as a new object
+        default_choice: The default choice
+
+    Returns:
+        The selected choice index or None for new object
+
+    Raises:
+        RuntimeError: If no choice is made.
+    """
+    table = Table(
+        title=f"Available {object_type}",
+        show_header=True,
+        border_style=None,
+        expand=True,
+        show_lines=True,
+    )
+    table.add_column("Choice", justify="left", width=1)
+    for h in headers:
+        table.add_column(
+            h.replace("_", " ").capitalize(), justify="left", width=10
+        )
+
+    i_shift = 0
+    if allow_zero_be_a_new_object:
+        i_shift = 1
+        table.add_row(
+            "[0]",
+            *([f"Create a new {object_type}"] * len(headers)),
+        )
+    for i, one_choice in enumerate(choices):
+        table.add_row(f"[{i+i_shift}]", *[str(x) for x in one_choice])
+    Console().print(table)
+
+    selected = Prompt.ask(
+        prompt_text,
+        choices=[str(i) for i in range(0, len(choices) + 1)],
+        default=default_choice,
+        show_choices=False,
+    )
+    if selected is None:
+        raise RuntimeError(f"No {object_type} was selected")
+
+    if selected == "0" and allow_zero_be_a_new_object:
+        return None
+    else:
+        return int(selected) - i_shift
+
+
+def requires_mac_env_var_warning() -> bool:
+    """Checks if a warning needs to be shown for a local Mac server.
+
+    This is for the case where a user is on a macOS system, trying to run a
+    local server but is missing the `OBJC_DISABLE_INITIALIZE_FORK_SAFETY`
+    environment variable.
+
+    Returns:
+        bool: True if a warning needs to be shown, False otherwise.
+    """
+    if sys.platform == "darwin":
+        mac_version_tuple = tuple(map(int, platform.release().split(".")[:2]))
+        return not os.getenv(
+            "OBJC_DISABLE_INITIALIZE_FORK_SAFETY"
+        ) and mac_version_tuple >= (10, 13)
+    return False

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,
@@ -1153,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.
 
@@ -1160,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.
@@ -1184,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)
@@ -1227,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,
@@ -1279,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]]]
@@ -1289,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.
@@ -1343,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,
@@ -2439,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:
@@ -2579,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:
@@ -3942,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(
@@ -4151,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 "

zenml/config/docker_settings.py

@@ -102,10 +102,11 @@ class DockerSettings(BaseSettings):
     Depending on the configuration of this object, requirements will be
     installed in the following order (each step optional):
     - The packages installed in your local python environment
-    - The packages specified via the `requirements` attribute
-    - The packages specified via the `required_integrations` and potentially
-      stack requirements
     - The packages specified via the `required_hub_plugins` attribute
+    - The packages required by the stack unless this is disabled by setting
+      `install_stack_requirements=False`.
+    - The packages specified via the `required_integrations`
+    - The packages specified via the `requirements` attribute
 
     Attributes:
         parent_image: Full name of the Docker image that should be
@@ -141,7 +142,9 @@ class DockerSettings(BaseSettings):
         target_repository: Name of the Docker repository to which the
             image should be pushed. This repository will be appended to the
             registry URI of the container registry of your stack and should
-            therefore **not** include any registry.
+            therefore **not** include any registry. If not specified, the
+            default repository name configured in the container registry
+            stack component settings will be used.
         python_package_installer: The package installer to use for python
             packages.
         python_package_installer_args: Arguments to pass to the python package
@@ -200,7 +203,7 @@ class DockerSettings(BaseSettings):
     build_options: Dict[str, Any] = {}
     parent_image_build_config: Optional[DockerBuildConfig] = None
     skip_build: bool = False
-    target_repository: str = "zenml"
+    target_repository: Optional[str] = None
     python_package_installer: PythonPackageInstaller = (
         PythonPackageInstaller.PIP
     )

zenml/constants.py

@@ -173,6 +173,7 @@ ENV_ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRES_MINUTES"
 )
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
+ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
@@ -263,6 +264,7 @@ DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS = 3
 DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT = 60 * 5  # 5 minutes
 DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING = 5  # seconds
 DEFAULT_HTTP_TIMEOUT = 30
+SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT = 120  # seconds
 ZENML_API_KEY_PREFIX = "ZENKEY_"
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
@@ -337,6 +339,7 @@ ARTIFACT_VISUALIZATIONS = "/artifact_visualizations"
 CODE_REFERENCES = "/code_references"
 CODE_REPOSITORIES = "/code_repositories"
 COMPONENT_TYPES = "/component-types"
+CONFIG = "/config"
 CURRENT_USER = "/current-user"
 DEACTIVATE = "/deactivate"
 DEVICES = "/devices"
@@ -346,6 +349,7 @@ EMAIL_ANALYTICS = "/email-opt-in"
 EVENT_FLAVORS = "/event-flavors"
 EVENT_SOURCES = "/event-sources"
 FLAVORS = "/flavors"
+FULL_STACK = "/full-stack"
 GET_OR_CREATE = "/get-or-create"
 GRAPH = "/graph"
 HEALTH = "/health"
@@ -372,13 +376,15 @@ SERVICE_CONNECTOR_CLIENT = "/client"
 SERVICE_CONNECTOR_RESOURCES = "/resources"
 SERVICE_CONNECTOR_TYPES = "/service_connector_types"
 SERVICE_CONNECTOR_VERIFY = "/verify"
-SERVICE_CONNECTOR_RESOURCES = "/resources"
+SERVICE_CONNECTOR_FULL_STACK = "/full_stack_resources"
 MODELS = "/models"
 MODEL_VERSIONS = "/model_versions"
 MODEL_VERSION_ARTIFACTS = "/model_version_artifacts"
 MODEL_VERSION_PIPELINE_RUNS = "/model_version_pipeline_runs"
 SERVICES = "/services"
 SERVICE_CONNECTORS = "/service_connectors"
+STACK = "/stack"
+STACK_DEPLOYMENT = "/stack-deployment"
 STACKS = "/stacks"
 STACK_COMPONENTS = "/components"
 STATISTICS = "/statistics"
@@ -388,7 +394,9 @@ STEPS = "/steps"
 TAGS = "/tags"
 TRIGGERS = "/triggers"
 TRIGGER_EXECUTIONS = "/trigger_executions"
+ONBOARDING_STATE = "/onboarding_state"
 USERS = "/users"
+URL = "/url"
 VERSION_1 = "/v1"
 VISUALIZE = "/visualize"
 WEBHOOKS = "/webhooks"
@@ -458,6 +466,7 @@ MLSTACKS_SUPPORTED_STACK_COMPONENTS = [
     "step_operator",
 ]
 
+
 # Parameters for internal ZenML Models
 TEXT_FIELD_MAX_LENGTH = 65535
 STR_FIELD_MAX_LENGTH = 255
@@ -480,3 +489,6 @@ FINISHED_ONBOARDING_SURVEY_KEY = "awareness_channels"
 
 # Name validation
 BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
+
+
+STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours

zenml/container_registries/base_container_registry.py

@@ -41,6 +41,7 @@ class BaseContainerRegistryConfig(AuthenticationConfigMixin):
     """
 
     uri: str
+    default_repository: Optional[str] = None
 
     @field_validator("uri")
     @classmethod

zenml/enums.py

@@ -384,3 +384,26 @@ class PluginSubType(StrEnum):
     WEBHOOK = "webhook"
     # Action Subtypes
     PIPELINE_RUN = "pipeline_run"
+
+
+class OnboardingStep(StrEnum):
+    """All onboarding steps."""
+
+    DEVICE_VERIFIED = "device_verified"
+    PIPELINE_RUN = "pipeline_run"
+    STARTER_SETUP_COMPLETED = "starter_setup_completed"
+    STACK_WITH_REMOTE_ORCHESTRATOR_CREATED = (
+        "stack_with_remote_orchestrator_created"
+    )
+    PIPELINE_RUN_WITH_REMOTE_ORCHESTRATOR = (
+        "pipeline_run_with_remote_orchestrator"
+    )
+    PRODUCTION_SETUP_COMPLETED = "production_setup_completed"
+
+
+class StackDeploymentProvider(StrEnum):
+    """All possible stack deployment providers."""
+
+    AWS = "aws"
+    GCP = "gcp"
+    AZURE = "azure"

zenml/event_hub/event_hub.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base class for all the Event Hub."""
 
-from functools import partial
 from typing import List
 
 from pydantic import ValidationError
@@ -125,13 +124,11 @@ class InternalEventHub(BaseEventHub):
         """
         # get all event sources configured for this flavor
         triggers: List[TriggerResponse] = depaginate(
-            partial(
-                self.zen_store.list_triggers,
-                trigger_filter_model=TriggerFilter(
-                    event_source_id=event_source.id, is_active=True
-                ),
-                hydrate=True,
-            )
+            self.zen_store.list_triggers,
+            trigger_filter_model=TriggerFilter(
+                event_source_id=event_source.id, is_active=True
+            ),
+            hydrate=True,
         )
 
         trigger_list: List[TriggerResponse] = []

zenml/integrations/__init__.py

@@ -23,6 +23,7 @@ from zenml.integrations.aws import AWSIntegration  # noqa
 from zenml.integrations.azure import AzureIntegration  # noqa
 from zenml.integrations.bentoml import BentoMLIntegration  # noqa
 from zenml.integrations.bitbucket import BitbucketIntegration  # noqa
+from zenml.integrations.databricks import DatabricksIntegration  # noqa
 from zenml.integrations.comet import CometIntegration  # noqa
 from zenml.integrations.deepchecks import DeepchecksIntegration  # noqa
 from zenml.integrations.discord import DiscordIntegration  # noqa

zenml/integrations/aws/__init__.py

@@ -43,6 +43,7 @@ class AWSIntegration(Integration):
         "kubernetes",
         "aws-profile-manager",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
     @staticmethod
     def activate() -> None:

zenml/integrations/azure/__init__.py

@@ -42,12 +42,13 @@ class AzureIntegration(Integration):
         "adlfs>=2021.10.0",
         "azure-keyvault-keys",
         "azure-keyvault-secrets",
-        "azure-identity==1.10.0",
-        "azureml-core==1.54.0.post1",
+        "azure-identity",
+        "azureml-core==1.56.0",
         "azure-mgmt-containerservice>=20.0.0",
         "azure-storage-blob==12.17.0",  # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056
         "kubernetes",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
     @staticmethod
     def activate() -> None:

zenml/integrations/constants.py

@@ -22,6 +22,7 @@ BENTOML = "bentoml"
 BITBUCKET = "bitbucket"
 COMET = "comet"
 DASH = "dash"
+DATABRICKS = "databricks"
 DEEPCHECKS = "deepchecks"
 DISCORD = "discord"
 EVIDENTLY = "evidently"

zenml/integrations/databricks/__init__.py

@@ -0,0 +1,52 @@
+#  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.
+"""Initialization of the Databricks integration for ZenML."""
+
+from typing import List, Type
+
+from zenml.integrations.constants import (
+    DATABRICKS,
+)
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+DATABRICKS_ORCHESTRATOR_FLAVOR = "databricks"
+DATABRICKS_MODEL_DEPLOYER_FLAVOR = "databricks"
+DATABRICKS_SERVICE_ARTIFACT = "databricks_deployment_service"
+
+
+class DatabricksIntegration(Integration):
+    """Definition of Databricks Integration for ZenML."""
+
+    NAME = DATABRICKS
+    REQUIREMENTS = ["databricks-sdk==0.28.0"]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Databricks integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.databricks.flavors import (
+            DatabricksOrchestratorFlavor,
+            DatabricksModelDeployerFlavor,
+        )
+
+        return [
+            DatabricksOrchestratorFlavor,
+            DatabricksModelDeployerFlavor,
+        ]
+
+DatabricksIntegration.check_installation()