zenml-nightly 0.61.0.dev20240712__py3-none-any.whl → 0.62.0.dev20240727__py3-none-any.whl

zenml/integrations/databricks/services/databricks_deployment.py

@@ -0,0 +1,407 @@
+#  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 Deployment service."""
+
+import time
+from typing import TYPE_CHECKING, Any, Dict, Generator, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import requests
+from databricks.sdk import WorkspaceClient as DatabricksClient
+from databricks.sdk.service.serving import (
+    EndpointCoreConfigInput,
+    EndpointStateConfigUpdate,
+    EndpointStateReady,
+    EndpointTag,
+    ServedModelInput,
+    ServingEndpointDetailed,
+)
+from pydantic import Field
+
+from zenml.client import Client
+from zenml.integrations.databricks.flavors.databricks_model_deployer_flavor import (
+    DatabricksBaseConfig,
+)
+from zenml.integrations.databricks.utils.databricks_utils import (
+    sanitize_labels,
+)
+from zenml.logger import get_logger
+from zenml.services import ServiceState, ServiceStatus, ServiceType
+from zenml.services.service import BaseDeploymentService, ServiceConfig
+
+logger = get_logger(__name__)
+
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+POLLING_TIMEOUT = 1200
+UUID_SLICE_LENGTH: int = 8
+
+
+class DatabricksDeploymentConfig(DatabricksBaseConfig, ServiceConfig):
+    """Databricks service configurations."""
+
+    model_uri: Optional[str] = Field(
+        None,
+        description="URI of the model to deploy. This can be a local path or a cloud storage path.",
+    )
+    host: Optional[str] = Field(
+        None, description="Databricks host URL for the deployment."
+    )
+
+    def get_databricks_deployment_labels(self) -> Dict[str, str]:
+        """Generate labels for the Databricks deployment from the service configuration.
+
+        These labels are attached to the Databricks deployment resource
+        and may be used as label selectors in lookup operations.
+
+        Returns:
+            The labels for the Databricks deployment.
+        """
+        labels = {}
+        if self.pipeline_name:
+            labels["zenml_pipeline_name"] = self.pipeline_name
+        if self.pipeline_step_name:
+            labels["zenml_pipeline_step_name"] = self.pipeline_step_name
+        if self.model_name:
+            labels["zenml_model_name"] = self.model_name
+        if self.model_uri:
+            labels["zenml_model_uri"] = self.model_uri
+        sanitize_labels(labels)
+        return labels
+
+
+class DatabricksServiceStatus(ServiceStatus):
+    """Databricks service status."""
+
+
+class DatabricksDeploymentService(BaseDeploymentService):
+    """Databricks model deployment service.
+
+    Attributes:
+        SERVICE_TYPE: a service type descriptor with information describing
+            the Databricks deployment service class
+        config: service configuration
+    """
+
+    SERVICE_TYPE = ServiceType(
+        name="databricks-deployment",
+        type="model-serving",
+        flavor="databricks",
+        description="Databricks inference endpoint prediction service",
+    )
+    config: DatabricksDeploymentConfig
+    status: DatabricksServiceStatus = Field(
+        default_factory=lambda: DatabricksServiceStatus()
+    )
+
+    def __init__(self, config: DatabricksDeploymentConfig, **attrs: Any):
+        """Initialize the Databricks deployment service.
+
+        Args:
+            config: service configuration
+            attrs: additional attributes to set on the service
+        """
+        super().__init__(config=config, **attrs)
+
+    def get_client_id_and_secret(self) -> Tuple[str, str, str]:
+        """Get the Databricks client id and secret.
+
+        Raises:
+            ValueError: If client id and secret are not found.
+
+        Returns:
+            Databricks client id and secret.
+        """
+        client = Client()
+        client_id = None
+        client_secret = None
+        host = None
+        from zenml.integrations.databricks.model_deployers.databricks_model_deployer import (
+            DatabricksModelDeployer,
+        )
+
+        model_deployer = client.active_stack.model_deployer
+        if not isinstance(model_deployer, DatabricksModelDeployer):
+            raise ValueError(
+                "DatabricksModelDeployer is not active in the stack."
+            )
+        host = model_deployer.config.host
+        self.config.host = host
+        if model_deployer.config.secret_name:
+            secret = client.get_secret(model_deployer.config.secret_name)
+            client_id = secret.secret_values["client_id"]
+            client_secret = secret.secret_values["client_secret"]
+
+        else:
+            client_id = model_deployer.config.client_id
+            client_secret = model_deployer.config.client_secret
+        if not client_id:
+            raise ValueError("Client id not found.")
+        if not client_secret:
+            raise ValueError("Client secret not found.")
+        if not host:
+            raise ValueError("Host not found.")
+        return host, client_id, client_secret
+
+    def _get_databricks_deployment_labels(self) -> Dict[str, str]:
+        """Generate the labels for the Databricks deployment from the service configuration.
+
+        Returns:
+            The labels for the Databricks deployment.
+        """
+        labels = self.config.get_databricks_deployment_labels()
+        labels["zenml_service_uuid"] = str(self.uuid)
+        sanitize_labels(labels)
+        return labels
+
+    @property
+    def databricks_client(self) -> DatabricksClient:
+        """Get the deployed Databricks inference endpoint.
+
+        Returns:
+            databricks inference endpoint.
+        """
+        return DatabricksClient(
+            host=self.get_client_id_and_secret()[0],
+            client_id=self.get_client_id_and_secret()[1],
+            client_secret=self.get_client_id_and_secret()[2],
+        )
+
+    @property
+    def databricks_endpoint(self) -> ServingEndpointDetailed:
+        """Get the deployed Hugging Face inference endpoint.
+
+        Returns:
+            Databricks inference endpoint.
+        """
+        return self.databricks_client.serving_endpoints.get(
+            name=self._generate_an_endpoint_name(),
+        )
+
+    @property
+    def prediction_url(self) -> Optional[str]:
+        """The prediction URI exposed by the prediction service.
+
+        Returns:
+            The prediction URI exposed by the prediction service, or None if
+            the service is not yet ready.
+        """
+        return f"{self.config.host}/serving-endpoints/{self._generate_an_endpoint_name()}/invocations"
+
+    def provision(self) -> None:
+        """Provision or update remote Databricks deployment instance."""
+        from databricks.sdk.service.serving import (
+            ServedModelInputWorkloadSize,
+            ServedModelInputWorkloadType,
+        )
+
+        tags = []
+        for key, value in self._get_databricks_deployment_labels().items():
+            tags.append(EndpointTag(key=key, value=value))
+        # Attempt to create and wait for the inference endpoint
+        served_model = ServedModelInput(
+            model_name=self.config.model_name,
+            model_version=self.config.model_version,
+            scale_to_zero_enabled=self.config.scale_to_zero_enabled,
+            workload_type=ServedModelInputWorkloadType(
+                self.config.workload_type
+            ),
+            workload_size=ServedModelInputWorkloadSize(
+                self.config.workload_size
+            ),
+        )
+
+        databricks_endpoint = (
+            self.databricks_client.serving_endpoints.create_and_wait(
+                name=self._generate_an_endpoint_name(),
+                config=EndpointCoreConfigInput(
+                    served_models=[served_model],
+                ),
+                tags=tags,
+            )
+        )
+        # Check if the endpoint URL is available after provisioning
+        if databricks_endpoint.endpoint_url:
+            logger.info(
+                f"Databricks inference endpoint successfully deployed and available. Endpoint URL: {databricks_endpoint.endpoint_url}"
+            )
+        else:
+            logger.error(
+                "Failed to start Databricks inference endpoint service: No URL available, please check the Databricks console for more details."
+            )
+
+    def check_status(self) -> Tuple[ServiceState, str]:
+        """Check the the current operational state of the Databricks deployment.
+
+        Returns:
+            The operational state of the Databricks deployment and a message
+            providing additional information about that state (e.g. a
+            description of the error, if one is encountered).
+        """
+        try:
+            status = self.databricks_endpoint.state or None
+            if (
+                status
+                and status.ready
+                and status.ready == EndpointStateReady.READY
+            ):
+                return (ServiceState.ACTIVE, "")
+            elif (
+                status
+                and status.config_update
+                and status.config_update
+                == EndpointStateConfigUpdate.UPDATE_FAILED
+            ):
+                return (
+                    ServiceState.ERROR,
+                    "Databricks Inference Endpoint deployment update failed",
+                )
+            elif (
+                status
+                and status.config_update
+                and status.config_update
+                == EndpointStateConfigUpdate.IN_PROGRESS
+            ):
+                return (ServiceState.PENDING_STARTUP, "")
+            return (ServiceState.PENDING_STARTUP, "")
+        except Exception as e:
+            return (
+                ServiceState.INACTIVE,
+                f"Databricks Inference Endpoint deployment is inactive or not found: {e}",
+            )
+
+    def deprovision(self, force: bool = False) -> None:
+        """Deprovision the remote Databricks deployment instance.
+
+        Args:
+            force: if True, the remote deployment instance will be
+                forcefully deprovisioned.
+        """
+        try:
+            self.databricks_client.serving_endpoints.delete(
+                name=self._generate_an_endpoint_name()
+            )
+        except Exception:
+            logger.error(
+                "Databricks Inference Endpoint is deleted or cannot be found."
+            )
+
+    def predict(
+        self, request: Union["NDArray[Any]", pd.DataFrame]
+    ) -> "NDArray[Any]":
+        """Make a prediction using the service.
+
+        Args:
+            request: The input data for the prediction.
+
+        Returns:
+            The prediction result.
+
+        Raises:
+            Exception: if the service is not running
+            ValueError: if the endpoint secret name is not provided.
+        """
+        if not self.is_running:
+            raise Exception(
+                "Databricks endpoint inference service is not running. "
+                "Please start the service before making predictions."
+            )
+        if self.prediction_url is not None:
+            if not self.config.endpoint_secret_name:
+                raise ValueError(
+                    "No endpoint secret name is provided for prediction."
+                )
+            databricks_token = Client().get_secret(
+                self.config.endpoint_secret_name
+            )
+            if not databricks_token.secret_values["token"]:
+                raise ValueError("No databricks token found.")
+            headers = {
+                "Authorization": f"Bearer {databricks_token.secret_values['token']}",
+                "Content-Type": "application/json",
+            }
+            if isinstance(request, pd.DataFrame):
+                response = requests.post(  # nosec
+                    self.prediction_url,
+                    json={"instances": request.to_dict("records")},
+                    headers=headers,
+                )
+            else:
+                response = requests.post(  # nosec
+                    self.prediction_url,
+                    json={"instances": request.tolist()},
+                    headers=headers,
+                )
+        else:
+            raise ValueError("No endpoint known for prediction.")
+        response.raise_for_status()
+
+        return np.array(response.json()["predictions"])
+
+    def get_logs(
+        self, follow: bool = False, tail: Optional[int] = None
+    ) -> Generator[str, bool, None]:
+        """Retrieve the service logs.
+
+        Args:
+            follow: if True, the logs will be streamed as they are written
+            tail: only retrieve the last NUM lines of log output.
+
+        Yields:
+            A generator that can be accessed to get the service logs.
+        """
+        logger.info(
+            "Databricks Endpoints provides access to the logs of your Endpoints through the UI in the `Logs` tab of your Endpoint"
+        )
+
+        def log_generator() -> Generator[str, bool, None]:
+            last_log_count = 0
+            while True:
+                logs = self.databricks_client.serving_endpoints.logs(
+                    name=self._generate_an_endpoint_name(),
+                    served_model_name=self.config.model_name,
+                )
+
+                log_lines = logs.logs.split("\n")
+
+                # Apply tail if specified and it's the first iteration
+                if tail is not None and last_log_count == 0:
+                    log_lines = log_lines[-tail:]
+
+                # Yield only new lines
+                for line in log_lines[last_log_count:]:
+                    yield line
+
+                last_log_count = len(log_lines)
+
+                if not follow:
+                    break
+
+                # Add a small delay to avoid excessive API calls
+                time.sleep(1)
+
+        yield from log_generator()
+
+    def _generate_an_endpoint_name(self) -> str:
+        """Generate a unique name for the Databricks Inference Endpoint.
+
+        Returns:
+            A unique name for the Databricks Inference Endpoint.
+        """
+        return (
+            f"{self.config.service_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}"
+        )

zenml/integrations/databricks/utils/__init__.py

@@ -0,0 +1,14 @@
+#  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.
+"""Utilities for Databricks integration."""

zenml/integrations/databricks/utils/databricks_utils.py

@@ -0,0 +1,87 @@
+#  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:
+#
+#       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 utilities."""
+
+import re
+from typing import Dict, List, Optional
+
+from databricks.sdk.service.compute import Library, PythonPyPiLibrary
+from databricks.sdk.service.jobs import PythonWheelTask, TaskDependency
+from databricks.sdk.service.jobs import Task as DatabricksTask
+
+from zenml import __version__
+
+
+def convert_step_to_task(
+    task_name: str,
+    command: str,
+    arguments: List[str],
+    libraries: Optional[List[str]] = None,
+    depends_on: Optional[List[str]] = None,
+    zenml_project_wheel: Optional[str] = None,
+    job_cluster_key: Optional[str] = None,
+) -> DatabricksTask:
+    """Convert a ZenML step to a Databricks task.
+
+    Args:
+        task_name: Name of the task.
+        command: Command to run.
+        arguments: Arguments to pass to the command.
+        libraries: List of libraries to install.
+        depends_on: List of tasks to depend on.
+        zenml_project_wheel: Path to the ZenML project wheel.
+        job_cluster_key: ID of the Databricks job_cluster_key.
+
+    Returns:
+        Databricks task.
+    """
+    db_libraries = []
+    if libraries:
+        for library in libraries:
+            db_libraries.append(Library(pypi=PythonPyPiLibrary(library)))
+    db_libraries.append(Library(whl=zenml_project_wheel))
+    db_libraries.append(
+        Library(pypi=PythonPyPiLibrary(f"zenml=={__version__}"))
+    )
+    return DatabricksTask(
+        task_key=task_name,
+        job_cluster_key=job_cluster_key,
+        libraries=db_libraries,
+        python_wheel_task=PythonWheelTask(
+            package_name="zenml",
+            entry_point=command,
+            parameters=arguments,
+        ),
+        depends_on=[TaskDependency(task) for task in depends_on]
+        if depends_on
+        else None,
+    )
+
+
+def sanitize_labels(labels: Dict[str, str]) -> None:
+    """Update the label values to be valid Kubernetes labels.
+
+    See:
+    https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
+
+    Args:
+        labels: the labels to sanitize.
+    """
+    for key, value in labels.items():
+        # Kubernetes labels must be alphanumeric, no longer than
+        # 63 characters, and must begin and end with an alphanumeric
+        # character ([a-z0-9A-Z])
+        labels[key] = re.sub(r"[^0-9a-zA-Z-_\.]+", "_", value)[:63].strip(
+            "-_."
+        )

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -363,7 +363,6 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 pipeline_func
             """
             step_name_to_dynamic_component: Dict[str, Any] = {}
-            node_selector_constraint: Optional[Tuple[str, str]] = None
 
             for step_name, step in deployment.step_configurations.items():
                 image = self.get_image(
@@ -410,23 +409,17 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                             "Volume mounts are set but not supported in "
                             "Vertex with Kubeflow Pipelines 2.x. Ignoring..."
                         )
-
-                    # apply pod settings
-                    if (
-                        GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
-                        in pod_settings.node_selectors.keys()
-                    ):
-                        node_selector_constraint = (
-                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
-                            pod_settings.node_selectors[
-                                GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
-                            ],
-                        )
-                    elif step_settings.node_selector_constraint:
-                        node_selector_constraint = (
-                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
-                            step_settings.node_selector_constraint[1],
-                        )
+                    for key in pod_settings.node_selectors:
+                        if (
+                            key
+                            != GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                        ):
+                            logger.warning(
+                                "Vertex only allows the %s node selector, "
+                                "ignoring the node selector %s.",
+                                GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                                key,
+                            )
 
                 step_name_to_dynamic_component[step_name] = dynamic_component
 
@@ -460,10 +453,33 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                         )
                         .after(*upstream_step_components)
                     )
+
+                    step_settings = cast(
+                        VertexOrchestratorSettings, self.get_settings(step)
+                    )
+                    pod_settings = step_settings.pod_settings
+
+                    node_selector_constraint: Optional[Tuple[str, str]] = None
+                    if pod_settings and (
+                        GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                        in pod_settings.node_selectors.keys()
+                    ):
+                        node_selector_constraint = (
+                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                            pod_settings.node_selectors[
+                                GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                            ],
+                        )
+                    elif step_settings.node_selector_constraint:
+                        node_selector_constraint = (
+                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                            step_settings.node_selector_constraint[1],
+                        )
+
                     self._configure_container_resources(
-                        task,
-                        step.config.resource_settings,
-                        node_selector_constraint,
+                        dynamic_component=task,
+                        resource_settings=step.config.resource_settings,
+                        node_selector_constraint=node_selector_constraint,
                     )
 
             return dynamic_pipeline
@@ -731,20 +747,20 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         )
 
         if node_selector_constraint:
-            (constraint_label, value) = node_selector_constraint
+            _, value = node_selector_constraint
             if gpu_limit is not None and gpu_limit > 0:
                 dynamic_component = (
                     dynamic_component.set_accelerator_type(value)
                     .set_accelerator_limit(gpu_limit)
                     .set_gpu_limit(gpu_limit)
                 )
-            elif (
-                constraint_label
-                == GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
-                and gpu_limit == 0
-            ):
+            else:
                 logger.warning(
-                    "GPU limit is set to 0 but a GPU type is specified. Ignoring GPU settings."
+                    "Accelerator type %s specified, but the GPU limit is not "
+                    "set or set to 0. The accelerator type will be ignored. "
+                    "To fix this warning, either remove the specified "
+                    "accelerator type or set the `gpu_count` using the "
+                    "ResourceSettings (https://docs.zenml.io/how-to/training-with-gpus#specify-resource-requirements-for-steps)."
                 )
 
         return dynamic_component

zenml/integrations/great_expectations/data_validators/ge_data_validator.py

@@ -284,14 +284,18 @@ class GreatExpectationsDataValidator(BaseDataValidator):
                         store_name=store_name,
                         store_config=store_config,
                     )
-                for site_name, site_config in zenml_context_config[
-                    "data_docs_sites"
-                ].items():
-                    self._context.config.data_docs_sites[site_name] = (
-                        site_config
-                    )
-
-            if self.config.configure_local_docs:
+                if self._context.config.data_docs_sites is not None:
+                    for site_name, site_config in zenml_context_config[
+                        "data_docs_sites"
+                    ].items():
+                        self._context.config.data_docs_sites[site_name] = (
+                            site_config
+                        )
+
+            if (
+                self.config.configure_local_docs
+                and self._context.config.data_docs_sites is not None
+            ):
                 client = Client()
                 artifact_store = client.active_stack.artifact_store
                 if artifact_store.flavor != "local":