content: #  Copyright (c) ZenML GmbH 2022. 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.
  """MLflow experiment tracker flavor."""
  
  from typing import TYPE_CHECKING, Any, Dict, Optional, Type
  
  from pydantic import Field, model_validator
  
  from zenml.config.base_settings import BaseSettings
  from zenml.experiment_trackers.base_experiment_tracker import (
      BaseExperimentTrackerConfig,
      BaseExperimentTrackerFlavor,
  )
  from zenml.integrations.mlflow import MLFLOW_MODEL_EXPERIMENT_TRACKER_FLAVOR
  from zenml.utils.secret_utils import SecretField
  
  if TYPE_CHECKING:
      from zenml.integrations.mlflow.experiment_trackers import (
          MLFlowExperimentTracker,
      )
  
  
  def is_remote_mlflow_tracking_uri(tracking_uri: str) -> bool:
      """Checks whether the given tracking uri is remote or not.
  
      Args:
          tracking_uri: The tracking uri to check.
  
      Returns:
          `True` if the tracking uri is remote, `False` otherwise.
      """
      return any(
          tracking_uri.startswith(prefix) for prefix in ["http://", "https://"]
      ) or is_databricks_tracking_uri(tracking_uri)
  
  
  def is_databricks_tracking_uri(tracking_uri: str) -> bool:
      """Checks whether the given tracking uri is a Databricks tracking uri.
  
      Args:
          tracking_uri: The tracking uri to check.
  
      Returns:
          `True` if the tracking uri is a Databricks tracking uri, `False`
          otherwise.
      """
      return tracking_uri == "databricks"
  
  
  class MLFlowExperimentTrackerSettings(BaseSettings):
      """Settings for the MLflow experiment tracker."""
  
      experiment_name: Optional[str] = Field(
          None,
          description="The MLflow experiment name to use for tracking runs.",
      )
      nested: bool = Field(
          False,
          description="If `True`, will create a nested sub-run for the step.",
      )
      tags: Dict[str, Any] = Field(
          default_factory=dict,
          description="Tags to attach to the MLflow run for categorization and filtering.",
      )
  
  
  class MLFlowExperimentTrackerConfig(
      BaseExperimentTrackerConfig, MLFlowExperimentTrackerSettings
  ):
      """Config for the MLflow experiment tracker."""
  
      tracking_uri: Optional[str] = Field(
          None,
          description="The URI of the MLflow tracking server. If no URI is set, "
          "your stack must contain a LocalArtifactStore and ZenML will point "
          "MLflow to a subdirectory of your artifact store instead.",
      )
      tracking_username: Optional[str] = SecretField(
          default=None,
          description="Username for authenticating with the MLflow tracking server. "
          "Required when using a remote tracking URI along with tracking_password.",
      )
      tracking_password: Optional[str] = SecretField(
          default=None,
          description="Password for authenticating with the MLflow tracking server. "
          "Required when using a remote tracking URI along with tracking_username.",
      )
      tracking_token: Optional[str] = SecretField(
          default=None,
          description="Token for authenticating with the MLflow tracking server. "
          "Alternative to username/password authentication for remote tracking URIs.",
      )
      tracking_insecure_tls: bool = Field(
          False,
          description="Skips verification of TLS connection to the MLflow tracking "
          "server if set to `True`. Use with caution in production environments.",
      )
      databricks_host: Optional[str] = Field(
          None,
          description="The host of the Databricks workspace with the MLflow managed "
          "server to connect to. Required when tracking_uri is set to 'databricks'.",
      )
      enable_unity_catalog: bool = Field(
          False,
          description="If `True`, will enable the Databricks Unity Catalog for "
          "logging and registering models.",
      )
  
      @model_validator(mode="after")
      def _ensure_authentication_if_necessary(
          self,
      ) -> "MLFlowExperimentTrackerConfig":
          """Ensures that credentials or a token for authentication exist.
  
          We make this check when running MLflow tracking with a remote backend.
  
          Returns:
              The validated values.
  
          Raises:
               ValueError: If neither credentials nor a token are provided.
          """
          if self.tracking_uri:
              if is_databricks_tracking_uri(self.tracking_uri):
                  # If the tracking uri is "databricks", then we need the
                  # databricks host to be set.
                  if not self.databricks_host:
                      raise ValueError(
                          "MLflow experiment tracking with a Databricks MLflow "
                          "managed tracking server requires the "
                          "`databricks_host` to be set in your stack component. "
                          "To update your component, run "
                          "`zenml experiment-tracker update "
                          "<NAME> --databricks_host=DATABRICKS_HOST` "
                          "and specify the hostname of your Databricks workspace."
                      )
  
              if is_remote_mlflow_tracking_uri(self.tracking_uri):
                  # we need either username + password or a token to authenticate
                  # to the remote backend
                  basic_auth = self.tracking_username and self.tracking_password
  
                  if not (basic_auth or self.tracking_token):
                      raise ValueError(
                          f"MLflow experiment tracking with a remote backend "
                          f"{self.tracking_uri} is only possible when specifying "
                          f"either username and password or an authentication "
                          f"token in your stack component. To update your "
                          f"component, run the following command: "
                          f"`zenml experiment-tracker update "
                          f"<NAME> --tracking_username=MY_USERNAME "
                          f"--tracking_password=MY_PASSWORD "
                          f"--tracking_token=MY_TOKEN` and specify either your "
                          f"username and password or token."
                      )
  
          return self
  
      @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.
          """
          if not self.tracking_uri or not is_remote_mlflow_tracking_uri(
              self.tracking_uri
          ):
              return True
          return False
  
  
  class MLFlowExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
      """Class for the `MLFlowExperimentTrackerFlavor`."""
  
      @property
      def name(self) -> str:
          """Name of the flavor.
  
          Returns:
              The name of the flavor.
          """
          return MLFLOW_MODEL_EXPERIMENT_TRACKER_FLAVOR
  
      @property
      def display_name(self) -> str:
          """Display name of the flavor.
  
          Returns:
              The display name of the flavor.
          """
          return "MLflow"
  
      @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/experiment_tracker/mlflow.png"
  
      @property
      def config_class(self) -> Type[MLFlowExperimentTrackerConfig]:
          """Returns `MLFlowExperimentTrackerConfig` config class.
  
          Returns:
                  The config class.
          """
          return MLFlowExperimentTrackerConfig
  
      @property
      def implementation_class(self) -> Type["MLFlowExperimentTracker"]:
          """Implementation class for this flavor.
  
          Returns:
              The implementation class.
          """
          from zenml.integrations.mlflow.experiment_trackers import (
              MLFlowExperimentTracker,
          )
  
          return MLFlowExperimentTracker
  