optimum/neuron/utils/patching.py

# coding=utf-8
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities related to patching."""

import functools
import importlib
import inspect
import sys
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Type, Union


if TYPE_CHECKING:
    from transformers import PreTrainedModel


class BasePatcher(ABC):
    """
    Base abstract class providing the core features for efficient context manager based patching.
    """

    def __init__(
        self, patching_specs: Optional[List[Tuple[Any, ...]]] = None, ignore_missing_attributes: bool = False
    ):
        self.patching_specs = self.process_patching_specs(
            patching_specs, ignore_missing_attributes=ignore_missing_attributes
        )
        self.already_patched = False

    @abstractmethod
    def process_patching_specs(
        self, patching_specs: Optional[List[Tuple[Any, Any]]] = None, ignore_missing_attributes: bool = False
    ) -> List[Tuple[Any, str, Any, Any, bool]]:
        pass

    def patch(self):
        if self.already_patched:
            return
        for module, attribute_name, _, patch, _ in self.patching_specs:
            setattr(module, attribute_name, patch)
        self.already_patched = True

    def restore(self):
        if not self.already_patched:
            return
        for module, attribute_name, orig, _, should_delete_attribute_at_restore in self.patching_specs:
            if should_delete_attribute_at_restore:
                delattr(module, attribute_name)
            else:
                setattr(module, attribute_name, orig)
        self.already_patched = False

    def __enter__(self):
        return self.patch()

    def __exit__(self, exc_type, exc_value, traceback):
        return self.restore()


class DynamicPatch:
    """
    Wrapper around a patch function.
    This can be used when the patch to apply is a function of the attribute it patches.
    """

    def __init__(self, patch_function: Callable[[Any], Any]):
        self.patch_function = patch_function

    def __call__(self, attribute: Any) -> Any:
        return self.patch_function(attribute)


class Patcher(BasePatcher):
    """
    Context manager that patches attributes of a module under its scope and restores everything after exit.
    """

    def process_patching_specs(
        self, patching_specs: Optional[List[Tuple[str, Any]]] = None, ignore_missing_attributes: bool = False
    ):
        processed_patching_specs = []
        for orig, patch in patching_specs or []:
            module_qualified_name, attribute_name = orig.rsplit(".", maxsplit=1)
            try:
                module = importlib.import_module(module_qualified_name)
            except ModuleNotFoundError as e:
                module_qualified_name, module_attribute_containing_attribute_name = module_qualified_name.rsplit(
                    ".", maxsplit=1
                )
                module = importlib.import_module(module_qualified_name)
                try:
                    module = getattr(module, module_attribute_containing_attribute_name)
                except AttributeError:
                    raise e

            module_has_attr = hasattr(module, attribute_name)
            if module_has_attr:
                attribute = getattr(module, attribute_name)
            elif ignore_missing_attributes and not isinstance(patch, DynamicPatch):
                attribute = None
            elif isinstance(patch, DynamicPatch):
                raise ValueError("Cannot ignore missing attribute with a DynamicPatch.")
            else:
                raise AttributeError(
                    f"Attribute {attribute_name} does not exist in {module}, set `ignore_missing_attributes=True` "
                    "to allow not failing when an attribute does not exist."
                )
            if isinstance(patch, DynamicPatch):
                patch = patch(attribute)
            processed_patching_specs.append((module, attribute_name, attribute, patch, not module_has_attr))
        return processed_patching_specs


class ModelPatcher(BasePatcher):
    """
    Context manager that patches attributes of a model under its scope and restores everything after exit.
    """

    def process_patching_specs(
        self,
        patching_specs: Optional[List[Tuple["PreTrainedModel", str, Any]]] = None,
        ignore_missing_attributes: bool = False,
    ):
        processed_patching_specs = []
        for model, attribute_qualified_name, patch in patching_specs or []:
            module_names = attribute_qualified_name.split(".")
            attribute_name = module_names.pop(-1)
            module = model
            for name in module_names:
                module = getattr(module, name)

            module_has_attr = hasattr(module, attribute_name)
            if module_has_attr:
                attribute = getattr(module, attribute_name)
            elif ignore_missing_attributes and not isinstance(patch, DynamicPatch):
                attribute = None
            elif isinstance(patch, DynamicPatch):
                raise ValueError("Cannot ignore missing attribute with a DynamicPatch.")
            else:
                raise AttributeError(
                    f"Attribute {attribute_name} does not exist in {module}, set `ignore_missing_attributes=True` "
                    "to allow not failing when an attribute does not exist."
                )

            if isinstance(patch, DynamicPatch):
                patch = patch(attribute)

            if inspect.ismethod(attribute):
                patch = patch.__get__(model)

            processed_patching_specs.append((module, attribute_name, attribute, patch, not module_has_attr))

        return processed_patching_specs


def patch_within_function(
    patching_specs: Union[List[Tuple[str, Any]], Tuple[str, Any]], ignore_missing_attributes: bool = False
):
    """
    Decorator that patches attributes of a module during the lifetime of the decorated function.

    Args:
        patching_specs (`Union[List[Tuple[str, Any]], Tuple[str, Any]]`):
            The specifications of what to patch.
        ignore_missing_attributes (`bool`, defaults to `False`):
            Whether or not the patch should fail if the attribute to patch does not exist.

    Returns:
        `Callable`: A patched version of the function.
    """
    if isinstance(patching_specs, tuple) and len(patching_specs) == 2:
        patching_specs = [patching_specs]

    patcher = Patcher(patching_specs, ignore_missing_attributes=ignore_missing_attributes)

    def decorator(func):
        is_bound = hasattr(func, "__self__")

        @functools.wraps(func.__func__ if is_bound else func)
        def wrapper(*args, **kwargs):
            with patcher:
                if is_bound:
                    args = args[1:]
                return func(*args, **kwargs)

        if is_bound:
            wrapper = wrapper.__get__(getattr(func, "__self__"))

        return wrapper

    return decorator


@functools.lru_cache()
def patch_everywhere(attribute_name: str, patch: Any, module_name_prefix: Optional[str] = None):
    """
    Finds all occurences of `attribute_name` in the loaded modules and patches them with `patch`.

    Args:
        attribute_name (`str`):
            The name of attribute to patch.
        patch (`Any`):
            The patch for the attribute.
        module_name_prefix (`Optional[str]`, defaults to `None`):
            If set, only module names starting with this prefix will be considered for patching.
    """
    for name, module in dict(sys.modules).items():
        if module_name_prefix is not None and not name.startswith(module_name_prefix):
            continue
        if hasattr(module, attribute_name):
            setattr(module, attribute_name, patch)


def replace_class_in_inheritance_hierarchy(obj: Any, orig_cls: Type, replacement_cls: Type):
    """
    Inspects the inheritance hierarchy of `obj` and replace `orig_cls` by `replacement_cls` if found.
    """
    to_visit = [obj.__class__]
    should_stop = False
    while to_visit and not should_stop:
        cls = to_visit.pop(0)
        if cls is object:
            continue
        bases = cls.__bases__
        new_bases = []
        for base in bases:
            to_visit.append(base)
            if base == orig_cls:
                new_bases.append(replacement_cls)
                should_stop = True
            elif base == replacement_cls:
                should_stop = True
                new_bases.append(base)
            else:
                new_bases.append(base)
        cls.__bases__ = tuple(new_bases)