genesis-flow 1.0.0__py3-none-any.whl

mlflow/legacy_databricks_cli/configure/provider.py

@@ -0,0 +1,482 @@
+# This module is copied from legacy databricks CLI python library
+# module `databricks_cli.configure.provider`,
+# but with some modification to make `EnvironmentVariableConfigProvider` supporting
+# 'DATABRICKS_CLIENT_ID' and 'DATABRICKS_CLIENT_SECRET' environmental variables,
+# and make ProfileConfigProvider supporting 'databricks-cli' authentication way,
+# 'databricks-cli' authentication way is for supporting U2M authentication.
+#
+# This is the original legacy databricks CLI python library provider module code:
+# https://github.com/databricks/databricks-cli/blob/0.18.0/databricks_cli/configure/provider.py
+#
+# The latest Databricks Runtime does not contain legacy databricks CLI
+# but MLflow still depends on it.
+
+import logging
+import os
+import sys
+import time
+from abc import ABCMeta, abstractmethod
+from configparser import ConfigParser
+from os.path import expanduser, join
+
+_logger = logging.getLogger(__name__)
+
+_home = expanduser("~")
+CONFIG_FILE_ENV_VAR = "DATABRICKS_CONFIG_FILE"
+HOST = "host"
+USERNAME = "username"
+PASSWORD = "password"
+TOKEN = "token"
+REFRESH_TOKEN = "refresh_token"
+INSECURE = "insecure"
+JOBS_API_VERSION = "jobs-api-version"
+DEFAULT_SECTION = "DEFAULT"
+CLIENT_ID = "client_id"
+CLIENT_SECRET = "client_secret"
+AUTH_TYPE = "auth_type"
+
+# User-provided override for the DatabricksConfigProvider
+_config_provider = None
+
+
+class InvalidConfigurationError(RuntimeError):
+    @staticmethod
+    def for_profile(profile):
+        if profile is None:
+            return InvalidConfigurationError(
+                "You haven't configured the CLI yet! "
+                f"Please configure by entering `{sys.argv[0]} configure`"
+            )
+        return InvalidConfigurationError(
+            f"You haven't configured the CLI yet for the profile {profile}! "
+            "Please configure by entering "
+            f"`{sys.argv[0]} configure --profile {profile}`"
+        )
+
+
+def _get_path():
+    return os.environ.get(CONFIG_FILE_ENV_VAR, join(_home, ".databrickscfg"))
+
+
+def _fetch_from_fs():
+    raw_config = ConfigParser()
+    raw_config.read(_get_path())
+    return raw_config
+
+
+def _create_section_if_absent(raw_config, profile):
+    if not raw_config.has_section(profile) and profile != DEFAULT_SECTION:
+        raw_config.add_section(profile)
+
+
+def _get_option_if_exists(raw_config, profile, option):
+    if profile == DEFAULT_SECTION:
+        # We must handle the DEFAULT_SECTION differently since it is not in the _sections property
+        # of raw config.
+        return raw_config.get(profile, option) if raw_config.has_option(profile, option) else None
+    # Check if option is defined in the profile.
+    elif option not in raw_config._sections.get(profile, {}).keys():
+        return None
+    return raw_config.get(profile, option)
+
+
+def _set_option(raw_config, profile, option, value):
+    if value:
+        raw_config.set(profile, option, value)
+    else:
+        raw_config.remove_option(profile, option)
+
+
+def _overwrite_config(raw_config):
+    config_path = _get_path()
+    # Create config file with owner only rw permissions
+    if not os.path.exists(config_path):
+        file_descriptor = os.open(config_path, os.O_CREAT | os.O_RDWR, 0o600)
+        os.close(file_descriptor)
+
+    # Change file permissions to owner only rw if that's not the case
+    if not os.stat(config_path).st_mode == 0o100600:
+        os.chmod(config_path, 0o600)
+
+    with open(config_path, "w") as cfg:
+        raw_config.write(cfg)
+
+
+def update_and_persist_config(profile, databricks_config):
+    """
+    Takes a DatabricksConfig and adds the in memory contents to the persisted version of the
+    config. This will overwrite any other config that was persisted to the file system under the
+    same profile.
+
+    Args:
+        profile: str
+        databricks_config: DatabricksConfig
+    """
+    profile = profile if profile else DEFAULT_SECTION
+    raw_config = _fetch_from_fs()
+    _create_section_if_absent(raw_config, profile)
+    _set_option(raw_config, profile, HOST, databricks_config.host)
+    _set_option(raw_config, profile, USERNAME, databricks_config.username)
+    _set_option(raw_config, profile, PASSWORD, databricks_config.password)
+    _set_option(raw_config, profile, TOKEN, databricks_config.token)
+    _set_option(raw_config, profile, REFRESH_TOKEN, databricks_config.refresh_token)
+    _set_option(raw_config, profile, INSECURE, databricks_config.insecure)
+    _set_option(raw_config, profile, JOBS_API_VERSION, databricks_config.jobs_api_version)
+    _overwrite_config(raw_config)
+
+
+def get_config():
+    """
+    Returns a DatabricksConfig containing the hostname and authentication used to talk to
+    the Databricks API. By default, we leverage the DefaultConfigProvider to get
+    this config, but this behavior may be overridden by calling 'set_config_provider'
+
+    If no DatabricksConfig can be found, an InvalidConfigurationError will be raised.
+    """
+    if _config_provider:
+        config = _config_provider.get_config()
+        if config:
+            return config
+        raise InvalidConfigurationError(
+            f"Custom provider returned no DatabricksConfig: {_config_provider}"
+        )
+
+    config = DefaultConfigProvider().get_config()
+    if config:
+        return config
+    raise InvalidConfigurationError.for_profile(None)
+
+
+def get_config_for_profile(profile):
+    """
+    [Deprecated] Reads from the filesystem and gets a DatabricksConfig for the
+    specified profile. If it does not exist, then return a DatabricksConfig with fields set
+    to None.
+
+    Internal callers should prefer get_config() to use user-specified overrides, and
+    to return appropriate error messages as opposited to invalid configurations.
+
+    If you want to read from a specific profile, please instead use
+    'ProfileConfigProvider(profile).get_config()'.
+
+    This method is maintained for backwards-compatibility. It may be removed in future versions.
+
+    Returns:
+        DatabricksConfig
+    """
+    profile = profile if profile else DEFAULT_SECTION
+    config = EnvironmentVariableConfigProvider().get_config()
+    if config and config.is_valid:
+        return config
+
+    config = ProfileConfigProvider(profile).get_config()
+    if config:
+        return config
+    return DatabricksConfig.empty()
+
+
+def set_config_provider(provider):
+    """
+    Sets a DatabricksConfigProvider that will be used for all future calls to get_config(),
+    used by the Databricks CLI code to discover the user's credentials.
+    """
+    global _config_provider
+    if provider and not isinstance(provider, DatabricksConfigProvider):
+        raise Exception(f"Must be instance of DatabricksConfigProvider: {_config_provider}")
+    _config_provider = provider
+
+
+def get_config_provider():
+    """
+    Returns the current DatabricksConfigProvider.
+    If None, the DefaultConfigProvider will be used.
+    """
+    return _config_provider
+
+
+class DatabricksConfigProvider:
+    """
+    Responsible for providing hostname and authentication information to make
+    API requests against the Databricks REST API.
+    This method should generally return None if it cannot provide credentials, in order
+    to facilitate chanining of providers.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def get_config(self):
+        pass
+
+
+class DefaultConfigProvider(DatabricksConfigProvider):
+    """Look for credentials in a chain of default locations."""
+
+    def __init__(self):
+        # The order of providers here will be used to determine
+        # the precedence order for the config provider used in `get_config`
+        self._providers = (
+            SparkTaskContextConfigProvider(),
+            EnvironmentVariableConfigProvider(),
+            ProfileConfigProvider(),
+            DatabricksModelServingConfigProvider(),
+        )
+
+    def get_config(self):
+        for provider in self._providers:
+            config = provider.get_config()
+            if config is not None and config.is_valid:
+                return config
+        return None
+
+
+class SparkTaskContextConfigProvider(DatabricksConfigProvider):
+    """Loads credentials from Spark TaskContext if running in a Spark Executor."""
+
+    @staticmethod
+    def _get_spark_task_context_or_none():
+        try:
+            from pyspark import TaskContext
+
+            return TaskContext.get()
+        except ImportError:
+            return None
+
+    @staticmethod
+    def set_insecure(x):
+        from pyspark import SparkContext
+
+        new_val = "True" if x else None
+        SparkContext._active_spark_context.setLocalProperty("spark.databricks.ignoreTls", new_val)
+
+    def get_config(self):
+        context = self._get_spark_task_context_or_none()
+        if context is not None:
+            host = context.getLocalProperty("spark.databricks.api.url")
+            token = context.getLocalProperty("spark.databricks.token")
+            insecure = context.getLocalProperty("spark.databricks.ignoreTls")
+            config = DatabricksConfig.from_token(
+                host=host, token=token, refresh_token=None, insecure=insecure, jobs_api_version=None
+            )
+            if config.is_valid:
+                return config
+        return None
+
+
+class EnvironmentVariableConfigProvider(DatabricksConfigProvider):
+    """Loads from system environment variables."""
+
+    def get_config(self):
+        host = os.environ.get("DATABRICKS_HOST")
+        username = os.environ.get("DATABRICKS_USERNAME")
+        password = os.environ.get("DATABRICKS_PASSWORD")
+        token = os.environ.get("DATABRICKS_TOKEN")
+        refresh_token = os.environ.get("DATABRICKS_REFRESH_TOKEN")
+        insecure = os.environ.get("DATABRICKS_INSECURE")
+        jobs_api_version = os.environ.get("DATABRICKS_JOBS_API_VERSION")
+        client_id = os.environ.get("DATABRICKS_CLIENT_ID")
+        client_secret = os.environ.get("DATABRICKS_CLIENT_SECRET")
+
+        config = DatabricksConfig(
+            host,
+            username,
+            password,
+            token,
+            refresh_token,
+            insecure,
+            jobs_api_version,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
+        if config.is_valid:
+            return config
+        return None
+
+
+class ProfileConfigProvider(DatabricksConfigProvider):
+    """Loads from the databrickscfg file."""
+
+    def __init__(self, profile=None):
+        self.profile = profile or DEFAULT_SECTION
+
+    def get_config(self):
+        raw_config = _fetch_from_fs()
+        host = _get_option_if_exists(raw_config, self.profile, HOST)
+        username = _get_option_if_exists(raw_config, self.profile, USERNAME)
+        password = _get_option_if_exists(raw_config, self.profile, PASSWORD)
+        token = _get_option_if_exists(raw_config, self.profile, TOKEN)
+        refresh_token = _get_option_if_exists(raw_config, self.profile, REFRESH_TOKEN)
+        insecure = _get_option_if_exists(raw_config, self.profile, INSECURE)
+        jobs_api_version = _get_option_if_exists(raw_config, self.profile, JOBS_API_VERSION)
+        client_id = _get_option_if_exists(raw_config, self.profile, CLIENT_ID)
+        client_secret = _get_option_if_exists(raw_config, self.profile, CLIENT_SECRET)
+        auth_type = _get_option_if_exists(raw_config, self.profile, AUTH_TYPE)
+        config = DatabricksConfig(
+            host,
+            username,
+            password,
+            token,
+            refresh_token,
+            insecure,
+            jobs_api_version,
+            client_id=client_id,
+            client_secret=client_secret,
+            auth_type=auth_type,
+        )
+        if config.is_valid:
+            return config
+        return None
+
+
+class DatabricksModelServingConfigProvider(DatabricksConfigProvider):
+    """Loads from OAuth credentials in the Databricks Model Serving environment."""
+
+    def get_config(self):
+        from mlflow.utils.databricks_utils import should_fetch_model_serving_environment_oauth
+
+        try:
+            if should_fetch_model_serving_environment_oauth():
+                config = DatabricksModelServingConfigProvider._get_databricks_model_serving_config()
+                if config.is_valid:
+                    return config
+            else:
+                return None
+        except Exception as e:
+            _logger.warning("Unexpected error resolving Databricks Model Serving config: %s", e)
+
+    @staticmethod
+    def _get_databricks_model_serving_config():
+        from mlflow.utils.databricks_utils import get_model_dependency_oauth_token
+
+        # Since we do not record OAuth expiration time in OAuth file, perform periodic refresh
+        # of OAuth environment variable cache here. As currently configured (02/24) OAuth token
+        # in model serving environment is guaranteed to have at least 30 min remaining on TTL
+        # at any point in time but refresh at higher rate of every 5 min here to be safe
+        # and conform with refresh logic for Brickstore tables.
+        OAUTH_CACHE_REFRESH_DURATION_SEC = 5 * 60
+        OAUTH_CACHE_ENV_VAR = "DB_DEPENDENCY_OAUTH_CACHE"
+        OAUTH_CACHE_EXPIRATION_ENV_VAR = "DB_DEPENDENCY_OAUTH_CACHE_EXPIRY_TS"
+        MODEL_SERVING_HOST_ENV_VAR = "DATABRICKS_MODEL_SERVING_HOST_URL"
+        DB_MODEL_SERVING_HOST_ENV_VAR = "DB_MODEL_SERVING_HOST_URL"
+
+        # read from DB_MODEL_SERVING_HOST_ENV_VAR if available otherwise MODEL_SERVING_HOST_ENV_VAR
+        host = os.environ.get(DB_MODEL_SERVING_HOST_ENV_VAR) or os.environ.get(
+            MODEL_SERVING_HOST_ENV_VAR
+        )
+
+        # check if dependency is cached in env var before reading from file
+        oauth_token = ""
+        if (
+            OAUTH_CACHE_ENV_VAR in os.environ
+            and OAUTH_CACHE_EXPIRATION_ENV_VAR in os.environ
+            and float(os.environ[OAUTH_CACHE_EXPIRATION_ENV_VAR]) > time.time()
+        ):
+            oauth_token = os.environ[OAUTH_CACHE_ENV_VAR]
+        else:
+            oauth_token = get_model_dependency_oauth_token()
+            os.environ[OAUTH_CACHE_ENV_VAR] = oauth_token
+            os.environ[OAUTH_CACHE_EXPIRATION_ENV_VAR] = str(
+                time.time() + OAUTH_CACHE_REFRESH_DURATION_SEC
+            )
+
+        return DatabricksConfig(
+            host=host,
+            token=oauth_token,
+            username=None,
+            password=None,
+            refresh_token=None,
+            insecure=None,
+            jobs_api_version=None,
+        )
+
+
+class DatabricksConfig:
+    def __init__(
+        self,
+        host,
+        username,
+        password,
+        token,
+        refresh_token=None,
+        insecure=None,
+        jobs_api_version=None,
+        client_id=None,
+        client_secret=None,
+        auth_type=None,
+    ):
+        self.host = host
+        self.username = username
+        self.password = password
+        self.token = token
+        self.refresh_token = refresh_token
+        self.insecure = insecure
+        self.jobs_api_version = jobs_api_version
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.auth_type = auth_type
+
+    @classmethod
+    def from_token(cls, host, token, refresh_token=None, insecure=None, jobs_api_version=None):
+        return DatabricksConfig(
+            host=host,
+            username=None,
+            password=None,
+            token=token,
+            refresh_token=refresh_token,
+            insecure=insecure,
+            jobs_api_version=jobs_api_version,
+        )
+
+    @classmethod
+    def from_password(cls, host, username, password, insecure=None, jobs_api_version=None):
+        return DatabricksConfig(
+            host=host,
+            username=username,
+            password=password,
+            token=None,
+            refresh_token=None,
+            insecure=insecure,
+            jobs_api_version=jobs_api_version,
+        )
+
+    @classmethod
+    def empty(cls):
+        return DatabricksConfig(
+            host=None,
+            username=None,
+            password=None,
+            token=None,
+            refresh_token=None,
+            insecure=None,
+            jobs_api_version=None,
+        )
+
+    @property
+    def is_valid_with_token(self):
+        return self.host is not None and self.token is not None
+
+    @property
+    def is_valid_with_password(self):
+        return self.host is not None and self.username is not None and self.password is not None
+
+    @property
+    def is_valid_with_client_id_secret(self):
+        return self.host and self.client_id and self.client_secret
+
+    @property
+    def is_databricks_cli_auth_type(self):
+        return self.auth_type == "databricks-cli"
+
+    @property
+    def is_azure_cli_auth_type(self):
+        return self.auth_type == "azure-cli"
+
+    @property
+    def is_valid(self):
+        return (
+            self.is_valid_with_token
+            or self.is_valid_with_password
+            or self.is_valid_with_client_id_secret
+            or self.is_databricks_cli_auth_type
+            or self.is_azure_cli_auth_type
+        )

mlflow/litellm/__init__.py

@@ -0,0 +1,175 @@
+import importlib.metadata
+import logging
+from contextlib import contextmanager
+from threading import Thread
+
+from packaging.version import Version
+
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import (
+    autologging_integration,
+    safe_patch,
+)
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+FLAVOR_NAME = "litellm"
+
+_logger = logging.getLogger(__name__)
+
+
+@experimental(version="2.19.0")
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from LiteLLM to MLflow. Currently, MLflow
+    only supports autologging for tracing.
+
+    Args:
+        log_traces: If ``True``, traces are logged for LiteLLM calls. If ``False``,
+            no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the LiteLLM autologging integration. If ``False``,
+            enables the LiteLLM autologging integration.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during LiteLLM
+            autologging. If ``False``, show all events and warnings.
+    """
+    import litellm
+
+    # This needs to be called before doing any safe-patching (otherwise safe-patch will be no-op).
+    # TODO: since this implementation is inconsistent, explore a universal way to solve the issue.
+    _autolog(log_traces=log_traces, disable=disable, silent=silent)
+
+    try:
+        from litellm.integrations.mlflow import MlflowLogger
+    except ImportError:
+        _logger.warning(
+            "MLflow LiteLLM integration is not supported for the installed LiteLLM version. "
+            "Please upgrade to a newer version to enable MLflow LiteLLM autologging."
+        )
+        return
+
+    if log_traces and not disable:
+        litellm.success_callback = _append_mlflow_callbacks(litellm.success_callback)
+        litellm.failure_callback = _append_mlflow_callbacks(litellm.failure_callback)
+
+        # Workaround for https://github.com/BerriAI/litellm/issues/8013
+        # TODO: Add upper bound version check when the issue is fixed.
+        if Version(importlib.metadata.version("litellm")) >= Version("1.59.4"):
+            litellm.failure_callback = [
+                cb if cb != "mlflow" else MlflowLogger() for cb in litellm.failure_callback
+            ]
+
+        if is_in_databricks_runtime():
+            # Patch main APIs e.g. completion to inject custom handling for threading.
+            # By default, those API will start a new thread when calling log_success_event()
+            # handler of the logging callbacks and never wait for it to finish. This is
+            # problematic in Databricks notebook, because the inline trace UI display
+            # assumes that the trace is generated synchronously. If the trace is generated
+            # asynchronously, it will be displayed in different later cells.
+            # To workaround this issue, we monkey-patch these APIs to wait for the logging
+            # threads to finish before returning the result.
+            # This is not required for OSS environment where we don't show inline trace UI.
+            for func in [
+                "completion",
+                "embedding",
+                "text_completion",
+                "image_generation",
+                "transcription",
+                "speech",
+            ]:
+                _patch_threading_in_function(litellm, func)
+
+            # For streaming case, we need to patch the iterator because traces are generated
+            # when consuming the generator, not when calling the main APIs.
+            _patch_threading_in_function(litellm.utils.CustomStreamWrapper, "__next__")
+
+            # NB: We don't need to patch async function because Databricks notebook waits
+            # for the async task to finish before finishing the cell.
+    else:
+        litellm.success_callback = _remove_mlflow_callbacks(litellm.success_callback)
+        litellm.failure_callback = _remove_mlflow_callbacks(litellm.failure_callback)
+        # Callback also needs to be removed from 'callbacks' as litellm adds
+        # success/failure callbacks to there as well.
+        litellm.callbacks = _remove_mlflow_callbacks(litellm.callbacks)
+
+
+# This is required by mlflow.autolog()
+autolog.integration_name = FLAVOR_NAME
+
+
+# NB: The @autologging_integration annotation must be applied here, and the callback injection
+# needs to happen outside the annotated function. This is because the annotated function is NOT
+# executed when disable=True is passed. This prevents us from removing our callback and patching
+# when autologging is turned off.
+@autologging_integration(FLAVOR_NAME)
+def _autolog(
+    log_traces: bool,
+    disable: bool = False,
+    silent: bool = False,
+):
+    pass
+
+
+def _patch_threading_in_function(target, function_name: str):
+    """
+    Apply the threading patch to a synchronous function.
+
+    We capture the threads started by the function using the _patch_thread_start context manager,
+    then join them to ensure they are finished before the notebook cell finishes executing.
+    """
+
+    def _patch_fn(original, *args, **kwargs):
+        with _patch_thread_start() as logging_threads:
+            result = original(*args, **kwargs)
+        for thread in logging_threads:
+            thread.join()
+        return result
+
+    safe_patch(FLAVOR_NAME, target, function_name, _patch_fn)
+
+
+@contextmanager
+def _patch_thread_start():
+    """
+    A context manager to collect threads started for logging handlers.
+    This is done by monkey-patching the start() method of threading.Thread.
+    Note that failure handlers are executed synchronously, so we don't need to patch them.
+    """
+    original = Thread.start
+    logging_threads = []
+
+    def patched_thread(self, *args, **kwargs):
+        target = getattr(self, "_target", None)
+        # success_handler is for normal request, and run_success_... is for streaming
+        # - https://github.com/BerriAI/litellm/blob/4f8a3fd4cfc20cf43b38379928b41c2691c85d36/litellm/utils.py#L946
+        # - https://github.com/BerriAI/litellm/blob/4f8a3fd4cfc20cf43b38379928b41c2691c85d36/litellm/utils.py#L7526
+        if target and target.__name__ in [
+            "success_handler",
+            "run_success_logging_and_cache_storage",
+        ]:
+            logging_threads.append(self)
+        return original(self, *args, **kwargs)
+
+    Thread.start = patched_thread
+    try:
+        yield logging_threads
+    finally:
+        Thread.start = original
+
+
+def _append_mlflow_callbacks(callbacks):
+    from litellm.integrations.mlflow import MlflowLogger
+
+    # MLflow callback can be stored as a string or the actual logger object
+    if not any(cb == "mlflow" or isinstance(cb, MlflowLogger) for cb in callbacks):
+        return callbacks + ["mlflow"]
+
+    return callbacks
+
+
+def _remove_mlflow_callbacks(callbacks):
+    from litellm.integrations.mlflow import MlflowLogger
+
+    return [cb for cb in callbacks if not (cb == "mlflow" or isinstance(cb, MlflowLogger))]

mlflow/llama_index/__init__.py

@@ -0,0 +1,22 @@
+from mlflow.llama_index.autolog import autolog
+from mlflow.llama_index.constant import FLAVOR_NAME
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+__all__ = ["autolog", "FLAVOR_NAME"]
+
+# Import model logging APIs only if mlflow skinny or full package is installed,
+# i.e., skip if only mlflow-tracing package is installed.
+if not IS_TRACING_SDK_ONLY:
+    from mlflow.llama_index.model import (
+        _load_pyfunc,
+        load_model,
+        log_model,
+        save_model,
+    )
+
+    __all__ += [
+        "load_model",
+        "log_model",
+        "save_model",
+        "_load_pyfunc",
+    ]