star

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)