zenml-nightly 0.61.0.dev20240713__py3-none-any.whl → 0.62.0.dev20240717__py3-none-any.whl

README.md

@@ -327,7 +327,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.61.0 is out. Check out the release notes
+    🎉 Version 0.62.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.

RELEASE_NOTES.md

@@ -1,4 +1,44 @@
 <!-- markdown-link-check-disable -->
+# 0.62.0
+
+Building on top of the last release, this release adds a new and easy way to deploy a GCP ZenML stack from the dashboard and the CLI. Give it a try by going to the `Stacks` section in the dashboard or running the `zenml stack deploy` command! For more information on this new feature, please do check out [the video and blog](https://www.zenml.io/blog/easy-mlops-pipelines) from our previous release.
+
+We also [updated our Hugging Face integration](https://github.com/zenml-io/zenml/pull/2851) to support the automatic display of an embedded `datasets` preview pane in the ZenML Dashboard whenever you return a `Dataset` from a step. This was recently released by the Hugging Face datasets team and it allows you to easily visualize and inspect your data from the comfort of the dashboard.
+
+## What's Changed
+
+* Fix release action docker limit by @schustmi in https://github.com/zenml-io/zenml/pull/2837
+* Upgrade ruff and yamlfix to latest versions before running formatting by @christianversloot in https://github.com/zenml-io/zenml/pull/2577
+* Fixed edge-case where step run is stored incompletely by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2827
+* Docs for stack registration + deployment wizards by @htahir1 in https://github.com/zenml-io/zenml/pull/2814
+* Make upgrade checks in formatting script optional by @avishniakov in https://github.com/zenml-io/zenml/pull/2839
+* Enable migration testing for version 0.61.0 by @schustmi in https://github.com/zenml-io/zenml/pull/2836
+* One-click GCP stack deployments by @stefannica in https://github.com/zenml-io/zenml/pull/2833
+* Only login to docker for PRs with secret access by @schustmi in https://github.com/zenml-io/zenml/pull/2842
+* Add GCP Stack creation Wizard (CLI) by @avishniakov in https://github.com/zenml-io/zenml/pull/2826
+* Update onboarding by @schustmi in https://github.com/zenml-io/zenml/pull/2794
+* Merged log files in Step Ops steps might be not available on main process, due to merge in the step op by @avishniakov in https://github.com/zenml-io/zenml/pull/2795
+* Fix some broken links, copy paste commands, and made secrets more visible  by @htahir1 in https://github.com/zenml-io/zenml/pull/2848
+* Update stack deployment docs and other small fixes by @stefannica in https://github.com/zenml-io/zenml/pull/2846
+* Improved the `StepInterfaceError` message for missing inputs by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2849
+* add image pull secrets to k8s pod settings by @wjayesh in https://github.com/zenml-io/zenml/pull/2847
+* Include apt installation of libgomp1 for docker images with lightgbm by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2813
+* Patch filter mflow by stage by @whoknowsB in https://github.com/zenml-io/zenml/pull/2798
+* Bump mlflow to version 2.14.2 by @christianversloot in https://github.com/zenml-io/zenml/pull/2825
+* Fix Accelerate string arguments passing by @avishniakov in https://github.com/zenml-io/zenml/pull/2845
+* Fix CI by @schustmi in https://github.com/zenml-io/zenml/pull/2850
+* Added some visualizations for the HF dataset by @htahir1 in https://github.com/zenml-io/zenml/pull/2851
+* Fix skypilot versioning for the lambda integration by @wjayesh in https://github.com/zenml-io/zenml/pull/2853
+* Improve custom visualization docs by @htahir1 in https://github.com/zenml-io/zenml/pull/2855
+* Fix list typo by @htahir1 in https://github.com/zenml-io/zenml/pull/2856
+* Endpoint to get existing and prospective resources for service connector by @avishniakov in https://github.com/zenml-io/zenml/pull/2854
+* Databricks integrations by @safoinme in https://github.com/zenml-io/zenml/pull/2823
+
+## New Contributors
+* @whoknowsB made their first contribution in https://github.com/zenml-io/zenml/pull/2798
+
+**Full Changelog**: https://github.com/zenml-io/zenml/compare/0.61.0...0.62.0
+
 # 0.61.0
 
 This release comes with a new and easy way to deploy an AWS ZenML stack from the dashboard and the CLI. Give it a try by going to the `Stacks` section in the dashboard or running the `zenml stack deploy` command!

zenml/VERSION

@@ -1 +1 @@
-0.61.0.dev20240713
+0.62.0.dev20240717

zenml/__init__.py

@@ -56,6 +56,7 @@ from zenml.new.pipelines.pipeline_decorator import pipeline
 from zenml.new.steps.step_decorator import step
 from zenml.new.steps.step_context import get_step_context
 from zenml.steps.utils import log_step_metadata
+from zenml.entrypoints import entrypoint
 
 __all__ = [
     "ArtifactConfig",
@@ -74,4 +75,5 @@ __all__ = [
     "save_artifact",
     "show",
     "step",
+    "entrypoint",
 ]

zenml/cli/stack.py

@@ -83,9 +83,10 @@ from zenml.models.v2.misc.full_stack import (
     ComponentInfo,
     FullStackRequest,
     ServiceConnectorInfo,
+    ServiceConnectorResourcesInfo,
 )
-from zenml.models.v2.misc.service_connector_type import (
-    ServiceConnectorTypedResourcesModel,
+from zenml.service_connectors.service_connector_utils import (
+    get_resources_options_from_resource_model_for_full_stack,
 )
 from zenml.utils.dashboard_utils import get_component_url, get_stack_url
 from zenml.utils.io_utils import create_dir_recursive_if_not_exists
@@ -329,6 +330,13 @@ def register_stack(
                 auto_configure=True,
                 verify=False,
             )
+        except NotImplementedError:
+            cli_utils.warning(
+                f"The {provider.upper()} service connector libraries are not "
+                "installed properly. Please run `zenml integration install "
+                f"{provider}` and try again to enable auto-discovery of the "
+                "connection configuration."
+            )
         except Exception:
             pass
         if service_connector_response:
@@ -390,8 +398,21 @@ def register_stack(
         labels["zenml:wizard"] = "true"
         if provider:
             labels["zenml:provider"] = provider
-        service_connector_resource_model = None
-        can_generate_long_tokens = False
+        resources_info = None
+        # explore the service connector
+        with console.status(
+            "Exploring resources available to the service connector...\n"
+        ):
+            resources_info = (
+                get_resources_options_from_resource_model_for_full_stack(
+                    connector_details=service_connector
+                )
+            )
+        if resources_info is None:
+            cli_utils.error(
+                f"Failed to fetch service connector resources information for {service_connector}..."
+            )
+
         # create components
         needed_components = (
             (StackComponentType.ARTIFACT_STORE, artifact_store),
@@ -408,21 +429,37 @@ def register_stack(
             else:
                 if isinstance(service_connector, UUID):
                     # find existing components under same connector
-                    existing_components = client.list_stack_components(
-                        type=component_type.value,
-                        connector_id=service_connector,
-                        size=100,
-                    )
-                    # if some existing components are found - prompt user what to do
-                    component_selected: Optional[int] = None
-                    if existing_components.total > 0:
+                    if (
+                        component_type
+                        in resources_info.components_resources_info
+                    ):
+                        existing_components = [
+                            existing_response
+                            for res_info in resources_info.components_resources_info[
+                                component_type
+                            ]
+                            for existing_response in res_info.connected_through_service_connector
+                        ]
+
+                        # if some existing components are found - prompt user what to do
+                        component_selected: Optional[int] = None
                         component_selected = cli_utils.multi_choice_prompt(
                             object_type=component_type.value.replace("_", " "),
                             choices=[
-                                [component.name]
-                                for component in existing_components.items
+                                [
+                                    component.flavor,
+                                    component.name,
+                                    component.configuration or "",
+                                    component.connector_resource_id,
+                                ]
+                                for component in existing_components
+                            ],
+                            headers=[
+                                "Type",
+                                "Name",
+                                "Configuration",
+                                "Connected as",
                             ],
-                            headers=["Name"],
                             prompt_text=f"We found these {component_type.value.replace('_', ' ')} "
                             "connected using the current service connector. Do you "
                             "want to create a new one or use existing one?",
@@ -433,61 +470,17 @@ def register_stack(
                     component_selected = None
 
                 if component_selected is None:
-                    if service_connector_resource_model is None:
-                        with console.status(
-                            "Exploring resources available to the service connector...\n"
-                        ):
-                            if isinstance(service_connector, UUID):
-                                service_connector_resource_model = (
-                                    client.verify_service_connector(
-                                        service_connector
-                                    )
-                                )
-                                existing_service_connector_info = (
-                                    client.get_service_connector(
-                                        service_connector
-                                    )
-                                )
-                                can_generate_long_tokens = not existing_service_connector_info.configuration.get(
-                                    "generate_temporary_tokens", True
-                                )
-                            else:
-                                _, service_connector_resource_model = (
-                                    client.create_service_connector(
-                                        name=stack_name,
-                                        connector_type=service_connector.type,
-                                        auth_method=service_connector.auth_method,
-                                        configuration=service_connector.configuration,
-                                        register=False,
-                                    )
-                                )
-                                can_generate_long_tokens = True
-                        if service_connector_resource_model is None:
-                            cli_utils.error(
-                                f"Failed to validate service connector {service_connector}..."
-                            )
-                    if provider is None:
-                        if isinstance(
-                            service_connector_resource_model.connector_type,
-                            str,
-                        ):
-                            provider = (
-                                service_connector_resource_model.connector_type
-                            )
-                        else:
-                            provider = service_connector_resource_model.connector_type.connector_type
-
                     component_info = _get_stack_component_info(
                         component_type=component_type.value,
-                        cloud_provider=provider,
-                        service_connector_resource_models=service_connector_resource_model.resources,
+                        cloud_provider=provider
+                        or resources_info.connector_type,
+                        resources_info=resources_info,
                         service_connector_index=0,
-                        can_generate_long_tokens=can_generate_long_tokens,
                     )
                     component_name = stack_name
                     created_objects.add(component_type.value)
                 else:
-                    selected_component = existing_components.items[
+                    selected_component = existing_components[
                         component_selected
                     ]
                     component_info = selected_component.id
@@ -1752,12 +1745,12 @@ def deploy(
                 Markdown(
                     "## Configuration\n"
                     "You will be asked to provide the following configuration "
-                    "values during the deployment process:\n"
-                )
+                    "values during the deployment process:"
+                ),
+                "\n",
             )
 
             console.print(
-                "\n",
                 deployment_config.configuration,
                 no_wrap=True,
                 overflow="ignore",
@@ -2349,7 +2342,7 @@ def _get_service_connector_info(
             choices.append([value.name, required])
 
         selected_auth_idx = cli_utils.multi_choice_prompt(
-            object_type=f"authentication methods for {cloud_provider}",
+            object_type=f"authentication methods for {cloud_provider.upper()}",
             choices=choices,
             headers=headers,
             prompt_text="Please choose one of the authentication option above",
@@ -2393,10 +2386,7 @@ def _get_service_connector_info(
 def _get_stack_component_info(
     component_type: str,
     cloud_provider: str,
-    service_connector_resource_models: List[
-        ServiceConnectorTypedResourcesModel
-    ],
-    can_generate_long_tokens: bool,
+    resources_info: ServiceConnectorResourcesInfo,
     service_connector_index: Optional[int] = None,
 ) -> ComponentInfo:
     """Get a stack component info with given type and service connector.
@@ -2404,8 +2394,7 @@ def _get_stack_component_info(
     Args:
         component_type: The type of component to create.
         cloud_provider: The cloud provider to use.
-        service_connector_resource_models: The list of the available service connector resource models.
-        can_generate_long_tokens: Whether connector can generate long-living tokens.
+        resources_info: The resources info of the service connector.
         service_connector_index: The index of the service connector to use.
 
     Returns:
@@ -2420,62 +2409,30 @@ def _get_stack_component_info(
     if cloud_provider not in {"aws", "azure", "gcp"}:
         raise ValueError(f"Unknown cloud provider {cloud_provider}")
 
-    AWS_DOCS = (
-        "https://docs.zenml.io/how-to/auth-management/aws-service-connector"
-    )
-    GCP_DOCS = (
-        "https://docs.zenml.io/how-to/auth-management/gcp-service-connector"
-    )
-
     flavor = "undefined"
     service_connector_resource_id = None
     config = {}
+    choices = [
+        [cri.flavor, resource_id]
+        for cri in resources_info.components_resources_info[
+            StackComponentType(component_type)
+        ]
+        for resource_id in cri.accessible_by_service_connector
+    ]
     if component_type == "artifact_store":
-        available_storages: List[str] = []
-        if cloud_provider == "aws":
-            for each in service_connector_resource_models:
-                if each.resource_type == "s3-bucket":
-                    available_storages = each.resource_ids or []
-            flavor = "s3"
-            if not available_storages:
-                cli_utils.error(
-                    "We were unable to find any S3 buckets available "
-                    "to configured service connector. Please, verify "
-                    "that needed permission are granted for the "
-                    "service connector.\nDocumentation for the S3 "
-                    "Buckets configuration can be found at "
-                    f"{AWS_DOCS}#s3-bucket"
-                )
-        elif cloud_provider == "azure":
-            flavor = "azure"
-        elif cloud_provider == "gcp":
-            flavor = "gcp"
-            for each in service_connector_resource_models:
-                if each.resource_type == "gcs-bucket":
-                    available_storages = each.resource_ids or []
-            if not available_storages:
-                cli_utils.error(
-                    "We were unable to find any GCS buckets available "
-                    "to configured service connector. Please, verify "
-                    "that needed permission are granted for the "
-                    "service connector.\nDocumentation for the GCS "
-                    "Buckets configuration can be found at "
-                    f"{GCP_DOCS}#gcs-bucket"
-                )
-
         selected_storage_idx = cli_utils.multi_choice_prompt(
             object_type=f"{cloud_provider.upper()} storages",
-            choices=[[st] for st in available_storages],
-            headers=["Storage"],
+            choices=choices,
+            headers=["Artifact Store Type", "Storage"],
             prompt_text="Please choose one of the storages for the new artifact store:",
         )
         if selected_storage_idx is None:
             cli_utils.error("No storage selected.")
 
-        selected_storage = available_storages[selected_storage_idx]
+        selected_storage = choices[selected_storage_idx]
 
-        config = {"path": selected_storage}
-        service_connector_resource_id = selected_storage
+        config = {"path": selected_storage[1]}
+        service_connector_resource_id = selected_storage[1]
     elif component_type == "orchestrator":
 
         def query_gcp_region(compute_type: str) -> str:
@@ -2492,139 +2449,41 @@ def _get_stack_component_info(
             )
             return region
 
-        if cloud_provider == "aws":
-            available_orchestrators = []
-            for each in service_connector_resource_models:
-                types = []
-                if each.resource_type == "aws-generic":
-                    types = ["Sagemaker"]
-                    if can_generate_long_tokens:
-                        types.append("Skypilot (EC2)")
-                if each.resource_type == "kubernetes-cluster":
-                    types = ["Kubernetes"]
-
-                if each.resource_ids:
-                    for orchestrator in each.resource_ids:
-                        for t in types:
-                            available_orchestrators.append([t, orchestrator])
-            if not available_orchestrators:
-                cli_utils.error(
-                    "We were unable to find any orchestrator engines "
-                    "available to the service connector. Please, verify "
-                    "that needed permission are granted for the "
-                    "service connector.\nDocumentation for the Generic "
-                    "AWS resource configuration can be found at "
-                    f"{AWS_DOCS}#generic-aws-resource\n"
-                    "Documentation for the Kubernetes resource "
-                    "configuration can be found at "
-                    f"{AWS_DOCS}#eks-kubernetes-cluster"
-                )
-        elif cloud_provider == "gcp":
-            available_orchestrators = []
-            for each in service_connector_resource_models:
-                types = []
-                if each.resource_type == "gcp-generic":
-                    types = ["Vertex AI"]
-                    if can_generate_long_tokens:
-                        types.append("Skypilot (Compute)")
-                if each.resource_type == "kubernetes-cluster":
-                    types = ["Kubernetes"]
-
-                if each.resource_ids:
-                    for orchestrator in each.resource_ids:
-                        for t in types:
-                            available_orchestrators.append([t, orchestrator])
-            if not available_orchestrators:
-                cli_utils.error(
-                    "We were unable to find any orchestrator engines "
-                    "available to the service connector. Please, verify "
-                    "that needed permission are granted for the "
-                    "service connector.\nDocumentation for the Generic "
-                    "GCP resource configuration can be found at "
-                    f"{GCP_DOCS}#generic-gcp-resource\n"
-                    "Documentation for the GKE Kubernetes resource "
-                    "configuration can be found at "
-                    f"{GCP_DOCS}#gke-kubernetes-cluster"
-                )
-        elif cloud_provider == "azure":
-            pass
-
         selected_orchestrator_idx = cli_utils.multi_choice_prompt(
             object_type=f"orchestrators on {cloud_provider.upper()}",
-            choices=available_orchestrators,
-            headers=["Orchestrator Type", "Orchestrator details"],
+            choices=choices,
+            headers=["Orchestrator Type", "Details"],
             prompt_text="Please choose one of the orchestrators for the new orchestrator:",
         )
         if selected_orchestrator_idx is None:
             cli_utils.error("No orchestrator selected.")
 
-        selected_orchestrator = available_orchestrators[
-            selected_orchestrator_idx
-        ]
+        selected_orchestrator = choices[selected_orchestrator_idx]
 
         config = {}
-        if selected_orchestrator[0] == "Sagemaker":
-            flavor = "sagemaker"
+        flavor = selected_orchestrator[0]
+        if flavor == "sagemaker":
             execution_role = Prompt.ask("Enter an execution role ARN:")
             config["execution_role"] = execution_role
-        elif selected_orchestrator[0] == "Skypilot (EC2)":
-            flavor = "vm_aws"
+        elif flavor == "vm_aws":
             config["region"] = selected_orchestrator[1]
-        elif selected_orchestrator[0] == "Skypilot (Compute)":
-            flavor = "vm_gcp"
+        elif flavor == "vm_azure":
             config["region"] = query_gcp_region("Skypilot cluster")
-        elif selected_orchestrator[0] == "Vertex AI":
-            flavor = "vertex"
+        elif flavor == "vertex":
             config["location"] = query_gcp_region("Vertex AI job")
-        elif selected_orchestrator[0] == "Kubernetes":
-            flavor = "kubernetes"
-        else:
-            raise ValueError(
-                f"Unknown orchestrator type {selected_orchestrator[0]}"
-            )
         service_connector_resource_id = selected_orchestrator[1]
     elif component_type == "container_registry":
-
-        def _get_registries(registry_name: str, docs_link: str) -> List[str]:
-            available_registries: List[str] = []
-            for each in service_connector_resource_models:
-                if each.resource_type == "docker-registry":
-                    available_registries = each.resource_ids or []
-            if not available_registries:
-                cli_utils.error(
-                    "We were unable to find any container registries "
-                    "available to the service connector. Please, verify "
-                    "that needed permission are granted for the "
-                    f"service connector.\nDocumentation for the {registry_name} "
-                    "container registry resource configuration can "
-                    f"be found at {docs_link}"
-                )
-            return available_registries
-
-        if cloud_provider == "aws":
-            flavor = "aws"
-            available_registries = _get_registries(
-                "ECR", f"{AWS_DOCS}#ecr-container-registry"
-            )
-        if cloud_provider == "gcp":
-            flavor = "gcp"
-            available_registries = _get_registries(
-                "GCR", f"{GCP_DOCS}#gcr-container-registry"
-            )
-        if cloud_provider == "azure":
-            flavor = "azure"
-
         selected_registry_idx = cli_utils.multi_choice_prompt(
             object_type=f"{cloud_provider.upper()} registries",
-            choices=[[st] for st in available_registries],
-            headers=["Container Registry"],
+            choices=choices,
+            headers=["Container Registry Type", "Container Registry"],
             prompt_text="Please choose one of the registries for the new container registry:",
         )
         if selected_registry_idx is None:
             cli_utils.error("No container registry selected.")
-        selected_registry = available_registries[selected_registry_idx]
-        config = {"uri": selected_registry}
-        service_connector_resource_id = selected_registry
+        selected_registry = choices[selected_registry_idx]
+        config = {"uri": selected_registry[1]}
+        service_connector_resource_id = selected_registry[1]
     else:
         raise ValueError(f"Unknown component type {component_type}")
 

zenml/cli/stack_components.py

@@ -138,14 +138,16 @@ def generate_stack_component_describe_command(
 
             if component_.connector:
                 # We also need the flavor to get the connector requirements
-                flavor = client.get_flavor_by_name_and_type(
+                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=flavor.connector_requirements,
+                connector_requirements=connector_requirements,
             )
 
             print_model_url(get_component_url(component_))

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_"
@@ -375,6 +376,7 @@ SERVICE_CONNECTOR_CLIENT = "/client"
 SERVICE_CONNECTOR_RESOURCES = "/resources"
 SERVICE_CONNECTOR_TYPES = "/service_connector_types"
 SERVICE_CONNECTOR_VERIFY = "/verify"
+SERVICE_CONNECTOR_FULL_STACK = "/full_stack_resources"
 MODELS = "/models"
 MODEL_VERSIONS = "/model_versions"
 MODEL_VERSION_ARTIFACTS = "/model_version_artifacts"

zenml/entrypoints/entrypoint.py

@@ -15,6 +15,7 @@
 
 import argparse
 import logging
+import os
 import sys
 
 from zenml import constants
@@ -44,7 +45,8 @@ def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument(f"--{ENTRYPOINT_CONFIG_SOURCE_OPTION}", required=True)
     args, remaining_args = parser.parse_known_args()
-
+    if os.environ.get(constants.ENV_ZENML_CUSTOM_SOURCE_ROOT):
+        source_utils.set_custom_source_root(source_root=os.getcwd())
     entrypoint_config_class = source_utils.load_and_validate_class(
         args.entrypoint_config_source,
         expected_class=BaseEntrypointConfiguration,

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/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()