import abc import functools import inspect import itertools import typing import uuid from contextlib import asynccontextmanager, contextmanager from typing import 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.tracking.client import MlflowClient 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.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 mlflow.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 (`.`) 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.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 (`.`) 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 = mlflow.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 = ( mlflow.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 = mlflow.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 mlflow.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 = mlflow.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 = ( mlflow.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 = mlflow.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 mlflow.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) """ 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(typing.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: typing.Callable 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), )