genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/autologging_utils/safety.py

@@ -0,0 +1,1104 @@
+import abc
+import functools
+import inspect
+import itertools
+import uuid
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, Callable, NamedTuple, Optional
+
+import mlflow
+import mlflow.utils.autologging_utils
+from mlflow.entities.run_status import RunStatus
+from mlflow.environment_variables import _MLFLOW_AUTOLOGGING_TESTING
+from mlflow.exceptions import MlflowException
+from mlflow.utils import gorilla, is_iterator
+from mlflow.utils.autologging_utils import _logger
+from mlflow.utils.autologging_utils.events import AutologgingEventLoggerWrapper
+from mlflow.utils.autologging_utils.logging_and_warnings import (
+    MlflowEventsAndWarningsBehaviorGlobally,
+    NonMlflowWarningsBehaviorForCurrentThread,
+)
+from mlflow.utils.mlflow_tags import MLFLOW_AUTOLOGGING
+
+_AUTOLOGGING_PATCHES = {}
+
+
+# Function attribute used for testing purposes to verify that a given function
+# has been wrapped with the `exception_safe_function_for_class` and
+# `picklable_exception_safe_function` decorators
+_ATTRIBUTE_EXCEPTION_SAFE = "exception_safe"
+
+
+_ERROR_MSG = "Encountered unexpected error during {} autologging: {}"
+
+
+def exception_safe_function_for_class(function):
+    """
+    Wraps the specified function with broad exception handling to guard
+    against unexpected errors during autologging.
+    Note this function creates an unpicklable function as `safe_function` is locally defined,
+    but a class instance containing methods decorated by this function should be pickalable,
+    because pickle only saves instance attributes, not methods.
+    See https://docs.python.org/3/library/pickle.html#pickling-class-instances for more details.
+    """
+    if is_testing():
+        setattr(function, _ATTRIBUTE_EXCEPTION_SAFE, True)
+
+    def safe_function(*args, **kwargs):
+        try:
+            return function(*args, **kwargs)
+        except Exception as e:
+            if is_testing():
+                raise
+            else:
+                _logger.warning("Encountered unexpected error during autologging: %s", e)
+
+    return update_wrapper_extended(safe_function, function)
+
+
+def _safe_function(function, *args, **kwargs):
+    try:
+        return function(*args, **kwargs)
+    except Exception as e:
+        if is_testing():
+            raise
+        else:
+            _logger.warning("Encountered unexpected error during autologging: %s", e)
+
+
+def picklable_exception_safe_function(function):
+    """
+    Wraps the specified function with broad exception handling to guard
+    against unexpected errors during autologging while preserving picklability.
+    """
+    if is_testing():
+        setattr(function, _ATTRIBUTE_EXCEPTION_SAFE, True)
+
+    return update_wrapper_extended(functools.partial(_safe_function, function), function)
+
+
+def _exception_safe_class_factory(base_class):
+    """
+    Creates an exception safe metaclass that inherits from `base_class`.
+    """
+
+    class _ExceptionSafeClass(base_class):
+        """
+        Metaclass that wraps all functions defined on the specified class with broad error handling
+        logic to guard against unexpected errors during autlogging.
+
+        Rationale: Patched autologging functions commonly pass additional class instances as
+        arguments to their underlying original training routines; for example, Keras autologging
+        constructs a subclass of `keras.callbacks.Callback` and forwards it to `Model.fit()`.
+        To prevent errors encountered during method execution within such classes from disrupting
+        model training, this metaclass wraps all class functions in a broad try / catch statement.
+
+        Note: `ExceptionSafeClass` does not handle exceptions in class methods or static methods,
+        as these are not always Python callables and are difficult to wrap
+        """
+
+        def __new__(cls, name, bases, dct):
+            for m in dct:
+                # class methods or static methods are not callable.
+                if callable(dct[m]):
+                    dct[m] = exception_safe_function_for_class(dct[m])
+            return base_class.__new__(cls, name, bases, dct)
+
+    return _ExceptionSafeClass
+
+
+ExceptionSafeClass = _exception_safe_class_factory(type)
+
+# `ExceptionSafeClass` causes an error when used with an abstract class.
+#
+# ```
+# class AbstractClass(abc.ABC):
+#    ...
+#
+# class DerivedClass(AbstractClass, metaclass=ExceptionSafeClass):
+#    ...
+# ```
+#
+# This raises:
+#
+# ```
+# TypeError: metaclass conflict: the metaclass of a derived class must be
+#            a (non-strict) subclass of the metaclasses of all its bases.
+# ```
+#
+# To avoid this error, create `ExceptionSafeAbstractClass` that is based on `abc.ABCMeta`.
+ExceptionSafeAbstractClass = _exception_safe_class_factory(abc.ABCMeta)
+
+
+def with_managed_run(autologging_integration, patch_function, tags=None):
+    """Given a `patch_function`, returns an `augmented_patch_function` that wraps the execution of
+    `patch_function` with an active MLflow run. The following properties apply:
+
+        - An MLflow run is only created if there is no active run present when the
+          patch function is executed
+
+        - If an active run is created by the `augmented_patch_function`, it is terminated
+          with the `FINISHED` state at the end of function execution
+
+        - If an active run is created by the `augmented_patch_function`, it is terminated
+          with the `FAILED` if an unhandled exception is thrown during function execution
+
+    Note that, if nested runs or non-fluent runs are created by `patch_function`, `patch_function`
+    is responsible for terminating them by the time it terminates
+    (or in the event of an exception).
+
+    Args:
+        autologging_integration: The autologging integration associated
+            with the `patch_function`.
+        patch_function: A function object compatible with `safe_patch`.
+        tags: A dictionary of string tags to set on each managed run created during the
+            execution of `patch_function`.
+    """
+    from mlflow.tracking.fluent import active_run
+    from mlflow.utils.autologging_utils import _has_active_training_session
+
+    def create_managed_run():
+        managed_run = mlflow.start_run(tags=tags)
+        _logger.info(
+            "Created MLflow autologging run with ID '%s', which will track hyperparameters,"
+            " performance metrics, model artifacts, and lineage information for the"
+            " current %s workflow",
+            managed_run.info.run_id,
+            autologging_integration,
+        )
+        return managed_run
+
+    def patch_with_managed_run(original, *args, **kwargs):
+        managed_run = None
+        # If there is an active training session but there is no active run
+        # in current thread, it means the thread is spawned by `estimator.fit`
+        # as a worker thread, we should disable autologging in
+        # these worker threads, so skip creating managed run.
+        if not active_run() and not _has_active_training_session():
+            managed_run = create_managed_run()
+
+        try:
+            result = patch_function(original, *args, **kwargs)
+        except (Exception, KeyboardInterrupt):
+            # In addition to standard Python exceptions, handle keyboard interrupts to ensure
+            # that runs are terminated if a user prematurely interrupts training execution
+            # (e.g. via sigint / ctrl-c)
+            if managed_run:
+                mlflow.end_run(RunStatus.to_string(RunStatus.FAILED))
+            raise
+        else:
+            if managed_run:
+                mlflow.end_run(RunStatus.to_string(RunStatus.FINISHED))
+            return result
+
+    return patch_with_managed_run
+
+
+def is_testing():
+    """
+    Indicates whether or not autologging functionality is running in test mode (as determined
+    by the `MLFLOW_AUTOLOGGING_TESTING` environment variable). Test mode performs additional
+    validation during autologging, including:
+
+        - Checks for the exception safety of arguments passed to model training functions
+          (i.e. all additional arguments should be "exception safe" functions or classes)
+        - Disables exception handling for patched function logic, ensuring that patch code
+          executes without errors during testing
+    """
+    return _MLFLOW_AUTOLOGGING_TESTING.get()
+
+
+def _resolve_extra_tags(autologging_integration, extra_tags):
+    tags = {MLFLOW_AUTOLOGGING: autologging_integration}
+    if extra_tags:
+        if isinstance(extra_tags, dict):
+            if MLFLOW_AUTOLOGGING in extra_tags:
+                extra_tags.pop(MLFLOW_AUTOLOGGING)
+                _logger.warning(
+                    f"Tag `{MLFLOW_AUTOLOGGING}` is ignored as it is a reserved tag by MLflow "
+                    f"autologging."
+                )
+            tags.update(extra_tags)
+        else:
+            raise mlflow.exceptions.MlflowException.invalid_parameter_value(
+                f"Invalid `extra_tags` type: expecting dictionary, "
+                f"received `{type(extra_tags).__name__}`"
+            )
+    return tags
+
+
+def safe_patch(
+    autologging_integration,
+    destination,
+    function_name,
+    patch_function,
+    manage_run=False,
+    extra_tags=None,
+):
+    """Patches the specified `function_name` on the specified `destination` class for autologging
+    purposes, preceding its implementation with an error-safe copy of the specified patch
+    `patch_function` with the following error handling behavior:
+        - Exceptions thrown from the underlying / original function
+          (`<destination>.<function_name>`) are propagated to the caller.
+        - Exceptions thrown from other parts of the patched implementation (`patch_function`)
+          are caught and logged as warnings.
+
+    Args:
+        autologging_integration: The name of the autologging integration associated with the
+            patch.
+        destination: The Python class on which the patch is being defined.
+        function_name: The name of the function to patch on the specified `destination` class.
+        patch_function: The patched function code to apply. The first argument should be reserved
+            for an `original` argument representing the underlying / original function. Subsequent
+            arguments should be identical to those of the original function being patched.
+        manage_run: If `True`, applies the `with_managed_run` wrapper to the specified
+            `patch_function`, which automatically creates & terminates an MLflow
+            active run during patch code execution if necessary. If `False`,
+            does not apply the `with_managed_run` wrapper to the specified
+            `patch_function`.
+        extra_tags: A dictionary of extra tags to set on each managed run created by autologging.
+    """
+    from mlflow.tracking.fluent import active_run
+    from mlflow.utils.autologging_utils import autologging_is_disabled, get_autologging_config
+
+    # NB: Checking the signature of the patch function rather than original, so that we don't
+    # accidentally change the behavior of existing patches that may use sync patch function
+    # for async original functions (e.g. LangChain).
+    is_async_function = inspect.iscoroutinefunction(patch_function)
+
+    if manage_run:
+        if is_async_function:
+            raise MlflowException("manage_run parameter is not supported for async functions.")
+
+        tags = _resolve_extra_tags(autologging_integration, extra_tags)
+        patch_function = with_managed_run(
+            autologging_integration,
+            patch_function,
+            tags=tags,
+        )
+
+    original_fn = gorilla.get_original_attribute(
+        destination, function_name, bypass_descriptor_protocol=False
+    )
+    # Retrieve raw attribute while bypassing the descriptor protocol
+    raw_original_obj = gorilla.get_original_attribute(
+        destination, function_name, bypass_descriptor_protocol=True
+    )
+    if original_fn != raw_original_obj:
+        raise RuntimeError(f"Unsupported patch on {destination}.{function_name}")
+    elif isinstance(original_fn, property):
+        if is_async_function:
+            raise MlflowException("Patching async property methods is not supported.")
+
+        is_property_method = True
+
+        # For property decorated methods (a kind of method delegation), e.g.
+        # class A:
+        #   @property
+        #   def f1(self):
+        #     ...
+        #     return delegated_f1
+        #
+        # suppose `a1` is an instance of class `A`,
+        # `A.f1.fget` will get the original `def f1(self)` method,
+        # and `A.f1.fget(a1)` will be equivalent to `a1.f1()` and
+        # its return value will be the `delegated_f1` function.
+        # So using the `property.fget` we can construct the (delegated) "original_fn"
+        def original(self, *args, **kwargs):
+            # the `original_fn.fget` will get the original method decorated by `property`
+            # the `original_fn.fget(self)` will get the delegated function returned by the
+            # property decorated method.
+            bound_delegate_method = original_fn.fget(self)
+            return bound_delegate_method(*args, **kwargs)
+
+    else:
+        original = original_fn
+        is_property_method = False
+
+    def safe_patch_function(*args, **kwargs):
+        """
+        A safe wrapper around the specified `patch_function` implementation designed to
+        handle exceptions thrown during the execution of `patch_function`. This wrapper
+        distinguishes exceptions thrown from the underlying / original function
+        (`<destination>.<function_name>`) from exceptions thrown from other parts of
+        `patch_function`. This distinction is made by passing an augmented version of the
+        underlying / original function to `patch_function` that uses nonlocal state to track
+        whether or not it has been executed and whether or not it threw an exception.
+        Exceptions thrown from the underlying / original function are propagated to the caller,
+        while exceptions thrown from other parts of `patch_function` are caught and logged as
+        warnings.
+
+        NB: PLEASE BE SUPER CAREFUL WHEN MODIFYING THIS FUNCTION. IT IS USED IN A WIDE VARIETY
+        OF CONTEXTX AND CRITICAL PATH IN DBR/MLR BY DEFAULT. ANY BUG HERE CAN BREAK USERS'
+        WORKLOAD WITHOUT THEM TAKING ANY ACTION.
+        """
+        # Reroute warnings encountered during the patch function implementation to an MLflow event
+        # logger, and enforce silent mode if applicable (i.e. if the corresponding autologging
+        # integration was called with `silent=True`), hiding MLflow event logging statements and
+        # hiding all warnings in the autologging preamble and postamble (i.e. the code surrounding
+        # the user's original / underlying ML function). Non-MLflow warnings are enabled during the
+        # execution of the original / underlying ML function
+        #
+        # Note that we've opted *not* to apply this context manager as a decorator on
+        # `safe_patch_function` because the context-manager-as-decorator pattern uses
+        # `contextlib.ContextDecorator`, which creates generator expressions that cannot be pickled
+        # during model serialization by ML frameworks such as scikit-learn
+        is_silent_mode = get_autologging_config(autologging_integration, "silent", False)
+        with (
+            MlflowEventsAndWarningsBehaviorGlobally(
+                # MLflow warnings emitted during autologging training sessions are likely not
+                # actionable and result from the autologging implementation invoking another MLflow
+                # API. Accordingly, we reroute these warnings to the MLflow event logger with level
+                # 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_event_logs=is_silent_mode,
+                disable_warnings=is_silent_mode,
+            ),
+            NonMlflowWarningsBehaviorForCurrentThread(
+                # non-MLflow Warnings emitted during the autologging preamble (before the original /
+                # underlying ML function is called) and postamble (after the original / underlying
+                # ML function is called) are likely not actionable and result from the autologging
+                # implementation invoking an API from a dependent library. Accordingly, we reroute
+                # these warnings to the MLflow event logger with level 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,
+            ),
+        ):
+            if is_testing():
+                preexisting_run_for_testing = active_run()
+
+            # Whether or not to exclude autologged content from user-created fluent runs
+            # (i.e. runs created manually via `mlflow.start_run()`)
+            exclusive = get_autologging_config(autologging_integration, "exclusive", False)
+            user_created_fluent_run_is_active = (
+                active_run() and not _AutologgingSessionManager.active_session()
+            )
+            active_session_failed = (
+                _AutologgingSessionManager.active_session() is not None
+                and _AutologgingSessionManager.active_session().state == "failed"
+            )
+
+            if (
+                active_session_failed
+                or autologging_is_disabled(autologging_integration)
+                or (user_created_fluent_run_is_active and exclusive)
+                or (
+                    mlflow.utils.autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED
+                    and autologging_integration
+                )
+            ):
+                # If the autologging integration associated with this patch is disabled,
+                # or if the current autologging integration is in exclusive mode and a user-created
+                # fluent run is active, call the original function and return. Restore the original
+                # warning behavior during original function execution, since autologging is being
+                # skipped
+                with NonMlflowWarningsBehaviorForCurrentThread(
+                    disable_warnings=False,
+                    reroute_warnings=False,
+                ):
+                    return original(*args, **kwargs)
+
+            # Whether or not the original / underlying function has been called during the
+            # execution of patched code
+            original_has_been_called = False
+            # The value returned by the call to the original / underlying function during
+            # the execution of patched code
+            original_result = None
+            # Whether or not an exception was raised from within the original / underlying function
+            # during the execution of patched code
+            failed_during_original = False
+            # The active MLflow run (if any) associated with patch code execution
+            patch_function_run_for_testing = None
+            # The exception raised during executing patching function
+            patch_error = None
+
+            with _AutologgingSessionManager.start_session(autologging_integration) as session:
+                event_logger = AutologgingEventLoggerWrapper(session, destination, function_name)
+
+                def call_original_fn_with_event_logging(original_fn, og_args, og_kwargs):
+                    try:
+                        event_logger.log_original_function_start(og_args, og_kwargs)
+
+                        original_fn_result = original_fn(*og_args, **og_kwargs)
+
+                        event_logger.log_original_function_success(og_args, og_kwargs)
+                        return original_fn_result
+                    except Exception as e:
+                        event_logger.log_original_function_error(og_args, og_kwargs, e)
+
+                        nonlocal failed_during_original
+                        failed_during_original = True
+                        raise
+
+                try:
+
+                    def call_original(*og_args, **og_kwargs):
+                        def _original_fn(*_og_args, **_og_kwargs):
+                            if is_testing():
+                                _validate_args(
+                                    autologging_integration,
+                                    function_name,
+                                    args,
+                                    kwargs,
+                                    og_args,
+                                    og_kwargs,
+                                )
+                                # By the time `original` is called by the patch implementation, we
+                                # assume that either: 1. the patch implementation has already
+                                # created an MLflow run or 2. the patch code will not create an
+                                # MLflow run during the current execution. Here, we capture a
+                                # reference to the active run, which we will use later on to
+                                # determine whether or not the patch implementation created
+                                # a run and perform validation if necessary
+                                nonlocal patch_function_run_for_testing
+                                patch_function_run_for_testing = active_run()
+
+                            nonlocal original_has_been_called
+                            original_has_been_called = True
+
+                            nonlocal original_result
+                            # Show all non-MLflow warnings as normal (i.e. not as event logs)
+                            # during original function execution, even if silent mode is enabled
+                            # (`silent=True`), since these warnings originate from the ML framework
+                            # or one of its dependencies and are likely relevant to the caller
+                            with NonMlflowWarningsBehaviorForCurrentThread(
+                                disable_warnings=False,
+                                reroute_warnings=False,
+                            ):
+                                original_result = original(*_og_args, **_og_kwargs)
+                                return original_result
+
+                        return call_original_fn_with_event_logging(_original_fn, og_args, og_kwargs)
+
+                    # Apply the name, docstring, and signature of `original` to `call_original`.
+                    # This is important because several autologging patch implementations inspect
+                    # the signature of the `original` argument during execution
+                    call_original = update_wrapper_extended(call_original, original)
+
+                    event_logger.log_patch_function_start(args, kwargs)
+
+                    patch_function(call_original, *args, **kwargs)
+
+                    session.state = "succeeded"
+                    event_logger.log_patch_function_success(args, kwargs)
+
+                except Exception as e:
+                    session.state = "failed"
+                    patch_error = e
+                    # Exceptions thrown during execution of the original function should be
+                    # propagated to the caller. Additionally, exceptions encountered during test
+                    # mode should be reraised to detect bugs in autologging implementations
+                    if failed_during_original or is_testing():
+                        raise
+
+                if is_testing() and not preexisting_run_for_testing:
+                    # If an MLflow run was created during the execution of patch code, verify that
+                    # it is no longer active and that it contains expected autologging tags
+                    assert not active_run(), (
+                        f"Autologging integration {autologging_integration} leaked an active run"
+                    )
+                    if patch_function_run_for_testing:
+                        _validate_autologging_run(
+                            autologging_integration, patch_function_run_for_testing.info.run_id
+                        )
+                try:
+                    if original_has_been_called:
+                        return original_result
+                    else:
+                        return call_original_fn_with_event_logging(original, args, kwargs)
+                finally:
+                    # If original function succeeds, but `patch_function_exception` exists,
+                    # it represent patching code unexpected failure, so we call
+                    # `log_patch_function_error` in this case.
+                    # If original function failed, we don't call `log_patch_function_error`
+                    # even if `patch_function_exception` exists, because original function failure
+                    # means there's some error in user code (e.g. user provide wrong arguments)
+                    if patch_error is not None and not failed_during_original:
+                        event_logger.log_patch_function_error(args, kwargs, patch_error)
+                        _logger.warning(_ERROR_MSG.format(autologging_integration, patch_error))
+
+    async def async_safe_patch_function(*args, **kwargs):
+        """
+        Async version of safe_patch_function.
+
+        This code brainlessly copies the synchronous version of the function, but with async
+        context managers and async functions. This is done to avoid the risk of introducing
+        any bugs or regressions in the async version of the function. Note that we need to
+        be really careful here, because autologging is enabled by-default in DBR/MLR, hence
+        any bug here can break users' workload without them taking any action.
+
+        That said, some long comments are omitted in this version to avoid redundancy. If
+        you want to understand the context of the code better, please refer to the
+        synchronous version as well.
+        """
+        is_silent_mode = get_autologging_config(autologging_integration, "silent", False)
+        async with (
+            MlflowEventsAndWarningsBehaviorGlobally(
+                reroute_warnings=True,
+                disable_event_logs=is_silent_mode,
+                disable_warnings=is_silent_mode,
+            ),
+            NonMlflowWarningsBehaviorForCurrentThread(
+                disable_warnings=is_silent_mode,
+                reroute_warnings=True,
+            ),
+        ):
+            if is_testing():
+                preexisting_run_for_testing = active_run()
+
+            # Whether or not to exclude autologged content from user-created fluent runs
+            # (i.e. runs created manually via `mlflow.start_run()`)
+            exclusive = get_autologging_config(autologging_integration, "exclusive", False)
+            user_created_fluent_run_is_active = (
+                active_run() and not _AutologgingSessionManager.active_session()
+            )
+            active_session_failed = (
+                _AutologgingSessionManager.active_session() is not None
+                and _AutologgingSessionManager.active_session().state == "failed"
+            )
+
+            if (
+                active_session_failed
+                or autologging_is_disabled(autologging_integration)
+                or (user_created_fluent_run_is_active and exclusive)
+                or (
+                    mlflow.utils.autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED
+                    and autologging_integration
+                )
+            ):
+                async with NonMlflowWarningsBehaviorForCurrentThread(False, False):
+                    return await original(*args, **kwargs)
+
+            original_has_been_called = False
+            original_result = None
+            failed_during_original = False
+            patch_function_run_for_testing = None
+            patch_error = None
+
+            async with _AutologgingSessionManager.astart_session(
+                autologging_integration
+            ) as session:
+                event_logger = AutologgingEventLoggerWrapper(session, destination, function_name)
+
+                async def call_original_fn_with_event_logging(original_fn, og_args, og_kwargs):
+                    try:
+                        event_logger.log_original_function_start(og_args, og_kwargs)
+                        original_fn_result = await original_fn(*og_args, **og_kwargs)
+                        event_logger.log_original_function_success(og_args, og_kwargs)
+                        return original_fn_result
+                    except Exception as e:
+                        event_logger.log_original_function_error(og_args, og_kwargs, e)
+                        nonlocal failed_during_original
+                        failed_during_original = True
+                        raise
+
+                try:
+
+                    async def call_original(*og_args, **og_kwargs):
+                        async def _original_fn(*_og_args, **_og_kwargs):
+                            if is_testing():
+                                _validate_args(
+                                    autologging_integration,
+                                    function_name,
+                                    args,
+                                    kwargs,
+                                    og_args,
+                                    og_kwargs,
+                                )
+                                nonlocal patch_function_run_for_testing
+                                patch_function_run_for_testing = active_run()
+
+                            nonlocal original_has_been_called
+                            original_has_been_called = True
+
+                            nonlocal original_result
+                            async with NonMlflowWarningsBehaviorForCurrentThread(False, False):
+                                original_result = await original(*_og_args, **_og_kwargs)
+                                return original_result
+
+                        return await call_original_fn_with_event_logging(
+                            _original_fn, og_args, og_kwargs
+                        )
+
+                    # Apply the name, docstring, and signature of `original` to `call_original`.
+                    # This is important because several autologging patch implementations inspect
+                    # the signature of the `original` argument during execution
+                    call_original = update_wrapper_extended(call_original, original)
+
+                    event_logger.log_patch_function_start(args, kwargs)
+
+                    await patch_function(call_original, *args, **kwargs)
+
+                    session.state = "succeeded"
+                    event_logger.log_patch_function_success(args, kwargs)
+
+                except Exception as e:
+                    session.state = "failed"
+                    patch_error = e
+                    # Exceptions thrown during execution of the original function should be
+                    # propagated to the caller. Additionally, exceptions encountered during test
+                    # mode should be reraised to detect bugs in autologging implementations
+                    if failed_during_original or is_testing():
+                        raise
+
+                if is_testing() and not preexisting_run_for_testing:
+                    # If an MLflow run was created during the execution of patch code, verify that
+                    # it is no longer active and that it contains expected autologging tags
+                    assert not active_run(), (
+                        f"Autologging integration {autologging_integration} leaked an active run"
+                    )
+                    if patch_function_run_for_testing:
+                        _validate_autologging_run(
+                            autologging_integration, patch_function_run_for_testing.info.run_id
+                        )
+                try:
+                    if original_has_been_called:
+                        return original_result
+                    else:
+                        return await call_original_fn_with_event_logging(original, args, kwargs)
+                finally:
+                    if patch_error is not None and not failed_during_original:
+                        event_logger.log_patch_function_error(args, kwargs, patch_error)
+                        _logger.warning(_ERROR_MSG.format(autologging_integration, patch_error))
+
+    if is_property_method:
+        # Create a patched function (also property decorated)
+        # like:
+        #
+        # class A:
+        # @property
+        # def get_bound_safe_patch_fn(self):
+        #   original_fn.fget(self) # do availability check
+        #   return bound_safe_patch_fn
+        #
+        # Suppose `a1` is instance of class A,
+        # then `a1.get_bound_safe_patch_fn(*args, **kwargs)` will be equivalent to
+        # `bound_safe_patch_fn(*args, **kwargs)`
+        def get_bound_safe_patch_fn(self):
+            # This `original_fn.fget` call is for availability check, if it raise error
+            # then `hasattr(obj, {func_name})` will return False
+            # so it mimic the original property behavior.
+            original_fn.fget(self)
+
+            def bound_safe_patch_fn(*args, **kwargs):
+                return safe_patch_function(self, *args, **kwargs)
+
+            # Make bound method `instance.target_method` keep the same doc and signature.
+            # Here return the bound safe patch function because user call property decorated
+            # method will like `instance.property_decorated_method(...)`, and internally it will
+            # call the `bound_safe_patch_fn`, the argument list don't include the `self` argument,
+            # so return bound function here.
+            return update_wrapper_extended(bound_safe_patch_fn, original_fn.fget)
+
+        # Make unbound method `class.target_method` keep the same doc and signature
+        get_bound_safe_patch_fn = update_wrapper_extended(get_bound_safe_patch_fn, original_fn.fget)
+        safe_patch_obj = property(get_bound_safe_patch_fn)
+    elif is_async_function:
+        safe_patch_obj = update_wrapper_extended(async_safe_patch_function, original)
+    else:
+        safe_patch_obj = update_wrapper_extended(safe_patch_function, original)
+
+    new_patch = _wrap_patch(destination, function_name, safe_patch_obj)
+    _store_patch(autologging_integration, new_patch)
+
+
+def revert_patches(autologging_integration):
+    """Reverts all patches on the specified destination class for autologging disablement purposes.
+
+    Args:
+        autologging_integration: The name of the autologging integration associated with the
+            patch. Note: If called via fluent api (`autologging_integration="mlflow"`), then revert
+            all patches for all active autologging integrations.
+
+    """
+    for patch in _AUTOLOGGING_PATCHES.get(autologging_integration, []):
+        gorilla.revert(patch)
+
+    _AUTOLOGGING_PATCHES.pop(autologging_integration, None)
+
+
+# Represents an active autologging session using two fields:
+# - integration: the name of the autologging integration corresponding to the session
+# - id: a unique session identifier (e.g., a UUID)
+# - state: the state of AutologgingSession, will be one of running/succeeded/failed
+class AutologgingSession:
+    def __init__(self, integration, id_):
+        self.integration = integration
+        self.id = id_
+        self.state = "running"
+
+
+class _AutologgingSessionManager:
+    _session = None
+
+    @classmethod
+    @contextmanager
+    def start_session(cls, integration):
+        try:
+            prev_session = cls._session
+            if prev_session is None:
+                session_id = uuid.uuid4().hex
+                cls._session = AutologgingSession(integration, session_id)
+            yield cls._session
+        finally:
+            # Only end the session upon termination of the context if we created
+            # the session; otherwise, leave the session open for later termination
+            # by its creator
+            if prev_session is None:
+                cls._end_session()
+
+    @classmethod
+    @asynccontextmanager
+    async def astart_session(cls, integration):
+        try:
+            prev_session = cls._session
+            if prev_session is None:
+                session_id = uuid.uuid4().hex
+                cls._session = AutologgingSession(integration, session_id)
+            yield cls._session
+        finally:
+            if prev_session is None:
+                cls._end_session()
+
+    @classmethod
+    def active_session(cls):
+        return cls._session
+
+    @classmethod
+    def _end_session(cls):
+        cls._session = None
+
+
+def update_wrapper_extended(wrapper, wrapped):
+    """Update a `wrapper` function to look like the `wrapped` function. This is an extension of
+    `functools.update_wrapper` that applies the docstring *and* signature of `wrapped` to
+    `wrapper`, producing a new function.
+
+    Returns:
+        A new function with the same implementation as `wrapper` and the same docstring
+        & signature as `wrapped`.
+    """
+    updated_wrapper = functools.update_wrapper(wrapper, wrapped)
+    # Assign the signature of the `wrapped` function to the updated wrapper function.
+    # Certain frameworks may disallow signature inspection, causing `inspect.signature()` to throw.
+    # One such example is the `tensorflow.estimator.Estimator.export_savedmodel()` function
+    try:
+        updated_wrapper.__signature__ = inspect.signature(wrapped)
+    except Exception:
+        _logger.debug("Failed to restore original signature for wrapper around %s", wrapped)
+    return updated_wrapper
+
+
+def _wrap_patch(destination, name, patch_obj, settings=None):
+    """Apply a patch.
+
+    Args:
+        destination: Patch destination.
+        name: Name of the attribute at the destination.
+        patch_obj: Patch object, it should be a function or a property decorated function
+            to be assigned to the patch point {destination}.{name}.
+        settings: Settings for gorilla.Patch.
+
+    """
+    if settings is None:
+        settings = gorilla.Settings(allow_hit=True, store_hit=True)
+
+    patch = gorilla.Patch(destination, name, patch_obj, settings=settings)
+    gorilla.apply(patch)
+    return patch
+
+
+def _store_patch(autologging_integration, patch):
+    """
+    Stores a patch for a specified autologging_integration class. Later to be used for being able
+    to revert the patch when disabling autologging.
+
+    Args:
+        autologging_integration: The name of the autologging integration associated with the
+            patch.
+        patch: The patch to be stored.
+    """
+    if autologging_integration in _AUTOLOGGING_PATCHES:
+        _AUTOLOGGING_PATCHES[autologging_integration].add(patch)
+    else:
+        _AUTOLOGGING_PATCHES[autologging_integration] = {patch}
+
+
+def _validate_autologging_run(autologging_integration, run_id):
+    """
+    For testing purposes, verifies that an MLflow run produced by an `autologging_integration`
+    satisfies the following properties:
+
+        - The run has an autologging tag whose value is the name of the autologging integration
+        - The run has a terminal status (e.g., KILLED, FAILED, FINISHED)
+    """
+    from mlflow.tracking.client import MlflowClient
+
+    client = MlflowClient()
+    run = client.get_run(run_id)
+    autologging_tag_value = run.data.tags.get(MLFLOW_AUTOLOGGING)
+    assert autologging_tag_value == autologging_integration, (
+        f"Autologging run with id {run_id} failed to set autologging tag with expected value. "
+        f"Expected: '{autologging_integration}', Actual: '{autologging_tag_value}'"
+    )
+    assert RunStatus.is_terminated(RunStatus.from_string(run.info.status)), (
+        f"Autologging run with id {run_id} has a non-terminal status '{run.info.status}'"
+    )
+
+
+class ValidationExemptArgument(NamedTuple):
+    """
+    A NamedTuple representing the properties of an argument that is exempt from validation
+
+    autologging_integration: The name of the autologging integration.
+    function_name: The name of the function that is being validated.
+    type_function: A Callable that accepts an object and returns True if the given object matches
+                   the argument type. Returns False otherwise.
+    positional_argument_index: The index of the argument in the function signature.
+    keyword_argument_name: The name of the argument in the function signature.
+    """
+
+    autologging_integration: str
+    function_name: str
+    type_function: Callable[..., Any]
+    positional_argument_index: Optional[int] = None
+    keyword_argument_name: Optional[str] = None
+
+    def matches(
+        self,
+        autologging_integration,
+        function_name,
+        value,
+        argument_index=None,
+        argument_name=None,
+    ):
+        """
+        This method checks if the properties provided through the function arguments matches the
+        properties defined in the NamedTuple.
+
+        Args:
+            autologging_integration: The name of an autologging integration.
+            function_name: The name of the function that is being matched.
+            value: The value of the argument.
+            argument_index: The index of the argument, if it is passed as a positional
+                argument. Otherwise it is None.
+            argument_name: The name of the argument, if it is passed as a keyword
+                argument. Otherwise it is None.
+
+        Returns:
+            Returns True if the given function properties matches the exempt argument's
+            properties. Returns False otherwise.
+        """
+        return (
+            self.autologging_integration == autologging_integration
+            and self.function_name == function_name
+            and (
+                self.positional_argument_index == argument_index
+                or self.keyword_argument_name == argument_name
+            )
+            and self.type_function(value)
+        )
+
+
+# WARNING: Exemptions should NOT be introduced unless absolutely necessary. If deemed necessary,
+#          clear reasons must be provided as comment in addition to thorough integration tests.
+_VALIDATION_EXEMPT_ARGUMENTS = [
+    # When extracting implicitly defined `batch_size` in the case that `x` is a generator or a
+    # generator class, we need to consume and restore the first element back to the generator to
+    # calculate the `batch_size`. This means that:
+    # 1. The type of `x` will become 'generator' regardless if user provided `x` as a generator or a
+    #    custom generator class.
+    # 2. The instance of `x` will be different, since we reconstructed the generator after consuming
+    #    the first element.
+    ValidationExemptArgument("tensorflow", "fit", is_iterator, 1, "x"),
+    ValidationExemptArgument("keras", "fit", is_iterator, 1, "x"),
+]
+
+
+def _is_arg_exempt_from_validation(
+    autologging_integration,
+    function_name,
+    argument,
+    argument_index=None,
+    argument_name=None,
+):
+    """This function is responsible for determining whether or not an argument is exempt from
+    autolog safety validations. This includes both type checking and immutable checking.
+
+    Args:
+        autologging_integration: The name of the autologging integration.
+        function_name: The name of the function that is being validated.
+        argument: The actual argument.
+        argument_index: The index of the argument, if it is passed as a positional
+            argument. Otherwise it is None.
+        argument_name: The name of the argument, if it is passed as a keyword argument.
+            Otherwise it is None.
+
+    Returns:
+        True or False
+    """
+    return any(
+        exemption.matches(
+            autologging_integration,
+            function_name,
+            argument,
+            argument_index,
+            argument_name,
+        )
+        for exemption in _VALIDATION_EXEMPT_ARGUMENTS
+    )
+
+
+def _validate_args(
+    autologging_integration,
+    function_name,
+    user_call_args,
+    user_call_kwargs,
+    autologging_call_args,
+    autologging_call_kwargs,
+):
+    """
+    Used for testing purposes to verify that, when a patched ML function calls its underlying
+    / original ML function, the following properties are satisfied:
+
+        - All arguments supplied to the patched ML function are forwarded to the
+          original ML function
+        - Any additional arguments supplied to the original function are exception safe (i.e.
+          they are either functions decorated with the `@exception_safe_function_for_class` or
+          `@pickalable_exception_safe_function` decorators, or classes / instances of classes with
+          type `ExceptionSafeClass`
+    """
+
+    def _validate_new_input(inp):
+        """
+        Validates a new input (arg or kwarg) introduced to the underlying / original ML function
+        call during the execution of a patched ML function. The new input is valid if:
+
+            - The new input is a function that has been decorated with
+              `exception_safe_function_for_class` or `pickalable_exception_safe_function`
+            - OR the new input is a class with the `ExceptionSafeClass` metaclass
+            - OR the new input is a list and each of its elements is valid according to the
+              these criteria
+        """
+        if type(inp) == list:
+            for item in inp:
+                _validate_new_input(item)
+        elif isinstance(inp, dict) and "callbacks" in inp:
+            _validate_new_input(inp["callbacks"])
+        elif callable(inp):
+            assert getattr(inp, _ATTRIBUTE_EXCEPTION_SAFE, False), (
+                f"New function argument '{inp}' passed to original function is not exception-safe."
+                " Please decorate the function with `exception_safe_function` or "
+                "`pickalable_exception_safe_function`"
+            )
+        else:
+            assert hasattr(inp, "__class__") and type(inp.__class__) in [
+                ExceptionSafeClass,
+                ExceptionSafeAbstractClass,
+            ], (
+                f"Invalid new input '{inp}'. New args / kwargs introduced to `original` function "
+                "calls by patched code must either be functions decorated with "
+                "`exception_safe_function_for_class`, instances of classes with the "
+                "`ExceptionSafeClass` or `ExceptionSafeAbstractClass` metaclass safe or lists of "
+                "such exception safe functions / classes."
+            )
+
+    def _assert_autologging_input_positional_args_are_superset(
+        autologging_call_input, user_call_input
+    ):
+        length_diff = len(autologging_call_input) - len(user_call_input)
+        assert length_diff >= 0, (
+            f"{length_diff} expected inputs are missing from the call to the original function."
+        )
+
+    def _assert_autologging_input_kwargs_are_superset(autologging_call_input, user_call_input):
+        assert set(user_call_input.keys()).issubset(set(autologging_call_input.keys())), (
+            "Keyword or dictionary arguments to original function omit"
+            " one or more expected keys: '{}'".format(
+                set(user_call_input.keys()) - set(autologging_call_input.keys())
+            )
+        )
+
+    def _validate(autologging_call_input, user_call_input=None):
+        """
+        Validates that the specified `autologging_call_input` and `user_call_input`
+        are compatible. If `user_call_input` is `None`, then `autologging_call_input`
+        is regarded as a new input added by autologging and is validated using
+        `_validate_new_input`. Otherwise, the following properties must hold:
+
+            - `autologging_call_input` and `user_call_input` must have the same type
+              (referred to as "input type")
+            - if the input type is a tuple, list or dictionary, then `autologging_call_input` must
+              be equivalent to `user_call_input` or be a superset of `user_call_input`
+            - for all other input types, `autologging_call_input` and `user_call_input`
+              must be equivalent by reference equality or by object equality
+
+        Args:
+            autologging_call_input: call input from autologging.
+            user_call_input: call input from user.
+        """
+
+        if user_call_input is None and autologging_call_input is not None:
+            _validate_new_input(autologging_call_input)
+            return
+
+        assert type(autologging_call_input) == type(user_call_input), (
+            "Type of input to original function '{}' does not match expected type '{}'".format(
+                type(autologging_call_input), type(user_call_input)
+            )
+        )
+
+        if type(autologging_call_input) in [list, tuple]:
+            _assert_autologging_input_positional_args_are_superset(
+                autologging_call_input, user_call_input
+            )
+            # If the autologging call input is longer than the user call input, we `zip_longest`
+            # will pad the user call input with `None` values to ensure that the subsequent calls
+            # to `_validate` identify new inputs added by the autologging call
+            for a, u in itertools.zip_longest(autologging_call_input, user_call_input):
+                _validate(a, u)
+        elif type(autologging_call_input) == dict:
+            _assert_autologging_input_kwargs_are_superset(autologging_call_input, user_call_input)
+            for key in autologging_call_input.keys():
+                _validate(autologging_call_input[key], user_call_input.get(key, None))
+
+        else:
+            assert (
+                autologging_call_input is user_call_input
+                or autologging_call_input == user_call_input
+            ), (
+                "Input to original function does not match expected input."
+                f" Original: '{autologging_call_input}'. Expected: '{user_call_input}'"
+            )
+
+    # Similar validation logic found in _validate, unraveling the list of arguments to exclude
+    # checks for any validation exempt positional arguments.
+    _assert_autologging_input_positional_args_are_superset(autologging_call_args, user_call_args)
+    for index, autologging_call_arg, user_call_arg in itertools.zip_longest(
+        range(len(user_call_args)), autologging_call_args, user_call_args
+    ):
+        if not _is_arg_exempt_from_validation(
+            autologging_integration,
+            function_name,
+            user_call_arg,
+            argument_index=index,
+        ):
+            _validate(autologging_call_arg, user_call_arg)
+
+    # Similar validation logic found in _validate, unraveling the dictionary of arguments to exclude
+    # checks for any validation exempt keyword arguments.
+    _assert_autologging_input_kwargs_are_superset(autologging_call_kwargs, user_call_kwargs)
+    for key in autologging_call_kwargs.keys():
+        if not _is_arg_exempt_from_validation(
+            autologging_integration,
+            function_name,
+            user_call_kwargs.get(key, None),
+            argument_name=key,
+        ):
+            _validate(
+                autologging_call_kwargs[key],
+                user_call_kwargs.get(key, None),
+            )