content: #  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.enums import ServiceState
  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.models.v2.misc.service import ServiceType
  from zenml.services import ServiceStatus
  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]}"
          )
  