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

zenml/integrations/databricks/flavors/__init__.py

@@ -0,0 +1,30 @@
+#  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.
+"""Databricks integration flavors."""
+
+from zenml.integrations.databricks.flavors.databricks_orchestrator_flavor import (
+    DatabricksOrchestratorConfig,
+    DatabricksOrchestratorFlavor,
+)
+from zenml.integrations.databricks.flavors.databricks_model_deployer_flavor import (
+    DatabricksModelDeployerConfig,
+    DatabricksModelDeployerFlavor,
+)
+
+__all__ = [
+    "DatabricksOrchestratorFlavor",
+    "DatabricksOrchestratorConfig",
+    "DatabricksModelDeployerFlavor",
+    "DatabricksModelDeployerConfig",
+]

zenml/integrations/databricks/flavors/databricks_model_deployer_flavor.py

@@ -0,0 +1,118 @@
+#  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.
+"""Databricks model deployer flavor."""
+
+from typing import TYPE_CHECKING, Dict, Optional, Type
+
+from pydantic import BaseModel
+
+from zenml.integrations.databricks import DATABRICKS_MODEL_DEPLOYER_FLAVOR
+from zenml.model_deployers.base_model_deployer import (
+    BaseModelDeployerConfig,
+    BaseModelDeployerFlavor,
+)
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (
+        DatabricksModelDeployer,
+    )
+
+
+class DatabricksBaseConfig(BaseModel):
+    """Databricks Inference Endpoint configuration."""
+
+    workload_size: str
+    scale_to_zero_enabled: bool = False
+    env_vars: Optional[Dict[str, str]] = None
+    workload_type: Optional[str] = None
+    endpoint_secret_name: Optional[str] = None
+
+
+class DatabricksModelDeployerConfig(BaseModelDeployerConfig):
+    """Configuration for the Databricks model deployer.
+
+    Attributes:
+        host: Databricks host.
+        secret_name: Secret name to use for authentication.
+        client_id: Databricks client id.
+        client_secret: Databricks client secret.
+    """
+
+    host: str
+    secret_name: Optional[str] = None
+    client_id: Optional[str] = SecretField(default=None)
+    client_secret: Optional[str] = SecretField(default=None)
+
+
+class DatabricksModelDeployerFlavor(BaseModelDeployerFlavor):
+    """Databricks Endpoint model deployer flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return DATABRICKS_MODEL_DEPLOYER_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/databricks.png"
+
+    @property
+    def config_class(self) -> Type[DatabricksModelDeployerConfig]:
+        """Returns `DatabricksModelDeployerConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return DatabricksModelDeployerConfig
+
+    @property
+    def implementation_class(self) -> Type["DatabricksModelDeployer"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (
+            DatabricksModelDeployer,
+        )
+
+        return DatabricksModelDeployer

zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py

@@ -0,0 +1,147 @@
+#  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.
+"""Databricks orchestrator base config and settings."""
+
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.databricks import DATABRICKS_ORCHESTRATOR_FLAVOR
+from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestratorConfig
+from zenml.orchestrators.base_orchestrator import BaseOrchestratorFlavor
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.databricks.orchestrators import (
+        DatabricksOrchestrator,
+    )
+
+
+logger = get_logger(__name__)
+
+
+class DatabricksOrchestratorSettings(BaseSettings):
+    """Databricks orchestrator base settings.
+
+    Attributes:
+        spark_version: Spark version.
+        num_workers: Number of workers.
+        node_type_id: Node type id.
+        policy_id: Policy id.
+        autotermination_minutes: Autotermination minutes.
+        autoscale: Autoscale.
+        single_user_name: Single user name.
+        spark_conf: Spark configuration.
+        spark_env_vars: Spark environment variables.
+        schedule_timezone: Schedule timezone.
+    """
+
+    # Resources
+    spark_version: Optional[str] = None
+    num_workers: Optional[int] = None
+    node_type_id: Optional[str] = None
+    policy_id: Optional[str] = None
+    autotermination_minutes: Optional[int] = None
+    autoscale: Tuple[int, int] = (2, 3)
+    single_user_name: Optional[str] = None
+    spark_conf: Optional[Dict[str, str]] = None
+    spark_env_vars: Optional[Dict[str, str]] = None
+    schedule_timezone: Optional[str] = None
+
+
+class DatabricksOrchestratorConfig(
+    BaseOrchestratorConfig, DatabricksOrchestratorSettings
+):
+    """Databricks orchestrator base config.
+
+    Attributes:
+        host: Databricks host.
+        client_id: Databricks client id.
+        client_secret: Databricks client secret.
+    """
+
+    host: str
+    client_id: str = SecretField(default=None)
+    client_secret: str = SecretField(default=None)
+
+    @property
+    def is_local(self) -> bool:
+        """Checks if this stack component is running locally.
+
+        Returns:
+            True if this config is for a local component, False otherwise.
+        """
+        return False
+
+
+class DatabricksOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Databricks orchestrator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return DATABRICKS_ORCHESTRATOR_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/databricks.png"
+
+    @property
+    def config_class(self) -> Type[DatabricksOrchestratorConfig]:
+        """Returns `KubeflowOrchestratorConfig` config class.
+
+        Returns:
+                The config class.
+        """
+        return DatabricksOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["DatabricksOrchestrator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.databricks.orchestrators import (
+            DatabricksOrchestrator,
+        )
+
+        return DatabricksOrchestrator

zenml/integrations/databricks/model_deployers/__init__.py

@@ -0,0 +1,20 @@
+#  Copyright (c) ZenML GmbH 2023. 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:
+#
+#       http://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 model deployers."""
+
+from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (  # noqa
+    DatabricksModelDeployer,
+)
+
+__all__ = ["DatabricksModelDeployer"]

zenml/integrations/databricks/model_deployers/databricks_model_deployer.py

@@ -0,0 +1,249 @@
+#  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.
+"""Implementation of the Databricks Model Deployer."""
+
+from typing import ClassVar, Dict, Optional, Tuple, Type, cast
+from uuid import UUID
+
+from zenml.analytics.enums import AnalyticsEvent
+from zenml.analytics.utils import track_handler
+from zenml.client import Client
+from zenml.integrations.databricks import DATABRICKS_SERVICE_ARTIFACT
+from zenml.integrations.databricks.flavors.databricks_model_deployer_flavor import (
+    DatabricksModelDeployerConfig,
+    DatabricksModelDeployerFlavor,
+)
+from zenml.integrations.databricks.services.databricks_deployment import (
+    DatabricksDeploymentConfig,
+    DatabricksDeploymentService,
+)
+from zenml.logger import get_logger
+from zenml.model_deployers import BaseModelDeployer
+from zenml.model_deployers.base_model_deployer import (
+    DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    BaseModelDeployerFlavor,
+)
+from zenml.services import BaseService, ServiceConfig
+from zenml.stack.stack import Stack
+from zenml.stack.stack_validator import StackValidator
+
+logger = get_logger(__name__)
+
+
+class DatabricksModelDeployer(BaseModelDeployer):
+    """Databricks endpoint model deployer."""
+
+    NAME: ClassVar[str] = "Databricks"
+    FLAVOR: ClassVar[Type[BaseModelDeployerFlavor]] = (
+        DatabricksModelDeployerFlavor
+    )
+
+    @property
+    def config(self) -> DatabricksModelDeployerConfig:
+        """Config class for the Databricks Model deployer settings class.
+
+        Returns:
+            The configuration.
+        """
+        return cast(DatabricksModelDeployerConfig, self._config)
+
+    @property
+    def validator(self) -> Optional[StackValidator]:
+        """Validates the stack.
+
+        Returns:
+            A validator that checks that the stack contains a remote artifact
+            store.
+        """
+
+        def _validate_if_secret_or_token_is_present(
+            stack: "Stack",
+        ) -> Tuple[bool, str]:
+            """Check if client id and client secret or secret name is present in the stack.
+
+            Args:
+                stack: The stack to validate.
+
+            Returns:
+                A tuple with a boolean indicating whether the stack is valid
+                and a message describing the validation result.
+            """
+            return bool(
+                (self.config.client_id and self.config.client_secret)
+                or self.config.secret_name
+            ), (
+                "The Databricks model deployer requires either a secret name"
+                " or a client id and client secret to be present in the stack."
+            )
+
+        return StackValidator(
+            custom_validation_function=_validate_if_secret_or_token_is_present,
+        )
+
+    def _create_new_service(
+        self, id: UUID, timeout: int, config: DatabricksDeploymentConfig
+    ) -> DatabricksDeploymentService:
+        """Creates a new DatabricksDeploymentService.
+
+        Args:
+            id: the UUID of the model to be deployed with Databricks model deployer.
+            timeout: the timeout in seconds to wait for the Databricks inference endpoint
+                to be provisioned and successfully started or updated.
+            config: the configuration of the model to be deployed with Databricks model deployer.
+
+        Returns:
+            The DatabricksDeploymentConfig object that can be used to interact
+            with the Databricks inference endpoint.
+        """
+        # create a new service for the new model
+        service = DatabricksDeploymentService(uuid=id, config=config)
+        logger.info(
+            f"Creating an artifact {DATABRICKS_SERVICE_ARTIFACT} with service instance attached as metadata."
+            " If there's an active pipeline and/or model this artifact will be associated with it."
+        )
+        service.start(timeout=timeout)
+        return service
+
+    def _clean_up_existing_service(
+        self,
+        timeout: int,
+        force: bool,
+        existing_service: DatabricksDeploymentService,
+    ) -> None:
+        """Stop existing services.
+
+        Args:
+            timeout: the timeout in seconds to wait for the Databricks
+                deployment to be stopped.
+            force: if True, force the service to stop
+            existing_service: Existing Databricks deployment service
+        """
+        # stop the older service
+        existing_service.stop(timeout=timeout, force=force)
+
+    def perform_deploy_model(
+        self,
+        id: UUID,
+        config: ServiceConfig,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    ) -> BaseService:
+        """Create a new Databricks deployment service or update an existing one.
+
+        This should serve the supplied model and deployment configuration.
+
+        Args:
+            id: the UUID of the model to be deployed with Databricks.
+            config: the configuration of the model to be deployed with Databricks.
+            timeout: the timeout in seconds to wait for the Databricks endpoint
+                to be provisioned and successfully started or updated. If set
+                to 0, the method will return immediately after the Databricks
+                server is provisioned, without waiting for it to fully start.
+
+        Returns:
+            The ZenML Databricks deployment service object that can be used to
+            interact with the remote Databricks inference endpoint server.
+        """
+        with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler:
+            config = cast(DatabricksDeploymentConfig, config)
+            # create a new DatabricksDeploymentService instance
+            service = self._create_new_service(
+                id=id, timeout=timeout, config=config
+            )
+            logger.info(
+                f"Creating a new Databricks inference endpoint service: {service}"
+            )
+            # Add telemetry with metadata that gets the stack metadata and
+            # differentiates between pure model and custom code deployments
+            stack = Client().active_stack
+            stack_metadata = {
+                component_type.value: component.flavor
+                for component_type, component in stack.components.items()
+            }
+            analytics_handler.metadata = {
+                "store_type": Client().zen_store.type.value,
+                **stack_metadata,
+            }
+
+        return service
+
+    def perform_stop_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+        force: bool = False,
+    ) -> BaseService:
+        """Method to stop a model server.
+
+        Args:
+            service: The service to stop.
+            timeout: Timeout in seconds to wait for the service to stop.
+            force: If True, force the service to stop.
+
+        Returns:
+            The stopped service.
+        """
+        service.stop(timeout=timeout, force=force)
+        return service
+
+    def perform_start_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    ) -> BaseService:
+        """Method to start a model server.
+
+        Args:
+            service: The service to start.
+            timeout: Timeout in seconds to wait for the service to start.
+
+        Returns:
+            The started service.
+        """
+        service.start(timeout=timeout)
+        return service
+
+    def perform_delete_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+        force: bool = False,
+    ) -> None:
+        """Method to delete all configuration of a model server.
+
+        Args:
+            service: The service to delete.
+            timeout: Timeout in seconds to wait for the service to stop.
+            force: If True, force the service to stop.
+        """
+        service = cast(DatabricksDeploymentService, service)
+        self._clean_up_existing_service(
+            existing_service=service, timeout=timeout, force=force
+        )
+
+    @staticmethod
+    def get_model_server_info(  # type: ignore[override]
+        service_instance: "DatabricksDeploymentService",
+    ) -> Dict[str, Optional[str]]:
+        """Return implementation specific information that might be relevant to the user.
+
+        Args:
+            service_instance: Instance of a DatabricksDeploymentService
+
+        Returns:
+            Model server information.
+        """
+        return {
+            "PREDICTION_URL": service_instance.get_prediction_url(),
+            "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(),
+        }

zenml/integrations/databricks/orchestrators/__init__.py

@@ -0,0 +1,20 @@
+#  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 ZenML orchestrator."""
+
+from zenml.integrations.databricks.orchestrators.databricks_orchestrator import (  # noqa
+    DatabricksOrchestrator,
+)
+
+__all__ = ["DatabricksOrchestrator"]