star

venv/lib/python3.9/site-packages/mlflow/utils/__init__.py

@@ -0,0 +1,271 @@
+import inspect
+import logging
+import socket
+import subprocess
+from contextlib import closing
+from itertools import islice
+from sys import version_info
+
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER  # noqa: F401
+
+PYTHON_VERSION = f"{version_info.major}.{version_info.minor}.{version_info.micro}"
+
+
+_logger = logging.getLogger(__name__)
+
+
+def get_major_minor_py_version(py_version):
+    return ".".join(py_version.split(".")[:2])
+
+
+def reraise(tp, value, tb=None):
+    # Taken from: https://github.com/benjaminp/six/blob/1.15.0/six.py#L694-L700
+    try:
+        if value is None:
+            value = tp()
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+    finally:
+        value = None
+        tb = None
+
+
+def chunk_list(l, chunk_size):
+    for i in range(0, len(l), chunk_size):
+        yield l[i : i + chunk_size]
+
+
+def _chunk_dict(d, chunk_size):
+    """
+    Splits a dictionary into chunks of the specified size.
+    Taken from: https://stackoverflow.com/a/22878842
+    """
+    it = iter(d)
+    for _ in range(0, len(d), chunk_size):
+        yield {k: d[k] for k in islice(it, chunk_size)}
+
+
+def _truncate_and_ellipsize(value, max_length):
+    """
+    Truncates the string representation of the specified value to the specified
+    maximum length, if necessary. The end of the string is ellipsized if truncation occurs
+    """
+    value = str(value)
+    if len(value) > max_length:
+        return value[: (max_length - 3)] + "..."
+    else:
+        return value
+
+
+def _truncate_dict(d, max_key_length=None, max_value_length=None):
+    """
+    Truncates keys and/or values in a dictionary to the specified maximum length.
+    Truncated items will be converted to strings and ellipsized.
+    """
+    key_is_none = max_key_length is None
+    val_is_none = max_value_length is None
+
+    if key_is_none and val_is_none:
+        raise ValueError("Must specify at least either `max_key_length` or `max_value_length`")
+
+    truncated = {}
+    for k, v in d.items():
+        should_truncate_key = (not key_is_none) and (len(str(k)) > max_key_length)
+        should_truncate_val = (not val_is_none) and (len(str(v)) > max_value_length)
+
+        new_k = _truncate_and_ellipsize(k, max_key_length) if should_truncate_key else k
+        if should_truncate_key:
+            # Use the truncated key for warning logs to avoid noisy printing to stdout
+            msg = f"Truncated the key `{new_k}`"
+            _logger.warning(msg)
+
+        new_v = _truncate_and_ellipsize(v, max_value_length) if should_truncate_val else v
+        if should_truncate_val:
+            # Use the truncated key and value for warning logs to avoid noisy printing to stdout
+            msg = f"Truncated the value of the key `{new_k}`. Truncated value: `{new_v}`"
+            _logger.warning(msg)
+
+        truncated[new_k] = new_v
+
+    return truncated
+
+
+def merge_dicts(dict_a, dict_b, raise_on_duplicates=True):
+    """This function takes two dictionaries and returns one singular merged dictionary.
+
+    Args:
+        dict_a: The first dictionary.
+        dict_b: The second dictionary.
+        raise_on_duplicates: If True, the function raises ValueError if there are duplicate keys.
+            Otherwise, duplicate keys in `dict_b` will override the ones in `dict_a`.
+
+    Returns:
+        A merged dictionary.
+
+    """
+    duplicate_keys = dict_a.keys() & dict_b.keys()
+    if raise_on_duplicates and len(duplicate_keys) > 0:
+        raise ValueError(f"The two merging dictionaries contains duplicate keys: {duplicate_keys}.")
+    return {**dict_a, **dict_b}
+
+
+def _get_fully_qualified_class_name(obj):
+    """
+    Obtains the fully qualified class name of the given object.
+    """
+    return obj.__class__.__module__ + "." + obj.__class__.__name__
+
+
+def _inspect_original_var_name(var, fallback_name):
+    """
+    Inspect variable name, will search above frames and fetch the same instance variable name
+    in the most outer frame.
+    If inspect failed, return fallback_name
+    """
+    if var is None:
+        return fallback_name
+    try:
+        original_var_name = fallback_name
+
+        frame = inspect.currentframe().f_back
+        while frame is not None:
+            arg_info = inspect.getargvalues(frame)
+
+            fixed_args = [arg_info.locals[arg_name] for arg_name in arg_info.args]
+            varlen_args = list(arg_info.locals[arg_info.varargs]) if arg_info.varargs else []
+            keyword_args = (
+                list(arg_info.locals[arg_info.keywords].values()) if arg_info.keywords else []
+            )
+
+            all_args = fixed_args + varlen_args + keyword_args
+
+            # check whether `var` is in arg list first. If yes, go to check parent frame.
+            if any(var is arg for arg in all_args):
+                # the var is passed in from caller, check parent frame.
+                frame = frame.f_back
+                continue
+
+            for var_name, var_val in frame.f_locals.items():
+                if var_val is var:
+                    original_var_name = var_name
+                    break
+
+            break
+
+        return original_var_name
+
+    except Exception:
+        return fallback_name
+
+
+def find_free_port():
+    """
+    Find free socket port on local machine.
+    """
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
+def check_port_connectivity():
+    port = find_free_port()
+    try:
+        with subprocess.Popen(
+            ["nc", "-l", "-p", str(port)],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        ) as server:
+            with subprocess.Popen(
+                ["nc", "-zv", "localhost", str(port)],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            ) as client:
+                client.wait()
+                server.terminate()
+                return client.returncode == 0
+    except Exception as e:
+        _logger.warning("Failed to check port connectivity: %s", e)
+        return False
+
+
+def is_iterator(obj):
+    """
+    Args:
+        obj: Any object.
+
+    Returns:
+        Boolean representing whether or not 'obj' is an iterator.
+    """
+    return (hasattr(obj, "__next__") or hasattr(obj, "next")) and hasattr(obj, "__iter__")
+
+
+def _is_in_ipython_notebook():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except Exception:
+        return False
+
+
+def get_results_from_paginated_fn(paginated_fn, max_results_per_page, max_results=None):
+    """Gets results by calling the ``paginated_fn`` until either no more results remain or
+    the specified ``max_results`` threshold has been reached.
+
+    Args:
+        paginated_fn: This function is expected to take in the number of results to retrieve
+            per page and a pagination token, and return a PagedList object.
+        max_results_per_page: The maximum number of results to retrieve per page.
+        max_results: The maximum number of results to retrieve overall. If unspecified,
+            all results will be retrieved.
+
+    Returns:
+        Returns a list of entities, as determined by the paginated_fn parameter, with no more
+        entities than specified by max_results.
+
+    """
+    all_results = []
+    next_page_token = None
+    returns_all = max_results is None
+    while returns_all or len(all_results) < max_results:
+        num_to_get = max_results_per_page if returns_all else max_results - len(all_results)
+        if num_to_get < max_results_per_page:
+            page_results = paginated_fn(num_to_get, next_page_token)
+        else:
+            page_results = paginated_fn(max_results_per_page, next_page_token)
+        all_results.extend(page_results)
+        if hasattr(page_results, "token") and page_results.token:
+            next_page_token = page_results.token
+        else:
+            break
+    return all_results
+
+
+class AttrDict(dict):
+    """
+    Dict-like object that exposes its keys as attributes.
+
+    Examples
+    --------
+    >>> d = AttrDict({"a": 1, "b": 2})
+    >>> d.a
+    1
+    >>> d = AttrDict({"a": 1, "b": {"c": 3, "d": 4}})
+    >>> d.b.c
+    3
+    """
+
+    def __getattr__(self, attr):
+        try:
+            value = self[attr]
+        except KeyError:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attr}'")
+        if isinstance(value, dict):
+            return AttrDict(value)
+        return value
+
+
+def get_parent_module(module):
+    return module[0 : module.rindex(".")]

venv/lib/python3.9/site-packages/mlflow/utils/_capture_modules.py

@@ -0,0 +1,256 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import argparse
+import builtins
+import functools
+import importlib
+import json
+import os
+import sys
+
+import mlflow
+from mlflow.models.model import MLMODEL_FILE_NAME, Model
+from mlflow.pyfunc import MAIN
+from mlflow.utils._spark_utils import _prepare_subprocess_environ_for_creating_local_spark_session
+from mlflow.utils.exception_utils import get_stacktrace
+from mlflow.utils.file_utils import write_to
+from mlflow.utils.requirements_utils import (
+    DATABRICKS_MODULES_TO_PACKAGES,
+    MLFLOW_MODULES_TO_PACKAGES,
+)
+
+
+def _get_top_level_module(full_module_name):
+    return full_module_name.split(".")[0]
+
+
+def _get_second_level_module(full_module_name):
+    return ".".join(full_module_name.split(".")[:2])
+
+
+class _CaptureImportedModules:
+    """
+    A context manager to capture imported modules by temporarily applying a patch to
+    `builtins.__import__` and `importlib.import_module`.
+
+    If `record_full_module` is set to `False`, it only captures top level modules
+    for inferring python package purpose.
+    If `record_full_module` is set to `True`, it captures full module name for all
+    imported modules and sub-modules. This is used in automatic model code path inference.
+    """
+
+    def __init__(self, record_full_module=False):
+        self.imported_modules = set()
+        self.original_import = None
+        self.original_import_module = None
+        self.record_full_module = record_full_module
+
+    def _wrap_import(self, original):
+        @functools.wraps(original)
+        def wrapper(name, globals=None, locals=None, fromlist=(), level=0):
+            is_absolute_import = level == 0
+            if not self.record_full_module and is_absolute_import:
+                self._record_imported_module(name)
+
+            result = original(name, globals, locals, fromlist, level)
+
+            if self.record_full_module:
+                if is_absolute_import:
+                    parent_modules = name.split(".")
+                else:
+                    parent_modules = globals["__name__"].split(".")
+                    if level > 1:
+                        parent_modules = parent_modules[: -(level - 1)]
+
+                if fromlist:
+                    for from_name in fromlist:
+                        full_modules = parent_modules + [from_name]
+                        full_module_name = ".".join(full_modules)
+                        if full_module_name in sys.modules:
+                            self._record_imported_module(full_module_name)
+                else:
+                    full_module_name = ".".join(parent_modules)
+                    self._record_imported_module(full_module_name)
+
+            return result
+
+        return wrapper
+
+    def _wrap_import_module(self, original):
+        @functools.wraps(original)
+        def wrapper(name, *args, **kwargs):
+            self._record_imported_module(name)
+            return original(name, *args, **kwargs)
+
+        return wrapper
+
+    def _record_imported_module(self, full_module_name):
+        if self.record_full_module:
+            self.imported_modules.add(full_module_name)
+            return
+
+        # If the module is an internal module (prefixed by "_") or is the "databricks"
+        # module, which is populated by many different packages, don't record it (specific
+        # module imports within the databricks namespace are still recorded and mapped to
+        # their corresponding packages)
+        if full_module_name.startswith("_") or full_module_name == "databricks":
+            return
+
+        top_level_module = _get_top_level_module(full_module_name)
+        second_level_module = _get_second_level_module(full_module_name)
+
+        if top_level_module == "databricks":
+            # Multiple packages populate the `databricks` module namespace on Databricks;
+            # to avoid bundling extraneous Databricks packages into model dependencies, we
+            # scope each module to its relevant package
+            if second_level_module in DATABRICKS_MODULES_TO_PACKAGES:
+                self.imported_modules.add(second_level_module)
+                return
+
+            for databricks_module in DATABRICKS_MODULES_TO_PACKAGES:
+                if full_module_name.startswith(databricks_module):
+                    self.imported_modules.add(databricks_module)
+                    return
+
+        # special casing for mlflow extras since they may not be required by default
+        if top_level_module == "mlflow":
+            if second_level_module in MLFLOW_MODULES_TO_PACKAGES:
+                self.imported_modules.add(second_level_module)
+                return
+
+        self.imported_modules.add(top_level_module)
+
+    def __enter__(self):
+        # Patch `builtins.__import__` and `importlib.import_module`
+        self.original_import = builtins.__import__
+        self.original_import_module = importlib.import_module
+        builtins.__import__ = self._wrap_import(self.original_import)
+        importlib.import_module = self._wrap_import_module(self.original_import_module)
+        return self
+
+    def __exit__(self, *_, **__):
+        # Revert the patches
+        builtins.__import__ = self.original_import
+        importlib.import_module = self.original_import_module
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--flavor", required=True)
+    parser.add_argument("--output-file", required=True)
+    parser.add_argument("--sys-path", required=True)
+    parser.add_argument("--module-to-throw", required=False)
+    parser.add_argument("--error-file", required=False)
+    parser.add_argument("--record-full-module", default=False, action="store_true")
+    return parser.parse_args()
+
+
+def store_imported_modules(
+    cap_cm, model_path, flavor, output_file, error_file=None, record_full_module=False
+):
+    # If `model_path` refers to an MLflow model directory, load the model using
+    # `mlflow.pyfunc.load_model`
+    if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(model_path):
+        mlflow_model = Model.load(model_path)
+        pyfunc_conf = mlflow_model.flavors.get(mlflow.pyfunc.FLAVOR_NAME)
+        input_example = mlflow_model.load_input_example(model_path)
+        params = mlflow_model.load_input_example_params(model_path)
+
+        def load_model_and_predict(original_load_fn, *args, **kwargs):
+            model = original_load_fn(*args, **kwargs)
+            if input_example is not None:
+                try:
+                    model.predict(input_example, params=params)
+                except Exception as e:
+                    if error_file:
+                        stack_trace = get_stacktrace(e)
+                        write_to(
+                            error_file,
+                            "Failed to run predict on input_example, dependencies "
+                            "introduced in predict are not captured.\n" + stack_trace,
+                        )
+                    else:
+                        raise e
+            return model
+
+        if record_full_module:
+            # Note: if we want to record all imported modules
+            # (for inferring code_paths purpose),
+            # The `importlib.import_module(pyfunc_conf[MAIN])` invocation
+            # must be wrapped with `cap_cm` context manager,
+            # because `pyfunc_conf[MAIN]` might also be a module loaded from
+            # code_paths.
+            with cap_cm:
+                # `mlflow.pyfunc.load_model` internally invokes
+                # `importlib.import_module(pyfunc_conf[MAIN])`
+                mlflow.pyfunc.load_model(model_path)
+        else:
+            loader_module = importlib.import_module(pyfunc_conf[MAIN])
+            original = loader_module._load_pyfunc
+
+            @functools.wraps(original)
+            def _load_pyfunc_patch(*args, **kwargs):
+                with cap_cm:
+                    return load_model_and_predict(original, *args, **kwargs)
+
+            loader_module._load_pyfunc = _load_pyfunc_patch
+            try:
+                mlflow.pyfunc.load_model(model_path)
+            finally:
+                loader_module._load_pyfunc = original
+    # Otherwise, load the model using `mlflow.<flavor>._load_pyfunc`.
+    # For models that don't contain pyfunc flavor (e.g. scikit-learn estimator
+    # that doesn't implement a `predict` method),
+    # we need to directly pass a model data path to this script.
+    else:
+        with cap_cm:
+            importlib.import_module(f"mlflow.{flavor}")._load_pyfunc(model_path)
+
+    # Store the imported modules in `output_file`
+    write_to(output_file, "\n".join(cap_cm.imported_modules))
+
+
+def main():
+    args = parse_args()
+    model_path = args.model_path
+    flavor = args.flavor
+    output_file = args.output_file
+    error_file = args.error_file
+    # Mirror `sys.path` of the parent process
+    sys.path = json.loads(args.sys_path)
+
+    if flavor == mlflow.spark.FLAVOR_NAME:
+        # Create a local spark environment within the subprocess
+        from mlflow.utils._spark_utils import _create_local_spark_session_for_loading_spark_model
+
+        _prepare_subprocess_environ_for_creating_local_spark_session()
+        _create_local_spark_session_for_loading_spark_model()
+
+    cap_cm = _CaptureImportedModules(record_full_module=args.record_full_module)
+    store_imported_modules(
+        cap_cm,
+        model_path,
+        flavor,
+        output_file,
+        error_file,
+        record_full_module=args.record_full_module,
+    )
+
+    # Clean up a spark session created by `mlflow.spark._load_pyfunc`
+    if flavor == mlflow.spark.FLAVOR_NAME:
+        from mlflow.utils._spark_utils import _get_active_spark_session
+
+        spark = _get_active_spark_session()
+        if spark:
+            try:
+                spark.stop()
+            except Exception:
+                # Swallow unexpected exceptions
+                pass
+
+
+if __name__ == "__main__":
+    main()

venv/lib/python3.9/site-packages/mlflow/utils/_capture_transformers_modules.py

@@ -0,0 +1,75 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import json
+import os
+import sys
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils._capture_modules import (
+    _CaptureImportedModules,
+    parse_args,
+    store_imported_modules,
+)
+
+
+class _CaptureImportedModulesForHF(_CaptureImportedModules):
+    """
+    A context manager to capture imported modules by temporarily applying a patch to
+    `builtins.__import__` and `importlib.import_module`.
+    Used for 'transformers' flavor only.
+    """
+
+    def __init__(self, module_to_throw, record_full_module=False):
+        super().__init__(record_full_module=record_full_module)
+        self.module_to_throw = module_to_throw
+
+    def _record_imported_module(self, full_module_name):
+        if full_module_name == self.module_to_throw or full_module_name.startswith(
+            f"{self.module_to_throw}."
+        ):
+            raise ImportError(f"Disabled package {full_module_name}")
+        return super()._record_imported_module(full_module_name)
+
+
+def main():
+    args = parse_args()
+    model_path = args.model_path
+    flavor = args.flavor
+    output_file = args.output_file
+    module_to_throw = args.module_to_throw
+    # Mirror `sys.path` of the parent process
+    sys.path = json.loads(args.sys_path)
+
+    if flavor != mlflow.transformers.FLAVOR_NAME:
+        raise MlflowException(
+            f"This script is only applicable to '{mlflow.transformers.FLAVOR_NAME}' flavor, "
+            "if you're applying other flavors, please use _capture_modules script.",
+        )
+
+    if module_to_throw == "":
+        raise MlflowException("Please specify the module to throw.")
+    elif module_to_throw == "tensorflow":
+        if os.environ.get("USE_TORCH", None) != "TRUE":
+            raise MlflowException(
+                "The environment variable USE_TORCH has to be set to TRUE to disable Tensorflow.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+    elif module_to_throw == "torch":
+        if os.environ.get("USE_TF", None) != "TRUE":
+            raise MlflowException(
+                "The environment variable USE_TF has to be set to TRUE to disable Pytorch.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    cap_cm = _CaptureImportedModulesForHF(
+        module_to_throw, record_full_module=args.record_full_module
+    )
+    store_imported_modules(cap_cm, model_path, flavor, output_file)
+
+
+if __name__ == "__main__":
+    main()

venv/lib/python3.9/site-packages/mlflow/utils/_spark_utils.py

@@ -0,0 +1,228 @@
+import contextlib
+import multiprocessing
+import os
+import shutil
+import tempfile
+import zipfile
+
+
+def _get_active_spark_session():
+    try:
+        from pyspark.sql import SparkSession
+    except ImportError:
+        # Return None if user doesn't have PySpark installed
+        return None
+    try:
+        # getActiveSession() only exists in Spark 3.0 and above
+        return SparkSession.getActiveSession()
+    except Exception:
+        # Fall back to this internal field for Spark 2.x and below.
+        return SparkSession._instantiatedSession
+
+
+# Suppose we have a parent process already initiate a spark session that connected to a spark
+# cluster, then the parent process spawns a child process, if child process directly creates
+# a local spark session, it does not work correctly, because of PYSPARK_GATEWAY_PORT and
+# PYSPARK_GATEWAY_SECRET are inherited from parent process and child process pyspark session
+# will try to connect to the port and cause error.
+# So the 2 lines here are to clear 'PYSPARK_GATEWAY_PORT' and 'PYSPARK_GATEWAY_SECRET' to
+# enforce launching a new pyspark JVM gateway.
+def _prepare_subprocess_environ_for_creating_local_spark_session():
+    from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+    if is_in_databricks_runtime():
+        os.environ["SPARK_DIST_CLASSPATH"] = "/databricks/jars/*"
+
+    os.environ.pop("PYSPARK_GATEWAY_PORT", None)
+    os.environ.pop("PYSPARK_GATEWAY_SECRET", None)
+
+
+def _get_spark_scala_version_from_spark_session(spark):
+    version = spark._jvm.scala.util.Properties.versionNumberString().split(".", 2)
+    return f"{version[0]}.{version[1]}"
+
+
+def _get_spark_scala_version_child_proc_target(result_queue):
+    from pyspark.sql import SparkSession
+
+    _prepare_subprocess_environ_for_creating_local_spark_session()
+    with SparkSession.builder.master("local[1]").getOrCreate() as spark_session:
+        scala_version = _get_spark_scala_version_from_spark_session(spark_session)
+        result_queue.put(scala_version)
+
+
+def _get_spark_scala_version():
+    from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+    if is_in_databricks_runtime() and "SPARK_SCALA_VERSION" in os.environ:
+        return os.environ["SPARK_SCALA_VERSION"]
+
+    if spark := _get_active_spark_session():
+        return _get_spark_scala_version_from_spark_session(spark)
+
+    result_queue = multiprocessing.Queue()
+
+    # If we need to create a new spark local session for reading scala version,
+    # we have to create the temporal spark session in a child process,
+    # if we create the temporal spark session in current process,
+    # after terminating the temporal spark session, creating another spark session
+    # with "spark.jars.packages" configuration doesn't work.
+    proc = multiprocessing.Process(
+        target=_get_spark_scala_version_child_proc_target, args=(result_queue,)
+    )
+    proc.start()
+    proc.join()
+    if proc.exitcode != 0:
+        raise RuntimeError("Failed to read scala version.")
+
+    return result_queue.get()
+
+
+def _create_local_spark_session_for_recipes():
+    """Create a sparksession to be used within an recipe step run in a subprocess locally."""
+
+    try:
+        from pyspark.sql import SparkSession
+    except ImportError:
+        # Return None if user doesn't have PySpark installed
+        return None
+
+    try:
+        spark_scala_version = _get_spark_scala_version()
+    except Exception as e:
+        raise RuntimeError("Failed to get spark scala version.") from e
+
+    _prepare_subprocess_environ_for_creating_local_spark_session()
+    return (
+        SparkSession.builder.master("local[*]")
+        .config("spark.jars.packages", f"io.delta:delta-spark_{spark_scala_version}:3.0.0")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+        )
+        .config("spark.sql.execution.arrow.pyspark.enabled", "true")
+        .getOrCreate()
+    )
+
+
+def _create_local_spark_session_for_loading_spark_model():
+    from pyspark.sql import SparkSession
+
+    return (
+        SparkSession.builder.config("spark.python.worker.reuse", "true")
+        # The config is a workaround for avoiding databricks delta cache issue when loading
+        # some specific model such as ALSModel.
+        .config("spark.databricks.io.cache.enabled", "false")
+        # In Spark 3.1 and above, we need to set this conf explicitly to enable creating
+        # a SparkSession on the workers
+        .config("spark.executor.allowSparkContext", "true")
+        # Binding "spark.driver.host" to 127.0.0.1 helps avoiding some local hostname
+        # related issues (e.g. https://github.com/mlflow/mlflow/issues/5733).
+        # Note that we should set "spark.driver.host" instead of "spark.driver.bindAddress",
+        # the latter one only set server binding host, but it doesn't set client side request
+        # destination host.
+        .config("spark.driver.host", "127.0.0.1")
+        .config("spark.executor.allowSparkContext", "true")
+        .config(
+            "spark.driver.extraJavaOptions",
+            "-Dlog4j.configuration=file:/usr/local/spark/conf/log4j.properties",
+        )
+        .master("local[1]")
+        .getOrCreate()
+    )
+
+
+_NFS_PATH_PREFIX = "nfs:"
+
+
+def _get_spark_distributor_nfs_cache_dir():
+    from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir  # avoid circular import
+
+    if (nfs_root_dir := get_nfs_cache_root_dir()) is not None:
+        cache_dir = os.path.join(nfs_root_dir, "mlflow_distributor_cache_dir")
+        os.makedirs(cache_dir, exist_ok=True)
+        return cache_dir
+    return None
+
+
+class _SparkDirectoryDistributor:
+    """Distribute spark directory from driver to executors."""
+
+    _extracted_dir_paths = {}
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def add_dir(spark, dir_path):
+        """Given a SparkSession and a model_path which refers to a pyfunc directory locally,
+        we will zip the directory up, enable it to be distributed to executors, and return
+        the "archive_path", which should be used as the path in get_or_load().
+        """
+        _, archive_basepath = tempfile.mkstemp()
+        # NB: We must archive the directory as Spark.addFile does not support non-DFS
+        # directories when recursive=True.
+        archive_path = shutil.make_archive(archive_basepath, "zip", dir_path)
+
+        if (nfs_cache_dir := _get_spark_distributor_nfs_cache_dir()) is not None:
+            # If NFS directory (shared by all spark nodes) is available, use NFS directory
+            # instead of `SparkContext.addFile` to distribute files.
+            # Because `SparkContext.addFile` is not secure, so it is not allowed to be called
+            # on a shared cluster.
+            dest_path = os.path.join(nfs_cache_dir, os.path.basename(archive_path))
+            shutil.copy(archive_path, dest_path)
+            return _NFS_PATH_PREFIX + dest_path
+
+        spark.sparkContext.addFile(archive_path)
+        return archive_path
+
+    @staticmethod
+    def get_or_extract(archive_path):
+        """Given a path returned by add_local_model(), this method will return a tuple of
+        (loaded_model, local_model_path).
+        If this Python process ever loaded the model before, we will reuse that copy.
+        """
+        from pyspark.files import SparkFiles
+
+        if archive_path in _SparkDirectoryDistributor._extracted_dir_paths:
+            return _SparkDirectoryDistributor._extracted_dir_paths[archive_path]
+
+        # BUG: Despite the documentation of SparkContext.addFile() and SparkFiles.get() in Scala
+        # and Python, it turns out that we actually need to use the basename as the input to
+        # SparkFiles.get(), as opposed to the (absolute) path.
+        if archive_path.startswith(_NFS_PATH_PREFIX):
+            local_path = archive_path[len(_NFS_PATH_PREFIX) :]
+        else:
+            archive_path_basename = os.path.basename(archive_path)
+            local_path = SparkFiles.get(archive_path_basename)
+        temp_dir = tempfile.mkdtemp()
+        zip_ref = zipfile.ZipFile(local_path, "r")
+        zip_ref.extractall(temp_dir)
+        zip_ref.close()
+
+        _SparkDirectoryDistributor._extracted_dir_paths[archive_path] = temp_dir
+        return _SparkDirectoryDistributor._extracted_dir_paths[archive_path]
+
+
+@contextlib.contextmanager
+def modified_environ(update):
+    """Temporarily updates the ``os.environ`` dictionary in-place.
+
+    The ``os.environ`` dictionary is updated in-place so that the modification
+    is sure to work in all situations.
+
+    Args:
+        update: Dictionary of environment variables and values to add/update.
+    """
+    update = update or {}
+    original_env = {k: os.environ.get(k) for k in update}
+
+    try:
+        os.environ.update(update)
+        yield
+    finally:
+        for k, v in original_env.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v

venv/lib/python3.9/site-packages/mlflow/utils/_unity_catalog_oss_utils.py

@@ -0,0 +1,97 @@
+import re
+
+from mlflow.entities.model_registry import (
+    ModelVersion,
+    ModelVersionSearch,
+    RegisteredModel,
+    RegisteredModelSearch,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    ModelVersionInfo,
+    ModelVersionStatus,
+    RegisteredModelInfo,
+)
+
+_STRING_TO_STATUS = {k: ModelVersionStatus.Value(k) for k in ModelVersionStatus.keys()}
+_STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+
+
+def get_registered_model_from_uc_oss_proto(uc_oss_proto: RegisteredModelInfo) -> RegisteredModel:
+    return RegisteredModel(
+        name=f"{uc_oss_proto.catalog_name}.{uc_oss_proto.schema_name}.{uc_oss_proto.name}",
+        creation_timestamp=uc_oss_proto.created_at,
+        last_updated_timestamp=uc_oss_proto.updated_at,
+        description=uc_oss_proto.comment,
+    )
+
+
+def get_model_version_from_uc_oss_proto(uc_oss_proto: ModelVersionInfo) -> ModelVersion:
+    return ModelVersion(
+        name=f"{uc_oss_proto.catalog_name}.{uc_oss_proto.schema_name}.{uc_oss_proto.model_name}",
+        version=uc_oss_proto.version,
+        creation_timestamp=uc_oss_proto.created_at,
+        last_updated_timestamp=uc_oss_proto.updated_at,
+        description=uc_oss_proto.comment,
+        source=uc_oss_proto.source,
+        run_id=uc_oss_proto.run_id,
+        status=uc_oss_model_version_status_to_string(uc_oss_proto.status),
+    )
+
+
+def get_registered_model_search_from_uc_oss_proto(
+    uc_oss_proto: RegisteredModelInfo,
+) -> RegisteredModelSearch:
+    return RegisteredModelSearch(
+        name=f"{uc_oss_proto.catalog_name}.{uc_oss_proto.schema_name}.{uc_oss_proto.name}",
+        creation_timestamp=uc_oss_proto.created_at,
+        last_updated_timestamp=uc_oss_proto.updated_at,
+        description=uc_oss_proto.comment,
+    )
+
+
+def get_model_version_search_from_uc_oss_proto(
+    uc_oss_proto: ModelVersionInfo,
+) -> ModelVersionSearch:
+    return ModelVersionSearch(
+        name=f"{uc_oss_proto.catalog_name}.{uc_oss_proto.schema_name}.{uc_oss_proto.model_name}",
+        version=uc_oss_proto.version,
+        creation_timestamp=uc_oss_proto.created_at,
+        last_updated_timestamp=uc_oss_proto.updated_at,
+        description=uc_oss_proto.comment,
+        source=uc_oss_proto.source,
+        run_id=uc_oss_proto.run_id,
+        status=uc_oss_model_version_status_to_string(uc_oss_proto.status),
+    )
+
+
+def uc_oss_model_version_status_to_string(status):
+    return _STATUS_TO_STRING[status]
+
+
+filter_pattern = re.compile(r"^name\s*=\s*'([^']+)'")
+
+
+def parse_model_name(filter):
+    trimmed_filter = filter.strip()
+    match = filter_pattern.match(trimmed_filter)
+    if match:
+        model_name_str = match.group(1)
+    elif trimmed_filter == "":
+        raise MlflowException(
+            "Missing filter: please specify a filter parameter in the format `name = 'model_name'`."
+        )
+    else:
+        raise MlflowException(
+            f"Unsupported filter query : `{trimmed_filter}`."
+            + " Please specify your filter parameter in "
+            + "the format `name = 'model_name'`."
+        )
+    parts = model_name_str.split(".")
+    if len(parts) != 3 or not all(parts):
+        raise MlflowException(
+            "Bad model name: please specify all three levels of the model in the"
+            "form `catalog_name.schema_name.model_name`"
+        )
+    catalog, schema, model = parts
+    return f"{catalog}.{schema}.{model}"

venv/lib/python3.9/site-packages/mlflow/utils/_unity_catalog_utils.py

@@ -0,0 +1,448 @@
+import logging
+from typing import Callable, Optional
+
+from mlflow.entities.model_registry import (
+    ModelVersion,
+    ModelVersionTag,
+    RegisteredModel,
+    RegisteredModelAlias,
+    RegisteredModelTag,
+)
+from mlflow.entities.model_registry.model_version_search import ModelVersionSearch
+from mlflow.entities.model_registry.registered_model_search import RegisteredModelSearch
+from mlflow.environment_variables import MLFLOW_USE_DATABRICKS_SDK_MODEL_ARTIFACTS_REPO_FOR_UC
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    EmitModelVersionLineageRequest,
+    EmitModelVersionLineageResponse,
+    IsDatabricksSdkModelsArtifactRepositoryEnabledRequest,
+    IsDatabricksSdkModelsArtifactRepositoryEnabledResponse,
+    ModelVersionLineageInfo,
+    SseEncryptionAlgorithm,
+    TemporaryCredentials,
+)
+from mlflow.protos.databricks_uc_registry_messages_pb2 import ModelVersion as ProtoModelVersion
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    ModelVersionStatus as ProtoModelVersionStatus,
+)
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    ModelVersionTag as ProtoModelVersionTag,
+)
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    RegisteredModel as ProtoRegisteredModel,
+)
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    RegisteredModelTag as ProtoRegisteredModelTag,
+)
+from mlflow.protos.databricks_uc_registry_service_pb2 import UcModelRegistryService
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    TemporaryCredentials as TemporaryCredentialsOSS,
+)
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+
+_logger = logging.getLogger(__name__)
+_METHOD_TO_INFO = extract_api_info_for_service(UcModelRegistryService, _REST_API_PATH_PREFIX)
+_STRING_TO_STATUS = {k: ProtoModelVersionStatus.Value(k) for k in ProtoModelVersionStatus.keys()}
+_STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+_ACTIVE_CATALOG_QUERY = "SELECT current_catalog() AS catalog"
+_ACTIVE_SCHEMA_QUERY = "SELECT current_database() AS schema"
+
+
+def uc_model_version_status_to_string(status):
+    return _STATUS_TO_STRING[status]
+
+
+def model_version_from_uc_proto(uc_proto: ProtoModelVersion) -> ModelVersion:
+    return ModelVersion(
+        name=uc_proto.name,
+        version=uc_proto.version,
+        creation_timestamp=uc_proto.creation_timestamp,
+        last_updated_timestamp=uc_proto.last_updated_timestamp,
+        description=uc_proto.description,
+        user_id=uc_proto.user_id,
+        source=uc_proto.source,
+        run_id=uc_proto.run_id,
+        status=uc_model_version_status_to_string(uc_proto.status),
+        status_message=uc_proto.status_message,
+        aliases=[alias.alias for alias in (uc_proto.aliases or [])],
+        tags=[ModelVersionTag(key=tag.key, value=tag.value) for tag in (uc_proto.tags or [])],
+    )
+
+
+def model_version_search_from_uc_proto(uc_proto: ProtoModelVersion) -> ModelVersionSearch:
+    return ModelVersionSearch(
+        name=uc_proto.name,
+        version=uc_proto.version,
+        creation_timestamp=uc_proto.creation_timestamp,
+        last_updated_timestamp=uc_proto.last_updated_timestamp,
+        description=uc_proto.description,
+        user_id=uc_proto.user_id,
+        source=uc_proto.source,
+        run_id=uc_proto.run_id,
+        status=uc_model_version_status_to_string(uc_proto.status),
+        status_message=uc_proto.status_message,
+        aliases=[],
+        tags=[],
+    )
+
+
+def registered_model_from_uc_proto(uc_proto: ProtoRegisteredModel) -> RegisteredModel:
+    return RegisteredModel(
+        name=uc_proto.name,
+        creation_timestamp=uc_proto.creation_timestamp,
+        last_updated_timestamp=uc_proto.last_updated_timestamp,
+        description=uc_proto.description,
+        aliases=[
+            RegisteredModelAlias(alias=alias.alias, version=alias.version)
+            for alias in (uc_proto.aliases or [])
+        ],
+        tags=[RegisteredModelTag(key=tag.key, value=tag.value) for tag in (uc_proto.tags or [])],
+    )
+
+
+def registered_model_search_from_uc_proto(uc_proto: ProtoRegisteredModel) -> RegisteredModelSearch:
+    return RegisteredModelSearch(
+        name=uc_proto.name,
+        creation_timestamp=uc_proto.creation_timestamp,
+        last_updated_timestamp=uc_proto.last_updated_timestamp,
+        description=uc_proto.description,
+        aliases=[],
+        tags=[],
+    )
+
+
+def uc_registered_model_tag_from_mlflow_tags(
+    tags: Optional[list[RegisteredModelTag]],
+) -> list[ProtoRegisteredModelTag]:
+    if tags is None:
+        return []
+    return [ProtoRegisteredModelTag(key=t.key, value=t.value) for t in tags]
+
+
+def uc_model_version_tag_from_mlflow_tags(
+    tags: Optional[list[ModelVersionTag]],
+) -> list[ProtoModelVersionTag]:
+    if tags is None:
+        return []
+    return [ProtoModelVersionTag(key=t.key, value=t.value) for t in tags]
+
+
+def get_artifact_repo_from_storage_info(
+    storage_location: str,
+    scoped_token: TemporaryCredentials,
+    base_credential_refresh_def: Callable[[], TemporaryCredentials],
+    is_oss: bool = False,
+) -> ArtifactRepository:
+    """
+    Get an ArtifactRepository instance capable of reading/writing to a UC model version's
+    file storage location
+
+    Args:
+        storage_location: Storage location of the model version
+        scoped_token: Protobuf scoped token to use to authenticate to blob storage
+        base_credential_refresh_def: Function that returns temporary credentials for accessing blob
+            storage. It is first used to determine the type of blob storage and to access it. It is
+            then passed to the relevant ArtifactRepository implementation to refresh credentials as
+            needed.
+        is_oss: Whether the user is using the OSS version of Unity Catalog
+    """
+    try:
+        if is_oss:
+            return _get_artifact_repo_from_storage_info_oss(
+                storage_location=storage_location,
+                scoped_token=scoped_token,
+                base_credential_refresh_def=base_credential_refresh_def,
+            )
+        else:
+            return _get_artifact_repo_from_storage_info(
+                storage_location=storage_location,
+                scoped_token=scoped_token,
+                base_credential_refresh_def=base_credential_refresh_def,
+            )
+    except ImportError as e:
+        raise MlflowException(
+            "Unable to import necessary dependencies to access model version files in "
+            "Unity Catalog. Please ensure you have the necessary dependencies installed, "
+            "e.g. by running 'pip install mlflow[databricks]' or "
+            "'pip install mlflow-skinny[databricks]'"
+        ) from e
+
+
+def _get_artifact_repo_from_storage_info(
+    storage_location: str,
+    scoped_token: TemporaryCredentials,
+    base_credential_refresh_def: Callable[[], TemporaryCredentials],
+) -> ArtifactRepository:
+    credential_type = scoped_token.WhichOneof("credentials")
+    if credential_type == "aws_temp_credentials":
+        # Verify upfront that boto3 is importable
+        import boto3  # noqa: F401
+
+        from mlflow.store.artifact.optimized_s3_artifact_repo import OptimizedS3ArtifactRepository
+
+        aws_creds = scoped_token.aws_temp_credentials
+        s3_upload_extra_args = _parse_aws_sse_credential(scoped_token)
+
+        def aws_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_aws_creds = new_scoped_token.aws_temp_credentials
+            new_s3_upload_extra_args = _parse_aws_sse_credential(new_scoped_token)
+            return {
+                "access_key_id": new_aws_creds.access_key_id,
+                "secret_access_key": new_aws_creds.secret_access_key,
+                "session_token": new_aws_creds.session_token,
+                "s3_upload_extra_args": new_s3_upload_extra_args,
+            }
+
+        return OptimizedS3ArtifactRepository(
+            artifact_uri=storage_location,
+            access_key_id=aws_creds.access_key_id,
+            secret_access_key=aws_creds.secret_access_key,
+            session_token=aws_creds.session_token,
+            credential_refresh_def=aws_credential_refresh,
+            s3_upload_extra_args=s3_upload_extra_args,
+        )
+    elif credential_type == "azure_user_delegation_sas":
+        from azure.core.credentials import AzureSasCredential
+
+        from mlflow.store.artifact.azure_data_lake_artifact_repo import (
+            AzureDataLakeArtifactRepository,
+        )
+
+        sas_token = scoped_token.azure_user_delegation_sas.sas_token
+
+        def azure_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_sas_token = new_scoped_token.azure_user_delegation_sas.sas_token
+            return {
+                "credential": AzureSasCredential(new_sas_token),
+            }
+
+        return AzureDataLakeArtifactRepository(
+            artifact_uri=storage_location,
+            credential=AzureSasCredential(sas_token),
+            credential_refresh_def=azure_credential_refresh,
+        )
+
+    elif credential_type == "gcp_oauth_token":
+        from google.cloud.storage import Client
+        from google.oauth2.credentials import Credentials
+
+        from mlflow.store.artifact.gcs_artifact_repo import GCSArtifactRepository
+
+        credentials = Credentials(scoped_token.gcp_oauth_token.oauth_token)
+
+        def gcp_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_gcp_creds = new_scoped_token.gcp_oauth_token
+            return {
+                "oauth_token": new_gcp_creds.oauth_token,
+            }
+
+        client = Client(project="mlflow", credentials=credentials)
+        return GCSArtifactRepository(
+            artifact_uri=storage_location,
+            client=client,
+            credential_refresh_def=gcp_credential_refresh,
+        )
+    elif credential_type == "r2_temp_credentials":
+        from mlflow.store.artifact.r2_artifact_repo import R2ArtifactRepository
+
+        r2_creds = scoped_token.r2_temp_credentials
+
+        def r2_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_r2_creds = new_scoped_token.r2_temp_credentials
+            return {
+                "access_key_id": new_r2_creds.access_key_id,
+                "secret_access_key": new_r2_creds.secret_access_key,
+                "session_token": new_r2_creds.session_token,
+            }
+
+        return R2ArtifactRepository(
+            artifact_uri=storage_location,
+            access_key_id=r2_creds.access_key_id,
+            secret_access_key=r2_creds.secret_access_key,
+            session_token=r2_creds.session_token,
+            credential_refresh_def=r2_credential_refresh,
+        )
+    else:
+        raise MlflowException(
+            f"Got unexpected credential type {credential_type} when attempting to "
+            "access model version files in Unity Catalog. Try upgrading to the latest "
+            "version of the MLflow Python client."
+        )
+
+
+def _get_artifact_repo_from_storage_info_oss(
+    storage_location: str,
+    scoped_token: TemporaryCredentialsOSS,
+    base_credential_refresh_def: Callable[[], TemporaryCredentialsOSS],
+) -> ArtifactRepository:
+    # OSS Temp Credential doesn't have a oneof credential field
+    # So, we must check for the individual cloud credentials
+    if len(scoped_token.aws_temp_credentials.access_key_id) > 0:
+        # Verify upfront that boto3 is importable
+        import boto3  # noqa: F401
+
+        from mlflow.store.artifact.optimized_s3_artifact_repo import OptimizedS3ArtifactRepository
+
+        aws_creds = scoped_token.aws_temp_credentials
+
+        def aws_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_aws_creds = new_scoped_token.aws_temp_credentials
+            return {
+                "access_key_id": new_aws_creds.access_key_id,
+                "secret_access_key": new_aws_creds.secret_access_key,
+                "session_token": new_aws_creds.session_token,
+            }
+
+        return OptimizedS3ArtifactRepository(
+            artifact_uri=storage_location,
+            access_key_id=aws_creds.access_key_id,
+            secret_access_key=aws_creds.secret_access_key,
+            session_token=aws_creds.session_token,
+            credential_refresh_def=aws_credential_refresh,
+        )
+    elif len(scoped_token.azure_user_delegation_sas.sas_token) > 0:
+        from azure.core.credentials import AzureSasCredential
+
+        from mlflow.store.artifact.azure_data_lake_artifact_repo import (
+            AzureDataLakeArtifactRepository,
+        )
+
+        sas_token = scoped_token.azure_user_delegation_sas.sas_token
+
+        def azure_credential_refresh():
+            new_scoped_token = base_credential_refresh_def()
+            new_sas_token = new_scoped_token.azure_user_delegation_sas.sas_token
+            return {
+                "credential": AzureSasCredential(new_sas_token),
+            }
+
+        return AzureDataLakeArtifactRepository(
+            artifact_uri=storage_location,
+            credential=AzureSasCredential(sas_token),
+            credential_refresh_def=azure_credential_refresh,
+        )
+
+    elif len(scoped_token.gcp_oauth_token.oauth_token) > 0:
+        from google.cloud.storage import Client
+        from google.oauth2.credentials import Credentials
+
+        from mlflow.store.artifact.gcs_artifact_repo import GCSArtifactRepository
+
+        credentials = Credentials(scoped_token.gcp_oauth_token.oauth_token)
+        client = Client(project="mlflow", credentials=credentials)
+        return GCSArtifactRepository(artifact_uri=storage_location, client=client)
+    else:
+        raise MlflowException(
+            "Got no credential type when attempting to "
+            "access model version files in Unity Catalog. Try upgrading to the latest "
+            "version of the MLflow Python client."
+        )
+
+
+def _parse_aws_sse_credential(scoped_token: TemporaryCredentials):
+    encryption_details = scoped_token.encryption_details
+    if not encryption_details:
+        return {}
+
+    if encryption_details.WhichOneof("encryption_details_type") != "sse_encryption_details":
+        return {}
+
+    sse_encryption_details = encryption_details.sse_encryption_details
+
+    if sse_encryption_details.algorithm == SseEncryptionAlgorithm.AWS_SSE_S3:
+        return {
+            "ServerSideEncryption": "AES256",
+        }
+    if sse_encryption_details.algorithm == SseEncryptionAlgorithm.AWS_SSE_KMS:
+        key_id = sse_encryption_details.aws_kms_key_arn.split("/")[-1]
+        return {
+            "ServerSideEncryption": "aws:kms",
+            "SSEKMSKeyId": key_id,
+        }
+    else:
+        return {}
+
+
+def get_full_name_from_sc(name, spark) -> str:
+    """
+    Constructs the full name of a registered model using the active catalog and schema in a spark
+    session / context.
+
+    Args:
+        name: The model name provided by the user.
+        spark: The active spark session.
+    """
+    num_levels = len(name.split("."))
+    if num_levels >= 3 or spark is None:
+        return name
+    catalog = spark.sql(_ACTIVE_CATALOG_QUERY).collect()[0]["catalog"]
+    # return the user provided name if the catalog is the hive metastore default
+    if catalog in {"spark_catalog", "hive_metastore"}:
+        return name
+    if num_levels == 2:
+        return f"{catalog}.{name}"
+    schema = spark.sql(_ACTIVE_SCHEMA_QUERY).collect()[0]["schema"]
+    return f"{catalog}.{schema}.{name}"
+
+
+def is_databricks_sdk_models_artifact_repository_enabled(host_creds):
+    # Return early if the environment variable is set to use the SDK models artifact repository
+    if MLFLOW_USE_DATABRICKS_SDK_MODEL_ARTIFACTS_REPO_FOR_UC.defined:
+        return MLFLOW_USE_DATABRICKS_SDK_MODEL_ARTIFACTS_REPO_FOR_UC.get()
+
+    endpoint, method = _METHOD_TO_INFO[IsDatabricksSdkModelsArtifactRepositoryEnabledRequest]
+    req_body = message_to_json(IsDatabricksSdkModelsArtifactRepositoryEnabledRequest())
+    response_proto = IsDatabricksSdkModelsArtifactRepositoryEnabledResponse()
+
+    try:
+        resp = call_endpoint(
+            host_creds=host_creds,
+            endpoint=endpoint,
+            method=method,
+            json_body=req_body,
+            response_proto=response_proto,
+        )
+        return resp.is_databricks_sdk_models_artifact_repository_enabled
+    except Exception as e:
+        _logger.warning(
+            "Failed to confirm if DatabricksSDKModelsArtifactRepository should be used; "
+            f"falling back to default. Error: {e}"
+        )
+    return False
+
+
+def emit_model_version_lineage(host_creds, name, version, entities, direction):
+    endpoint, method = _METHOD_TO_INFO[EmitModelVersionLineageRequest]
+
+    req_body = message_to_json(
+        EmitModelVersionLineageRequest(
+            name=name,
+            version=version,
+            model_version_lineage_info=ModelVersionLineageInfo(
+                entities=entities,
+                direction=direction,
+            ),
+        )
+    )
+    response_proto = EmitModelVersionLineageResponse()
+    try:
+        call_endpoint(
+            host_creds=host_creds,
+            endpoint=endpoint,
+            method=method,
+            json_body=req_body,
+            response_proto=response_proto,
+        )
+    except Exception as e:
+        _logger.warning(f"Failed to emit best-effort model version lineage. Error: {e}")

venv/lib/python3.9/site-packages/mlflow/utils/annotations.py

@@ -0,0 +1,206 @@
+import inspect
+import re
+import types
+import warnings
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar, Union
+
+C = TypeVar("C", bound=Callable[..., Any])
+
+
+def _get_min_indent_of_docstring(docstring_str: str) -> str:
+    """
+    Get the minimum indentation string of a docstring, based on the assumption
+    that the closing triple quote for multiline comments must be on a new line.
+    Note that based on ruff rule D209, the closing triple quote for multiline
+    comments must be on a new line.
+
+    Args:
+        docstring_str: string with docstring
+
+    Returns:
+        Whitespace corresponding to the indent of a docstring.
+    """
+
+    if not docstring_str or "\n" not in docstring_str:
+        return ""
+
+    return re.match(r"^\s*", docstring_str.rsplit("\n", 1)[-1]).group()
+
+
+def experimental(api_or_type: Union[C, str]) -> C:
+    """Decorator / decorator creator for marking APIs experimental in the docstring.
+
+    Args:
+        api_or_type: An API to mark, or an API typestring for which to generate a decorator.
+
+    Returns:
+        Decorated API (if a ``api_or_type`` is an API) or a function that decorates
+        the specified API type (if ``api_or_type`` is a typestring).
+    """
+    if isinstance(api_or_type, str):
+
+        def f(api: C) -> C:
+            return _experimental(api=api, api_type=api_or_type)
+
+        return f
+    elif inspect.isclass(api_or_type):
+        return _experimental(api=api_or_type, api_type="class")
+    elif inspect.isfunction(api_or_type):
+        return _experimental(api=api_or_type, api_type="function")
+    elif isinstance(api_or_type, (property, types.MethodType)):
+        return _experimental(api=api_or_type, api_type="property")
+    else:
+        return _experimental(api=api_or_type, api_type=str(type(api_or_type)))
+
+
+def _experimental(api: C, api_type: str) -> C:
+    indent = _get_min_indent_of_docstring(api.__doc__) if api.__doc__ else ""
+    notice = (
+        indent + f".. Note:: Experimental: This {api_type} may change or "
+        "be removed in a future release without warning.\n\n"
+    )
+    if api_type == "property":
+        api.__doc__ = api.__doc__ + "\n\n" + notice if api.__doc__ else notice
+    else:
+        api.__doc__ = notice + api.__doc__ if api.__doc__ else notice
+    return api
+
+
+def developer_stable(func):
+    """
+    The API marked here as `@developer_stable` has certain protections associated with future
+    development work.
+    Classes marked with this decorator implicitly apply this status to all methods contained within
+    them.
+
+    APIs that are annotated with this decorator are guaranteed (except in cases of notes below) to:
+    - maintain backwards compatibility such that earlier versions of any MLflow client, cli, or
+      server will not have issues with any changes being made to them from an interface perspective.
+    - maintain a consistent contract with respect to existing named arguments such that
+      modifications will not alter or remove an existing named argument.
+    - maintain implied or declared types of arguments within its signature.
+    - maintain consistent behavior with respect to return types.
+
+    Note: Should an API marked as `@developer_stable` require a modification for enhanced feature
+      functionality, a deprecation warning will be added to the API well in advance of its
+      modification.
+
+    Note: Should an API marked as `@developer_stable` require patching for any security reason,
+      advanced notice is not guaranteed and the labeling of such API as stable will be ignored
+      for the sake of such a security patch.
+
+    """
+    return func
+
+
+_DEPRECATED_MARK_ATTR_NAME = "__deprecated"
+
+
+def mark_deprecated(func):
+    """
+    Mark a function as deprecated by setting a private attribute on it.
+    """
+    setattr(func, _DEPRECATED_MARK_ATTR_NAME, True)
+
+
+def is_marked_deprecated(func):
+    """
+    Is the function marked as deprecated.
+    """
+    return getattr(func, _DEPRECATED_MARK_ATTR_NAME, False)
+
+
+def deprecated(
+    alternative: Optional[str] = None, since: Optional[str] = None, impact: Optional[str] = None
+):
+    """Annotation decorator for marking APIs as deprecated in docstrings and raising a warning if
+    called.
+
+    Args:
+        alternative: The name of a superseded replacement function, method,
+            or class to use in place of the deprecated one.
+        since: A version designator defining during which release the function,
+            method, or class was marked as deprecated.
+        impact: Indication of whether the method, function, or class will be
+            removed in a future release.
+
+    Returns:
+        Decorated function or class.
+    """
+
+    def deprecated_decorator(obj):
+        since_str = f" since {since}" if since else ""
+        impact_str = impact if impact else "This method will be removed in a future release."
+
+        qual_name = f"{obj.__module__}.{obj.__qualname__}"
+        notice = f"``{qual_name}`` is deprecated{since_str}. {impact_str}"
+        if alternative and alternative.strip():
+            notice += f" Use ``{alternative}`` instead."
+
+        if inspect.isclass(obj):
+            original_init = obj.__init__
+
+            @wraps(original_init)
+            def new_init(self, *args, **kwargs):
+                warnings.warn(notice, category=FutureWarning, stacklevel=2)
+                original_init(self, *args, **kwargs)
+
+            obj.__init__ = new_init
+
+            if obj.__doc__:
+                obj.__doc__ = f".. Warning:: {notice}\n{obj.__doc__}"
+            else:
+                obj.__doc__ = f".. Warning:: {notice}"
+
+            mark_deprecated(obj)
+            return obj
+
+        elif isinstance(obj, (types.FunctionType, types.MethodType)):
+
+            @wraps(obj)
+            def deprecated_func(*args, **kwargs):
+                warnings.warn(notice, category=FutureWarning, stacklevel=2)
+                return obj(*args, **kwargs)
+
+            if obj.__doc__:
+                indent = _get_min_indent_of_docstring(obj.__doc__)
+                deprecated_func.__doc__ = f"{indent}.. Warning:: {notice}\n{obj.__doc__}"
+            else:
+                deprecated_func.__doc__ = f".. Warning:: {notice}"
+
+            mark_deprecated(deprecated_func)
+            return deprecated_func
+
+        else:
+            return obj
+
+    return deprecated_decorator
+
+
+def keyword_only(func):
+    """A decorator that forces keyword arguments in the wrapped method."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        if len(args) > 0:
+            raise TypeError(f"Method {func.__name__} only takes keyword arguments.")
+        return func(**kwargs)
+
+    indent = _get_min_indent_of_docstring(wrapper.__doc__) if wrapper.__doc__ else ""
+    notice = indent + ".. note:: This method requires all argument be specified by keyword.\n"
+    wrapper.__doc__ = notice + wrapper.__doc__ if wrapper.__doc__ else notice
+
+    return wrapper
+
+
+def filter_user_warnings_once(func):
+    """A decorator that filter user warnings to only show once in the wrapped method."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with warnings.catch_warnings():
+            warnings.simplefilter("once", category=UserWarning)
+            return func(*args, **kwargs)
+
+    return wrapper

venv/lib/python3.9/site-packages/mlflow/utils/arguments_utils.py

@@ -0,0 +1,16 @@
+import inspect
+
+
+def _get_arg_names(f):
+    """Get the argument names of a function.
+
+    Args:
+        f: A function.
+
+    Returns:
+        A list of argument names.
+
+    """
+    # `inspect.getargspec` or `inspect.getfullargspec` doesn't work properly for a wrapped function.
+    # See https://hynek.me/articles/decorators#mangled-signatures for details.
+    return list(inspect.signature(f).parameters.keys())

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/__init__.py

@@ -0,0 +1 @@
+from mlflow.utils.async_logging import run_operations  # noqa: F401

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/async_artifacts_logging_queue.py

@@ -0,0 +1,258 @@
+"""
+Defines an AsyncArtifactsLoggingQueue that provides async fashion artifact writes using
+queue based approach.
+"""
+
+import atexit
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, Callable, Union
+
+from mlflow.utils.async_logging.run_artifact import RunArtifact
+from mlflow.utils.async_logging.run_operations import RunOperations
+
+if TYPE_CHECKING:
+    import PIL.Image
+
+_logger = logging.getLogger(__name__)
+
+
+class AsyncArtifactsLoggingQueue:
+    """
+    This is a queue based run data processor that queue incoming data and process it using a single
+    worker thread. This class is used to process artifacts saving in async fashion.
+
+    Args:
+        logging_func: A callable function that takes in three arguments:
+            - filename: The name of the artifact file.
+            - artifact_path: The path to the artifact.
+            - artifact: The artifact to be logged.
+    """
+
+    def __init__(
+        self, artifact_logging_func: Callable[[str, str, Union["PIL.Image.Image"]], None]
+    ) -> None:
+        self._queue: Queue[RunArtifact] = Queue()
+        self._lock = threading.RLock()
+        self._artifact_logging_func = artifact_logging_func
+
+        self._stop_data_logging_thread_event = threading.Event()
+        self._is_activated = False
+
+    def _at_exit_callback(self) -> None:
+        """Callback function to be executed when the program is exiting.
+
+        Stops the data processing thread and waits for the queue to be drained. Finally, shuts down
+        the thread pools used for data logging and artifact processing status check.
+        """
+        try:
+            # Stop the data processing thread
+            self._stop_data_logging_thread_event.set()
+            # Waits till logging queue is drained.
+            self._artifact_logging_thread.join()
+            self._artifact_logging_worker_threadpool.shutdown(wait=True)
+            self._artifact_status_check_threadpool.shutdown(wait=True)
+        except Exception as e:
+            _logger.error(f"Encountered error while trying to finish logging: {e}")
+
+    def flush(self) -> None:
+        """Flush the async logging queue.
+
+        Calling this method will flush the queue to ensure all the data are logged.
+        """
+        # Stop the data processing thread.
+        self._stop_data_logging_thread_event.set()
+        # Waits till logging queue is drained.
+        self._artifact_logging_thread.join()
+        self._artifact_logging_worker_threadpool.shutdown(wait=True)
+        self._artifact_status_check_threadpool.shutdown(wait=True)
+
+        # Restart the thread to listen to incoming data after flushing.
+        self._stop_data_logging_thread_event.clear()
+        self._set_up_logging_thread()
+
+    def _logging_loop(self) -> None:
+        """
+        Continuously logs run data until `self._continue_to_process_data` is set to False.
+        If an exception occurs during logging, a `MlflowException` is raised.
+        """
+        try:
+            while not self._stop_data_logging_thread_event.is_set():
+                self._log_artifact()
+            # Drain the queue after the stop event is set.
+            while not self._queue.empty():
+                self._log_artifact()
+        except Exception as e:
+            from mlflow.exceptions import MlflowException
+
+            raise MlflowException(f"Exception inside the run data logging thread: {e}")
+
+    def _log_artifact(self) -> None:
+        """Process the run's artifacts in the running runs queues.
+
+        For each run in the running runs queues, this method retrieves the next artifact of run
+        from the queue and processes it by calling the `_artifact_logging_func` method with the run
+        ID and artifact. If the artifact is empty, it is skipped. After processing the artifact,
+        the processed watermark is updated and the artifact event is set.
+        If an exception occurs during processing, the exception is logged and the artifact event
+        is set with the exception. If the queue is empty, it is ignored.
+        """
+        try:
+            run_artifact = self._queue.get(timeout=1)
+        except Empty:
+            # Ignore empty queue exception
+            return
+
+        def logging_func(run_artifact):
+            try:
+                self._artifact_logging_func(
+                    filename=run_artifact.filename,
+                    artifact_path=run_artifact.artifact_path,
+                    artifact=run_artifact.artifact,
+                )
+
+                # Signal the artifact processing is done.
+                run_artifact.completion_event.set()
+
+            except Exception as e:
+                _logger.error(f"Failed to log artifact {run_artifact.filename}. Exception: {e}")
+                run_artifact.exception = e
+                run_artifact.completion_event.set()
+
+        self._artifact_logging_worker_threadpool.submit(logging_func, run_artifact)
+
+    def _wait_for_artifact(self, artifact: RunArtifact) -> None:
+        """Wait for given artifacts to be processed by the logging thread.
+
+        Args:
+            artifact: The artifact to wait for.
+
+        Raises:
+            Exception: If an exception occurred while processing the artifact.
+        """
+        artifact.completion_event.wait()
+        if artifact.exception:
+            raise artifact.exception
+
+    def __getstate__(self):
+        """Return the state of the object for pickling.
+
+        This method is called by the `pickle` module when the object is being pickled. It returns a
+        dictionary containing the object's state, with non-picklable attributes removed.
+
+        Returns:
+            dict: A dictionary containing the object's state.
+        """
+        state = self.__dict__.copy()
+        del state["_queue"]
+        del state["_lock"]
+        del state["_is_activated"]
+
+        if "_stop_data_logging_thread_event" in state:
+            del state["_stop_data_logging_thread_event"]
+        if "_artifact_logging_thread" in state:
+            del state["_artifact_logging_thread"]
+        if "_artifact_logging_worker_threadpool" in state:
+            del state["_artifact_logging_worker_threadpool"]
+        if "_artifact_status_check_threadpool" in state:
+            del state["_artifact_status_check_threadpool"]
+
+        return state
+
+    def __setstate__(self, state):
+        """Set the state of the object from a given state dictionary.
+
+        It pops back the removed non-picklable attributes from `self.__getstate__()`.
+
+        Args:
+            state (dict): A dictionary containing the state of the object.
+
+        Returns:
+            None
+        """
+        self.__dict__.update(state)
+        self._queue = Queue()
+        self._lock = threading.RLock()
+        self._is_activated = False
+        self._artifact_logging_thread = None
+        self._artifact_logging_worker_threadpool = None
+        self._artifact_status_check_threadpool = None
+        self._stop_data_logging_thread_event = threading.Event()
+
+    def log_artifacts_async(self, filename, artifact_path, artifact) -> RunOperations:
+        """Asynchronously logs runs artifacts.
+
+        Args:
+            filename: Filename of the artifact to be logged.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+            artifact: The artifact to be logged.
+
+        Returns:
+            mlflow.utils.async_utils.RunOperations: An object that encapsulates the
+                asynchronous operation of logging the artifact of run data.
+                The object contains a list of `concurrent.futures.Future` objects that can be used
+                to check the status of the operation and retrieve any exceptions
+                that occurred during the operation.
+        """
+        from mlflow import MlflowException
+
+        if not self._is_activated:
+            raise MlflowException("AsyncArtifactsLoggingQueue is not activated.")
+        artifact = RunArtifact(
+            filename=filename,
+            artifact_path=artifact_path,
+            artifact=artifact,
+            completion_event=threading.Event(),
+        )
+        self._queue.put(artifact)
+        operation_future = self._artifact_status_check_threadpool.submit(
+            self._wait_for_artifact, artifact
+        )
+        return RunOperations(operation_futures=[operation_future])
+
+    def is_active(self) -> bool:
+        return self._is_activated
+
+    def _set_up_logging_thread(self) -> None:
+        """Sets up the logging thread.
+
+        If the logging thread is already set up, this method does nothing.
+        """
+        with self._lock:
+            self._artifact_logging_thread = threading.Thread(
+                target=self._logging_loop,
+                name="MLflowAsyncArtifactsLoggingLoop",
+                daemon=True,
+            )
+            self._artifact_logging_worker_threadpool = ThreadPoolExecutor(
+                max_workers=5,
+                thread_name_prefix="MLflowArtifactsLoggingWorkerPool",
+            )
+
+            self._artifact_status_check_threadpool = ThreadPoolExecutor(
+                max_workers=5,
+                thread_name_prefix="MLflowAsyncArtifactsLoggingStatusCheck",
+            )
+            self._artifact_logging_thread.start()
+
+    def activate(self) -> None:
+        """Activates the async logging queue
+
+        1. Initializes queue draining thread.
+        2. Initializes threads for checking the status of logged artifact.
+        3. Registering an atexit callback to ensure that any remaining log data
+            is flushed before the program exits.
+
+        If the queue is already activated, this method does nothing.
+        """
+        with self._lock:
+            if self._is_activated:
+                return
+
+            self._set_up_logging_thread()
+            atexit.register(self._at_exit_callback)
+
+            self._is_activated = True

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/async_logging_queue.py

@@ -0,0 +1,366 @@
+"""
+Defines an AsyncLoggingQueue that provides async fashion logging of metrics/tags/params using
+queue based approach.
+"""
+
+import atexit
+import enum
+import logging
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from queue import Empty, Queue
+from typing import Callable
+
+from mlflow.entities.metric import Metric
+from mlflow.entities.param import Param
+from mlflow.entities.run_tag import RunTag
+from mlflow.environment_variables import (
+    MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS,
+    MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE,
+)
+from mlflow.utils.async_logging.run_batch import RunBatch
+from mlflow.utils.async_logging.run_operations import RunOperations
+
+_logger = logging.getLogger(__name__)
+
+
+ASYNC_LOGGING_WORKER_THREAD_PREFIX = "MLflowBatchLoggingWorkerPool"
+ASYNC_LOGGING_STATUS_CHECK_THREAD_PREFIX = "MLflowAsyncLoggingStatusCheck"
+
+
+class QueueStatus(enum.Enum):
+    """Status of the async queue"""
+
+    # The queue is listening to new data and logging enqueued data to MLflow.
+    ACTIVE = 1
+    # The queue is not listening to new data, but still logging enqueued data to MLflow.
+    TEAR_DOWN = 2
+    # The queue is neither listening to new data or logging enqueued data to MLflow.
+    IDLE = 3
+
+
+_MAX_ITEMS_PER_BATCH = 1000
+_MAX_PARAMS_PER_BATCH = 100
+_MAX_TAGS_PER_BATCH = 100
+
+
+class AsyncLoggingQueue:
+    """
+    This is a queue based run data processor that queues incoming batches and processes them using
+    single worker thread.
+    """
+
+    def __init__(
+        self, logging_func: Callable[[str, list[Metric], list[Param], list[RunTag]], None]
+    ) -> None:
+        """Initializes an AsyncLoggingQueue object.
+
+        Args:
+            logging_func: A callable function that takes in four arguments: a string
+                representing the run_id, a list of Metric objects,
+                a list of Param objects, and a list of RunTag objects.
+        """
+        self._queue = Queue()
+        self._lock = threading.RLock()
+        self._logging_func = logging_func
+
+        self._stop_data_logging_thread_event = threading.Event()
+        self._status = QueueStatus.IDLE
+
+    def _at_exit_callback(self) -> None:
+        """Callback function to be executed when the program is exiting.
+
+        Stops the data processing thread and waits for the queue to be drained. Finally, shuts down
+        the thread pools used for data logging and batch processing status check.
+        """
+        try:
+            # Stop the data processing thread
+            self._stop_data_logging_thread_event.set()
+            # Waits till logging queue is drained.
+            self._batch_logging_thread.join()
+            self._batch_logging_worker_threadpool.shutdown(wait=True)
+            self._batch_status_check_threadpool.shutdown(wait=True)
+        except Exception as e:
+            _logger.error(f"Encountered error while trying to finish logging: {e}")
+
+    def end_async_logging(self) -> None:
+        with self._lock:
+            # Stop the data processing thread.
+            self._stop_data_logging_thread_event.set()
+            # Waits till logging queue is drained.
+            self._batch_logging_thread.join()
+            # Set the status to tear down. The worker threads will still process
+            # the remaining data.
+            self._status = QueueStatus.TEAR_DOWN
+            # Clear the status to avoid blocking next logging.
+            self._stop_data_logging_thread_event.clear()
+
+    def shut_down_async_logging(self) -> None:
+        """
+        Shut down the async logging queue and wait for the queue to be drained.
+        Use this method if the async logging should be terminated.
+        """
+        self.end_async_logging()
+        self._batch_logging_worker_threadpool.shutdown(wait=True)
+        self._batch_status_check_threadpool.shutdown(wait=True)
+        self._status = QueueStatus.IDLE
+
+    def flush(self) -> None:
+        """
+        Flush the async logging queue and restart thread to listen
+        to incoming data after flushing.
+
+        Calling this method will flush the queue to ensure all the data are logged.
+        """
+        self.shut_down_async_logging()
+        # Reinitialize the logging thread and set the status to active.
+        self.activate()
+
+    def _logging_loop(self) -> None:
+        """
+        Continuously logs run data until `self._continue_to_process_data` is set to False.
+        If an exception occurs during logging, a `MlflowException` is raised.
+        """
+        try:
+            while not self._stop_data_logging_thread_event.is_set():
+                self._log_run_data()
+            # Drain the queue after the stop event is set.
+            while not self._queue.empty():
+                self._log_run_data()
+        except Exception as e:
+            from mlflow.exceptions import MlflowException
+
+            raise MlflowException(f"Exception inside the run data logging thread: {e}")
+
+    def _fetch_batch_from_queue(self) -> list[RunBatch]:
+        """Fetches a batch of run data from the queue.
+
+        Returns:
+            RunBatch: A batch of run data.
+        """
+        batches = []
+        if self._queue.empty():
+            return batches
+        queue_size = self._queue.qsize()  # Estimate the queue's size.
+        merged_batch = self._queue.get()
+        for i in range(queue_size - 1):
+            if self._queue.empty():
+                # `queue_size` is an estimate, so we need to check if the queue is empty.
+                break
+            batch = self._queue.get()
+
+            if (
+                merged_batch.run_id != batch.run_id
+                or (
+                    len(merged_batch.metrics + merged_batch.params + merged_batch.tags)
+                    + len(batch.metrics + batch.params + batch.tags)
+                )
+                >= _MAX_ITEMS_PER_BATCH
+                or len(merged_batch.params) + len(batch.params) >= _MAX_PARAMS_PER_BATCH
+                or len(merged_batch.tags) + len(batch.tags) >= _MAX_TAGS_PER_BATCH
+            ):
+                # Make a new batch if the run_id is different or the batch is full.
+                batches.append(merged_batch)
+                merged_batch = batch
+            else:
+                merged_batch.add_child_batch(batch)
+                merged_batch.params.extend(batch.params)
+                merged_batch.tags.extend(batch.tags)
+                merged_batch.metrics.extend(batch.metrics)
+
+        batches.append(merged_batch)
+        return batches
+
+    def _log_run_data(self) -> None:
+        """Process the run data in the running runs queues.
+
+        For each run in the running runs queues, this method retrieves the next batch of run data
+        from the queue and processes it by calling the `_processing_func` method with the run ID,
+        metrics, parameters, and tags in the batch. If the batch is empty, it is skipped. After
+        processing the batch, the processed watermark is updated and the batch event is set.
+        If an exception occurs during processing, the exception is logged and the batch event is set
+        with the exception. If the queue is empty, it is ignored.
+
+        Returns: None
+        """
+        async_logging_buffer_seconds = MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS.get()
+        try:
+            if async_logging_buffer_seconds:
+                self._stop_data_logging_thread_event.wait(async_logging_buffer_seconds)
+                run_batches = self._fetch_batch_from_queue()
+            else:
+                run_batches = [self._queue.get(timeout=1)]
+        except Empty:
+            # Ignore empty queue exception
+            return
+
+        def logging_func(run_batch):
+            try:
+                self._logging_func(
+                    run_id=run_batch.run_id,
+                    metrics=run_batch.metrics,
+                    params=run_batch.params,
+                    tags=run_batch.tags,
+                )
+            except Exception as e:
+                _logger.error(f"Run Id {run_batch.run_id}: Failed to log run data: Exception: {e}")
+                run_batch.exception = e
+            finally:
+                run_batch.complete()
+
+        for run_batch in run_batches:
+            try:
+                self._batch_logging_worker_threadpool.submit(logging_func, run_batch)
+            except Exception as e:
+                _logger.error(
+                    f"Failed to submit batch for logging: {e}. Usually this means you are not "
+                    "shutting down MLflow properly before exiting. Please make sure you are using "
+                    "context manager, e.g., `with mlflow.start_run():` or call `mlflow.end_run()`"
+                    "explicitly to terminate MLflow logging before exiting."
+                )
+                run_batch.exception = e
+                run_batch.complete()
+
+    def _wait_for_batch(self, batch: RunBatch) -> None:
+        """Wait for the given batch to be processed by the logging thread.
+
+        Args:
+            batch: The batch to wait for.
+
+        Raises:
+            Exception: If an exception occurred while processing the batch.
+        """
+        batch.completion_event.wait()
+        if batch.exception:
+            raise batch.exception
+
+    def __getstate__(self):
+        """Return the state of the object for pickling.
+
+        This method is called by the `pickle` module when the object is being pickled. It returns a
+        dictionary containing the object's state, with non-picklable attributes removed.
+
+        Returns:
+            dict: A dictionary containing the object's state.
+        """
+        state = self.__dict__.copy()
+        del state["_queue"]
+        del state["_lock"]
+        del state["_status"]
+
+        if "_run_data_logging_thread" in state:
+            del state["_run_data_logging_thread"]
+        if "_stop_data_logging_thread_event" in state:
+            del state["_stop_data_logging_thread_event"]
+        if "_batch_logging_thread" in state:
+            del state["_batch_logging_thread"]
+        if "_batch_logging_worker_threadpool" in state:
+            del state["_batch_logging_worker_threadpool"]
+        if "_batch_status_check_threadpool" in state:
+            del state["_batch_status_check_threadpool"]
+
+        return state
+
+    def __setstate__(self, state):
+        """Set the state of the object from a given state dictionary.
+
+        It pops back the removed non-picklable attributes from `self.__getstate__()`.
+
+        Args:
+            state (dict): A dictionary containing the state of the object.
+
+        Returns:
+            None
+        """
+        self.__dict__.update(state)
+        self._queue = Queue()
+        self._lock = threading.RLock()
+        self._status = QueueStatus.IDLE
+        self._batch_logging_thread = None
+        self._batch_logging_worker_threadpool = None
+        self._batch_status_check_threadpool = None
+        self._stop_data_logging_thread_event = threading.Event()
+
+    def log_batch_async(
+        self, run_id: str, params: list[Param], tags: list[RunTag], metrics: list[Metric]
+    ) -> RunOperations:
+        """Asynchronously logs a batch of run data (parameters, tags, and metrics).
+
+        Args:
+            run_id (str): The ID of the run to log data for.
+            params (list[mlflow.entities.Param]): A list of parameters to log for the run.
+            tags (list[mlflow.entities.RunTag]): A list of tags to log for the run.
+            metrics (list[mlflow.entities.Metric]): A list of metrics to log for the run.
+
+        Returns:
+            mlflow.utils.async_utils.RunOperations: An object that encapsulates the
+                asynchronous operation of logging the batch of run data.
+                The object contains a list of `concurrent.futures.Future` objects that can be used
+                to check the status of the operation and retrieve any exceptions
+                that occurred during the operation.
+        """
+        from mlflow import MlflowException
+
+        if not self.is_active():
+            raise MlflowException("AsyncLoggingQueue is not activated.")
+        batch = RunBatch(
+            run_id=run_id,
+            params=params,
+            tags=tags,
+            metrics=metrics,
+            completion_event=threading.Event(),
+        )
+        self._queue.put(batch)
+        operation_future = self._batch_status_check_threadpool.submit(self._wait_for_batch, batch)
+        return RunOperations(operation_futures=[operation_future])
+
+    def is_active(self) -> bool:
+        return self._status == QueueStatus.ACTIVE
+
+    def is_idle(self) -> bool:
+        return self._status == QueueStatus.IDLE
+
+    def _set_up_logging_thread(self) -> None:
+        """
+        Sets up the logging thread.
+
+        This method shouldn't be called directly without shutting down the async
+        logging first if an existing async logging exists, otherwise it might
+        hang the program.
+        """
+        with self._lock:
+            self._batch_logging_thread = threading.Thread(
+                target=self._logging_loop,
+                name="MLflowAsyncLoggingLoop",
+                daemon=True,
+            )
+            self._batch_logging_worker_threadpool = ThreadPoolExecutor(
+                max_workers=MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.get() or 10,
+                thread_name_prefix=ASYNC_LOGGING_WORKER_THREAD_PREFIX,
+            )
+
+            self._batch_status_check_threadpool = ThreadPoolExecutor(
+                max_workers=MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.get() or 10,
+                thread_name_prefix=ASYNC_LOGGING_STATUS_CHECK_THREAD_PREFIX,
+            )
+
+            self._batch_logging_thread.start()
+
+    def activate(self) -> None:
+        """Activates the async logging queue
+
+        1. Initializes queue draining thread.
+        2. Initializes threads for checking the status of logged batch.
+        3. Registering an atexit callback to ensure that any remaining log data
+            is flushed before the program exits.
+
+        If the queue is already activated, this method does nothing.
+        """
+        with self._lock:
+            if self.is_active():
+                return
+
+            self._set_up_logging_thread()
+            atexit.register(self._at_exit_callback)
+
+            self._status = QueueStatus.ACTIVE

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/run_artifact.py

@@ -0,0 +1,38 @@
+import threading
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import PIL
+
+
+class RunArtifact:
+    def __init__(
+        self,
+        filename: str,
+        artifact_path: str,
+        artifact: Union["PIL.Image.Image"],
+        completion_event: threading.Event,
+    ) -> None:
+        """Initializes an instance of `RunArtifacts`.
+
+        Args:
+            filename: Filename of the artifact to be logged
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+            artifact: The artifact to be logged.
+            completion_event: A threading.Event object.
+        """
+        self.filename = filename
+        self.artifact_path = artifact_path
+        self.artifact = artifact
+        self.completion_event = completion_event
+        self._exception = None
+
+    @property
+    def exception(self):
+        """Exception raised during logging the batch."""
+        return self._exception
+
+    @exception.setter
+    def exception(self, exception):
+        self._exception = exception

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/run_batch.py

@@ -0,0 +1,58 @@
+import threading
+from typing import Optional
+
+from mlflow.entities.metric import Metric
+from mlflow.entities.param import Param
+from mlflow.entities.run_tag import RunTag
+
+
+class RunBatch:
+    def __init__(
+        self,
+        run_id: str,
+        params: Optional[list["Param"]] = None,
+        tags: Optional[list["RunTag"]] = None,
+        metrics: Optional[list["Metric"]] = None,
+        completion_event: Optional[threading.Event] = None,
+    ):
+        """Initializes an instance of `RunBatch`.
+
+        Args:
+            run_id: The ID of the run.
+            params: A list of parameters. Default is None.
+            tags: A list of tags. Default is None.
+            metrics: A list of metrics. Default is None.
+            completion_event: A threading.Event object. Default is None.
+        """
+        self.run_id = run_id
+        self.params = params or []
+        self.tags = tags or []
+        self.metrics = metrics or []
+        self.completion_event = completion_event
+        self._exception = None
+        self.child_batches = []
+
+    @property
+    def exception(self):
+        """Exception raised during logging the batch."""
+        return self._exception
+
+    @exception.setter
+    def exception(self, exception):
+        self._exception = exception
+
+    def add_child_batch(self, child_batch):
+        """Add a child batch to the current batch.
+
+        This is useful when merging child batches into a parent batch. Child batches are kept so
+        that we can properly notify the system when child batches have been processed.
+        """
+        self.child_batches.append(child_batch)
+
+    def complete(self):
+        """Mark the batch as completed."""
+        if self.completion_event:
+            self.completion_event.set()
+
+        for child_batch in self.child_batches:
+            child_batch.complete()

venv/lib/python3.9/site-packages/mlflow/utils/async_logging/run_operations.py

@@ -0,0 +1,49 @@
+class RunOperations:
+    """Class that helps manage the futures of MLflow async logging."""
+
+    def __init__(self, operation_futures):
+        self._operation_futures = operation_futures or []
+
+    def wait(self):
+        """Blocks on completion of all futures."""
+        from mlflow.exceptions import MlflowException
+
+        failed_operations = []
+        for future in self._operation_futures:
+            try:
+                future.result()
+            except Exception as e:
+                failed_operations.append(e)
+
+        if len(failed_operations) > 0:
+            raise MlflowException(
+                "The following failures occurred while performing one or more async logging "
+                f"operations: {failed_operations}"
+            )
+
+
+def get_combined_run_operations(run_operations_list: list[RunOperations]) -> RunOperations:
+    """Combine a list of RunOperations objects into a single RunOperations object.
+
+    Given a list of `RunOperations`, returns a single `RunOperations` object that represents the
+    combined set of operations. If the input list is empty, returns None. If the input list
+    contains only one element, returns that element. Otherwise, creates a new `RunOperations`
+    object that combines the operation futures from each input RunOperations object.
+
+    Args:
+        run_operations_list: A list of `RunOperations` objects to combine.
+
+    Returns:
+        A single `RunOperations` object that represents the combined set of operations.
+    """
+    if not run_operations_list:
+        return None
+    if len(run_operations_list) == 1:
+        return run_operations_list[0]
+
+    if len(run_operations_list) > 1:
+        operation_futures = []
+        for run_operations in run_operations_list:
+            if run_operations and run_operations._operation_futures:
+                operation_futures.extend(run_operations._operation_futures)
+        return RunOperations(operation_futures)

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/__init__.py

@@ -0,0 +1,731 @@
+import contextlib
+import importlib
+import inspect
+import logging
+import threading
+import time
+from typing import Any, Callable, Optional
+
+import mlflow
+from mlflow.entities import Metric
+from mlflow.tracking.client import MlflowClient
+from mlflow.utils.validation import MAX_METRICS_PER_BATCH
+
+# Define the module-level logger for autologging utilities before importing utilities defined in
+# submodules (e.g., `safety`, `events`) that depend on the module-level logger. Add the `noqa: E402`
+# comment after each subsequent import to ignore "import not at top of file" code style errors
+_logger = logging.getLogger(__name__)
+
+# Import autologging utilities used by this module
+from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS, FLAVOR_TO_MODULE_NAME
+from mlflow.utils.autologging_utils.client import MlflowAutologgingQueueingClient  # noqa: F401
+from mlflow.utils.autologging_utils.events import AutologgingEventLogger
+from mlflow.utils.autologging_utils.logging_and_warnings import (
+    MlflowEventsAndWarningsBehaviorGlobally,
+    NonMlflowWarningsBehaviorForCurrentThread,
+)
+
+# Wildcard import other autologging utilities (e.g. safety utilities, event logging utilities) used
+# in autologging integration implementations, which reference them via the
+# `mlflow.utils.autologging_utils` module
+from mlflow.utils.autologging_utils.safety import (  # noqa: F401
+    ExceptionSafeAbstractClass,
+    ExceptionSafeClass,
+    exception_safe_function_for_class,
+    is_testing,
+    picklable_exception_safe_function,
+    revert_patches,
+    safe_patch,
+    update_wrapper_extended,
+    with_managed_run,
+)
+from mlflow.utils.autologging_utils.versioning import (
+    get_min_max_version_and_pip_release,
+    is_flavor_supported_for_associated_package_versions,
+)
+
+INPUT_EXAMPLE_SAMPLE_ROWS = 5
+ENSURE_AUTOLOGGING_ENABLED_TEXT = (
+    "please ensure that autologging is enabled before constructing the dataset."
+)
+
+# Flag indicating whether autologging is globally disabled for all integrations.
+_AUTOLOGGING_GLOBALLY_DISABLED = False
+
+# Autologging config key indicating whether or not a particular autologging integration
+# was configured (i.e. its various `log_models`, `disable`, etc. configuration options
+# were set) via a call to `mlflow.autolog()`, rather than via a call to the integration-specific
+# autologging method (e.g., `mlflow.tensorflow.autolog()`, ...)
+AUTOLOGGING_CONF_KEY_IS_GLOBALLY_CONFIGURED = "globally_configured"
+
+# Dict mapping integration name to its config.
+AUTOLOGGING_INTEGRATIONS = {}
+
+# When the library version installed in the user's environment is outside of the supported
+# version range declared in `ml-package-versions.yml`, a warning message is issued to the user.
+# However, some libraries releases versions very frequently, and our configuration (updated on
+# MLflow release) cannot keep up with the pace, resulting in false alarms. Therefore, we
+# suppress warnings for certain libraries that are known to have frequent releases.
+_AUTOLOGGING_SUPPORTED_VERSION_WARNING_SUPPRESS_LIST = [
+    "langchain",
+    "llama_index",
+    "litellm",
+    "openai",
+    "dspy",
+    "autogen",
+    "gemini",
+    "anthropic",
+    "crewai",
+    "bedrock",
+]
+
+# Global lock for turning on / off autologging
+# Note "RLock" is required instead of plain lock, for avoid dead-lock
+_autolog_conf_global_lock = threading.RLock()
+
+_logger = logging.getLogger(__name__)
+
+
+def autologging_conf_lock(fn):
+    """
+    Apply a global lock on functions that enable / disable autologging.
+    """
+
+    def wrapper(*args, **kwargs):
+        with _autolog_conf_global_lock:
+            return fn(*args, **kwargs)
+
+    return update_wrapper_extended(wrapper, fn)
+
+
+def get_mlflow_run_params_for_fn_args(fn, args, kwargs, unlogged=None):
+    """Given arguments explicitly passed to a function, generate a dictionary of MLflow Run
+    parameter key / value pairs.
+
+    Args:
+        fn: function whose parameters are to be logged.
+        args: arguments explicitly passed into fn. If `fn` is defined on a class,
+            `self` should not be part of `args`; the caller is responsible for
+            filtering out `self` before calling this function.
+        kwargs: kwargs explicitly passed into fn.
+        unlogged: parameters not to be logged.
+
+    Returns:
+        A dictionary of MLflow Run parameter key / value pairs.
+    """
+    unlogged = unlogged or []
+    param_spec = inspect.signature(fn).parameters
+    # Filter out `self` from the signature under the assumption that it is not contained
+    # within the specified `args`, as stipulated by the documentation
+    relevant_params = [param for param in param_spec.values() if param.name != "self"]
+
+    # Fetch the parameter names for specified positional arguments from the function
+    # signature & create a mapping from positional argument name to specified value
+    params_to_log = {
+        param_info.name: param_val
+        for param_info, param_val in zip(list(relevant_params)[: len(args)], args)
+    }
+    # Add all user-specified keyword arguments to the set of parameters to log
+    params_to_log.update(kwargs)
+    # Add parameters that were not explicitly specified by the caller to the mapping,
+    # using their default values
+    params_to_log.update(
+        {
+            param.name: param.default
+            for param in list(relevant_params)[len(args) :]
+            if param.name not in kwargs
+        }
+    )
+    # Filter out any parameters that should not be logged, as specified by the `unlogged` parameter
+    return {key: value for key, value in params_to_log.items() if key not in unlogged}
+
+
+def log_fn_args_as_params(fn, args, kwargs, unlogged=None):
+    """Log arguments explicitly passed to a function as MLflow Run parameters to the current active
+    MLflow Run.
+
+    Args:
+        fn: function whose parameters are to be logged
+        args: arguments explicitly passed into fn. If `fn` is defined on a class,
+            `self` should not be part of `args`; the caller is responsible for
+            filtering out `self` before calling this function.
+        kwargs: kwargs explicitly passed into fn
+        unlogged: parameters not to be logged
+
+    Returns:
+        None
+
+    """
+    params_to_log = get_mlflow_run_params_for_fn_args(fn, args, kwargs, unlogged)
+    mlflow.log_params(params_to_log)
+
+
+class InputExampleInfo:
+    """
+    Stores info about the input example collection before it is needed.
+
+    For example, in xgboost and lightgbm, an InputExampleInfo object is attached to the dataset,
+    where its value is read later by the train method.
+
+    Exactly one of input_example or error_msg should be populated.
+    """
+
+    def __init__(self, input_example=None, error_msg=None):
+        self.input_example = input_example
+        self.error_msg = error_msg
+
+
+def resolve_input_example_and_signature(
+    get_input_example, infer_model_signature, log_input_example, log_model_signature, logger
+):
+    """Handles the logic of calling functions to gather the input example and infer the model
+    signature.
+
+    Args:
+        get_input_example: Function which returns an input example, usually sliced from a
+            dataset. This function can raise an exception, its message will be
+            shown to the user in a warning in the logs.
+        infer_model_signature: Function which takes an input example and returns the signature
+            of the inputs and outputs of the model. This function can raise
+            an exception, its message will be shown to the user in a warning
+            in the logs.
+        log_input_example: Whether to log errors while collecting the input example, and if it
+            succeeds, whether to return the input example to the user. We collect
+            it even if this parameter is False because it is needed for inferring
+            the model signature.
+        log_model_signature: Whether to infer and return the model signature.
+        logger: The logger instance used to log warnings to the user during input example
+            collection and model signature inference.
+
+    Returns:
+        A tuple of input_example and signature. Either or both could be None based on the
+        values of log_input_example and log_model_signature.
+
+    """
+
+    input_example = None
+    input_example_user_msg = None
+    input_example_failure_msg = None
+    if log_input_example or log_model_signature:
+        try:
+            input_example = get_input_example()
+        except Exception as e:
+            input_example_failure_msg = str(e)
+            input_example_user_msg = "Failed to gather input example: " + str(e)
+
+    model_signature = None
+    model_signature_user_msg = None
+    if log_model_signature:
+        try:
+            if input_example is None:
+                raise Exception(
+                    "could not sample data to infer model signature: " + input_example_failure_msg
+                )
+            model_signature = infer_model_signature(input_example)
+        except Exception as e:
+            model_signature_user_msg = "Failed to infer model signature: " + str(e)
+
+    # disable input_example signature inference in model logging if `log_model_signature`
+    # is set to `False` or signature inference in autologging fails
+    if (
+        model_signature is None
+        and input_example is not None
+        and (not log_model_signature or model_signature_user_msg is not None)
+    ):
+        model_signature = False
+
+    if log_input_example and input_example_user_msg is not None:
+        logger.warning(input_example_user_msg)
+    if log_model_signature and model_signature_user_msg is not None:
+        logger.warning(model_signature_user_msg)
+
+    return input_example if log_input_example else None, model_signature
+
+
+class BatchMetricsLogger:
+    """
+    The BatchMetricsLogger will log metrics in batch against an mlflow run.
+    If run_id is passed to to constructor then all recording and logging will
+    happen against that run_id.
+    If no run_id is passed into constructor, then the run ID will be fetched
+    from `mlflow.active_run()` each time `record_metrics()` or `flush()` is called; in this
+    case, callers must ensure that an active run is present before invoking
+    `record_metrics()` or `flush()`.
+    """
+
+    def __init__(self, run_id=None, tracking_uri=None):
+        self.run_id = run_id
+        self.client = MlflowClient(tracking_uri)
+
+        # data is an array of Metric objects
+        self.data = []
+        self.total_training_time = 0
+        self.total_log_batch_time = 0
+        self.previous_training_timestamp = None
+
+    def flush(self):
+        """
+        The metrics accumulated by BatchMetricsLogger will be batch logged to an MLflow run.
+        """
+        self._timed_log_batch()
+        self.data = []
+
+    def _timed_log_batch(self):
+        # Retrieving run_id from active mlflow run when run_id is empty.
+        current_run_id = mlflow.active_run().info.run_id if self.run_id is None else self.run_id
+
+        start = time.time()
+        metrics_slices = [
+            self.data[i : i + MAX_METRICS_PER_BATCH]
+            for i in range(0, len(self.data), MAX_METRICS_PER_BATCH)
+        ]
+        for metrics_slice in metrics_slices:
+            self.client.log_batch(run_id=current_run_id, metrics=metrics_slice)
+        end = time.time()
+        self.total_log_batch_time += end - start
+
+    def _should_flush(self):
+        target_training_to_logging_time_ratio = 10
+        if (
+            self.total_training_time
+            >= self.total_log_batch_time * target_training_to_logging_time_ratio
+        ):
+            return True
+
+        return False
+
+    def record_metrics(self, metrics, step=None):
+        """
+        Submit a set of metrics to be logged. The metrics may not be immediately logged, as this
+        class will batch them in order to not increase execution time too much by logging
+        frequently.
+
+        Args:
+            metrics: Dictionary containing key, value pairs of metrics to be logged.
+            step: The training step that the metrics correspond to.
+        """
+        current_timestamp = time.time()
+        if self.previous_training_timestamp is None:
+            self.previous_training_timestamp = current_timestamp
+
+        training_time = current_timestamp - self.previous_training_timestamp
+
+        self.total_training_time += training_time
+
+        # log_batch() requires step to be defined. Therefore will set step to 0 if not defined.
+        if step is None:
+            step = 0
+
+        for key, value in metrics.items():
+            self.data.append(Metric(key, value, int(current_timestamp * 1000), step))
+
+        if self._should_flush():
+            self.flush()
+
+        self.previous_training_timestamp = current_timestamp
+
+
+@contextlib.contextmanager
+def batch_metrics_logger(run_id):
+    """
+    Context manager that yields a BatchMetricsLogger object, which metrics can be logged against.
+    The BatchMetricsLogger keeps metrics in a list until it decides they should be logged, at
+    which point the accumulated metrics will be batch logged. The BatchMetricsLogger ensures
+    that logging imposes no more than a 10% overhead on the training, where the training is
+    measured by adding up the time elapsed between consecutive calls to record_metrics.
+
+    If logging a batch fails, a warning will be emitted and subsequent metrics will continue to
+    be collected.
+
+    Once the context is closed, any metrics that have yet to be logged will be logged.
+
+    Args:
+        run_id: ID of the run that the metrics will be logged to.
+    """
+
+    batch_metrics_logger = BatchMetricsLogger(run_id)
+    yield batch_metrics_logger
+    batch_metrics_logger.flush()
+
+
+def gen_autologging_package_version_requirements_doc(integration_name):
+    """
+    Returns:
+        A document note string saying the compatibility for the specified autologging
+        integration's associated package versions.
+    """
+    min_ver, max_ver, pip_release = get_min_max_version_and_pip_release(integration_name)
+    required_pkg_versions = f"``{min_ver}`` <= ``{pip_release}`` <= ``{max_ver}``"
+
+    return (
+        "    .. Note:: Autologging is known to be compatible with the following package versions: "
+        + required_pkg_versions
+        + ". Autologging may not succeed when used with package versions outside of this range."
+        + "\n\n"
+    )
+
+
+def _check_and_log_warning_for_unsupported_package_versions(integration_name):
+    """
+    When autologging is enabled and `disable_for_unsupported_versions=False` for the specified
+    autologging integration, check whether the currently-installed versions of the integration's
+    associated package versions are supported by the specified integration. If the package versions
+    are not supported, log a warning message.
+    """
+    if (
+        integration_name in FLAVOR_TO_MODULE_NAME
+        and integration_name not in _AUTOLOGGING_SUPPORTED_VERSION_WARNING_SUPPRESS_LIST
+        and not get_autologging_config(integration_name, "disable", True)
+        and not get_autologging_config(integration_name, "disable_for_unsupported_versions", False)
+        and not is_flavor_supported_for_associated_package_versions(integration_name)
+    ):
+        min_var, max_var, pip_release = get_min_max_version_and_pip_release(integration_name)
+        module = importlib.import_module(FLAVOR_TO_MODULE_NAME[integration_name])
+        _logger.warning(
+            f"MLflow {integration_name} autologging is known to be compatible with "
+            f"{min_var} <= {pip_release} <= {max_var}, but the installed version is "
+            f"{module.__version__}. If you encounter errors during autologging, try upgrading "
+            f"/ downgrading {pip_release} to a compatible version, or try upgrading MLflow.",
+        )
+
+
+def autologging_integration(name):
+    """
+    **All autologging integrations should be decorated with this wrapper.**
+
+    Wraps an autologging function in order to store its configuration arguments. This enables
+    patch functions to broadly obey certain configurations (e.g., disable=True) without
+    requiring specific logic to be present in each autologging integration.
+    """
+
+    def validate_param_spec(param_spec):
+        if "disable" not in param_spec or param_spec["disable"].default is not False:
+            raise Exception(
+                f"Invalid `autolog()` function for integration '{name}'. `autolog()` functions"
+                " must specify a 'disable' argument with default value 'False'"
+            )
+        elif "silent" not in param_spec or param_spec["silent"].default is not False:
+            raise Exception(
+                f"Invalid `autolog()` function for integration '{name}'. `autolog()` functions"
+                " must specify a 'silent' argument with default value 'False'"
+            )
+
+    def wrapper(_autolog):
+        param_spec = inspect.signature(_autolog).parameters
+        validate_param_spec(param_spec)
+
+        AUTOLOGGING_INTEGRATIONS[name] = {}
+        default_params = {param.name: param.default for param in param_spec.values()}
+
+        @autologging_conf_lock
+        def autolog(*args, **kwargs):
+            config_to_store = dict(default_params)
+            config_to_store.update(
+                {param.name: arg for arg, param in zip(args, param_spec.values())}
+            )
+            config_to_store.update(kwargs)
+            AUTOLOGGING_INTEGRATIONS[name] = config_to_store
+
+            try:
+                # Pass `autolog()` arguments to `log_autolog_called` in keyword format to enable
+                # event loggers to more easily identify important configuration parameters
+                # (e.g., `disable`) without examining positional arguments. Passing positional
+                # arguments to `log_autolog_called` is deprecated in MLflow > 1.13.1
+                AutologgingEventLogger.get_logger().log_autolog_called(name, (), config_to_store)
+            except Exception:
+                pass
+
+            revert_patches(name)
+
+            # If disabling autologging using fluent api, then every active integration's autolog
+            # needs to be called with disable=True. So do not short circuit and let
+            # `mlflow.autolog()` invoke all active integrations with disable=True.
+            if name != "mlflow" and get_autologging_config(name, "disable", True):
+                return
+
+            is_silent_mode = get_autologging_config(name, "silent", False)
+            # Reroute non-MLflow warnings encountered during autologging enablement to an
+            # MLflow event logger, and enforce silent mode if applicable (i.e. if the corresponding
+            # autologging integration was called with `silent=True`)
+            with (
+                MlflowEventsAndWarningsBehaviorGlobally(
+                    # MLflow warnings emitted during autologging setup / enablement are likely
+                    # actionable and relevant to the user, so they should be emitted as normal
+                    # when `silent=False`. For reference, see recommended warning and event logging
+                    # behaviors from https://docs.python.org/3/howto/logging.html#when-to-use-logging
+                    reroute_warnings=False,
+                    disable_event_logs=is_silent_mode,
+                    disable_warnings=is_silent_mode,
+                ),
+                NonMlflowWarningsBehaviorForCurrentThread(
+                    # non-MLflow warnings emitted during autologging setup / enablement are not
+                    # actionable for the user, as they are a byproduct of the autologging
+                    # implementation. Accordingly, they should be rerouted to `logger.warning()`.
+                    # For reference, see recommended warning and event logging
+                    # behaviors from https://docs.python.org/3/howto/logging.html#when-to-use-logging
+                    reroute_warnings=True,
+                    disable_warnings=is_silent_mode,
+                ),
+            ):
+                _check_and_log_warning_for_unsupported_package_versions(name)
+
+                return _autolog(*args, **kwargs)
+
+        wrapped_autolog = update_wrapper_extended(autolog, _autolog)
+        # Set the autologging integration name as a function attribute on the wrapped autologging
+        # function, allowing the integration name to be extracted from the function. This is used
+        # during the execution of import hooks for `mlflow.autolog()`.
+        wrapped_autolog.integration_name = name
+
+        if name in FLAVOR_TO_MODULE_NAME:
+            wrapped_autolog.__doc__ = gen_autologging_package_version_requirements_doc(name) + (
+                wrapped_autolog.__doc__ or ""
+            )
+        return wrapped_autolog
+
+    return wrapper
+
+
+def get_autologging_config(flavor_name, config_key, default_value=None):
+    """
+    Returns a desired config value for a specified autologging integration.
+
+    Returns `None` if specified `flavor_name` has no recorded configs.
+    If `config_key` is not set on the config object, default value is returned.
+
+    Args:
+        flavor_name: An autologging integration flavor name.
+        config_key: The key for the desired config value.
+        default_value: The default_value to return.
+    """
+    config = AUTOLOGGING_INTEGRATIONS.get(flavor_name)
+    if config is not None:
+        return config.get(config_key, default_value)
+    else:
+        return default_value
+
+
+def autologging_is_disabled(integration_name):
+    """Returns a boolean flag of whether the autologging integration is disabled.
+
+    Args:
+        integration_name: An autologging integration flavor name.
+
+    """
+    explicit_disabled = get_autologging_config(integration_name, "disable", True)
+    if explicit_disabled:
+        return True
+
+    if (
+        integration_name in FLAVOR_TO_MODULE_NAME
+        and get_autologging_config(integration_name, "disable_for_unsupported_versions", False)
+        and not is_flavor_supported_for_associated_package_versions(integration_name)
+    ):
+        return True
+
+    return False
+
+
+def is_autolog_supported(integration_name: str) -> bool:
+    """
+    Whether the specified autologging integration is supported by the current environment.
+
+    Args:
+        integration_name: An autologging integration flavor name.
+    """
+    # NB: We don't check for the presence of autolog() function as it requires importing
+    #   the flavor module, which may cause import error or overhead.
+    return "autologging" in _ML_PACKAGE_VERSIONS.get(integration_name, {})
+
+
+def get_autolog_function(integration_name: str) -> Optional[Callable[..., Any]]:
+    """
+    Get the autolog() function for the specified integration.
+    Returns None if the flavor does not have an autolog() function.
+    """
+    flavor_module = importlib.import_module(f"mlflow.{integration_name}")
+    return getattr(flavor_module, "autolog", None)
+
+
+@contextlib.contextmanager
+def disable_autologging():
+    """
+    Context manager that temporarily disables autologging globally for all integrations upon
+    entry and restores the previous autologging configuration upon exit.
+    """
+    global _AUTOLOGGING_GLOBALLY_DISABLED
+    _AUTOLOGGING_GLOBALLY_DISABLED = True
+    try:
+        yield
+    finally:
+        _AUTOLOGGING_GLOBALLY_DISABLED = False
+
+
+@contextlib.contextmanager
+def disable_discrete_autologging(flavors_to_disable: list[str]) -> None:
+    """
+    Context manager for disabling specific autologging integrations temporarily while another
+    flavor's autologging is activated. This context wrapper is useful in the event that, for
+    example, a particular library calls upon another library within a training API that has a
+    current MLflow autologging integration.
+    For instance, the transformers library's Trainer class, when running metric scoring,
+    builds a sklearn model and runs evaluations as part of its accuracy scoring. Without this
+    temporary autologging disabling, a new run will be generated that contains a sklearn model
+    that holds no use for tracking purposes as it is only used during the metric evaluation phase
+    of training.
+
+    Args:
+        flavors_to_disable: A list of flavors that need to be temporarily disabled while
+            executing another flavor's autologging to prevent spurious run
+            logging of unrelated models, metrics, and parameters.
+    """
+    enabled_flavors = []
+    for flavor in flavors_to_disable:
+        if not autologging_is_disabled(flavor):
+            enabled_flavors.append(flavor)
+            autolog_func = getattr(mlflow, flavor)
+            autolog_func.autolog(disable=True)
+    yield
+    for flavor in enabled_flavors:
+        autolog_func = getattr(mlflow, flavor)
+        autolog_func.autolog(disable=False)
+
+
+_training_sessions = []
+
+
+def _get_new_training_session_class():
+    """
+    Returns a session manager class for nested autologging runs.
+
+    Examples
+    --------
+    >>> class Parent:
+    ...     pass
+    >>> class Child:
+    ...     pass
+    >>> class Grandchild:
+    ...     pass
+    >>>
+    >>> _TrainingSession = _get_new_training_session_class()
+    >>> with _TrainingSession(Parent, False) as p:
+    ...     with _SklearnTrainingSession(Child, True) as c:
+    ...         with _SklearnTrainingSession(Grandchild, True) as g:
+    ...             print(p.should_log(), c.should_log(), g.should_log())
+    True False False
+    >>>
+    >>> with _TrainingSession(Parent, True) as p:
+    ...     with _TrainingSession(Child, False) as c:
+    ...         with _TrainingSession(Grandchild, True) as g:
+    ...             print(p.should_log(), c.should_log(), g.should_log())
+    True True False
+    >>>
+    >>> with _TrainingSession(Child, True) as c1:
+    ...     with _TrainingSession(Child, True) as c2:
+    ...         print(c1.should_log(), c2.should_log())
+    True False
+    """
+
+    # NOTE: The current implementation doesn't guarantee thread-safety, but that's okay for now
+    # because:
+    # 1. We don't currently have any use cases for allow_children=True.
+    # 2. The list append & pop operations are thread-safe, so we will always clear the session stack
+    #    once all _TrainingSessions exit.
+    class _TrainingSession:
+        _session_stack = []
+
+        def __init__(self, estimator, allow_children=True):
+            """A session manager for nested autologging runs.
+
+            Args:
+                estimator: An estimator that this session originates from.
+                allow_children: If True, allows autologging in child sessions.
+                    If False, disallows autologging in all descendant sessions.
+
+            """
+            self.allow_children = allow_children
+            self.estimator = estimator
+            self._parent = None
+
+        def __enter__(self):
+            if len(_TrainingSession._session_stack) > 0:
+                self._parent = _TrainingSession._session_stack[-1]
+                self.allow_children = (
+                    _TrainingSession._session_stack[-1].allow_children and self.allow_children
+                )
+            _TrainingSession._session_stack.append(self)
+            return self
+
+        def __exit__(self, tp, val, traceback):
+            _TrainingSession._session_stack.pop()
+
+        def should_log(self):
+            """
+            Returns True when at least one of the following conditions satisfies:
+
+            1. This session is the root session.
+            2. The parent session allows autologging and its estimator differs from this session's
+               estimator.
+            """
+            for training_session in _TrainingSession._session_stack:
+                if training_session is self:
+                    break
+                elif training_session.estimator is self.estimator:
+                    return False
+
+            return self._parent is None or self._parent.allow_children
+
+        @staticmethod
+        def is_active():
+            return len(_TrainingSession._session_stack) != 0
+
+        @staticmethod
+        def get_current_session():
+            if _TrainingSession.is_active():
+                return _TrainingSession._session_stack[-1]
+            return None
+
+    _training_sessions.append(_TrainingSession)
+    return _TrainingSession
+
+
+def _has_active_training_session():
+    return any(s.is_active() for s in _training_sessions)
+
+
+def get_instance_method_first_arg_value(method, call_pos_args, call_kwargs):
+    """Get instance method first argument value (exclude the `self` argument).
+
+    Args:
+        method: A `cls.method` object which includes the `self` argument.
+        call_pos_args: positional arguments excluding the first `self` argument.
+        call_kwargs: keywords arguments.
+    """
+    if len(call_pos_args) >= 1:
+        return call_pos_args[0]
+    else:
+        param_sig = inspect.signature(method).parameters
+        first_arg_name = list(param_sig.keys())[1]
+        assert param_sig[first_arg_name].kind not in [
+            inspect.Parameter.VAR_KEYWORD,
+            inspect.Parameter.VAR_POSITIONAL,
+        ]
+        return call_kwargs.get(first_arg_name)
+
+
+def get_method_call_arg_value(arg_index, arg_name, default_value, call_pos_args, call_kwargs):
+    """Get argument value for a method call.
+
+    Args:
+        arg_index: The argument index in the function signature. Start from 0.
+        arg_name: The argument name in the function signature.
+        default_value: Default argument value.
+        call_pos_args: The positional argument values in the method call.
+        call_kwargs: The keyword argument values in the method call.
+    """
+    if arg_name in call_kwargs:
+        return call_kwargs[arg_name]
+    elif arg_index < len(call_pos_args):
+        return call_pos_args[arg_index]
+    else:
+        return default_value

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/client.py

@@ -0,0 +1,417 @@
+"""
+Defines an MlflowAutologgingQueueingClient developer API that provides batching, queueing, and
+asynchronous execution capabilities for a subset of MLflow Tracking logging operations used most
+frequently by autologging operations.
+
+TODO(dbczumar): Migrate request batching, queueing, and async execution support from
+MlflowAutologgingQueueingClient to MlflowClient in order to provide broader benefits to end users.
+Remove this developer API.
+"""
+
+import logging
+import os
+from collections import namedtuple
+from concurrent.futures import ThreadPoolExecutor
+from itertools import zip_longest
+from typing import Any, Optional, Union
+
+from mlflow.entities import Metric, Param, RunTag
+from mlflow.entities.dataset_input import DatasetInput
+from mlflow.exceptions import MlflowException
+from mlflow.tracking.client import MlflowClient
+from mlflow.utils import _truncate_dict, chunk_list
+from mlflow.utils.time import get_current_time_millis
+from mlflow.utils.validation import (
+    MAX_DATASETS_PER_BATCH,
+    MAX_ENTITIES_PER_BATCH,
+    MAX_ENTITY_KEY_LENGTH,
+    MAX_METRICS_PER_BATCH,
+    MAX_PARAM_VAL_LENGTH,
+    MAX_PARAMS_TAGS_PER_BATCH,
+    MAX_TAG_VAL_LENGTH,
+)
+
+_logger = logging.getLogger(__name__)
+
+_PendingCreateRun = namedtuple(
+    "_PendingCreateRun", ["experiment_id", "start_time", "tags", "run_name"]
+)
+_PendingSetTerminated = namedtuple("_PendingSetTerminated", ["status", "end_time"])
+
+
+class PendingRunId:
+    """
+    Serves as a placeholder for the ID of a run that does not yet exist, enabling additional
+    metadata (e.g. metrics, params, ...) to be enqueued for the run prior to its creation.
+    """
+
+
+class RunOperations:
+    """
+    Represents a collection of operations on one or more MLflow Runs, such as run creation
+    or metric logging.
+    """
+
+    def __init__(self, operation_futures):
+        self._operation_futures = operation_futures
+
+    def await_completion(self):
+        """
+        Blocks on completion of the MLflow Run operations.
+        """
+        failed_operations = []
+        for future in self._operation_futures:
+            try:
+                future.result()
+            except Exception as e:
+                failed_operations.append(e)
+
+        if len(failed_operations) > 0:
+            raise MlflowException(
+                message=(
+                    "The following failures occurred while performing one or more logging"
+                    f" operations: {failed_operations}"
+                )
+            )
+
+
+# Define a threadpool for use across `MlflowAutologgingQueueingClient` instances to ensure that
+# `MlflowAutologgingQueueingClient` instances can be pickled (ThreadPoolExecutor objects are not
+# pickleable and therefore cannot be assigned as instance attributes).
+#
+# We limit the number of threads used for run operations, using at most 8 threads or 2 * the number
+# of CPU cores available on the system (whichever is smaller)
+num_cpus = os.cpu_count() or 4
+num_logging_workers = min(num_cpus * 2, 8)
+_AUTOLOGGING_QUEUEING_CLIENT_THREAD_POOL = ThreadPoolExecutor(
+    max_workers=num_logging_workers,
+    thread_name_prefix="MlflowAutologgingQueueingClient",
+)
+
+
+class MlflowAutologgingQueueingClient:
+    """
+    Efficiently implements a subset of MLflow Tracking's  `MlflowClient` and fluent APIs to provide
+    automatic batching and async execution of run operations by way of queueing, as well as
+    parameter / tag truncation for autologging use cases. Run operations defined by this client,
+    such as `create_run` and `log_metrics`, enqueue data for future persistence to MLflow
+    Tracking. Data is not persisted until the queue is flushed via the `flush()` method, which
+    supports synchronous and asynchronous execution.
+
+    MlflowAutologgingQueueingClient is not threadsafe; none of its APIs should be called
+    concurrently.
+    """
+
+    def __init__(self, tracking_uri=None):
+        self._client = MlflowClient(tracking_uri)
+        self._pending_ops_by_run_id = {}
+
+    def __enter__(self):
+        """
+        Enables `MlflowAutologgingQueueingClient` to be used as a context manager with
+        synchronous flushing upon exit, removing the need to call `flush()` for use cases
+        where logging completion can be waited upon synchronously.
+
+        Run content is only flushed if the context exited without an exception.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc, traceback):
+        """
+        Enables `MlflowAutologgingQueueingClient` to be used as a context manager with
+        synchronous flushing upon exit, removing the need to call `flush()` for use cases
+        where logging completion can be waited upon synchronously.
+
+        Run content is only flushed if the context exited without an exception.
+        """
+        # NB: Run content is only flushed upon context exit to ensure that we don't elide the
+        # original exception thrown by the context (because `flush()` itself may throw). This
+        # is consistent with the behavior of a routine that calls `flush()` explicitly: content
+        # is not logged if an exception preempts the call to `flush()`
+        if exc is None and exc_type is None and traceback is None:
+            self.flush(synchronous=True)
+        else:
+            _logger.debug(
+                "Skipping run content logging upon MlflowAutologgingQueueingClient context because"
+                " an exception was raised within the context: %s",
+                exc,
+            )
+
+    def create_run(
+        self,
+        experiment_id: str,
+        start_time: Optional[int] = None,
+        tags: Optional[dict[str, Any]] = None,
+        run_name: Optional[str] = None,
+    ) -> PendingRunId:
+        """
+        Enqueues a CreateRun operation with the specified attributes, returning a `PendingRunId`
+        instance that can be used as input to other client logging APIs (e.g. `log_metrics`,
+        `log_params`, ...).
+
+        Returns:
+            A `PendingRunId` that can be passed as the `run_id` parameter to other client
+            logging APIs, such as `log_params` and `log_metrics`.
+        """
+        tags = tags or {}
+        tags = _truncate_dict(
+            tags, max_key_length=MAX_ENTITY_KEY_LENGTH, max_value_length=MAX_TAG_VAL_LENGTH
+        )
+        run_id = PendingRunId()
+        self._get_pending_operations(run_id).enqueue(
+            create_run=_PendingCreateRun(
+                experiment_id=experiment_id,
+                start_time=start_time,
+                tags=[RunTag(key, str(value)) for key, value in tags.items()],
+                run_name=run_name,
+            )
+        )
+        return run_id
+
+    def set_terminated(
+        self,
+        run_id: Union[str, PendingRunId],
+        status: Optional[str] = None,
+        end_time: Optional[int] = None,
+    ) -> None:
+        """
+        Enqueues an UpdateRun operation with the specified `status` and `end_time` attributes
+        for the specified `run_id`.
+        """
+        self._get_pending_operations(run_id).enqueue(
+            set_terminated=_PendingSetTerminated(status=status, end_time=end_time)
+        )
+
+    def log_params(self, run_id: Union[str, PendingRunId], params: dict[str, Any]) -> None:
+        """
+        Enqueues a collection of Parameters to be logged to the run specified by `run_id`.
+        """
+        params = _truncate_dict(
+            params, max_key_length=MAX_ENTITY_KEY_LENGTH, max_value_length=MAX_PARAM_VAL_LENGTH
+        )
+        params_arr = [Param(key, str(value)) for key, value in params.items()]
+        self._get_pending_operations(run_id).enqueue(params=params_arr)
+
+    def log_inputs(
+        self, run_id: Union[str, PendingRunId], datasets: Optional[list[DatasetInput]]
+    ) -> None:
+        """
+        Enqueues a collection of Dataset to be logged to the run specified by `run_id`.
+        """
+        if datasets is None or len(datasets) == 0:
+            return
+        self._get_pending_operations(run_id).enqueue(datasets=datasets)
+
+    def log_metrics(
+        self,
+        run_id: Union[str, PendingRunId],
+        metrics: dict[str, float],
+        step: Optional[int] = None,
+    ) -> None:
+        """
+        Enqueues a collection of Metrics to be logged to the run specified by `run_id` at the
+        step specified by `step`.
+        """
+        metrics = _truncate_dict(metrics, max_key_length=MAX_ENTITY_KEY_LENGTH)
+        timestamp_ms = get_current_time_millis()
+        metrics_arr = [
+            Metric(key, value, timestamp_ms, step or 0) for key, value in metrics.items()
+        ]
+        self._get_pending_operations(run_id).enqueue(metrics=metrics_arr)
+
+    def set_tags(self, run_id: Union[str, PendingRunId], tags: dict[str, Any]) -> None:
+        """
+        Enqueues a collection of Tags to be logged to the run specified by `run_id`.
+        """
+        tags = _truncate_dict(
+            tags, max_key_length=MAX_ENTITY_KEY_LENGTH, max_value_length=MAX_TAG_VAL_LENGTH
+        )
+        tags_arr = [RunTag(key, str(value)) for key, value in tags.items()]
+        self._get_pending_operations(run_id).enqueue(tags=tags_arr)
+
+    def flush(self, synchronous=True):
+        """
+        Flushes all queued run operations, resulting in the creation or mutation of runs
+        and run data.
+
+        Args:
+            synchronous: If `True`, run operations are performed synchronously, and a
+                `RunOperations` result object is only returned once all operations
+                are complete. If `False`, run operations are performed asynchronously,
+                and an `RunOperations` object is returned that represents the ongoing
+                run operations.
+
+        Returns:
+            A `RunOperations` instance representing the flushed operations. These operations
+            are already complete if `synchronous` is `True`. If `synchronous` is `False`, these
+            operations may still be inflight. Operation completion can be synchronously waited
+            on via `RunOperations.await_completion()`.
+        """
+        logging_futures = []
+        for pending_operations in self._pending_ops_by_run_id.values():
+            future = _AUTOLOGGING_QUEUEING_CLIENT_THREAD_POOL.submit(
+                self._flush_pending_operations,
+                pending_operations=pending_operations,
+            )
+            logging_futures.append(future)
+        self._pending_ops_by_run_id = {}
+
+        logging_operations = RunOperations(logging_futures)
+        if synchronous:
+            logging_operations.await_completion()
+        return logging_operations
+
+    def _get_pending_operations(self, run_id):
+        """
+        Returns:
+            A `_PendingRunOperations` containing all pending operations for the
+            specified `run_id`.
+        """
+        if run_id not in self._pending_ops_by_run_id:
+            self._pending_ops_by_run_id[run_id] = _PendingRunOperations(run_id=run_id)
+        return self._pending_ops_by_run_id[run_id]
+
+    def _try_operation(self, fn, *args, **kwargs):
+        """
+        Attempt to evaluate the specified function, `fn`, on the specified `*args` and `**kwargs`,
+        returning either the result of the function evaluation (if evaluation was successful) or
+        the exception raised by the function evaluation (if evaluation was unsuccessful).
+        """
+        try:
+            return fn(*args, **kwargs)
+        except Exception as e:
+            return e
+
+    def _flush_pending_operations(self, pending_operations):
+        """
+        Synchronously and sequentially flushes the specified list of pending run operations.
+
+        NB: Operations are not parallelized on a per-run basis because MLflow's File Store, which
+        is frequently used for local ML development, does not support threadsafe metadata logging
+        within a given run.
+        """
+        if pending_operations.create_run:
+            create_run_tags = pending_operations.create_run.tags
+            num_additional_tags_to_include_during_creation = MAX_ENTITIES_PER_BATCH - len(
+                create_run_tags
+            )
+            if num_additional_tags_to_include_during_creation > 0:
+                create_run_tags.extend(
+                    pending_operations.tags_queue[:num_additional_tags_to_include_during_creation]
+                )
+                pending_operations.tags_queue = pending_operations.tags_queue[
+                    num_additional_tags_to_include_during_creation:
+                ]
+
+            new_run = self._client.create_run(
+                experiment_id=pending_operations.create_run.experiment_id,
+                start_time=pending_operations.create_run.start_time,
+                tags={tag.key: tag.value for tag in create_run_tags},
+            )
+            pending_operations.run_id = new_run.info.run_id
+
+        run_id = pending_operations.run_id
+        assert not isinstance(run_id, PendingRunId), "Run ID cannot be pending for logging"
+
+        operation_results = []
+
+        param_batches_to_log = chunk_list(
+            pending_operations.params_queue,
+            chunk_size=MAX_PARAMS_TAGS_PER_BATCH,
+        )
+        tag_batches_to_log = chunk_list(
+            pending_operations.tags_queue,
+            chunk_size=MAX_PARAMS_TAGS_PER_BATCH,
+        )
+        for params_batch, tags_batch in zip_longest(
+            param_batches_to_log, tag_batches_to_log, fillvalue=[]
+        ):
+            metrics_batch_size = min(
+                MAX_ENTITIES_PER_BATCH - len(params_batch) - len(tags_batch),
+                MAX_METRICS_PER_BATCH,
+            )
+            metrics_batch_size = max(metrics_batch_size, 0)
+            metrics_batch = pending_operations.metrics_queue[:metrics_batch_size]
+            pending_operations.metrics_queue = pending_operations.metrics_queue[metrics_batch_size:]
+
+            operation_results.append(
+                self._try_operation(
+                    self._client.log_batch,
+                    run_id=run_id,
+                    metrics=metrics_batch,
+                    params=params_batch,
+                    tags=tags_batch,
+                )
+            )
+
+        for metrics_batch in chunk_list(
+            pending_operations.metrics_queue, chunk_size=MAX_METRICS_PER_BATCH
+        ):
+            operation_results.append(
+                self._try_operation(self._client.log_batch, run_id=run_id, metrics=metrics_batch)
+            )
+
+        for datasets_batch in chunk_list(
+            pending_operations.datasets_queue, chunk_size=MAX_DATASETS_PER_BATCH
+        ):
+            operation_results.append(
+                self._try_operation(self._client.log_inputs, run_id=run_id, datasets=datasets_batch)
+            )
+
+        if pending_operations.set_terminated:
+            operation_results.append(
+                self._try_operation(
+                    self._client.set_terminated,
+                    run_id=run_id,
+                    status=pending_operations.set_terminated.status,
+                    end_time=pending_operations.set_terminated.end_time,
+                )
+            )
+
+        failures = [result for result in operation_results if isinstance(result, Exception)]
+        if len(failures) > 0:
+            raise MlflowException(
+                message=(
+                    f"Failed to perform one or more operations on the run with ID {run_id}."
+                    f" Failed operations: {failures}"
+                )
+            )
+
+
+class _PendingRunOperations:
+    """
+    Represents a collection of queued / pending MLflow Run operations.
+    """
+
+    def __init__(self, run_id):
+        self.run_id = run_id
+        self.create_run = None
+        self.set_terminated = None
+        self.params_queue = []
+        self.tags_queue = []
+        self.metrics_queue = []
+        self.datasets_queue = []
+
+    def enqueue(
+        self,
+        params=None,
+        tags=None,
+        metrics=None,
+        datasets=None,
+        create_run=None,
+        set_terminated=None,
+    ):
+        """
+        Enqueues a new pending logging operation for the associated MLflow Run.
+        """
+        if create_run:
+            assert not self.create_run, "Attempted to create the same run multiple times"
+            self.create_run = create_run
+        if set_terminated:
+            assert not self.set_terminated, "Attempted to terminate the same run multiple times"
+            self.set_terminated = set_terminated
+
+        self.params_queue += params or []
+        self.tags_queue += tags or []
+        self.metrics_queue += metrics or []
+        self.datasets_queue += datasets or []

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/config.py

@@ -0,0 +1,39 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.utils.autologging_utils import AUTOLOGGING_INTEGRATIONS
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AutoLoggingConfig:
+    """
+    A dataclass to hold common autologging configuration options.
+    """
+
+    log_models: bool
+    log_input_examples: bool
+    log_model_signatures: bool
+    log_traces: bool
+    extra_tags: Optional[dict] = None
+
+    def should_log_optional_artifacts(self):
+        """
+        Check if any optional artifacts should be logged to MLflow.
+        """
+        return self.log_models or self.log_input_examples or self.log_model_signatures
+
+    @classmethod
+    def init(cls, flavor_name: str):
+        config_dict = AUTOLOGGING_INTEGRATIONS.get(flavor_name, {})
+        # NB: These defaults are only used when the autolog() function for the
+        # flavor does not specify the corresponding configuration option
+        return cls(
+            log_models=config_dict.get("log_models", False),
+            log_input_examples=config_dict.get("log_input_examples", False),
+            log_model_signatures=config_dict.get("log_model_signatures", False),
+            log_traces=config_dict.get("log_traces", True),
+            extra_tags=config_dict.get("extra_tags", None),
+        )

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/events.py

@@ -0,0 +1,294 @@
+import warnings
+from typing import Any
+
+from mlflow.utils.autologging_utils import _logger
+
+
+def _catch_exception(fn):
+    """A decorator that catches exceptions thrown by the wrapped function and logs them."""
+
+    def wrapper(*args):
+        try:
+            fn(*args)
+        except Exception as e:
+            _logger.debug(f"Failed to log autologging event via '{fn}'. Exception: {e}")
+
+    return wrapper
+
+
+class AutologgingEventLoggerWrapper:
+    """
+    A wrapper around AutologgingEventLogger for DRY:
+      - Store common arguments to avoid passing them to each logger method
+      - Catches exceptions thrown by the logger and logs them
+
+    NB: We could not modify the AutologgingEventLogger class directly because
+        it is used in Databricks code base as well.
+    """
+
+    def __init__(self, session, destination: Any, function_name: str):
+        self._session = session
+        self._destination = destination
+        self._function_name = function_name
+        self._logger = AutologgingEventLogger.get_logger()
+
+    @_catch_exception
+    def log_patch_function_start(self, args, kwargs):
+        self._logger.log_patch_function_start(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_patch_function_success(self, args, kwargs):
+        self._logger.log_patch_function_success(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_patch_function_error(self, args, kwargs, exception):
+        self._logger.log_patch_function_error(
+            self._session, self._destination, self._function_name, args, kwargs, exception
+        )
+
+    @_catch_exception
+    def log_original_function_start(self, args, kwargs):
+        self._logger.log_original_function_start(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_original_function_success(self, args, kwargs):
+        self._logger.log_original_function_success(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_original_function_error(self, args, kwargs, exception):
+        self._logger.log_original_function_error(
+            self._session, self._destination, self._function_name, args, kwargs, exception
+        )
+
+
+class AutologgingEventLogger:
+    """
+    Provides instrumentation hooks for important autologging lifecycle events, including:
+
+        - Calls to `mlflow.autolog()` APIs
+        - Calls to patched APIs with associated termination states
+          ("success" and "failure due to error")
+        - Calls to original / underlying APIs made by patched function code with
+          associated termination states ("success" and "failure due to error")
+
+    Default implementations are included for each of these hooks, which emit corresponding
+    DEBUG-level logging statements. Developers can provide their own hook implementations
+    by subclassing `AutologgingEventLogger` and calling the static
+    `AutologgingEventLogger.set_logger()` method to supply a new event logger instance.
+
+    Callers fetch the configured logger via `AutologgingEventLogger.get_logger()`
+    and invoke one or more hooks (e.g., `AutologgingEventLogger.get_logger().log_autolog_called()`).
+    """
+
+    _event_logger = None
+
+    @staticmethod
+    def get_logger():
+        """Fetches the configured `AutologgingEventLogger` instance for logging.
+
+        Returns:
+            The instance of `AutologgingEventLogger` specified via `set_logger`
+            (if configured) or the default implementation of `AutologgingEventLogger`
+            (if a logger was not configured via `set_logger`).
+
+        """
+        return AutologgingEventLogger._event_logger or AutologgingEventLogger()
+
+    @staticmethod
+    def set_logger(logger):
+        """Configures the `AutologgingEventLogger` instance for logging. This instance
+        is exposed via `AutologgingEventLogger.get_logger()` and callers use it to invoke
+        logging hooks (e.g., AutologgingEventLogger.get_logger().log_autolog_called()).
+
+        Args:
+            logger: The instance of `AutologgingEventLogger` to use when invoking logging hooks.
+
+        """
+        AutologgingEventLogger._event_logger = logger
+
+    def log_autolog_called(self, integration, call_args, call_kwargs):
+        """Called when the `autolog()` method for an autologging integration
+        is invoked (e.g., when a user invokes `mlflow.sklearn.autolog()`)
+
+        Args:
+            integration: The autologging integration for which `autolog()` was called.
+            call_args: **DEPRECATED** The positional arguments passed to the `autolog()` call.
+                This field is empty in MLflow > 1.13.1; all arguments are passed in
+                keyword form via `call_kwargs`.
+            call_kwargs: The arguments passed to the `autolog()` call in keyword form.
+                Any positional arguments should also be converted to keyword form
+                and passed via `call_kwargs`.
+        """
+        if len(call_args) > 0:
+            warnings.warn(
+                f"Received {len(call_args)} positional arguments via `call_args`. `call_args` is"
+                " deprecated in MLflow > 1.13.1, and all arguments should be passed"
+                " in keyword form via `call_kwargs`.",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+        _logger.debug(
+            "Called autolog() method for %s autologging with args '%s' and kwargs '%s'",
+            integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_start(self, session, patch_obj, function_name, call_args, call_kwargs):
+        """Called upon invocation of a patched API associated with an autologging integration
+        (e.g., `sklearn.linear_model.LogisticRegression.fit()`).
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+
+        """
+        _logger.debug(
+            "Invoked patched API '%s.%s' for %s autologging with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_success(self, session, patch_obj, function_name, call_args, call_kwargs):
+        """
+        Called upon successful termination of a patched API associated with an autologging
+        integration (e.g., `sklearn.linear_model.LogisticRegression.fit()`).
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+        """
+        _logger.debug(
+            "Patched API call '%s.%s' for %s autologging completed successfully. Patched ML"
+            " API was called with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_error(
+        self, session, patch_obj, function_name, call_args, call_kwargs, exception
+    ):
+        """Called when execution of a patched API associated with an autologging integration
+        (e.g., `sklearn.linear_model.LogisticRegression.fit()`) terminates with an exception.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+            exception: The exception that caused the patched API call to terminate.
+        """
+        _logger.debug(
+            "Patched API call '%s.%s' for %s autologging threw exception. Patched API was"
+            " called with args '%s' and kwargs '%s'. Exception: %s",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+            exception,
+        )
+
+    def log_original_function_start(
+        self, session, patch_obj, function_name, call_args, call_kwargs
+    ):
+        """
+        Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API is invoked. For example, this is called when
+        a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes
+        the original implementation of `sklearn.linear_model.LogisticRegression.fit()`.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+        """
+        _logger.debug(
+            "Original function invoked during execution of patched API '%s.%s' for %s"
+            " autologging. Original function was invoked with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_original_function_success(
+        self, session, patch_obj, function_name, call_args, call_kwargs
+    ):
+        """Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API invocation terminates successfully. For example,
+        when a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes the
+        original / underlying implementation of `LogisticRegression.fit()`, then this function is
+        called if the original / underlying implementation successfully completes.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+
+        """
+        _logger.debug(
+            "Original function invocation completed successfully during execution of patched API"
+            " call '%s.%s' for %s autologging. Original function was invoked with with"
+            " args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_original_function_error(
+        self, session, patch_obj, function_name, call_args, call_kwargs, exception
+    ):
+        """Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API invocation terminates with an error. For example,
+        when a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes the
+        original / underlying implementation of `LogisticRegression.fit()`, then this function is
+        called if the original / underlying implementation terminates with an exception.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+            exception: The exception that caused the original API call to terminate.
+        """
+        _logger.debug(
+            "Original function invocation threw exception during execution of patched"
+            " API call '%s.%s' for %s autologging. Original function was invoked with"
+            " args '%s' and kwargs '%s'. Exception: %s",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+            exception,
+        )

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/logging_and_warnings.py

@@ -0,0 +1,328 @@
+import os
+import warnings
+from pathlib import Path
+from threading import RLock
+from threading import get_ident as get_current_thread_id
+
+import mlflow
+from mlflow.utils import logging_utils
+
+ORIGINAL_SHOWWARNING = warnings.showwarning
+
+
+class _WarningsController:
+    """
+    Provides threadsafe utilities to modify warning behavior for MLflow autologging, including:
+
+    - Global disablement of MLflow warnings across all threads
+    - Global rerouting of MLflow warnings to an MLflow event logger (i.e. `logger.warning()`)
+      across all threads
+    - Disablement of non-MLflow warnings for the current thread
+    - Rerouting of non-MLflow warnings to an MLflow event logger for the current thread
+    """
+
+    def __init__(self):
+        self._mlflow_root_path = Path(os.path.dirname(mlflow.__file__)).resolve()
+        self._state_lock = RLock()
+
+        self._did_patch_showwarning = False
+
+        self._disabled_threads = set()
+        self._rerouted_threads = set()
+        self._mlflow_warnings_disabled_globally = False
+        self._mlflow_warnings_rerouted_to_event_logs = False
+
+    def _patched_showwarning(self, message, category, filename, lineno, *args, **kwargs):
+        """
+        A patched implementation of `warnings.showwarning` that enforces the warning configuration
+        options configured on the controller (e.g. rerouting or disablement of MLflow warnings,
+        disablement of all warnings for the current thread).
+
+        Note that reassigning `warnings.showwarning` is the standard / recommended approach for
+        modifying warning message display behaviors. For reference, see
+        https://docs.python.org/3/library/warnings.html#warnings.showwarning
+        """
+        # NB: We explicitly avoid blocking on the `self._state_lock` lock during `showwarning`
+        # to so that threads don't have to execute serially whenever they emit warnings with
+        # `warnings.warn()`. We only lock during configuration changes to ensure that
+        # `warnings.showwarning` is patched or unpatched at the correct times.
+
+        from mlflow.utils.autologging_utils import _logger
+
+        # If the warning's source file is contained within the MLflow package's base
+        # directory, it is an MLflow warning and should be emitted via `logger.warning`
+        warning_source_path = Path(filename).resolve()
+        is_mlflow_warning = self._mlflow_root_path in warning_source_path.parents
+        curr_thread = get_current_thread_id()
+
+        if (curr_thread in self._disabled_threads) or (
+            is_mlflow_warning and self._mlflow_warnings_disabled_globally
+        ):
+            return
+        elif (curr_thread in self._rerouted_threads and not is_mlflow_warning) or (
+            is_mlflow_warning and self._mlflow_warnings_rerouted_to_event_logs
+        ):
+            _logger.warning(
+                'MLflow autologging encountered a warning: "%s:%d: %s: %s"',
+                filename,
+                lineno,
+                category.__name__,
+                message,
+            )
+        else:
+            ORIGINAL_SHOWWARNING(message, category, filename, lineno, *args, **kwargs)
+
+    def _should_patch_showwarning(self):
+        return (
+            (len(self._disabled_threads) > 0)
+            or (len(self._rerouted_threads) > 0)
+            or self._mlflow_warnings_disabled_globally
+            or self._mlflow_warnings_rerouted_to_event_logs
+        )
+
+    def _modify_patch_state_if_necessary(self):
+        """
+        Patches or unpatches `warnings.showwarning` if necessary, as determined by:
+            - Whether or not `warnings.showwarning` is already patched
+            - Whether or not any custom warning state has been configured on the warnings
+              controller (i.e. disablement or rerouting of certain warnings globally or for a
+              particular thread)
+
+        Note that reassigning `warnings.showwarning` is the standard / recommended approach for
+        modifying warning message display behaviors. For reference, see
+        https://docs.python.org/3/library/warnings.html#warnings.showwarning
+        """
+        with self._state_lock:
+            if self._should_patch_showwarning() and not self._did_patch_showwarning:
+                # NB: guard to prevent patching an instance of a patch
+                if warnings.showwarning != self._patched_showwarning:
+                    warnings.showwarning = self._patched_showwarning
+                self._did_patch_showwarning = True
+            elif not self._should_patch_showwarning() and self._did_patch_showwarning:
+                # NB: only unpatch iff the patched function is active
+                if warnings.showwarning == self._patched_showwarning:
+                    warnings.showwarning = ORIGINAL_SHOWWARNING
+                self._did_patch_showwarning = False
+
+    def set_mlflow_warnings_disablement_state_globally(self, disabled=True):
+        """Disables (or re-enables) MLflow warnings globally across all threads.
+
+        Args:
+            disabled: If `True`, disables MLflow warnings globally across all threads.
+                If `False`, enables MLflow warnings globally across all threads.
+
+        """
+        with self._state_lock:
+            self._mlflow_warnings_disabled_globally = disabled
+            self._modify_patch_state_if_necessary()
+
+    def set_mlflow_warnings_rerouting_state_globally(self, rerouted=True):
+        """
+        Enables (or disables) rerouting of MLflow warnings to an MLflow event logger with level
+        WARNING (e.g. `logger.warning()`) globally across all threads.
+
+        Args:
+            rerouted: If `True`, enables MLflow warning rerouting globally across all threads.
+                      If `False`, disables MLflow warning rerouting globally across all threads.
+        """
+        with self._state_lock:
+            self._mlflow_warnings_rerouted_to_event_logs = rerouted
+            self._modify_patch_state_if_necessary()
+
+    def set_non_mlflow_warnings_disablement_state_for_current_thread(self, disabled=True):
+        """Disables (or re-enables) non-MLflow warnings for the current thread.
+
+        Args:
+            disabled: If `True`, disables non-MLflow warnings for the current thread. If `False`,
+                enables non-MLflow warnings for the current thread. non-MLflow warning
+                behavior in other threads is unaffected.
+
+        """
+        with self._state_lock:
+            if disabled:
+                self._disabled_threads.add(get_current_thread_id())
+            else:
+                self._disabled_threads.discard(get_current_thread_id())
+            self._modify_patch_state_if_necessary()
+
+    def set_non_mlflow_warnings_rerouting_state_for_current_thread(self, rerouted=True):
+        """Enables (or disables) rerouting of non-MLflow warnings to an MLflow event logger with
+        level WARNING (e.g. `logger.warning()`) for the current thread.
+
+        Args:
+            rerouted: If `True`, enables non-MLflow warning rerouting for the current thread.
+                If `False`, disables non-MLflow warning rerouting for the current thread.
+                non-MLflow warning behavior in other threads is unaffected.
+
+        """
+        with self._state_lock:
+            if rerouted:
+                self._rerouted_threads.add(get_current_thread_id())
+            else:
+                self._rerouted_threads.discard(get_current_thread_id())
+            self._modify_patch_state_if_necessary()
+
+    def get_warnings_disablement_state_for_current_thread(self):
+        """
+        Returns:
+            True if non-MLflow warnings are disabled for the current thread. False otherwise.
+        """
+        return get_current_thread_id() in self._disabled_threads
+
+    def get_warnings_rerouting_state_for_current_thread(self):
+        """
+        Returns:
+            True if non-MLflow warnings are rerouted to an MLflow event logger with level
+            WARNING for the current thread. False otherwise.
+        """
+        return get_current_thread_id() in self._rerouted_threads
+
+
+_WARNINGS_CONTROLLER = _WarningsController()
+
+
+class NonMlflowWarningsBehaviorForCurrentThread:
+    """
+    Context manager that modifies the behavior of non-MLflow warnings upon entry, according to the
+    specified parameters.
+
+    Args:
+        disable_warnings: If `True`, disable  (mutate & discard) non-MLflow warnings. If `False`,
+            do not disable non-MLflow warnings.
+        reroute_warnings: If `True`, reroute non-MLflow warnings to an MLflow event logger with
+            level WARNING. If `False`, do not reroute non-MLflow warnings.
+    """
+
+    def __init__(self, disable_warnings, reroute_warnings):
+        self._disable_warnings = disable_warnings
+        self._reroute_warnings = reroute_warnings
+        self._prev_disablement_state = None
+        self._prev_rerouting_state = None
+
+    def __enter__(self):
+        self._enter_impl()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        self._enter_impl()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        self._prev_disablement_state = (
+            _WARNINGS_CONTROLLER.get_warnings_disablement_state_for_current_thread()
+        )
+        self._prev_rerouting_state = (
+            _WARNINGS_CONTROLLER.get_warnings_rerouting_state_for_current_thread()
+        )
+        try:
+            _WARNINGS_CONTROLLER.set_non_mlflow_warnings_disablement_state_for_current_thread(
+                disabled=self._disable_warnings
+            )
+            _WARNINGS_CONTROLLER.set_non_mlflow_warnings_rerouting_state_for_current_thread(
+                rerouted=self._reroute_warnings
+            )
+        except Exception:
+            pass
+
+    def _exit_impl(self, *args, **kwargs):
+        _WARNINGS_CONTROLLER.set_non_mlflow_warnings_disablement_state_for_current_thread(
+            disabled=self._prev_disablement_state
+        )
+        _WARNINGS_CONTROLLER.set_non_mlflow_warnings_rerouting_state_for_current_thread(
+            rerouted=self._prev_rerouting_state
+        )
+
+
+class MlflowEventsAndWarningsBehaviorGlobally:
+    """
+    Threadsafe context manager that modifies the behavior of MLflow event logging statements
+    and MLflow warnings upon entry, according to the specified parameters. Modifications are
+    applied globally across all threads and are not reverted until all threads that have made
+    a particular modification have exited the context.
+
+    Args:
+        disable_event_logs: If `True`, disable (mute & discard) MLflow event logging statements.
+            If `False`, do not disable MLflow event logging statements.
+        disable_warnings: If `True`, disable  (mutate & discard) MLflow warnings. If `False`,
+            do not disable MLflow warnings.
+        reroute_warnings: If `True`, reroute MLflow warnings to an MLflow event logger with
+            level WARNING. If `False`, do not reroute MLflow warnings.
+
+    """
+
+    _lock = RLock()
+    _disable_event_logs_count = 0
+    _disable_warnings_count = 0
+    _reroute_warnings_count = 0
+
+    def __init__(self, disable_event_logs, disable_warnings, reroute_warnings):
+        self._disable_event_logs = disable_event_logs
+        self._disable_warnings = disable_warnings
+        self._reroute_warnings = reroute_warnings
+
+    def __enter__(self):
+        self._enter_impl()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        self._enter_impl()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        try:
+            with MlflowEventsAndWarningsBehaviorGlobally._lock:
+                if self._disable_event_logs:
+                    if MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count <= 0:
+                        logging_utils.disable_logging()
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count += 1
+
+                if self._disable_warnings:
+                    if MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count <= 0:
+                        _WARNINGS_CONTROLLER.set_mlflow_warnings_disablement_state_globally(
+                            disabled=True
+                        )
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count += 1
+
+                if self._reroute_warnings:
+                    if MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count <= 0:
+                        _WARNINGS_CONTROLLER.set_mlflow_warnings_rerouting_state_globally(
+                            rerouted=True
+                        )
+                    MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count += 1
+        except Exception:
+            pass
+
+    def _exit_impl(self, *args, **kwargs):
+        try:
+            with MlflowEventsAndWarningsBehaviorGlobally._lock:
+                if self._disable_event_logs:
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count -= 1
+                if self._disable_warnings:
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count -= 1
+                if self._reroute_warnings:
+                    MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count -= 1
+
+                if MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count <= 0:
+                    logging_utils.enable_logging()
+                if MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count <= 0:
+                    _WARNINGS_CONTROLLER.set_mlflow_warnings_disablement_state_globally(
+                        disabled=False
+                    )
+                if MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count <= 0:
+                    _WARNINGS_CONTROLLER.set_mlflow_warnings_rerouting_state_globally(
+                        rerouted=False
+                    )
+        except Exception:
+            pass

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/metrics_queue.py

@@ -0,0 +1,69 @@
+import concurrent.futures
+from threading import RLock
+
+from mlflow.entities import Metric
+from mlflow.tracking.client import MlflowClient
+
+_metrics_queue_lock = RLock()
+_metrics_queue = []
+_thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+
+_MAX_METRIC_QUEUE_SIZE = 500
+
+
+def _assoc_list_to_map(lst):
+    """
+    Convert an association list to a dictionary.
+    """
+    d = {}
+    for run_id, metric in lst:
+        d[run_id] = d[run_id] + [metric] if run_id in d else [metric]
+    return d
+
+
+def flush_metrics_queue():
+    """Flush the metric queue and log contents in batches to MLflow.
+
+    Queue is divided into batches according to run id.
+    """
+    try:
+        # Multiple queue flushes may be scheduled simultaneously on different threads
+        # (e.g., if the queue is at its flush threshold and several more items
+        # are added before a flush occurs). For correctness and efficiency, only one such
+        # flush operation should proceed; all others are redundant and should be dropped
+        acquired_lock = _metrics_queue_lock.acquire(blocking=False)
+        if acquired_lock:
+            client = MlflowClient()
+            # For thread safety and to avoid modifying a list while iterating over it, we record a
+            # separate list of the items being flushed and remove each one from the metric queue,
+            # rather than clearing the metric queue or reassigning it (clearing / reassigning is
+            # dangerous because we don't block threads from adding to the queue while a flush is
+            # in progress)
+            snapshot = _metrics_queue[:]
+            for item in snapshot:
+                _metrics_queue.remove(item)
+
+            metrics_by_run = _assoc_list_to_map(snapshot)
+            for run_id, metrics in metrics_by_run.items():
+                client.log_batch(run_id, metrics=metrics, params=[], tags=[])
+    finally:
+        if acquired_lock:
+            _metrics_queue_lock.release()
+
+
+def add_to_metrics_queue(key, value, step, time, run_id):
+    """Add a metric to the metric queue.
+
+    Flush the queue if it exceeds max size.
+
+    Args:
+        key: string, the metrics key,
+        value: float, the metrics value.
+        step: int, the step of current metric.
+        time: int, the timestamp of current metric.
+        run_id: string, the run id of the associated mlflow run.
+    """
+    met = Metric(key=key, value=value, timestamp=time, step=step)
+    _metrics_queue.append((run_id, met))
+    if len(_metrics_queue) > _MAX_METRIC_QUEUE_SIZE:
+        _thread_pool.submit(flush_metrics_queue)

venv/lib/python3.9/site-packages/mlflow/utils/autologging_utils/versioning.py

@@ -0,0 +1,95 @@
+import importlib
+import importlib.metadata
+import re
+from typing import Literal
+
+from packaging.version import InvalidVersion, Version
+
+from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS, FLAVOR_TO_MODULE_NAME
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+
+
+def _check_version_in_range(ver, min_ver, max_ver):
+    return Version(min_ver) <= Version(ver) <= Version(max_ver)
+
+
+def _check_spark_version_in_range(ver, min_ver, max_ver):
+    """
+    Utility function for allowing late addition release changes to PySpark minor version increments
+    to be accepted, provided that the previous minor version has been previously validated.
+    For example, if version 3.2.1 has been validated as functional with MLflow, an upgrade of
+    PySpark's minor version to 3.2.2 will still provide a valid version check.
+    """
+    parsed_ver = Version(ver)
+    if parsed_ver > Version(min_ver):
+        ver = f"{parsed_ver.major}.{parsed_ver.minor}"
+    return _check_version_in_range(ver, min_ver, max_ver)
+
+
+def _violates_pep_440(ver):
+    try:
+        Version(ver)
+        return False
+    except InvalidVersion:
+        return True
+
+
+def _is_pre_or_dev_release(ver):
+    v = Version(ver)
+    return v.is_devrelease or v.is_prerelease
+
+
+def _strip_dev_version_suffix(version):
+    return re.sub(r"(\.?)dev.*", "", version)
+
+
+def get_min_max_version_and_pip_release(
+    flavor_name: str, category: Literal["autologging", "models"] = "autologging"
+):
+    if flavor_name == "pyspark.ml":
+        # pyspark.ml is a special case of spark flavor
+        flavor_name = "spark"
+
+    min_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["minimum"]
+    max_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["maximum"]
+    pip_release = _ML_PACKAGE_VERSIONS[flavor_name]["package_info"]["pip_release"]
+    return min_version, max_version, pip_release
+
+
+def is_flavor_supported_for_associated_package_versions(flavor_name):
+    """
+    Returns:
+        True if the specified flavor is supported for the currently-installed versions of its
+        associated packages.
+    """
+    module_name = FLAVOR_TO_MODULE_NAME[flavor_name]
+
+    try:
+        actual_version = importlib.import_module(module_name).__version__
+    except AttributeError:
+        try:
+            # NB: Module name is not necessarily the same as the package name. However,
+            # we assume they are the same here for simplicity. If they are not the same,
+            # this will fail and fallback to 'True', which is not a disaster.
+            actual_version = importlib.metadata.version(module_name)
+        except importlib.metadata.PackageNotFoundError:
+            # Some package (e.g. dspy) do not publish version info in a standard format.
+            # For this case, we assume the package version is supported by MLflow.
+            return True
+
+    # In Databricks, treat 'pyspark 3.x.y.dev0' as 'pyspark 3.x.y'
+    if module_name == "pyspark" and is_in_databricks_runtime():
+        actual_version = _strip_dev_version_suffix(actual_version)
+
+    if _violates_pep_440(actual_version) or _is_pre_or_dev_release(actual_version):
+        return False
+    min_version, max_version, _ = get_min_max_version_and_pip_release(flavor_name)
+
+    if module_name == "pyspark" and is_in_databricks_runtime():
+        # MLflow 1.25.0 is known to be compatible with PySpark 3.3.0 on Databricks, despite the
+        # fact that PySpark 3.3.0 was not available in PyPI at the time of the MLflow 1.25.0 release
+        if Version(max_version) < Version("3.3.0"):
+            max_version = "3.3.0"
+        return _check_spark_version_in_range(actual_version, min_version, max_version)
+    else:
+        return _check_version_in_range(actual_version, min_version, max_version)

venv/lib/python3.9/site-packages/mlflow/utils/checkpoint_utils.py

@@ -0,0 +1,206 @@
+import logging
+import os
+import posixpath
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.utils.autologging_utils import (
+    ExceptionSafeAbstractClass,
+)
+from mlflow.utils.file_utils import TempDir
+from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
+
+_logger = logging.getLogger(__name__)
+
+
+_CHECKPOINT_DIR = "checkpoints"
+_CHECKPOINT_METRIC_FILENAME = "checkpoint_metrics.json"
+_CHECKPOINT_MODEL_FILENAME = "checkpoint"
+_LATEST_CHECKPOINT_PREFIX = "latest_"
+_CHECKPOINT_EPOCH_PREFIX = "epoch_"
+_CHECKPOINT_GLOBAL_STEP_PREFIX = "global_step_"
+_WEIGHT_ONLY_CHECKPOINT_SUFFIX = ".weights"
+
+
+class MlflowModelCheckpointCallbackBase(metaclass=ExceptionSafeAbstractClass):
+    """Callback base class for automatic model checkpointing to MLflow.
+
+    You must implement "save_checkpoint" method to save the model as the checkpoint file.
+    and you must call `check_and_save_checkpoint_if_needed` method in relevant
+    callback events to trigger automatic checkpointing.
+
+    Args:
+        checkpoint_file_suffix: checkpoint file suffix.
+        monitor: In automatic model checkpointing, the metric name to monitor if
+            you set `model_checkpoint_save_best_only` to True.
+        save_best_only: If True, automatic model checkpointing only saves when
+            the model is considered the "best" model according to the quantity
+            monitored and previous checkpoint model is overwritten.
+        mode: one of {"min", "max"}. In automatic model checkpointing,
+            if save_best_only=True, the decision to overwrite the current save file is made
+            based on either the maximization or the minimization of the monitored quantity.
+        save_weights_only: In automatic model checkpointing, if True, then
+            only the model’s weights will be saved. Otherwise, the optimizer states,
+            lr-scheduler states, etc are added in the checkpoint too.
+        save_freq: `"epoch"` or integer. When using `"epoch"`, the callback
+            saves the model after each epoch. When using integer, the callback
+            saves the model at end of this many batches. Note that if the saving isn't
+            aligned to epochs, the monitored metric may potentially be less reliable (it
+            could reflect as little as 1 batch, since the metrics get reset
+            every epoch). Defaults to `"epoch"`.
+    """
+
+    def __init__(
+        self,
+        checkpoint_file_suffix,
+        monitor,
+        mode,
+        save_best_only,
+        save_weights_only,
+        save_freq,
+    ):
+        self.checkpoint_file_suffix = checkpoint_file_suffix
+        self.monitor = monitor
+        self.mode = mode
+        self.save_best_only = save_best_only
+        self.save_weights_only = save_weights_only
+        self.save_freq = save_freq
+        self.last_monitor_value = None
+
+        self.mlflow_tracking_uri = mlflow.get_tracking_uri()
+
+        if self.save_best_only:
+            if self.monitor is None:
+                raise MlflowException(
+                    "If checkpoint 'save_best_only' config is set to True, you need to set "
+                    "'monitor' config as well."
+                )
+            if self.mode not in ["min", "max"]:
+                raise MlflowException(
+                    "If checkpoint 'save_best_only' config is set to True, you need to set "
+                    "'mode' config and available modes includes 'min' and 'max', but you set "
+                    f"'mode' to '{self.mode}'."
+                )
+
+    def _is_new_checkpoint_better(self, new_monitor_value):
+        if self.last_monitor_value is None:
+            return True
+
+        if self.mode == "min":
+            return new_monitor_value < self.last_monitor_value
+
+        return new_monitor_value > self.last_monitor_value
+
+    def save_checkpoint(self, filepath: str):
+        raise NotImplementedError()
+
+    def check_and_save_checkpoint_if_needed(self, current_epoch, global_step, metric_dict):
+        # For distributed model training, trainer workers need to use the driver process
+        # mlflow_tracking_uri.
+        # Note that `self.mlflow_tracking_uri` value is assigned in the driver process
+        # then it is pickled to trainer workers.
+        mlflow.set_tracking_uri(self.mlflow_tracking_uri)
+
+        if self.save_best_only:
+            if self.monitor not in metric_dict:
+                _logger.warning(
+                    "Checkpoint logging is skipped, because checkpoint 'save_best_only' config is "
+                    "True, it requires to compare the monitored metric value, but the provided "
+                    "monitored metric value is not available."
+                )
+                return
+
+            new_monitor_value = metric_dict[self.monitor]
+            if not self._is_new_checkpoint_better(new_monitor_value):
+                # Current checkpoint is worse than last saved checkpoint,
+                # so skip checkpointing.
+                return
+
+            self.last_monitor_value = new_monitor_value
+
+        suffix = self.checkpoint_file_suffix
+
+        if self.save_best_only:
+            if self.save_weights_only:
+                checkpoint_model_filename = (
+                    f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}"
+                    f"{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
+                )
+            else:
+                checkpoint_model_filename = (
+                    f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}{suffix}"
+                )
+            checkpoint_metrics_filename = (
+                f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_METRIC_FILENAME}"
+            )
+            checkpoint_artifact_dir = _CHECKPOINT_DIR
+        else:
+            if self.save_freq == "epoch":
+                sub_dir_name = f"{_CHECKPOINT_EPOCH_PREFIX}{current_epoch}"
+            else:
+                sub_dir_name = f"{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}"
+
+            if self.save_weights_only:
+                checkpoint_model_filename = (
+                    f"{_CHECKPOINT_MODEL_FILENAME}{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
+                )
+            else:
+                checkpoint_model_filename = f"{_CHECKPOINT_MODEL_FILENAME}{suffix}"
+            checkpoint_metrics_filename = _CHECKPOINT_METRIC_FILENAME
+            checkpoint_artifact_dir = f"{_CHECKPOINT_DIR}/{sub_dir_name}"
+
+        mlflow.set_tag(
+            LATEST_CHECKPOINT_ARTIFACT_TAG_KEY,
+            f"{checkpoint_artifact_dir}/{checkpoint_model_filename}",
+        )
+
+        mlflow.log_dict(
+            {**metric_dict, "epoch": current_epoch, "global_step": global_step},
+            f"{checkpoint_artifact_dir}/{checkpoint_metrics_filename}",
+        )
+
+        with TempDir() as tmp_dir:
+            tmp_model_save_path = os.path.join(tmp_dir.path(), checkpoint_model_filename)
+            self.save_checkpoint(tmp_model_save_path)
+            mlflow.log_artifact(tmp_model_save_path, checkpoint_artifact_dir)
+
+
+def download_checkpoint_artifact(run_id=None, epoch=None, global_step=None, dst_path=None):
+    from mlflow.client import MlflowClient
+    from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
+
+    client = MlflowClient()
+
+    if run_id is None:
+        run = mlflow.active_run()
+        if run is None:
+            raise MlflowException(
+                "There is no active run, please provide the 'run_id' argument for "
+                "'load_checkpoint' invocation."
+            )
+        run_id = run.info.run_id
+    else:
+        run = client.get_run(run_id)
+
+    latest_checkpoint_artifact_path = run.data.tags.get(LATEST_CHECKPOINT_ARTIFACT_TAG_KEY)
+    if latest_checkpoint_artifact_path is None:
+        raise MlflowException("There is no logged checkpoint artifact in the current run.")
+
+    checkpoint_filename = posixpath.basename(latest_checkpoint_artifact_path)
+
+    if epoch is not None and global_step is not None:
+        raise MlflowException(
+            "Only one of 'epoch' and 'global_step' can be set for 'load_checkpoint'."
+        )
+    elif global_step is not None:
+        checkpoint_artifact_path = (
+            f"{_CHECKPOINT_DIR}/{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}/{checkpoint_filename}"
+        )
+    elif epoch is not None:
+        checkpoint_artifact_path = (
+            f"{_CHECKPOINT_DIR}/{_CHECKPOINT_EPOCH_PREFIX}{epoch}/{checkpoint_filename}"
+        )
+    else:
+        checkpoint_artifact_path = latest_checkpoint_artifact_path
+
+    return client.download_artifacts(run_id, checkpoint_artifact_path, dst_path=dst_path)

venv/lib/python3.9/site-packages/mlflow/utils/class_utils.py

@@ -0,0 +1,6 @@
+import importlib
+
+
+def _get_class_from_string(fully_qualified_class_name):
+    module, class_name = fully_qualified_class_name.rsplit(".", maxsplit=1)
+    return getattr(importlib.import_module(module), class_name)

venv/lib/python3.9/site-packages/mlflow/utils/cli_args.py

@@ -0,0 +1,257 @@
+"""
+Definitions of click options shared by several CLI commands.
+"""
+
+import warnings
+
+import click
+
+from mlflow.environment_variables import MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING
+from mlflow.utils import env_manager as _EnvManager
+
+MODEL_PATH = click.option(
+    "--model-path",
+    "-m",
+    default=None,
+    metavar="PATH",
+    required=True,
+    help="Path to the model. The path is relative to the run with the given "
+    "run-id or local filesystem path without run-id.",
+)
+
+_model_uri_help_string = (
+    "URI to the model. A local path, a 'runs:/' URI, or a"
+    " remote storage URI (e.g., an 's3://' URI). For more information"
+    " about supported remote URIs for model artifacts, see"
+    " https://mlflow.org/docs/latest/tracking.html#artifact-stores"
+)
+
+MODEL_URI_BUILD_DOCKER = click.option(
+    "--model-uri",
+    "-m",
+    metavar="URI",
+    default=None,
+    required=False,
+    help="[Optional] " + _model_uri_help_string,
+)
+
+MODEL_URI = click.option(
+    "--model-uri",
+    "-m",
+    metavar="URI",
+    required=True,
+    help=_model_uri_help_string,
+)
+
+MLFLOW_HOME = click.option(
+    "--mlflow-home",
+    default=None,
+    metavar="PATH",
+    help="Path to local clone of MLflow project. Use for development only.",
+)
+
+RUN_ID = click.option(
+    "--run-id",
+    "-r",
+    default=None,
+    required=False,
+    metavar="ID",
+    help="ID of the MLflow run that generated the referenced content.",
+)
+
+
+def _resolve_env_manager(_, __, env_manager):
+    if env_manager is not None:
+        _EnvManager.validate(env_manager)
+        if env_manager == _EnvManager.CONDA and not MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING.get():
+            warnings.warn(
+                (
+                    "Use of conda is discouraged. If you use it, please ensure that your use of "
+                    "conda complies with Anaconda's terms of service "
+                    "(https://legal.anaconda.com/policies/en/?name=terms-of-service). "
+                    "virtualenv is the recommended tool for environment reproducibility. "
+                    f"To suppress this warning, set the {MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING} "
+                    "environment variable to 'TRUE'."
+                ),
+                UserWarning,
+                stacklevel=2,
+            )
+        return env_manager
+
+    return None
+
+
+def _create_env_manager_option(help_string, default=None):
+    return click.option(
+        "--env-manager",
+        default=default,
+        type=click.UNPROCESSED,
+        callback=_resolve_env_manager,
+        help=help_string,
+    )
+
+
+ENV_MANAGER = _create_env_manager_option(
+    default=_EnvManager.VIRTUALENV,
+    # '\b' prevents rewrapping text:
+    # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
+    help_string="""
+If specified, create an environment for MLmodel using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, default to virtualenv.
+""",
+)
+
+ENV_MANAGER_PROJECTS = _create_env_manager_option(
+    help_string="""
+If specified, create an environment for MLproject using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, the appropriate environment manager is automatically selected based on
+the project configuration. For example, if `MLproject.yaml` contains a `python_env` key,
+virtualenv is used.
+""",
+)
+
+ENV_MANAGER_DOCKERFILE = _create_env_manager_option(
+    default=None,
+    # '\b' prevents rewrapping text:
+    # https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
+    help_string="""
+If specified, create an environment for MLmodel using the specified
+environment manager. The following values are supported:
+
+\b
+- local: use the local environment
+- virtualenv: use virtualenv (and pyenv for Python version management)
+- conda: use conda
+
+If unspecified, default to None, then MLflow will automatically pick the env manager
+based on the model's flavor configuration.
+If model-uri is specified: if python version is specified in the flavor configuration
+and no java installation is required, then we use local environment. Otherwise we use virtualenv.
+If no model-uri is provided, we use virtualenv.
+""",
+)
+
+
+INSTALL_MLFLOW = click.option(
+    "--install-mlflow",
+    is_flag=True,
+    default=False,
+    help="If specified and there is a conda or virtualenv environment to be activated "
+    "mlflow will be installed into the environment after it has been "
+    "activated. The version of installed mlflow will be the same as "
+    "the one used to invoke this command.",
+)
+
+HOST = click.option(
+    "--host",
+    "-h",
+    envvar="MLFLOW_HOST",
+    metavar="HOST",
+    default="127.0.0.1",
+    help="The network address to listen on (default: 127.0.0.1). "
+    "Use 0.0.0.0 to bind to all addresses if you want to access the tracking "
+    "server from other machines.",
+)
+
+PORT = click.option(
+    "--port",
+    "-p",
+    envvar="MLFLOW_PORT",
+    default=5000,
+    help="The port to listen on (default: 5000).",
+)
+
+TIMEOUT = click.option(
+    "--timeout",
+    "-t",
+    envvar="MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT",
+    default=60,
+    help="Timeout in seconds to serve a request (default: 60).",
+)
+
+# We use None to disambiguate manually selecting "4"
+WORKERS = click.option(
+    "--workers",
+    "-w",
+    envvar="MLFLOW_WORKERS",
+    default=None,
+    help="Number of gunicorn worker processes to handle requests (default: 4).",
+)
+
+MODELS_WORKERS = click.option(
+    "--workers",
+    "-w",
+    envvar="MLFLOW_MODELS_WORKERS",
+    default=None,
+    help="Number of uvicorn workers to handle requests when serving mlflow models (default: 1).",
+)
+
+ENABLE_MLSERVER = click.option(
+    "--enable-mlserver",
+    is_flag=True,
+    default=False,
+    help=(
+        "Enable serving with MLServer through the v2 inference protocol. "
+        "You can use environment variables to configure MLServer. "
+        "(See https://mlserver.readthedocs.io/en/latest/reference/settings.html)"
+    ),
+)
+
+ARTIFACTS_DESTINATION = click.option(
+    "--artifacts-destination",
+    envvar="MLFLOW_ARTIFACTS_DESTINATION",
+    metavar="URI",
+    default="./mlartifacts",
+    help=(
+        "The base artifact location from which to resolve artifact upload/download/list requests "
+        "(e.g. 's3://my-bucket'). Defaults to a local './mlartifacts' directory. This option only "
+        "applies when the tracking server is configured to stream artifacts and the experiment's "
+        "artifact root location is http or mlflow-artifacts URI."
+    ),
+)
+
+SERVE_ARTIFACTS = click.option(
+    "--serve-artifacts/--no-serve-artifacts",
+    envvar="MLFLOW_SERVE_ARTIFACTS",
+    is_flag=True,
+    default=True,
+    help="Enables serving of artifact uploads, downloads, and list requests "
+    "by routing these requests to the storage location that is specified by "
+    "'--artifacts-destination' directly through a proxy. The default location that "
+    "these requests are served from is a local './mlartifacts' directory which can be "
+    "overridden via the '--artifacts-destination' argument. To disable artifact serving, "
+    "specify `--no-serve-artifacts`. Default: True",
+)
+
+NO_CONDA = click.option(
+    "--no-conda",
+    is_flag=True,
+    help="If specified, use local environment.",
+)
+
+INSTALL_JAVA = click.option(
+    "--install-java",
+    is_flag=False,
+    flag_value=True,
+    default=None,
+    type=bool,
+    help="Installs Java in the image if needed. Default is None, "
+    "allowing MLflow to determine installation. Flavors requiring "
+    "Java, such as Spark, enable this automatically. "
+    "Note: This option only works with the UBUNTU base image; "
+    "Python base images do not support Java installation.",
+)

venv/lib/python3.9/site-packages/mlflow/utils/conda.py

@@ -0,0 +1,354 @@
+import hashlib
+import json
+import logging
+import os
+
+import yaml
+
+from mlflow.environment_variables import MLFLOW_CONDA_CREATE_ENV_CMD, MLFLOW_CONDA_HOME
+from mlflow.exceptions import ExecutionException
+from mlflow.utils import process
+from mlflow.utils.environment import Environment
+from mlflow.utils.os import is_windows
+
+_logger = logging.getLogger(__name__)
+
+CONDA_EXE = "CONDA_EXE"
+
+
+def get_conda_command(conda_env_name):
+    #  Checking for newer conda versions
+    if not is_windows() and (CONDA_EXE in os.environ or MLFLOW_CONDA_HOME.defined):
+        conda_path = get_conda_bin_executable("conda")
+        activate_conda_env = [f"source {os.path.dirname(conda_path)}/../etc/profile.d/conda.sh"]
+        activate_conda_env += [f"conda activate {conda_env_name} 1>&2"]
+    else:
+        activate_path = get_conda_bin_executable("activate")
+        # in case os name is not 'nt', we are not running on windows. It introduces
+        # bash command otherwise.
+        if not is_windows():
+            return [f"source {activate_path} {conda_env_name} 1>&2"]
+        else:
+            return [f"conda activate {conda_env_name}"]
+    return activate_conda_env
+
+
+def get_conda_bin_executable(executable_name):
+    """
+    Return path to the specified executable, assumed to be discoverable within the 'bin'
+    subdirectory of a conda installation.
+
+    The conda home directory (expected to contain a 'bin' subdirectory) is configurable via the
+    ``mlflow.projects.MLFLOW_CONDA_HOME`` environment variable. If
+    ``mlflow.projects.MLFLOW_CONDA_HOME`` is unspecified, this method simply returns the passed-in
+    executable name.
+    """
+    if conda_home := MLFLOW_CONDA_HOME.get():
+        return os.path.join(conda_home, f"bin/{executable_name}")
+    # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
+    if conda_exe := os.getenv(CONDA_EXE):
+        conda_bin_dir = os.path.dirname(conda_exe)
+        return os.path.join(conda_bin_dir, executable_name)
+    return executable_name
+
+
+def _get_conda_env_name(conda_env_path, env_id=None, env_root_dir=None):
+    if conda_env_path:
+        with open(conda_env_path) as f:
+            conda_env_contents = f.read()
+    else:
+        conda_env_contents = ""
+
+    if env_id:
+        conda_env_contents += env_id
+
+    env_name = "mlflow-{}".format(
+        hashlib.sha1(conda_env_contents.encode("utf-8"), usedforsecurity=False).hexdigest()
+    )
+    if env_root_dir:
+        env_root_dir = os.path.normpath(env_root_dir)
+        # Generate env name with format "mlflow-{conda_env_contents_hash}-{env_root_dir_hash}"
+        # hashing `conda_env_contents` and `env_root_dir` separately helps debugging
+        env_name += "-{}".format(
+            hashlib.sha1(env_root_dir.encode("utf-8"), usedforsecurity=False).hexdigest()
+        )
+
+    return env_name
+
+
+def _get_conda_executable_for_create_env():
+    """
+    Returns the executable that should be used to create environments. This is "conda"
+    by default, but it can be set to something else by setting the environment variable
+
+    """
+
+    return get_conda_bin_executable(MLFLOW_CONDA_CREATE_ENV_CMD.get())
+
+
+def _list_conda_environments(extra_env=None):
+    """Return a list of names of conda environments.
+
+    Args:
+        extra_env: Extra environment variables for running "conda env list" command.
+
+    """
+    prc = process._exec_cmd(
+        [get_conda_bin_executable("conda"), "env", "list", "--json"], extra_env=extra_env
+    )
+    return list(map(os.path.basename, json.loads(prc.stdout).get("envs", [])))
+
+
+_CONDA_ENVS_DIR = "conda_envs"
+_CONDA_CACHE_PKGS_DIR = "conda_cache_pkgs"
+_PIP_CACHE_DIR = "pip_cache_pkgs"
+
+
+def _create_conda_env(
+    conda_env_path,
+    conda_env_create_path,
+    project_env_name,
+    conda_extra_env_vars,
+    capture_output,
+):
+    if conda_env_path:
+        process._exec_cmd(
+            [
+                conda_env_create_path,
+                "env",
+                "create",
+                "-n",
+                project_env_name,
+                "--file",
+                conda_env_path,
+            ],
+            extra_env=conda_extra_env_vars,
+            capture_output=capture_output,
+        )
+    else:
+        process._exec_cmd(
+            [
+                conda_env_create_path,
+                "create",
+                "--channel",
+                "conda-forge",
+                "--yes",
+                "--override-channels",
+                "-n",
+                project_env_name,
+                "python",
+            ],
+            extra_env=conda_extra_env_vars,
+            capture_output=capture_output,
+        )
+
+    return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
+
+
+def _create_conda_env_retry(
+    conda_env_path, conda_env_create_path, project_env_name, conda_extra_env_vars, _capture_output
+):
+    """
+    `conda env create` command can fail due to network issues such as `ConnectionResetError`
+    while collecting package metadata. This function retries the command up to 3 times.
+    """
+    num_attempts = 3
+    for attempt in range(num_attempts):
+        try:
+            return _create_conda_env(
+                conda_env_path,
+                conda_env_create_path,
+                project_env_name,
+                conda_extra_env_vars,
+                capture_output=True,
+            )
+        except process.ShellCommandException as e:
+            error_str = str(e)
+            if (num_attempts - attempt - 1) > 0 and (
+                "ConnectionResetError" in error_str or "ChunkedEncodingError" in error_str
+            ):
+                _logger.warning("Conda env creation failed due to network issue. Retrying...")
+                continue
+            raise
+
+
+def _get_conda_extra_env_vars(env_root_dir=None):
+    """
+    Given the `env_root_dir` (See doc of PyFuncBackend constructor argument `env_root_dir`),
+    return a dict of environment variables which are used to config conda to generate envs
+    under the expected `env_root_dir`.
+    """
+    if env_root_dir is None:
+        return None
+
+    # Create isolated conda package cache dir "conda_pkgs" under the env_root_dir
+    # for each python process.
+    # Note: shared conda package cache dir causes race condition issues:
+    # See https://github.com/conda/conda/issues/8870
+    # See https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-environment-directories-envs-dirs
+    # and https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-package-directories-pkgs-dirs
+
+    conda_envs_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR)
+    conda_pkgs_path = os.path.join(env_root_dir, _CONDA_CACHE_PKGS_DIR)
+    pip_cache_dir = os.path.join(env_root_dir, _PIP_CACHE_DIR)
+
+    os.makedirs(conda_envs_path, exist_ok=True)
+    os.makedirs(conda_pkgs_path, exist_ok=True)
+    os.makedirs(pip_cache_dir, exist_ok=True)
+
+    return {
+        "CONDA_ENVS_PATH": conda_envs_path,
+        "CONDA_PKGS_DIRS": conda_pkgs_path,
+        "PIP_CACHE_DIR": pip_cache_dir,
+        # PIP_NO_INPUT=1 makes pip run in non-interactive mode,
+        # otherwise pip might prompt "yes or no" and ask stdin input
+        "PIP_NO_INPUT": "1",
+    }
+
+
+def get_or_create_conda_env(
+    conda_env_path,
+    env_id=None,
+    capture_output=False,
+    env_root_dir=None,
+    pip_requirements_override=None,
+    extra_envs=None,
+):
+    """Given a `Project`, creates a conda environment containing the project's dependencies if such
+    a conda environment doesn't already exist. Returns the name of the conda environment.
+
+    Args:
+        conda_env_path: Path to a conda yaml file.
+        env_id: Optional string that is added to the contents of the yaml file before
+            calculating the hash. It can be used to distinguish environments that have the
+            same conda dependencies but are supposed to be different based on the context.
+            For example, when serving the model we may install additional dependencies to the
+            environment after the environment has been activated.
+        capture_output: Specify the capture_output argument while executing the
+            "conda env create" command.
+        env_root_dir: See doc of PyFuncBackend constructor argument `env_root_dir`.
+        pip_requirements_override: If specified, install the specified python dependencies to
+            the environment (upgrade if already installed).
+        extra_envs: If specified, a dictionary of extra environment variables will be passed to the
+            model inference environment.
+
+    Returns:
+        The name of the conda environment.
+
+    """
+
+    conda_path = get_conda_bin_executable("conda")
+    conda_env_create_path = _get_conda_executable_for_create_env()
+
+    try:
+        # Checks if Conda executable exists
+        process._exec_cmd([conda_path, "--help"], throw_on_error=False, extra_env=extra_envs)
+    except OSError:
+        raise ExecutionException(
+            f"Could not find Conda executable at {conda_path}. "
+            "Ensure Conda is installed as per the instructions at "
+            "https://conda.io/projects/conda/en/latest/"
+            "user-guide/install/index.html. "
+            "You can also configure MLflow to look for a specific "
+            f"Conda executable by setting the {MLFLOW_CONDA_HOME} environment variable "
+            "to the path of the Conda executable"
+        )
+
+    try:
+        # Checks if executable for environment creation exists
+        process._exec_cmd(
+            [conda_env_create_path, "--help"], throw_on_error=False, extra_env=extra_envs
+        )
+    except OSError:
+        raise ExecutionException(
+            f"You have set the env variable {MLFLOW_CONDA_CREATE_ENV_CMD}, but "
+            f"{conda_env_create_path} does not exist or it is not working properly. "
+            f"Note that {conda_env_create_path} and the conda executable need to be "
+            "in the same conda environment. You can change the search path by"
+            f"modifying the env variable {MLFLOW_CONDA_HOME}"
+        )
+
+    conda_extra_env_vars = _get_conda_extra_env_vars(env_root_dir)
+    if extra_envs:
+        conda_extra_env_vars.update(extra_envs)
+
+    # Include the env_root_dir hash in the project_env_name,
+    # this is for avoid conda env name conflicts between different CONDA_ENVS_PATH.
+    project_env_name = _get_conda_env_name(conda_env_path, env_id=env_id, env_root_dir=env_root_dir)
+    if env_root_dir is not None:
+        project_env_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR, project_env_name)
+    else:
+        project_env_path = project_env_name
+
+    if project_env_name in _list_conda_environments(conda_extra_env_vars):
+        _logger.info("Conda environment %s already exists.", project_env_path)
+        return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
+
+    _logger.info("=== Creating conda environment %s ===", project_env_path)
+    try:
+        _create_conda_env_func = (
+            # Retry conda env creation in a pytest session to avoid flaky test failures
+            _create_conda_env_retry if "PYTEST_CURRENT_TEST" in os.environ else _create_conda_env
+        )
+        conda_env = _create_conda_env_func(
+            conda_env_path,
+            conda_env_create_path,
+            project_env_name,
+            conda_extra_env_vars,
+            capture_output,
+        )
+
+        if pip_requirements_override:
+            _logger.info(
+                "Installing additional dependencies specified"
+                f"by pip_requirements_override: {pip_requirements_override}"
+            )
+            cmd = [
+                conda_path,
+                "install",
+                "-n",
+                project_env_name,
+                "--yes",
+                *pip_requirements_override,
+            ]
+            process._exec_cmd(cmd, extra_env=conda_extra_env_vars, capture_output=capture_output)
+
+        return conda_env
+
+    except Exception:
+        try:
+            if project_env_name in _list_conda_environments(conda_extra_env_vars):
+                _logger.warning(
+                    "Encountered unexpected error while creating conda environment. Removing %s.",
+                    project_env_path,
+                )
+                process._exec_cmd(
+                    [
+                        conda_path,
+                        "remove",
+                        "--yes",
+                        "--name",
+                        project_env_name,
+                        "--all",
+                    ],
+                    extra_env=conda_extra_env_vars,
+                    capture_output=False,
+                )
+        except Exception as e:
+            _logger.warning(
+                "Removing conda environment %s failed (error: %s)",
+                project_env_path,
+                repr(e),
+            )
+        raise
+
+
+def _get_conda_dependencies(conda_yaml_path):
+    """Extracts conda dependencies from a conda yaml file.
+
+    Args:
+        conda_yaml_path: Conda yaml file path.
+    """
+    with open(conda_yaml_path) as f:
+        conda_yaml = yaml.safe_load(f)
+        return [d for d in conda_yaml.get("dependencies", []) if isinstance(d, str)]

venv/lib/python3.9/site-packages/mlflow/utils/credentials.py

@@ -0,0 +1,231 @@
+import configparser
+import getpass
+import logging
+import os
+from typing import NamedTuple, Optional
+
+from mlflow.environment_variables import (
+    MLFLOW_TRACKING_AUTH,
+    MLFLOW_TRACKING_AWS_SIGV4,
+    MLFLOW_TRACKING_CLIENT_CERT_PATH,
+    MLFLOW_TRACKING_INSECURE_TLS,
+    MLFLOW_TRACKING_PASSWORD,
+    MLFLOW_TRACKING_SERVER_CERT_PATH,
+    MLFLOW_TRACKING_TOKEN,
+    MLFLOW_TRACKING_USERNAME,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.utils.rest_utils import MlflowHostCreds
+
+_logger = logging.getLogger(__name__)
+
+
+class MlflowCreds(NamedTuple):
+    username: Optional[str]
+    password: Optional[str]
+
+
+def _get_credentials_path() -> str:
+    return os.path.expanduser("~/.mlflow/credentials")
+
+
+def _read_mlflow_creds_from_file() -> tuple[Optional[str], Optional[str]]:
+    path = _get_credentials_path()
+    if not os.path.exists(path):
+        return None, None
+
+    config = configparser.ConfigParser()
+    config.read(path)
+    if "mlflow" not in config:
+        return None, None
+
+    mlflow_cfg = config["mlflow"]
+    username_key = MLFLOW_TRACKING_USERNAME.name.lower()
+    password_key = MLFLOW_TRACKING_PASSWORD.name.lower()
+    return mlflow_cfg.get(username_key), mlflow_cfg.get(password_key)
+
+
+def _read_mlflow_creds_from_env() -> tuple[Optional[str], Optional[str]]:
+    return MLFLOW_TRACKING_USERNAME.get(), MLFLOW_TRACKING_PASSWORD.get()
+
+
+def read_mlflow_creds() -> MlflowCreds:
+    username_file, password_file = _read_mlflow_creds_from_file()
+    username_env, password_env = _read_mlflow_creds_from_env()
+    return MlflowCreds(
+        username=username_env or username_file,
+        password=password_env or password_file,
+    )
+
+
+def get_default_host_creds(store_uri):
+    creds = read_mlflow_creds()
+    return MlflowHostCreds(
+        host=store_uri,
+        username=creds.username,
+        password=creds.password,
+        token=MLFLOW_TRACKING_TOKEN.get(),
+        aws_sigv4=MLFLOW_TRACKING_AWS_SIGV4.get(),
+        auth=MLFLOW_TRACKING_AUTH.get(),
+        ignore_tls_verification=MLFLOW_TRACKING_INSECURE_TLS.get(),
+        client_cert_path=MLFLOW_TRACKING_CLIENT_CERT_PATH.get(),
+        server_cert_path=MLFLOW_TRACKING_SERVER_CERT_PATH.get(),
+    )
+
+
+def login(backend: str = "databricks", interactive: bool = True) -> None:
+    """Configure MLflow server authentication and connect MLflow to tracking server.
+
+    This method provides a simple way to connect MLflow to its tracking server. Currently only
+    Databricks tracking server is supported. Users will be prompted to enter the credentials if no
+    existing Databricks profile is found, and the credentials will be saved to `~/.databrickscfg`.
+
+    Args:
+        backend: string, the backend of the tracking server. Currently only "databricks" is
+            supported.
+
+        interactive: bool, controls request for user input on missing credentials. If true, user
+            input will be requested if no credentials are found, otherwise an exception will be
+            raised if no credentials are found.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        mlflow.login()
+        with mlflow.start_run():
+            mlflow.log_param("p", 0)
+    """
+    from mlflow.tracking import set_tracking_uri
+
+    if backend == "databricks":
+        _databricks_login(interactive)
+        set_tracking_uri("databricks")
+    else:
+        raise MlflowException(
+            f"Currently only 'databricks' backend is supported, received `backend={backend}`."
+        )
+
+
+def _validate_databricks_auth():
+    # Check if databricks credentials are valid.
+    try:
+        from databricks.sdk import WorkspaceClient
+    except ImportError:
+        raise ImportError(
+            "Databricks SDK is not installed. To use `mlflow.login()`, please install "
+            "databricks-sdk by `pip install databricks-sdk`."
+        )
+
+    try:
+        w = WorkspaceClient()
+        if "community" in w.config.host:
+            # Databricks community edition cannot use `w.current_user.me()` for auth validation.
+            w.clusters.list_zones()
+        else:
+            # If credentials are invalid, `w.current_user.me()` will throw an error.
+            w.current_user.me()
+        _logger.info(
+            f"Successfully connected to MLflow hosted tracking server! Host: {w.config.host}."
+        )
+    except Exception as e:
+        raise MlflowException(f"Failed to validate databricks credentials: {e}")
+
+
+def _overwrite_or_create_databricks_profile(
+    file_name,
+    profile,
+    profile_name="DEFAULT",
+):
+    """Overwrite or create a profile in the databricks config file.
+
+    Args:
+        file_name: string, the file name of the databricks config file, usually `~/.databrickscfg`.
+        profile: dict, contains the authentiacation profile information.
+        profile_name: string, the name of the profile to be overwritten or created.
+    """
+    profile_name = f"[{profile_name}]"
+    lines = []
+    # Read `file_name` if the file exists, otherwise `lines=[]`.
+    if os.path.exists(file_name):
+        with open(file_name) as file:
+            lines = file.readlines()
+    start_index = -1
+    end_index = -1
+    # Find the start and end indices of the profile to overwrite.
+    for i in range(len(lines)):
+        if lines[i].strip() == profile_name:
+            start_index = i
+            break
+
+    if start_index != -1:
+        for i in range(start_index + 1, len(lines)):
+            # Reach an empty line or a new profile.
+            if lines[i].strip() == "" or lines[i].startswith("["):
+                end_index = i
+                break
+        end_index = end_index if end_index != -1 else len(lines)
+        del lines[start_index : end_index + 1]
+
+    # Write the new profile to the top of the file.
+    new_profile = []
+    new_profile.append(profile_name + "\n")
+    new_profile.append(f"host = {profile['host']}\n")
+    if "token" in profile:
+        new_profile.append(f"token = {profile['token']}\n")
+    else:
+        new_profile.append(f"username = {profile['username']}\n")
+        new_profile.append(f"password = {profile['password']}\n")
+    new_profile.append("\n")
+    lines = new_profile + lines
+
+    # Write back the modified lines to the file.
+    with open(file_name, "w") as file:
+        file.writelines(lines)
+
+
+def _databricks_login(interactive):
+    """Set up databricks authentication."""
+    try:
+        # Failed validation will throw an error.
+        _validate_databricks_auth()
+        return
+    except Exception:
+        if interactive:
+            _logger.info("No valid Databricks credentials found, please enter your credentials...")
+        else:
+            raise MlflowException(
+                "No valid Databricks credentials found while running in non-interactive mode."
+            )
+
+    while True:
+        host = input("Databricks Host (should begin with https://): ")
+        if not host.startswith("https://"):
+            _logger.error("Invalid host: {host}, host must begin with https://, please retry.")
+        break
+
+    profile = {"host": host}
+    if "community" in host:
+        # Databricks community edition requires username and password for authentication.
+        username = input("Username: ")
+        password = getpass.getpass("Password: ")
+        profile["username"] = username
+        profile["password"] = password
+    else:
+        # Production or staging Databricks requires personal token for authentication.
+        token = getpass.getpass("Token: ")
+        profile["token"] = token
+
+    file_name = os.environ.get(
+        "DATABRICKS_CONFIG_FILE", f"{os.path.expanduser('~')}/.databrickscfg"
+    )
+    profile_name = os.environ.get("DATABRICKS_CONFIG_PROFILE", "DEFAULT")
+    _overwrite_or_create_databricks_profile(file_name, profile, profile_name)
+
+    try:
+        # Failed validation will throw an error.
+        _validate_databricks_auth()
+    except Exception as e:
+        # If user entered invalid auth, we will raise an error and ask users to retry.
+        raise MlflowException(f"`mlflow.login()` failed with error: {e}")

venv/lib/python3.9/site-packages/mlflow/utils/data_utils.py

@@ -0,0 +1,17 @@
+import urllib.parse
+
+
+def parse_s3_uri(uri):
+    """Parse an S3 URI, returning (bucket, path)"""
+    parsed = urllib.parse.urlparse(uri)
+    if parsed.scheme != "s3":
+        raise Exception(f"Not an S3 URI: {uri}")
+    path = parsed.path
+    if path.startswith("/"):
+        path = path[1:]
+    return parsed.netloc, path
+
+
+def is_uri(string):
+    parsed_uri = urllib.parse.urlparse(string)
+    return len(parsed_uri.scheme) > 0

venv/lib/python3.9/site-packages/mlflow/utils/docstring_utils.py

@@ -0,0 +1,475 @@
+import textwrap
+import warnings
+
+from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
+from mlflow.utils.autologging_utils.versioning import (
+    get_min_max_version_and_pip_release,
+)
+
+
+def _create_placeholder(key: str):
+    return "{{ " + key + " }}"
+
+
+def _replace_keys_with_placeholders(d: dict) -> dict:
+    return {_create_placeholder(k): v for k, v in d.items()}
+
+
+def _get_indentation_of_key(line: str, placeholder: str) -> str:
+    index = line.find(placeholder)
+    return (index * " ") if index != -1 else ""
+
+
+def _indent(text: str, indent: str) -> str:
+    """Indent everything but first line in text."""
+    lines = text.splitlines()
+    if len(lines) <= 1:
+        return text
+
+    else:
+        first_line = lines[0]
+        subsequent_lines = "\n".join(list(lines[1:]))
+        indented_subsequent_lines = textwrap.indent(subsequent_lines, indent)
+        return first_line + "\n" + indented_subsequent_lines
+
+
+def _replace_all(text: str, replacements: dict[str, str]) -> str:
+    """
+    Replace all instances of replacements.keys() with their corresponding
+    values in text. The replacements will be inserted on the same line
+    with wrapping to the same level of indentation, for example:
+
+    ```
+    Args:
+        param_1: {{ key }}
+    ```
+
+    will become...
+
+    ```
+    Args:
+        param_1: replaced_value_at same indentation as prior
+                 and if there are more lines they will also
+                 have the same indentation.
+    ```
+    """
+    for key, value in replacements.items():
+        if key in text:
+            indent = _get_indentation_of_key(text, key)
+            indented_value = _indent(value, indent)
+            text = text.replace(key, indented_value)
+    return text
+
+
+class ParamDocs(dict):
+    """
+    Represents a set of parameter documents in the docstring.
+    """
+
+    def __repr__(self):
+        return f"ParamDocs({super().__repr__()})"
+
+    def format(self, **kwargs):
+        """
+        Formats values to be substituted in via the format_docstring() method.
+
+        Args:
+            kwargs: A `dict` in the form of `{"< placeholder name >": "< value >"}`.
+
+        Returns:
+            A new `ParamDocs` instance with the formatted param docs.
+
+        .. code-block:: text
+            :caption: Example
+
+            >>> pd = ParamDocs(p1="{{ doc1 }}", p2="{{ doc2 }}")
+            >>> pd.format(doc1="foo", doc2="bar")
+            ParamDocs({'p1': 'foo', 'p2': 'bar'})
+        """
+        replacements = _replace_keys_with_placeholders(kwargs)
+        return ParamDocs({k: _replace_all(v, replacements) for k, v in self.items()})
+
+    def format_docstring(self, docstring: str) -> str:
+        """
+        Formats placeholders in `docstring`.
+
+        Args:
+            docstring: A docstring with placeholders to be replaced.
+                If provided with None, will return None.
+
+        .. code-block:: text
+            :caption: Example
+
+            >>> pd = ParamDocs(p1="doc1", p2="doc2
+            doc2 second line")
+            >>> docstring = '''
+            ... Args:
+            ...     p1: {{ p1 }}
+            ...     p2: {{ p2 }}
+            ... '''.strip()
+            >>> print(pd.format_docstring(docstring))
+        """
+        if docstring is None:
+            return None
+
+        replacements = _replace_keys_with_placeholders(self)
+        lines = docstring.splitlines()
+        for i, line in enumerate(lines):
+            lines[i] = _replace_all(line, replacements)
+
+        return "\n".join(lines)
+
+
+def format_docstring(param_docs):
+    """
+    Returns a decorator that replaces param doc placeholders (e.g. '{{ param_name }}') in the
+    docstring of the decorated function.
+
+    Args:
+        param_docs: A `ParamDocs` instance or `dict`.
+
+    Returns:
+        A decorator to apply the formatting.
+
+    .. code-block:: text
+        :caption: Example
+
+        >>> param_docs = {"p1": "doc1", "p2": "doc2
+        doc2 second line"}
+        >>> @format_docstring(param_docs)
+        ... def func(p1, p2):
+        ...     '''
+        ...     Args:
+        ...         p1: {{ p1 }}
+        ...         p2: {{ p2 }}
+        ...     '''
+        >>> import textwrap
+        >>> print(textwrap.dedent(func.__doc__).strip())
+
+        Args:
+            p1: doc1
+            p2: doc2
+                doc2 second line
+    """
+    param_docs = ParamDocs(param_docs)
+
+    def decorator(func):
+        func.__doc__ = param_docs.format_docstring(func.__doc__)
+        return func
+
+    return decorator
+
+
+# `{{ ... }}` represents a placeholder.
+LOG_MODEL_PARAM_DOCS = ParamDocs(
+    {
+        "conda_env": (
+            """Either a dictionary representation of a Conda environment or the path to a conda
+environment yaml file. If provided, this describes the environment this model should be run in.
+At a minimum, it should specify the dependencies contained in :func:`get_default_conda_env()`.
+If ``None``, a conda environment with pip requirements inferred by
+:func:`mlflow.models.infer_pip_requirements` is added
+to the model. If the requirement inference fails, it falls back to using
+:func:`get_default_pip_requirements`. pip requirements from ``conda_env`` are written to a pip
+``requirements.txt`` file and the full conda environment is written to ``conda.yaml``.
+The following is an *example* dictionary representation of a conda environment::
+
+    {
+        "name": "mlflow-env",
+        "channels": ["conda-forge"],
+        "dependencies": [
+            "python=3.8.15",
+            {
+                "pip": [
+                    "{{ package_name }}==x.y.z"
+                ],
+            },
+        ],
+    }"""
+        ),
+        "pip_requirements": (
+            """Either an iterable of pip requirement strings
+(e.g. ``["{{ package_name }}", "-r requirements.txt", "-c constraints.txt"]``) or the string path to
+a pip requirements file on the local filesystem (e.g. ``"requirements.txt"``). If provided, this
+describes the environment this model should be run in. If ``None``, a default list of requirements
+is inferred by :func:`mlflow.models.infer_pip_requirements` from the current software environment.
+If the requirement inference fails, it falls back to using :func:`get_default_pip_requirements`.
+Both requirements and constraints are automatically parsed and written to ``requirements.txt`` and
+``constraints.txt`` files, respectively, and stored as part of the model. Requirements are also
+written to the ``pip`` section of the model's conda environment (``conda.yaml``) file."""
+        ),
+        "extra_pip_requirements": (
+            """Either an iterable of pip
+requirement strings
+(e.g. ``["pandas", "-r requirements.txt", "-c constraints.txt"]``) or the string path to
+a pip requirements file on the local filesystem (e.g. ``"requirements.txt"``). If provided, this
+describes additional pip requirements that are appended to a default set of pip requirements
+generated automatically based on the user's current software environment. Both requirements and
+constraints are automatically parsed and written to ``requirements.txt`` and ``constraints.txt``
+files, respectively, and stored as part of the model. Requirements are also written to the ``pip``
+section of the model's conda environment (``conda.yaml``) file.
+
+.. warning::
+    The following arguments can't be specified at the same time:
+
+    - ``conda_env``
+    - ``pip_requirements``
+    - ``extra_pip_requirements``
+
+`This example <https://github.com/mlflow/mlflow/blob/master/examples/pip_requirements/pip_requirements.py>`_ demonstrates how to specify pip requirements using
+``pip_requirements`` and ``extra_pip_requirements``."""  # noqa: E501
+        ),
+        "signature": (
+            """an instance of the :py:class:`ModelSignature <mlflow.models.ModelSignature>`
+class that describes the model's inputs and outputs. If not specified but an
+``input_example`` is supplied, a signature will be automatically inferred
+based on the supplied input example and model. To disable automatic signature
+inference when providing an input example, set ``signature`` to ``False``.
+To manually infer a model signature, call
+:py:func:`infer_signature() <mlflow.models.infer_signature>` on datasets
+with valid model inputs, such as a training dataset with the target column
+omitted, and valid model outputs, like model predictions made on the training
+dataset, for example:
+
+.. code-block:: python
+
+    from mlflow.models import infer_signature
+
+    train = df.drop_column("target_label")
+    predictions = ...  # compute model predictions
+    signature = infer_signature(train, predictions)
+"""
+        ),
+        "metadata": (
+            "Custom metadata dictionary passed to the model and stored in the MLmodel file."
+        ),
+        "input_example": (
+            """one or several instances of valid model input. The input example is used
+as a hint of what data to feed the model. It will be converted to a Pandas
+DataFrame and then serialized to json using the Pandas split-oriented
+format, or a numpy array where the example will be serialized to json
+by converting it to a list. Bytes are base64-encoded. When the ``signature`` parameter is
+``None``, the input example is used to infer a model signature.
+"""
+        ),
+        "example_no_conversion": (
+            """This parameter is deprecated and will be removed in a future release.
+It's no longer used and can be safely removed. Input examples are not converted anymore.
+"""
+        ),
+        "prompt_template": (
+            """A string that, if provided, will be used to format the user's input prior
+to inference. The string should contain a single placeholder, ``{prompt}``, which will be
+replaced with the user's input. For example: ``"Answer the following question. Q: {prompt} A:"``.
+
+Currently, only the following pipeline types are supported:
+
+- `feature-extraction <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.FeatureExtractionPipeline>`_
+- `fill-mask <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.FillMaskPipeline>`_
+- `summarization <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.SummarizationPipeline>`_
+- `text2text-generation <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.Text2TextGenerationPipeline>`_
+- `text-generation <https://huggingface.co/transformers/main_classes/pipelines.html#transformers.TextGenerationPipeline>`_
+"""
+        ),
+        "code_paths": (
+            """A list of local filesystem paths to Python file dependencies (or directories
+containing file dependencies). These files are *prepended* to the system path when the model
+is loaded. Files declared as dependencies for a given model should have relative
+imports declared from a common root path if multiple files are defined with import dependencies
+between them to avoid import errors when loading the model.
+
+For a detailed explanation of ``code_paths`` functionality, recommended usage patterns and
+limitations, see the
+`code_paths usage guide <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_.
+"""
+        ),
+        # Only pyfunc flavor supports `infer_code_paths`.
+        "code_paths_pyfunc": (
+            """A list of local filesystem paths to Python file dependencies (or directories
+containing file dependencies). These files are *prepended* to the system path when the model
+is loaded. Files declared as dependencies for a given model should have relative
+imports declared from a common root path if multiple files are defined with import dependencies
+between them to avoid import errors when loading the model.
+
+You can leave ``code_paths`` argument unset but set ``infer_code_paths`` to ``True`` to let MLflow
+infer the model code paths. See ``infer_code_paths`` argument doc for details.
+
+For a detailed explanation of ``code_paths`` functionality, recommended usage patterns and
+limitations, see the
+`code_paths usage guide <https://mlflow.org/docs/latest/model/dependencies.html?highlight=code_paths#saving-extra-code-with-an-mlflow-model>`_.
+"""
+        ),
+        "infer_code_paths": (
+            """If set to ``True``, MLflow automatically infers model code paths. The inferred
+            code path files only include necessary python module files. Only python code files
+            under current working directory are automatically inferable. Default value is
+            ``False``.
+
+.. warning::
+    Please ensure that the custom python module code does not contain sensitive data such as
+    credential token strings, otherwise they might be included in the automatic inferred code
+    path files and be logged to MLflow artifact repository.
+
+    If your custom python module depends on non-python files (e.g. a JSON file) with a relative
+    path to the module code file path, the non-python files can't be automatically inferred as the
+    code path file. To address this issue, you should put all used non-python files outside
+    your custom code directory.
+
+    If a python code file is loaded as the python ``__main__`` module, then this code file can't be
+    inferred as the code path file. If your model depends on classes / functions defined in
+    ``__main__`` module, you should use `cloudpickle` to dump your model instance in order to pickle
+    classes / functions in ``__main__``.
+
+.. Note:: Experimental: This parameter may change or be removed in a future release without warning.
+"""
+        ),
+        "save_pretrained": (
+            """If set to ``False``, MLflow will not save the Transformer model weight files,
+instead only saving the reference to the HuggingFace Hub model repository and its commit hash.
+This is useful when you load the pretrained model from HuggingFace Hub and want to log or save
+it to MLflow without modifying the model weights. In such case, specifying this flag to
+``False`` will save the storage space and reduce time to save the model. Please refer to the
+`Storage-Efficient Model Logging
+<../../llms/transformers/large-models.html#transformers-save-pretrained-guide>`_ for more detailed
+usage.
+
+
+.. warning::
+
+    If the model is saved with ``save_pretrained`` set to ``False``, the model cannot be
+    registered to the MLflow Model Registry. In order to convert the model to the one that
+    can be registered, you can use :py:func:`mlflow.transformers.persist_pretrained_model()`
+    to download the model weights from the HuggingFace Hub and save it in the existing model
+    artifacts. Please refer to `Transformers flavor documentation
+    <../../llms/transformers/large-models.html#persist-pretrained-guide>`_
+    for more detailed usage.
+
+    .. code-block:: python
+
+        import mlflow.transformers
+
+        model_uri = "YOUR_MODEL_URI_LOGGED_WITH_SAVE_PRETRAINED_FALSE"
+        model = mlflow.transformers.persist_pretrained_model(model_uri)
+        mlflow.register_model(model_uri, "model_name")
+
+.. important::
+
+    When you save the `PEFT <https://huggingface.co/docs/peft/en/index>`_ model, MLflow will
+    override the `save_pretrained` flag to `False` and only store the PEFT adapter weights. The
+    base model weights are not saved but the reference to the HuggingFace repository and
+    its commit hash are logged instead.
+"""
+        ),
+        "auth_policy": (
+            """Specifies the authentication policy for the model, which includes two key components.
+            Note that only one of `auth_policy` or `resources` should be defined.
+
+                - **System Auth Policy**: A list of resources required to serve this model.
+                - **User Auth Policy**: A minimal list of scopes that the user should have access to
+                    ,in order to invoke this model.
+
+    .. Note::
+        Experimental: This parameter may change or be removed in a future release without warning.
+            """
+        ),
+        "prompts": """\
+A list of prompt URIs registered in the MLflow Prompt Registry, to be associated with the model.
+Each prompt URI should be in the form ``prompt:/<name>/<version>``. The prompts should be
+registered in the MLflow Prompt Registry before being associated with the model.
+
+This will create a mutual link between the model and the prompt. The associated prompts can be
+seen in the model's metadata stored in the MLmodel file. From the Prompt Registry UI, you can
+navigate to the model as well.
+
+.. code-block:: python
+
+    import mlflow
+
+    prompt_template = "Hi, {name}! How are you doing today?"
+
+    # Register a prompt in the MLflow Prompt Registry
+    mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")
+
+    # Log a model with the registered prompt
+    with mlflow.start_run():
+        model_info = mlflow.pyfunc.log_model(
+            MyModel(),
+            artifact_path="model",
+            prompts=["prompt:/my_prompt/1"]
+        )
+
+    print(model_info.prompts)
+    # Output: ['prompt:/my_prompt/1']
+
+    # Load the prompt
+    prompt = mlflow.load_prompt(model_info.prompts[0])
+""",
+    }
+)
+
+
+def get_module_min_and_max_supported_ranges(flavor_name):
+    """
+    Extracts the minimum and maximum supported package versions from the provided module name.
+    The version information is provided via the yaml-to-python-script generation script in
+    dev/update_ml_package_versions.py which writes a python file to the importable namespace of
+    mlflow.ml_package_versions
+
+    Args:
+        flavor_name: The flavor name registered in ml_package_versions.py
+
+    Returns:
+        tuple of module name, minimum supported version, maximum supported version as strings.
+    """
+    if flavor_name == "pyspark.ml":
+        # pyspark.ml is a special case of spark flavor
+        flavor_name = "spark"
+
+    module_name = _ML_PACKAGE_VERSIONS[flavor_name]["package_info"].get("module_name", flavor_name)
+    versions = _ML_PACKAGE_VERSIONS[flavor_name]["models"]
+    min_version = versions["minimum"]
+    max_version = versions["maximum"]
+    return module_name, min_version, max_version
+
+
+def _do_version_compatibility_warning(msg: str):
+    """
+    Isolate the warn call to show the warning only once.
+    """
+    warnings.warn(msg, category=UserWarning, stacklevel=2)
+
+
+def docstring_version_compatibility_warning(integration_name):
+    """
+    Generates a docstring that can be applied as a note stating a version compatibility range for
+    a given flavor and optionally raises a warning if the installed version is outside of the
+    supported range.
+
+    Args:
+        integration_name: The name of the module as stored within ml-package-versions.yml
+
+    Returns:
+        The wrapped function with the additional docstring header applied
+    """
+
+    def annotated_func(func):
+        # NB: if using this decorator, ensure the package name to module name reference is
+        # updated with the flavor's `save` and `load` functions being used within
+        # ml-package-version.yml file.
+        min_ver, max_ver, pip_release = get_min_max_version_and_pip_release(
+            integration_name, "models"
+        )
+        notice = (
+            f"The '{integration_name}' MLflow Models integration is known to be compatible with "
+            f"``{min_ver}`` <= ``{pip_release}`` <= ``{max_ver}``. "
+            f"MLflow Models integrations with {integration_name} may not succeed when used with "
+            "package versions outside of this range."
+        )
+
+        func.__doc__ = (
+            "    .. Note:: " + notice + "\n" * 2 + func.__doc__ if func.__doc__ else notice
+        )
+
+        return func
+
+    return annotated_func

venv/lib/python3.9/site-packages/mlflow/utils/doctor.py

@@ -0,0 +1,133 @@
+import os
+import platform
+
+import click
+import importlib_metadata
+import yaml
+from packaging.requirements import Requirement
+
+import mlflow
+from mlflow.utils.databricks_utils import get_databricks_runtime_version
+
+
+def doctor(mask_envs=False):
+    """Prints out useful information for debugging issues with MLflow.
+
+    Args:
+        mask_envs: If True, mask the MLflow environment variable values
+            (e.g. `"MLFLOW_ENV_VAR": "***"`) in the output to prevent leaking sensitive
+            information.
+
+    .. warning::
+
+        - This API should only be used for debugging purposes.
+        - The output may contain sensitive information such as a database URI containing a password.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        with mlflow.start_run():
+            mlflow.doctor()
+
+    .. code-block:: text
+        :caption: Output
+
+        System information: Linux #58~20.04.1-Ubuntu SMP Thu Oct 13 13:09:46 UTC 2022
+        Python version: 3.8.13
+        MLflow version: 2.0.1
+        MLflow module location: /usr/local/lib/python3.8/site-packages/mlflow/__init__.py
+        Tracking URI: sqlite:///mlflow.db
+        Registry URI: sqlite:///mlflow.db
+        MLflow environment variables:
+          MLFLOW_TRACKING_URI: sqlite:///mlflow.db
+        MLflow dependencies:
+          Flask: 2.2.2
+          Jinja2: 3.0.3
+          alembic: 1.8.1
+          click: 8.1.3
+          cloudpickle: 2.2.0
+          databricks-cli: 0.17.4.dev0
+          docker: 6.0.0
+          entrypoints: 0.4
+          gitpython: 3.1.29
+          gunicorn: 20.1.0
+          importlib-metadata: 5.0.0
+          markdown: 3.4.1
+          matplotlib: 3.6.1
+          numpy: 1.23.4
+          packaging: 21.3
+          pandas: 1.5.1
+          protobuf: 3.19.6
+          pyarrow: 9.0.0
+          pytz: 2022.6
+          pyyaml: 6.0
+          querystring-parser: 1.2.4
+          requests: 2.28.1
+          scikit-learn: 1.1.3
+          scipy: 1.9.3
+          shap: 0.41.0
+          sqlalchemy: 1.4.42
+          sqlparse: 0.4.3
+    """
+    items = [
+        ("System information", " ".join((platform.system(), platform.version()))),
+        ("Python version", platform.python_version()),
+        ("MLflow version", mlflow.__version__),
+        ("MLflow module location", mlflow.__file__),
+        ("Tracking URI", mlflow.get_tracking_uri()),
+        ("Registry URI", mlflow.get_registry_uri()),
+    ]
+
+    if (runtime := get_databricks_runtime_version()) is not None:
+        items.append(("Databricks runtime version", runtime))
+
+    active_run = mlflow.active_run()
+    if active_run:
+        items.extend(
+            [
+                ("Active experiment ID", active_run.info.experiment_id),
+                ("Active run ID", active_run.info.run_id),
+                ("Active run artifact URI", active_run.info.artifact_uri),
+            ]
+        )
+
+    mlflow_envs = {
+        k: ("***" if mask_envs else v) for k, v in os.environ.items() if k.startswith("MLFLOW_")
+    }
+    if mlflow_envs:
+        items.append(
+            (
+                "MLflow environment variables",
+                yaml.dump({"_": mlflow_envs}, indent=2).replace("'", "").lstrip("_:").rstrip("\n"),
+            )
+        )
+
+    try:
+        requires = importlib_metadata.requires("mlflow")
+    except importlib_metadata.PackageNotFoundError:
+        requires = importlib_metadata.requires("mlflow-skinny")
+
+    mlflow_dependencies = {}
+    for req in requires:
+        req = Requirement(req)
+        try:
+            dist = importlib_metadata.distribution(req.name)
+        except importlib_metadata.PackageNotFoundError:
+            continue
+        else:
+            mlflow_dependencies[req.name] = dist.version
+
+    items.append(
+        (
+            "MLflow dependencies",
+            yaml.dump({"_": mlflow_dependencies}, indent=2)
+            .replace("'", "")
+            .lstrip("_:")
+            .rstrip("\n"),
+        )
+    )
+    for key, val in items:
+        click.secho(key, fg="blue", nl=False)
+        click.echo(f": {val}")

venv/lib/python3.9/site-packages/mlflow/utils/download_cloud_file_chunk.py

@@ -0,0 +1,43 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import argparse
+import importlib.util
+import json
+import os
+import sys
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--range-start", required=True, type=int)
+    parser.add_argument("--range-end", required=True, type=int)
+    parser.add_argument("--headers", required=True, type=str)
+    parser.add_argument("--download-path", required=True, type=str)
+    parser.add_argument("--http-uri", required=True, type=str)
+    return parser.parse_args()
+
+
+def main():
+    file_path = os.path.join(os.path.dirname(__file__), "request_utils.py")
+    module_name = "mlflow.utils.request_utils"
+
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    download_chunk = module.download_chunk
+
+    args = parse_args()
+    download_chunk(
+        range_start=args.range_start,
+        range_end=args.range_end,
+        headers=json.loads(args.headers),
+        download_path=args.download_path,
+        http_uri=args.http_uri,
+    )
+
+
+if __name__ == "__main__":
+    main()

venv/lib/python3.9/site-packages/mlflow/utils/env_manager.py

@@ -0,0 +1,16 @@
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+LOCAL = "local"
+CONDA = "conda"
+VIRTUALENV = "virtualenv"
+UV = "uv"
+
+
+def validate(env_manager):
+    allowed_values = [LOCAL, CONDA, VIRTUALENV, UV]
+    if env_manager not in allowed_values:
+        raise MlflowException(
+            f"Invalid value for `env_manager`: {env_manager}. Must be one of {allowed_values}",
+            error_code=INVALID_PARAMETER_VALUE,
+        )