zenml-nightly 0.61.0.dev20240714__py3-none-any.whl → 0.62.0.dev20240719__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/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":

zenml/integrations/huggingface/materializers/huggingface_datasets_materializer.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  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.
@@ -16,12 +16,21 @@
 import os
 from collections import defaultdict
 from tempfile import TemporaryDirectory, mkdtemp
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 from datasets import Dataset, load_from_disk
 from datasets.dataset_dict import DatasetDict
 
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.pandas_materializer import PandasMaterializer
@@ -33,6 +42,31 @@ if TYPE_CHECKING:
 DEFAULT_DATASET_DIR = "hf_datasets"
 
 
+def extract_repo_name(checksum_str: str) -> Optional[str]:
+    """Extracts the repo name from the checksum string.
+
+    An example of a checksum_str is:
+    "hf://datasets/nyu-mll/glue@bcdcba79d07bc864c1c254ccfcedcce55bcc9a8c/mrpc/train-00000-of-00001.parquet"
+    and the expected output is "nyu-mll/glue".
+
+    Args:
+        checksum_str: The checksum_str to extract the repo name from.
+
+    Returns:
+        str: The extracted repo name.
+    """
+    dataset = None
+    try:
+        parts = checksum_str.split("/")
+        if len(parts) >= 4:
+            # Case: nyu-mll/glue
+            dataset = f"{parts[3]}/{parts[4].split('@')[0]}"
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+    return dataset
+
+
 class HFDatasetMaterializer(BaseMaterializer):
     """Materializer to read data to and from huggingface datasets."""
 
@@ -103,3 +137,54 @@ class HFDatasetMaterializer(BaseMaterializer):
                     metadata[key][dataset_name] = value
             return dict(metadata)
         raise ValueError(f"Unsupported type {type(ds)}")
+
+    def save_visualizations(
+        self, ds: Union[Dataset, DatasetDict]
+    ) -> Dict[str, VisualizationType]:
+        """Save visualizations for the dataset.
+
+        Args:
+            ds: The Dataset or DatasetDict to visualize.
+
+        Returns:
+            A dictionary mapping visualization paths to their types.
+
+        Raises:
+            ValueError: If the given object is not a `Dataset` or `DatasetDict`.
+        """
+        visualizations = {}
+
+        if isinstance(ds, Dataset):
+            datasets = {"default": ds}
+        elif isinstance(ds, DatasetDict):
+            datasets = ds
+        else:
+            raise ValueError(f"Unsupported type {type(ds)}")
+
+        for name, dataset in datasets.items():
+            # Generate a unique identifier for the dataset
+            if dataset.info.download_checksums:
+                dataset_id = extract_repo_name(
+                    [x for x in dataset.info.download_checksums.keys()][0]
+                )
+                if dataset_id:
+                    # Create the iframe HTML
+                    html = f"""
+                    <iframe
+                    src="https://huggingface.co/datasets/{dataset_id}/embed/viewer"
+                    frameborder="0"
+                    width="100%"
+                    height="560px"
+                    ></iframe>
+                    """
+
+                    # Save the HTML to a file
+                    visualization_path = os.path.join(
+                        self.uri, f"{name}_viewer.html"
+                    )
+                    with fileio.open(visualization_path, "w") as f:
+                        f.write(html)
+
+                    visualizations[visualization_path] = VisualizationType.HTML
+
+        return visualizations

zenml/integrations/huggingface/steps/accelerate_runner.py

@@ -111,15 +111,9 @@ def run_with_accelerate(
                     if isinstance(v, bool):
                         if v:
                             commands.append(f"--{k}")
-                    elif isinstance(v, str):
-                        commands += [f"--{k}", '"{v}"']
                     elif type(v) in (list, tuple, set):
                         for each in v:
-                            commands.append(f"--{k}")
-                            if isinstance(each, str):
-                                commands.append(f'"{each}"')
-                            else:
-                                commands.append(f"{each}")
+                            commands += [f"--{k}", f"{each}"]
                     else:
                         commands += [f"--{k}", f"{v}"]
 

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -139,10 +139,17 @@ def build_pod_manifest(
         ],
         security_context=security_context,
     )
+    image_pull_secrets = []
+    if pod_settings:
+        image_pull_secrets = [
+            k8s_client.V1LocalObjectReference(name=name)
+            for name in pod_settings.image_pull_secrets
+        ]
 
     pod_spec = k8s_client.V1PodSpec(
         containers=[container_spec],
         restart_policy="Never",
+        image_pull_secrets=image_pull_secrets,
     )
 
     if service_account_name is not None: